#!/usr/bin/env python3
"""
气候季节起始日分析脚本
依据：《气候季节划分指标》标准（GB/T 41538-2022）
- 常年法：基于历史多年平均气温，计算常年春/夏/秋/冬起始日
- 当年法：基于当年气温数据，结合常年参考进行初判与二次判断
"""

import csv
import json
from collections import defaultdict
from datetime import datetime
from typing import Optional


# ---------------------------------------------------------------------------
# 工具函数
# ---------------------------------------------------------------------------

def month_day_to_doy(month: int, day: int) -> int:
    """月份+日期 -> 一年中的第几天（1-based，按非闰年算）"""
    days_in_month = [0, 31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]
    return sum(days_in_month[:month]) + day


def doy_to_date(doy: int, year: int = 2023) -> str:
    """日序 -> YYYY-MM-DD"""
    days_in_month = [0, 31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]
    m = 1
    while doy > days_in_month[m]:
        doy -= days_in_month[m]
        m += 1
    return f"{year}-{m:02d}-{doy:02d}"


def day_diff(date1: str, date2: str) -> int:
    """两个 YYYY-MM-DD 字符串的天数差（date2 - date1）"""
    d1 = datetime.strptime(date1, "%Y-%m-%d")
    d2 = datetime.strptime(date2, "%Y-%m-%d")
    return (d2 - d1).days


# ---------------------------------------------------------------------------
# 1. 数据加载
# ---------------------------------------------------------------------------

def load_temperature_data(csv_path: str) -> dict[str, dict[str, float]]:
    """
    读取 CSV，过滤掉 mean_temp 为空的行。
    返回 {city: {date_str: mean_temp}}
    """
    data: dict[str, dict[str, float]] = defaultdict(dict)
    with open(csv_path, encoding="utf-8-sig") as f:
        reader = csv.reader(f)
        next(reader)  # skip header
        for row in reader:
            if len(row) < 3 or not row[2].strip():
                continue
            data[row[1]][row[0]] = float(row[2])
    return dict(data)


# ---------------------------------------------------------------------------
# 2. 常年日平均气温序列（公式1）
# ---------------------------------------------------------------------------

def compute_clim_daily_avg(
    city_data: dict[str, float],
    years: list[str],
) -> list[float]:
    """
    对选定年份逐日（日序 1~365）求平均，得常年日平均气温序列。
    返回列表索引 j 对应日序 j+1（0-based）。
    """
    totals = [0.0] * 366
    counts = [0] * 366

    for year in years:
        yd = {k: v for k, v in city_data.items() if k.startswith(year)}
        for date_str, temp in yd.items():
            m = int(date_str[5:7])
            d = int(date_str[8:10])
            doy = month_day_to_doy(m, d)
            if doy == 60:  # skip leap Feb 29
                continue
            totals[doy] += temp
            counts[doy] += 1

    return [
        round(totals[j] / counts[j], 2) if counts[j] > 0 else None
        for j in range(1, 366)
    ]


# ---------------------------------------------------------------------------
# 3. 5日滑动平均（公式2，首尾相接）
# ---------------------------------------------------------------------------

def moving_avg_5d(series: list[float]) -> list[float]:
    """
    计算5日滑动平均序列。
    边界处用相邻年份数据补齐（首尾相接），保留2位小数。
    返回列表与输入等长，索引 j 对应日序 j+1。
    """
    n = len(series)
    # 前补2个、后补4个（需要4个来覆盖最后两天的右边界窗口）
    padded = series[-2:] + series + series[:4]
    avg = []
    for i in range(n):
        vals = [padded[i + 2], padded[i + 3], padded[i + 4],
                padded[i + 5], padded[i + 6]]
        non_none = [v for v in vals if v is not None]
        if non_none:
            avg.append(round(sum(non_none) / len(non_none), 2))
        else:
            avg.append(None)
    return avg


# ---------------------------------------------------------------------------
# 4. 找到首次连续W天满足条件的起始索引（0-based）
# ---------------------------------------------------------------------------

