"""
环境数据 API 路由

提供空气污染物时间序列和滞后相关性分析接口
"""

import logging
from pathlib import Path
from typing import Optional

import pandas as pd
from fastapi import APIRouter, HTTPException, Query
from pydantic import BaseModel
from scipy.stats import pearsonr

logger = logging.getLogger("cbpoa.environment")

PROJECT_ROOT = Path(__file__).parent.parent.parent

router = APIRouter(prefix="/api/environment", tags=["environment"])

# 污染物列表（CO 无基值列，仅存在于 lag_features.parquet 的滞后列中）
POLLUTANTS = ["AQI", "PM2.5", "PM10", "SO2", "NO2", "O3"]
LAGS = [1, 2, 3, 5, 7, 14]

# station_daily 列名映射（PM25 无点号 -> PM2.5 带点号）
_STATION_COL_MAP: dict[str, str] | None = None
# lag_features 基值列映射
_LAG_BASE_MAP: dict[str, str] | None = None


# ============== Response Models ==============


class LagCorrelationItem(BaseModel):
    pollutant: str  # "AQI", "PM2.5", "PM10", "SO2", "NO2", "O3"
    lag_days: int  # 1, 2, 3, 5, 7, 14
    correlation: float  # Pearson r


class LagCorrelationResponse(BaseModel):
    correlations: list[LagCorrelationItem]
    data_note: str  # "CO excluded - no base column in lag_features.parquet"


class PollutantPoint(BaseModel):
    date: str
    AQI: float
    PM25: float
    PM10: float
    SO2: float
    NO2: float
    O3: float
    CO: float


class PollutantResponse(BaseModel):
    data: list[PollutantPoint]
    station_count: int
    date_range: dict  # {start, end}


# ============== Helper Functions ==============


def _get_station_col_map() -> dict[str, str]:
    """返回 station_daily parquet 中实际列名到标准名称的映射。

    station_daily 文件中 PM2.5 列名为 "PM25"（无点号），
    需要映射到前端期望的 "PM2.5"。
    """
    global _STATION_COL_MAP
    if _STATION_COL_MAP is not None:
        return _STATION_COL_MAP

    path = PROJECT_ROOT / "processed" / "weather" / "station_daily_2022.parquet"
    df = pd.read_parquet(path)
    cols = set(df.columns)

    col_map = {}
    for standard in POLLUTANTS:
        if standard in cols:
            col_map[standard] = standard
        elif "PM25" in cols and standard == "PM2.5":
            col_map[standard] = "PM25"
        else:
            col_map[standard] = standard  # fallback

    # CO is in station_daily but not in POLLUTANTS
    if "CO" in cols:
        col_map["CO"] = "CO"

    _STATION_COL_MAP = col_map
    return _STATION_COL_MAP


def _get_lag_base_map() -> dict[str, str]:
    """返回 lag_features parquet 中基值列名到标准名称的映射。

    lag_features 中 PM2.5 列名为 "PM2.5"（带点号），与标准名称一致。
    此函数在运行时验证实际列名。
    """
    global _LAG_BASE_MAP
    if _LAG_BASE_MAP is not None:
        return _LAG_BASE_MAP

    path = PROJECT_ROOT / "processed" / "weather" / "lag_features.parquet"
    df = pd.read_parquet(path)
    cols = set(df.columns)

    base_map = {}
    for p in POLLUTANTS:
        if p in cols:
            base_map[p] = p
        elif p == "PM2.5" and "PM25" in cols:
            base_map[p] = "PM25"
        else:
            base_map[p] = p  # will be checked later

    _LAG_BASE_MAP = base_map
    return _LAG_BASE_MAP


def _load_lag_features() -> pd.DataFrame:
    """加载 lag_features.parquet 并转换日期列。"""
    path = PROJECT_ROOT / "processed" / "weather" / "lag_features.parquet"
    if not path.exists():
        raise FileNotFoundError(f"lag_features.parquet not found at {path}")
    df = pd.read_parquet(path)
    df["date"] = pd.to_datetime(df["date"])
    return df


def _load_cases_daily() -> pd.DataFrame:
    """加载 cases_combined.parquet 并按日期汇总每日总病例数。"""
    path = PROJECT_ROOT / "processed" / "cases_combined.parquet"
    if not path.exists():
        raise FileNotFoundError(f"cases_combined.parquet not found at {path}")
    df = pd.read_parquet(path)
    df["date"] = pd.to_datetime(df["date"])
    daily = df.groupby("date").size().reset_index(name="total_cases")
    return daily


def _load_station_daily() -> pd.DataFrame:
    """加载并合并 station_daily_2022.parquet 和 station_daily_2023.parquet。"""
    dfs = []
    for year in [2022, 2023]:
        path = PROJECT_ROOT / "processed" / "weather" / f"station_daily_{year}.parquet"
        if not path.exists():
            logger.warning("station_daily_%s.parquet not found at %s", year, path)
            continue
        df = pd.read_parquet(path)
        dfs.append(df)
    if not dfs:
        raise FileNotFoundError("No station_daily parquet files found")
    combined = pd.concat(dfs, ignore_index=True)
    combined["date"] = pd.to_datetime(combined["date"])
    return combined


# ============== Endpoints ==============


