"""
   @Author      : hejiangming
   @Description : 计算 ABA 月度词频热度 + 同比 + 主题标签
   @SourceTable :
                  ①dwt_aba_st_analytics
                  ②keywords_match_theme（PG）
                  ③match_character_dict（MySQL）
   @SinkTable   : dws_aba_word_heat
   @CreateTime  : 2026/06/11 18:27
   @UpdateTime  : 2026/06/16 16:55
"""

import os
import re
import sys

import inflect
import pandas as pd

sys.path.append(os.path.dirname(sys.path[0]))
from utils.templates import Templates
from utils.common_util import CommonUtil
from utils.word_normalize import to_base_form, clean_word  # 归一/清洗统一公共模块
from utils.hdfs_utils import HdfsUtils
from utils.spark_util import SparkUtil
from utils.db_util import DBUtil
from pyspark.sql import functions as F


class DwsAbaWordHeat(Templates):

    def __init__(self, site_name="us", date_type="month", date_info="2026-05"):
        super().__init__()
        self.site_name = site_name
        self.date_type = date_type
        self.date_info = date_info
        self.db_save = "dws_aba_word_heat"
        self.spark = self.create_spark_object(
            app_name=f"{self.db_save}: {self.site_name},{self.date_type},{self.date_info}")

        # 写入、分区初始化
        self.df_save = self.spark.sql("select 1+1;")
        self.partitions_by = ['site_name', 'date_type', 'date_info']
        self.reset_partitions(partitions_num=10)

        # 去年同月（同比基准）
        self.last_year_month = CommonUtil.get_month_offset(self.date_info, -12)

        # 初始化全局 df 和 sp_symbols
        self.sp_symbols = []
        self.df_curr = self.spark.sql("select 1+1;")          # 当月 search_term + rank
        self.df_last = self.spark.sql("select 1+1;")          # 去年同月 search_term + rank
        self.df_theme_agg = self.spark.sql("select 1+1;")     # 词典 word + theme_ch_list（聚合后）
        self.df_word_to_base = self.spark.sql("select 1+1;")  # variant_word → base_form（全量映射）

    def read_data(self):
        # ========== 1) 读当月 dwt_aba_st_analytics ==========
        # 提前用 take(1) 校验有数据，避免空分区导致全表静默错误
        # take(1) 而非 count()：LZO 文本格式下 count 会扫全表，take 读到一行就返回
        has_curr = self.spark.sql(f"""
            SELECT 1 FROM dwt_aba_st_analytics
            WHERE site_name='{self.site_name}'
              AND date_type='{self.date_type}'
              AND date_info='{self.date_info}'
            LIMIT 1
        """).take(1)
        assert len(has_curr) > 0, (
            f"上游 dwt_aba_st_analytics 当月分区无数据 "
            f"site_name={self.site_name}/date_type={self.date_type}/date_info={self.date_info}，"
            f"请先确认上游调度是否完成"
        )

        sql_curr = f"""
            select search_term, rank
            from dwt_aba_st_analytics
            where site_name = '{self.site_name}'
              and date_type = '{self.date_type}'
              and date_info = '{self.date_info}'
              and rank > 0
              and (st_bsr_cate_1_id_new not in ('audible','books','digital-text','dmusic','mobile-apps','movies-tv','music','software','videogames')
                   or st_bsr_cate_1_id_new is null)
              and (is_self_max_num_asin <> 1 or is_self_max_num_asin is null)
        """
        self.df_curr = self.spark.sql(sql_curr).repartition(40, 'search_term').cache()
        print("self.df_curr:")
        self.df_curr.show(10, truncate=True)

        # ========== 2) 读去年同月 dwt_aba_st_analytics（同比基准） ==========
        # 业务确认：us/uk/de 三站点 2025 年都有数据，所以不做"分区不存在"兜底
        sql_last = f"""
            select search_term, rank
            from dwt_aba_st_analytics
            where site_name = '{self.site_name}'
              and date_type = '{self.date_type}'
              and date_info = '{self.last_year_month}'
              and rank > 0
              and (st_bsr_cate_1_id_new not in ('audible','books','digital-text','dmusic','mobile-apps','movies-tv','music','software','videogames')
                   or st_bsr_cate_1_id_new is null)
              and (is_self_max_num_asin <> 1 or is_self_max_num_asin is null)
        """
        self.df_last = self.spark.sql(sql_last).repartition(40, 'search_term').cache()
        print(f"self.df_last (去年同月 date_info={self.last_year_month}):")
        self.df_last.show(10, truncate=True)

        # ========== 3) 读 PG 词典 keywords_match_theme ==========
        # 21821 更新：词典源从 MySQL selection.aba_match_theme 改为 PG 集群 public.keywords_match_theme
        # 与 dws_st_theme 同源，单词级 word → theme_ch 映射
        sql_theme = """
            select label_en_lower as word, theme_ch
            from public.keywords_match_theme
            where label_ch is not null
              and label_en_lower is not null
              and theme_ch is not null
        """
        pg_conn = DBUtil.get_connection_info("postgresql_cluster", "us")
        df_theme_raw = SparkUtil.read_jdbc_query(
            session=self.spark, url=pg_conn["url"],
            pwd=pg_conn["pwd"], username=pg_conn["username"],
            query=sql_theme
        )
        # 词典 theme_ch 可能是逗号双值（如 "主题,风格"），且中英文逗号都可能出现。
        # 先把全角逗号统一成半角，下游 _build_base_theme 的 split(theme_ch_list, ',') 才能把双值拆成原子主题；
        # 否则中文逗号 "主题，风格" split 不开 → 整块变成一个废主题 token（与 dws_st_theme 同一个病）。
        df_theme_raw = df_theme_raw.withColumn(
            "theme_ch", F.regexp_replace("theme_ch", "，", ",")
        )
        # 同一 word 在词典里可能有多行主题（如某词既是材质也是细分人群）
        # 这里先按 word 聚合（sort_array 保证幂等）；后续 _build_base_theme 还会再做一次 base 维度聚合
        self.df_theme_agg = df_theme_raw.groupBy('word').agg(
            F.concat_ws(',', F.sort_array(F.collect_set('theme_ch'))).alias('theme_ch_list')
        ).cache()
        print("self.df_theme_agg:")
        self.df_theme_agg.show(10, truncate=False)

        # ========== 4) 读 MySQL 特殊字符列表 sp_symbols（当前不启用过滤） ==========
        # 前端按空格拆分时保留孤立特殊字符（?/!/& 这种）作为"词"，大数据保持口径一致
        # 代码保留：后续如果业务要求开启过滤，去掉 _calc_word_heat 里 sp_symbols 那段的注释即可
        sql_sp = """
            select character_name
            from match_character_dict
            where match_type = '特殊字符'
        """
        mysql_conn = DBUtil.get_connection_info("mysql", "us")
        sp_df = SparkUtil.read_jdbc_query(
            session=self.spark, url=mysql_conn["url"],
            pwd=mysql_conn["pwd"], username=mysql_conn["username"],
            query=sql_sp
        )
        self.sp_symbols = [row['character_name'] for row in sp_df.collect()]
        print(f"sp_symbols（当前未启用过滤）: {self.sp_symbols}")

        # ========== 5) 构建全量 word → base_form 映射表 ==========
        # 来源：拆词侧 distinct + 词典侧 word，union 后 distinct，保证拆词和词典走同一份 base
        # 这样拆词的 base 和词典的 base 一定对得齐，theme 查找不会因归一不一致而 miss
        self.df_word_to_base = self._build_word_to_base_mapping()

    def _build_word_to_base_mapping(self):
        """
        构建权威 word → base_form 映射表（全量，含拆词侧 + 词典侧），返回 broadcast df。
        保证拆词的 base 和词典的 base 一定对得齐，防止同一个词在两侧归一结果不一致导致 theme 查不到。
        """
        # ============================================================
        # Step 1：拆词侧 + 词典侧 union 后 distinct
        # 为什么 union：要保证拆词的 base 和词典的 base 用同一份映射规则
        # 量级（us 站点）：200w 搜索词 × 平均 3-4 词 ≈ 600-800w explode 行 → distinct 后约 50-100w 词
        # ============================================================
        df_curr_words = self.df_curr.select(
            F.explode(F.split(F.lower('search_term'), ' ')).alias('word')
        )
        df_last_words = self.df_last.select(
            F.explode(F.split(F.lower('search_term'), ' ')).alias('word')
        )
        df_theme_words = self.df_theme_agg.select('word')

        # cache 原因：下面会跑两个 action（count + toPandas），不 cache 会触发两次完整的 distinct 计算
        # 200w 搜索词 union distinct 比较重，cache 能减半 shuffle
        df_all = df_curr_words.union(df_last_words).union(df_theme_words) \
            .filter((F.col('word') != '') & F.col('word').isNotNull()) \
            .distinct() \
            .cache()

        cnt = df_all.count()
        print(f"distinct word 总数（拆词+词典 union）: {cnt}")

        # ============================================================
        # Step 2：toPandas 拉到 driver，单线程跑 to_base_form 归一
        # 为什么在 driver 跑：inflect/WordNet 没有 spark UDF 版本，每个词 ~1ms，100w 词约 70 秒可接受
        # 跑完即 unpersist，释放 executor 内存
        # ============================================================
        pdf = df_all.toPandas()
        df_all.unpersist()

        p = inflect.engine()
        # 先 clean_word 清洗（None=丢弃，脏词不进映射 → join 时被丢），再归一 base_form
        mappings = []
        for w in pdf['word']:
            cw = clean_word(w)
            if cw is None:
                continue
            mappings.append((w, cw, to_base_form(cw, p)))

        # ============================================================
        # Step 3：打印归一示例方便调试（只用 pandas 做筛选展示，不参与 createDataFrame）
        # 跑批日志里能看到前 20 条"被归一的"词，业务/工程师快速抽查
        # ============================================================
        pdf_map = pd.DataFrame(mappings, columns=['word', 'clean_word', 'base_form'])
        diff_sample = pdf_map[
            (pdf_map['word'] != pdf_map['clean_word']) | (pdf_map['clean_word'] != pdf_map['base_form'])
        ].head(20)
        print("word → clean_word → base_form 示例（前 20 条有变化的）：")
        print(diff_sample.to_string(index=False))

        # ============================================================
        # Step 4：list of tuples + 显式 schema 创建 DataFrame，再 broadcast
        # 为什么不走 createDataFrame(pdf)：pandas 2.x 移除了 DataFrame.iteritems()，
        #   PySpark 3.1.x 内部还在调用，会报 "'DataFrame' object has no attribute 'iteritems'"
        #   走 list + schema 方式与 pandas 版本完全解耦
        # 为什么 broadcast：后续 4 处 join（拆词归一 ×2 + 词典归一 + 反向展开）都用这份映射
        #   小表（10w 级别，几 MB），broadcast 无 shuffle
        # ============================================================
        df_map = self.spark.createDataFrame(mappings, schema=['word', 'clean_word', 'base_form'])
        return F.broadcast(df_map)

    def _calc_word_heat(self, df_aba):
        """
        拆词 + 归一到 base_form + 按 base_form 聚合热度。
        热度公式：sum(100000 / rank)，单复数会在 base_form 维度合并（girl/girls 共享同一份热度）。
        """
        # ============================================================
        # Step 1：拆词（lower + split + explode）
        # 与前端拆词口径完全一致（前端就是按空格拆 + 转小写）
        # 不去重叠词：'shoe shoe rack' 让 'shoe' 累加两次热度（前端词频也会算两次）
        # ============================================================
        df_words = df_aba.select(
            F.explode(F.split(F.lower('search_term'), ' ')).alias('word'),
            F.col('rank')
        ).filter((F.col('word') != '') & F.col('word').isNotNull())

        # 当前前端不过滤孤立特殊字符（?/!/&），大数据保持一致
        # 后续如需启用打开下面注释，给词典维护者控制
        # if self.sp_symbols:
        #     df_words = df_words.filter(~F.col('word').isin(self.sp_symbols))

        # ============================================================
        # Step 2：broadcast join 归一到 base_form（inner）
        # 为什么 broadcast：word_to_base 是 10w 行级小表，避免大表 shuffle
        # 映射含拆词侧 + 词典侧所有 distinct word，拆词的 word 几乎都能 join 上；
        #   只有 clean_word=None 的脏词（中韩 / 纯符号）没进映射 → inner join 在这里把它们丢弃
        # 例：'sweaters' → base_form='sweater'（归一）；中韩 / 纯符号词不在映射 → 被丢弃
        # ============================================================
        df_normalized = df_words.join(
            self.df_word_to_base, on='word', how='inner'
        )

        # ============================================================
        # Step 3：按 base_form 聚合热度
        # girl/girls 都被归一到 base='girl' → 在这里合并热度
        # 公式：sum(100000 / rank)
        # 例：'girl' 出现在 rank=10 和 rank=100 → heat = 100000/10 + 100000/100 = 11000
        # ============================================================
        return df_normalized.groupBy('base_form').agg(
            F.sum(F.lit(100000.0) / F.col('rank')).alias('word_heat')
        )

    def _build_base_theme(self):
        """
        词典侧按 base_form 聚合 theme_ch_list。
        修复早期 bug：同 base 多词典行（如 girl→'人群' + girls→'女性'）的 theme 会互相覆盖。
        """
        # ============================================================
        # Step 1：把聚合好的 theme_ch_list 重新 explode 成单个 theme 一行
        # 为什么 explode：方便后续按 base_form 维度重新 collect_set 合并
        # 输入：df_theme_agg(word='girl', theme_ch_list='人群,女孩气')
        # explode 后：(word='girl', theme='人群'), (word='girl', theme='女孩气')
        # ============================================================
        df_exploded = self.df_theme_agg \
            .withColumn('theme', F.explode(F.split('theme_ch_list', ','))) \
            .filter(F.col('theme') != '')

        # ============================================================
        # Step 2：broadcast join 把 word 归一到 base_form
        # 跟拆词侧用同一份映射，保证 base 对齐
        # 例：word='girls' → base_form='girl'，word='girl' → base_form='girl'
        # ============================================================
        df_normalized = df_exploded \
            .join(self.df_word_to_base, on='word', how='inner')

        # ============================================================
        # Step 3：按 base_form 聚合所有 theme tokens
        # 关键修复点：collect_set 自然合并不同词典行的 theme（不会再覆盖）
        #   例：词典 A 行：girl→'人群'，B 行：girls→'女孩气'
        #   归一后两者 base_form 都是 'girl'，collect_set 得到 {'人群', '女孩气'}
        #   早期 bug 版本用 (base, variant) 去重，B 行的 '女孩气' 会被跳过丢失
        # sort_array 保证幂等：concat_ws 拼接结果稳定，重跑不会因顺序不同导致 diff
        # ============================================================
        return df_normalized.groupBy('base_form').agg(
            F.concat_ws(',', F.sort_array(F.collect_set('theme'))).alias('theme_ch_list')
        )

    def handle_data(self):
        # ========== 1) 算当月 + 去年同月的 base_form 维度热度 ==========
        df_curr_heat = self._calc_word_heat(self.df_curr)
        df_last_heat = self._calc_word_heat(self.df_last)

        # ========== 2) 同比 = (当月-去年)/去年；新词→1000（与 dwt_aba_last_change_rate 惯例一致） ==========
        # 左 join：保留所有当月有热度的 base；去年没有的 → last_heat 为 NULL → 同比填 1000
        df_heat = df_curr_heat.withColumnRenamed('word_heat', 'curr_heat').join(
            df_last_heat.withColumnRenamed('word_heat', 'last_heat'),
            on='base_form', how='left'
        ).withColumn(
            'word_heat_change_rate',
            F.when(
                F.col('last_heat').isNull(), F.lit(1000.0)
            ).otherwise(
                F.round((F.col('curr_heat') - F.col('last_heat')) / F.col('last_heat'), 4)
            )
        ).select(
            'base_form',
            F.col('curr_heat').alias('word_heat'),
            'word_heat_change_rate'
        )

        # ========== 3) 词典按 base_form 聚合 theme（修复 theme 互相覆盖 bug） ==========
        df_base_theme = self._build_base_theme()

        # ========== 4) 反向展开：从当月拆词侧收集所有 (word, base_form) variant 组合 ==========
        # 只展开"实际出现过的 variant"——前端按 word 查询，拆不到的词无需输出
        # 例：词典录 girls 但当月 ABA 只出现 girl → 只输出 word='girl'，不输出 word='girls'
        # ★口径(C1)：word 列存【原始词】(不是 clean_word)，让前端用原始拆词(towel®/men's)能直接查到；
        #   inner join 仍会丢掉 clean=None 的脏词(中韩/纯符号)；
        #   热度/同比/主题按 base_form(clean_word 归一后)合并 → towel® 这行显示的是所有 towel 的热度
        df_curr_variants = self.df_curr.select(
            F.explode(F.split(F.lower('search_term'), ' ')).alias('word')
        ).filter((F.col('word') != '') & F.col('word').isNotNull()) \
         .distinct() \
         .join(self.df_word_to_base, on='word', how='inner') \
         .select('word', 'base_form') \
         .distinct()

        # ========== 5) 组装最终结果 ==========
        # df_heat(base 维) → join theme(base 维) → join variants(base→多 variant 展开) → 每 variant 一行
        # inner join variants：df_curr_variants 的 base ⊆ df_heat 的 base，inner 与 left 等价但更直观
        df_final = df_heat \
            .join(df_base_theme, on='base_form', how='left') \
            .fillna({'theme_ch_list': '-1'}) \
            .join(df_curr_variants, on='base_form', how='inner') \
            .select(
                F.col('word'),
                F.round('word_heat', 4).alias('word_heat'),
                F.col('word_heat_change_rate'),
                F.col('theme_ch_list')
            )

        # ========== 6) 入库前字段处理（精度对齐 PG numeric(20,4)，附 site/date 分区字段） ==========
        self.df_save = df_final.select(
            F.col('word'),
            F.col('word_heat'),
            F.col('word_heat_change_rate'),
            F.col('theme_ch_list'),
            F.date_format(F.current_timestamp(), 'yyyy-MM-dd HH:mm:SS').alias('created_time'),
            F.date_format(F.current_timestamp(), 'yyyy-MM-dd HH:mm:SS').alias('updated_time'),
            F.lit(self.site_name).alias('site_name'),
            F.lit(self.date_type).alias('date_type'),
            F.lit(self.date_info).alias('date_info')
        )

    def save_data(self):
        # 清 HDFS 分区目录后写入（与项目其他 dws/dwt 脚本保持一致）
        # 先清是为了避免 saveAsTable append 模式下旧文件残留导致重复
        hdfs_path = CommonUtil.build_hdfs_path(
            self.db_save,
            partition_dict={
                'site_name': self.site_name,
                'date_type': self.date_type,
                'date_info': self.date_info
            }
        )
        print(f"清除 HDFS 目录: {hdfs_path}")
        HdfsUtils.delete_file_in_folder(hdfs_path)

        self.df_save = self.df_save.repartition(10)
        print(f"当前存储的表名为：{self.db_save}, 分区为 {self.partitions_by}")
        self.df_save.write.saveAsTable(
            name=self.db_save, format='hive', mode='append', partitionBy=self.partitions_by
        )

        # 注意：unpersist 必须放在 saveAsTable 之后
        # 原版放在 handle_data 末尾的 bug：handle_data 只构造 lazy DataFrame，saveAsTable 才触发计算
        # 提前 unpersist 会导致 saveAsTable 阶段 cache 已被标记释放 → 重新扫上游 → cache 失效
        self.df_curr.unpersist()
        self.df_last.unpersist()
        self.df_theme_agg.unpersist()

        print("success")


if __name__ == '__main__':
    site_name = CommonUtil.get_sys_arg(1, None)
    date_type = CommonUtil.get_sys_arg(2, None)
    date_info = CommonUtil.get_sys_arg(3, None)
    assert site_name is not None, "site_name 不能为空"
    assert date_type is not None, "date_type 不能为空"
    assert date_info is not None, "date_info 不能为空"
    obj = DwsAbaWordHeat(site_name=site_name, date_type=date_type, date_info=date_info)
    obj.run()