def find_season_start(
    avg_list: list[float],
    threshold: float,
    direction: str,
    search_start_doy: int = 1,
) -> Optional[int]:
    """
    在序列中找到首次连续5天满足条件的起始日序。
    direction='above': >= threshold（春/夏）
    direction='below': < threshold  （秋/冬）
    search_start_doy: 从该日序开始搜索（避免跨年边界误判），默认从1开始。
    返回日序（1-based）或 None。
    """
    streak = 0
    for j, v in enumerate(avg_list):
        doy = j + 1  # 1-based day of year
        if doy < search_start_doy:
            streak = 0
            continue
        if v is None:
            streak = 0
            continue
        if direction == "above" and v >= threshold:
            streak += 1
        elif direction == "below" and v < threshold:
            streak += 1
        else:
            streak = 0
        if streak == 5:
            return doy - 4  # convert to 1-based day-of-year of first day
    return None


# ---------------------------------------------------------------------------
# 5. 常年气候季节分析
# ---------------------------------------------------------------------------

def analyze_climate_season(
    city_data: dict[str, float],
    hist_years: list[str],
) -> dict:
    daily_avg = compute_clim_daily_avg(city_data, hist_years)
    avg_5d = moving_avg_5d(daily_avg)

    # 秋季搜索从夏季起始日后开始，冬季搜索从秋季起始日后开始
    autumn_start = (climate_ref_summer_doy or 1) + 1 if False else None
    # 由于我们还没有计算 climate_ref，用保守估计：
    # 秋季从夏季日后30天开始搜索（roughly June+30 = July），冬季从秋季日后30天开始
    summer_doy_val = find_season_start(avg_5d, 22.0, "above")
    autumn_doy_val = find_season_start(avg_5d, 22.0, "below", search_start_doy=(summer_doy_val or 1) + 30)
    winter_doy_val = find_season_start(avg_5d, 10.0, "below", search_start_doy=(autumn_doy_val or 1) + 30)

    return {
        "daily_avg": daily_avg,
        "avg_5d":    avg_5d,
        "spring_doy":   find_season_start(avg_5d, 10.0, "above"),
        "summer_doy":   summer_doy_val,
        "autumn_doy":   autumn_doy_val,
        "winter_doy":   winter_doy_val,
    }


# ---------------------------------------------------------------------------
# 6. 当年季节起始日分析（含初判 & 二次判断）
# ---------------------------------------------------------------------------

def compute_current_year_avg_5d(
    city_data: dict[str, float],
    target_year: str,
) -> tuple[list[str], list[float]]:
    """
    计算目标年每日5日滑动平均。
    返回 (sorted_dates, avg_5d_list)，avg_5d_list 与 sorted_dates 一一对应。
    """
    yd = {k: v for k, v in city_data.items() if k.startswith(target_year)}
    sorted_dates = sorted(yd.keys())
    temps = [yd[d] for d in sorted_dates]
    n = len(temps)
    # 首尾相接补齐
    padded = temps[-2:] + temps + temps[:4]
    avg_5d = []
    for i in range(n):
        vals = [padded[i + 2], padded[i + 3], padded[i + 4],
                padded[i + 5], padded[i + 6]]
        non_none = [v for v in vals if v is not None]
        if non_none:
            avg_5d.append(round(sum(non_none) / len(non_none), 2))
        else:
            avg_5d.append(None)
    return sorted_dates, avg_5d


def find_all_season_starts(
    avg_list: list[float],
    threshold: float,
    direction: str,
) -> list[int]:
    """
    找出所有首次连续5日满足条件的起始日期（0-based 索引）。
    用于收集初判及所有二次/多次候选。
    """
    candidates = []
    streak = 0
    for j, v in enumerate(avg_list):
        if v is None:
            streak = 0
            continue
        if direction == "above" and v >= threshold:
            streak += 1
        elif direction == "below" and v < threshold:
            streak += 1
        else:
            streak = 0
        if streak == 5:
            candidates.append(j - 4)  # 0-based index of first day
            streak = 0  # 不重叠，继续搜索
    return candidates