@router.get("/lag-correlations", response_model=LagCorrelationResponse)
async def get_lag_correlations():
    """获取污染物滞后相关性分析数据。

    计算各污染物在不同滞后天数（1, 2, 3, 5, 7, 14 天）下
    与每日病例总数之间的 Pearson 相关系数。

    Returns:
        LagCorrelationResponse: 包含 36 个相关系数（6 种污染物 × 6 个滞后天数）
    """
    try:
        lag_df = _load_lag_features()
        cases_daily = _load_cases_daily()
    except FileNotFoundError as e:
        logger.warning("Data file not found for lag-correlations: %s", e)
        return LagCorrelationResponse(
            correlations=[],
            data_note="CO excluded - no base column in lag_features.parquet",
        )

    # 计算每日全市均值（按日期聚合，对 23 个站点取平均）
    base_map = _get_lag_base_map()
    mean_cols = {p: base_map.get(p, p) for p in POLLUTANTS}
    daily_mean = lag_df.groupby("date")[list(mean_cols.values())].mean().reset_index()

    # 重命名列为标准名称以便一致访问
    rename_map = {v: k for k, v in mean_cols.items() if v != k}
    if rename_map:
        daily_mean = daily_mean.rename(columns=rename_map)

    # 对每种污染物计算每日均值
    # 对齐污染物时间序列与病例数据
    merged = daily_mean.merge(cases_daily, on="date", how="inner")
    merged = merged.sort_values("date")

    results: list[LagCorrelationItem] = []

    for pollutant in POLLUTANTS:
        # 构建滞后列映射（lag_features 文件名用点号 "PM2.5"）
        lag_base_name = base_map.get(pollutant, pollutant)

        for lag in LAGS:
            lag_col = f"{lag_base_name}_lag{lag}"
            if lag_col not in lag_df.columns:
                logger.debug("Lag column %s not found, skipping", lag_col)
                continue

            # 从原始 lag_features 提取该污染物的滞后数据（按日期取全市均值）
            lag_series = lag_df.groupby("date")[lag_col].mean().reset_index()
            lag_series = lag_series.rename(columns={lag_col: f"{pollutant}_lag{lag}"})

            # 将滞后污染物数据与病例数据对齐
            # 滞后列的值代表的是 t-lag 时刻的污染物，病例是 t 时刻
            # 所以将 lag 列的时间向后平移 lag 天，使其与病例时间对齐
            lag_series["align_date"] = lag_series["date"] + pd.Timedelta(days=lag)
            combined = lag_series.merge(
                cases_daily, left_on="align_date", right_on="date", how="inner"
            )

            if len(combined) < 10:
                logger.debug(
                    "Insufficient data for %s lag%d: %d rows, skipping",
                    pollutant,
                    lag,
                    len(combined),
                )
                continue

            # 计算 Pearson 相关系数
            r, _ = pearsonr(
                combined[f"{pollutant}_lag{lag}"], combined["total_cases"]
            )
            results.append(
                LagCorrelationItem(
                    pollutant=pollutant, lag_days=lag, correlation=round(float(r), 4)
                )
            )

    return LagCorrelationResponse(
        correlations=results,
        data_note="CO excluded - no base column in lag_features.parquet",
    )


@router.get("/pollutants", response_model=PollutantResponse)
async def get_pollutants(
    days: Optional[int] = Query(default=30, ge=1, le=730, description="返回最近 N 天的数据"),
    start_date: Optional[str] = Query(default=None, description="开始日期 YYYY-MM-DD"),
    end_date: Optional[str] = Query(default=None, description="结束日期 YYYY-MM-DD"),
):
    """获取每日全市均值污染物时间序列。

    合并 2022/2023 两个年度的站点日数据，按日期聚合所有站点取均值。

    可通过 ?days=N（默认 30）或 ?start_date=YYYY-MM-DD&end_date=YYYY-MM-DD 筛选日期范围。
    """
    try:
        df = _load_station_daily()
    except FileNotFoundError as e:
        logger.warning("Data file not found for pollutants: %s", e)
        return PollutantResponse(data=[], station_count=0, date_range={})

    station_count = df["station_id"].nunique()

    # 获取列名映射
    col_map = _get_station_col_map()

    # 选择污染物列
    poll_cols = [col_map.get(p, p) for p in POLLUTANTS] + ["CO"]
    # 确保需要的列都存在
    available_cols = [c for c in poll_cols if c in df.columns]

    # 按日期聚合取均值
    daily_mean = df.groupby("date")[available_cols].mean().reset_index()

    # 重命名为标准名称
    rename_map = {}
    for std_name, actual_name in col_map.items():
        if actual_name != std_name and actual_name in daily_mean.columns:
            rename_map[actual_name] = std_name
    if rename_map:
        daily_mean = daily_mean.rename(columns=rename_map)

    # 日期筛选
    if start_date and end_date:
        start_dt = pd.to_datetime(start_date)
        end_dt = pd.to_datetime(end_date)
        date_mask = (daily_mean["date"] >= start_dt) & (
            daily_mean["date"] <= end_dt
        )
        daily_mean = daily_mean[date_mask].copy()
    else:
        daily_mean = daily_mean.sort_values("date").tail(days)

    daily_mean = daily_mean.sort_values("date")

    # 构建响应
    data: list[PollutantPoint] = []
    for _, row in daily_mean.iterrows():
        pt = PollutantPoint(
            date=row["date"].strftime("%Y-%m-%d"),
            AQI=round(float(row.get("AQI", 0)), 2),
            PM25=round(float(row.get("PM2.5", 0)), 2),
            PM10=round(float(row.get("PM10", 0)), 2),
            SO2=round(float(row.get("SO2", 0)), 2),
            NO2=round(float(row.get("NO2", 0)), 2),
            O3=round(float(row.get("O3", 0)), 2),
            CO=round(float(row.get("CO", 0)), 2),
        )
        data.append(pt)

    date_range = {}
    if daily_mean.shape[0] > 0:
        date_range = {
            "start": daily_mean["date"].iloc[0].strftime("%Y-%m-%d"),
            "end": daily_mean["date"].iloc[-1].strftime("%Y-%m-%d"),
        }

    return PollutantResponse(
        data=data, station_count=station_count, date_range=date_range
    )