"""
 年度：dws_aba_word_heat_last365.py
计算 ABA 年度滚动词频热度 + 同比 + 主题标签，写入 Hive dws_aba_word_heat_last365 表

业务场景：
  前端"词频统计"页面支持"年度"切换：展示该单词在过去 12 个月的平均热度 + 年度同比
  本表是滚动型：每跑一次基准月，输出"基于该基准月往前 12 个月汇总"的结果
  仅 us 站点（业务要求）

聚合算法：
  年度热度 = sum(12个月 word_heat) / count(distinct 实际出现的月份)
  - 动态分母（不是固定除以 12）：照顾季节性词（如 christmas 只在 11-12 月出现）
  - 不除空缺月份，反映该词在出现的月份里的平均热度

年度同比：
  (今年年度热度 - 去年年度热度) / 去年年度热度
  - 去年年度热度同算法，从 dws_aba_word_heat 的去年 12 个月分区聚出来
  - 新词（去年 12 个月都没该词）→ 1000（与月度同款占位）

theme_ch_list 取 12 个月合并去重：
  跟 dwt_aba_last365 处理 st_attribute_label 同思路——把 12 个月的 theme_ch_list 全部合并去重
  原因：原来只取最新月口径,但基准月该词没出现就丢标签(假阴性,前端筛选漏数据)
  改造：collect_set + flatten + array_distinct + sort_array,保留所有历史月命中的主题

"""

import os
import sys

sys.path.append(os.path.dirname(sys.path[0]))
from utils.templates import Templates
from utils.common_util import CommonUtil
from utils.hdfs_utils import HdfsUtils
from pyspark.sql import functions as F