def analyze_current_year(
    city_data: dict[str, float],
    target_year: str,
    climate_info: dict,
) -> dict:
    sorted_dates, avg_5d_curr = compute_current_year_avg_5d(city_data, target_year)
    n = len(sorted_dates)

    spring_candidates = find_all_season_starts(avg_5d_curr, 10.0, "above")
    summer_candidates = find_all_season_starts(avg_5d_curr, 22.0, "above")

    # 初判
    initial_spring = sorted_dates[spring_candidates[0]] if spring_candidates else None
    initial_summer = sorted_dates[summer_candidates[0]] if summer_candidates else None

    # 二次/多次候选（初判之后的所有候选）
    secondary_spring = [sorted_dates[i] for i in spring_candidates[1:]]
    secondary_summer = [sorted_dates[i] for i in summer_candidates[1:]]

    # 与常年对比
    ref_spring_doy = climate_info["spring_doy"]
    ref_summer_doy = climate_info["summer_doy"]

    spring_diff_days: Optional[int] = None
    summer_diff_days: Optional[int] = None
    if initial_spring and ref_spring_doy:
        ref_date = doy_to_date(ref_spring_doy, int(target_year))
        spring_diff_days = day_diff(ref_date, initial_spring)
    if initial_summer and ref_summer_doy:
        ref_date = doy_to_date(ref_summer_doy, int(target_year))
        summer_diff_days = day_diff(ref_date, initial_summer)

    # 最终判定（按 6.1.2 规则）
    final_spring = initial_spring
    final_summer = initial_summer
    spring_note = ""
    summer_note = ""

    if initial_spring and ref_spring_doy:
        if spring_diff_days > 30:
            if secondary_spring:
                final_spring = secondary_spring[0]
                spring_note = "（二次判断）"
            else:
                spring_note = "（无二次候选，保留初判）"

    if initial_summer and ref_summer_doy:
        if summer_diff_days > 30:
            if secondary_summer:
                final_summer = secondary_summer[0]
                summer_note = "（二次判断）"
            else:
                summer_note = "（无二次候选，保留初判）"

    # 数据完整性提示
    warnings = []
    last_date = sorted_dates[-1] if sorted_dates else ""
    if initial_spring is None and last_date <= f"{target_year}-08-31":
        warnings.append("数据截至8月底，春季尚未完成判定")
    if initial_summer is None and last_date <= f"{target_year}-08-31":
        warnings.append("数据截至8月底，夏季尚未完成判定")

    return {
        "target_year":      target_year,
        "initial_spring":   initial_spring,
        "initial_summer":   initial_summer,
        "secondary_spring": secondary_spring,
        "secondary_summer": secondary_summer,
        "final_spring":     final_spring,
        "final_summer":     final_summer,
        "spring_diff_days": spring_diff_days,
        "summer_diff_days": summer_diff_days,
        "spring_note":      spring_note,
        "summer_note":      summer_note,
        "warnings":         warnings,
    }


# ---------------------------------------------------------------------------
# 7. 输出
# ---------------------------------------------------------------------------

def print_results(all_cities: dict, hist_years: list[str], target_year: str):
    print("=" * 70)
    print(f"  气候季节起始日分析  |  历史年: {', '.join(hist_years)}  |  目标年: {target_year}")
    print("=" * 70)

    for city in sorted(all_cities):
        city_data = all_cities[city]

        # 常年
        climate = analyze_climate_season(city_data, hist_years)
        c_spring  = doy_to_date(climate["spring_doy"], 2023) if climate["spring_doy"]  else "--"
        c_summer  = doy_to_date(climate["summer_doy"], 2023) if climate["summer_doy"]  else "--"
        c_autumn  = doy_to_date(climate["autumn_doy"], 2023) if climate["autumn_doy"]  else "--"
        c_winter  = doy_to_date(climate["winter_doy"], 2023) if climate["winter_doy"]  else "--"

        # 当年
        current = analyze_current_year(city_data, target_year, climate)

        print(f"\n【{city}】")
        print(f"  常年参考（{','.join(hist_years)} 平均）:")
        print(f"    春季起始日: {c_spring}   夏季起始日: {c_summer}")
        print(f"    秋季起始日: {c_autumn}   冬季起始日: {c_winter}")

        init_s = current["initial_spring"] or "--"
        init_sm = current["initial_summer"] or "--"

        # 初判信息
        diff_s_str = ""
        if current["spring_diff_days"] is not None:
            if current["spring_diff_days"] < 0:
                diff_s_str = f"  (比常年早{abs(current['spring_diff_days'])}天，触发二次判断)"
            elif current["spring_diff_days"] <= 30:
                diff_s_str = f"  (比常年晚{current['spring_diff_days']}天，初判有效)"
            else:
                diff_s_str = f"  (比常年晚{current['spring_diff_days']}天，需进一步判断)"
        else:
            diff_s_str = "  (无常年参考)"

        diff_sm_str = ""
        if current["summer_diff_days"] is not None:
            if current["summer_diff_days"] < 0:
                diff_sm_str = f"  (比常年早{abs(current['summer_diff_days'])}天，触发二次判断)"
            elif current["summer_diff_days"] <= 30:
                diff_sm_str = f"  (比常年晚{current['summer_diff_days']}天，初判有效)"
            else:
                diff_sm_str = f"  (比常年晚{current['summer_diff_days']}天，需进一步判断)"
        else:
            diff_sm_str = "  (无常年参考)"

        print(f"  当年法判定（{target_year}）:")
        print(f"    初判春季起始日: {init_s}{diff_s_str}")
        sec_spring_list = current["secondary_spring"]
        if sec_spring_list:
            shown = ', '.join(sec_spring_list[:5])
            suffix = f' 等{len(sec_spring_list)}个' if len(sec_spring_list) > 5 else ''
            print(f"    二次候选春季: {shown}{suffix}")
        else:
            print(f"    二次候选春季: (无)")

        print(f"    初判夏季起始日: {init_sm}{diff_sm_str}")
        sec_summer_list = current["secondary_summer"]
        if sec_summer_list:
            shown = ', '.join(sec_summer_list[:5])
            suffix = f' 等{len(sec_summer_list)}个' if len(sec_summer_list) > 5 else ''
            print(f"    二次候选夏季: {shown}{suffix}")
        else:
            print(f"    二次候选夏季: (无)")

        print(f"    >> 最终判定 春季起始日: 【{current['final_spring'] or '--'}】{current['spring_note']}")
        print(f"    >> 最终判定 夏季起始日: 【{current['final_summer'] or '--'}】{current['summer_note']}")

        for w in current["warnings"]:
            print(f"    提示: {w}")

    print("\n" + "=" * 70)


# ---------------------------------------------------------------------------
# 8. 导出常年参考（JSON）
# ---------------------------------------------------------------------------

def export_climate_reference(
    all_cities: dict,
    hist_years: list[str],
    out_path: str,
):
    result = {}
    for city, city_data in all_cities.items():
        climate = analyze_climate_season(city_data, hist_years)
        result[city] = {
            "spring_doy":   climate["spring_doy"],
            "summer_doy":   climate["summer_doy"],
            "autumn_doy":   climate["autumn_doy"],
            "winter_doy":   climate["winter_doy"],
        }
    with open(out_path, "w", encoding="utf-8") as f:
        json.dump(result, f, ensure_ascii=False, indent=2)
    print(f"常年参考已导出至: {out_path}")


# ---------------------------------------------------------------------------
# 主程序
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    CSV_PATH  = "/home/proc/new_mage/china_city_temperature_2020_2026.csv"
    HIST_YEARS = ["2020", "2021", "2022", "2023", "2024", "2025"]
    TARGET_YEAR = "2026"
    CLIMATE_OUT = "/home/proc/new_mage/climate_reference.json"

    print("正在加载数据...")
    all_cities = load_temperature_data(CSV_PATH)
    print(f"已加载 {len(all_cities)} 个城市\n")
    for c in sorted(all_cities):
        print(f"  {c}: {len(all_cities[c])} 条记录")

    export_climate_reference(all_cities, HIST_YEARS, CLIMATE_OUT)
    print()
    print_results(all_cities, HIST_YEARS, TARGET_YEAR)
    print("\n完成。")