class DwsAbaWordHeatLast365(Templates):

    def __init__(self, site_name="us", date_type="month", date_info="2026-05"):
        super().__init__()
        # 业务确认：本表仅 us 站点
        assert site_name == 'us', "dws_aba_word_heat_last365 仅 us 站点跑"
        self.site_name = site_name
        self.date_type = date_type
        self.date_info = date_info
        self.db_save = "dws_aba_word_heat_last365"
        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)

        # 过去12个月 list：基准月含本月，往前 11 个月
        # 例：date_info='2026-05' → [2026-05, 2026-04, ..., 2025-06]
        self.last_12_month = [CommonUtil.get_month_offset(self.date_info, -i) for i in range(12)]
        # 再往前12个月 list（用于年度同比）
        # 例：[2025-05, 2025-04, ..., 2024-06]
        self.prev_12_month = [CommonUtil.get_month_offset(self.date_info, -12 - i) for i in range(12)]

        # 初始化全局 df
        self.df_curr_year = self.spark.sql("select 1+1;")   # 当年 12 个月明细
        self.df_prev_year = self.spark.sql("select 1+1;")   # 去年 12 个月明细
        self.df_theme_latest = self.spark.sql("select 1+1;")  # 基准月 theme_ch_list

    def read_data(self):
        # 1) has_data assert 基准月（必须有数据）
        # 历史月份缺失允许（影响分母但不阻断），基准月缺失直接报错
        has_curr = self.spark.sql(f"""
            SELECT 1 FROM dws_aba_word_heat
            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"上游 dws_aba_word_heat 基准月分区无数据 "
            f"site_name={self.site_name}/date_type={self.date_type}/date_info={self.date_info}，"
            f"请先确认 dws_aba_word_heat 月度调度是否完成"
        )

        # 2) 读当年 12 个月 dws_aba_word_heat（用于算年度平均）
        sql_curr = f"""
            select word, word_heat, date_info
            from dws_aba_word_heat
            where site_name='{self.site_name}'
              and date_type='{self.date_type}'
              and date_info in ({CommonUtil.list_to_insql(self.last_12_month)})
        """
        self.df_curr_year = self.spark.sql(sql_curr).repartition(40, 'word').cache()
        print(f"self.df_curr_year (当年 12 个月，{self.last_12_month[-1]} ~ {self.last_12_month[0]}):")
        self.df_curr_year.show(10, truncate=True)

        # 3) 读去年 12 个月 dws_aba_word_heat（用于算年度同比）
        # 启动期前 24 个月数据可能不齐，缺失月份不参与聚合，影响精度但不阻断
        sql_prev = f"""
            select word, word_heat, date_info
            from dws_aba_word_heat
            where site_name='{self.site_name}'
              and date_type='{self.date_type}'
              and date_info in ({CommonUtil.list_to_insql(self.prev_12_month)})
        """
        self.df_prev_year = self.spark.sql(sql_prev).repartition(40, 'word').cache()
        print(f"self.df_prev_year (去年 12 个月，{self.prev_12_month[-1]} ~ {self.prev_12_month[0]}):")
        self.df_prev_year.show(10, truncate=True)

        # 4) 读 12 个月 theme_ch_list,后续在 handle_data 里按 word 合并去重
        # 原来只读基准月(最新月口径),但基准月该词没出现就丢标签(假阴性,前端筛选漏数据)
        # 改成 12 月合并去重,保留所有历史月命中的主题(与 dwt_aba_last365 的 st_attribute_label 同步改造)
        sql_theme = f"""
            select word, theme_ch_list
            from dws_aba_word_heat
            where site_name='{self.site_name}'
              and date_type='{self.date_type}'
              and date_info in ({CommonUtil.list_to_insql(self.last_12_month)})
        """
        self.df_theme_latest = self.spark.sql(sql_theme).repartition(40, 'word').cache()
        print(f"self.df_theme_latest (12 个月,{self.last_12_month[-1]} ~ {self.last_12_month[0]}):")
        self.df_theme_latest.show(10, truncate=False)

    def handle_data(self):
        # 1) 当年年度热度 = sum(word_heat) / count(distinct 实际出现的月份)
        # 用 count(distinct date_info) 做动态分母，照顾季节性词
        df_curr_agg = self.df_curr_year.groupBy('word').agg(
            F.sum('word_heat').alias('curr_total'),
            F.countDistinct('date_info').alias('curr_months')
        ).withColumn(
            'curr_avg_heat',
            F.round(F.col('curr_total') / F.col('curr_months'), 4)
        ).select('word', 'curr_avg_heat')

        # 2) 去年年度热度（同算法）
        df_prev_agg = self.df_prev_year.groupBy('word').agg(
            F.sum('word_heat').alias('prev_total'),
            F.countDistinct('date_info').alias('prev_months')
        ).withColumn(
            'prev_avg_heat',
            F.round(F.col('prev_total') / F.col('prev_months'), 4)
        ).select('word', 'prev_avg_heat')

        # 3) 当年 left join 去年算同比
        # 新词（去年 12 个月都没该词）→ word_heat_change_rate=1000
        df_heat = df_curr_agg.join(
            df_prev_agg, on='word', how='left'
        ).withColumn(
            'word_heat_change_rate',
            F.when(
                F.col('prev_avg_heat').isNull(), F.lit(1000.0)
            ).otherwise(
                F.round((F.col('curr_avg_heat') - F.col('prev_avg_heat')) / F.col('prev_avg_heat'), 4)
            )
        ).withColumnRenamed('curr_avg_heat', 'word_heat').drop('prev_avg_heat')

        # 4) 按 word 合并 12 个月所有非 -1 主题,去重后重新拼串(同 dwt_aba_last365 的 st_attribute_label 处理)
        # 步骤: 单月逗号串 split → flatten 跨月合并 → array_distinct 去重 → sort_array 排序 → concat_ws 拼串
        # 过滤 '-1' 占位,防止合并进有效结果;12 月全 -1 时输出空串,nullif 转 null 后由 fillna 转 '-1'
        self.df_theme_latest = self.df_theme_latest.groupBy('word').agg(
            F.expr("""
                nullif(
                    concat_ws(',',
                        sort_array(array_distinct(flatten(
                            collect_set(
                                case when theme_ch_list != '-1' then split(theme_ch_list, ',') end
                            )
                        )))
                    ),
                    ''
                )
            """).alias('theme_ch_list')
        )
        df_heat = df_heat.join(
            self.df_theme_latest, on='word', how='left'
        ).fillna({'theme_ch_list': '-1'})

        # 5) 入库前字段处理（精度按 PG numeric(20,4) 对齐）
        self.df_save = df_heat.select(
            F.col('word'),
            F.round('word_heat', 4).alias('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')
        )

        # 释放上游 cache
        self.df_curr_year.unpersist()
        self.df_prev_year.unpersist()
        self.df_theme_latest.unpersist()

    def save_data(self):
        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
        )
        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 = DwsAbaWordHeatLast365(site_name=site_name, date_type=date_type, date_info=date_info)
    obj.run()
