"""
09_time_varying_weights.py
时变权重模块：5 年滚动窗口 PCA 权重重估
Time-Varying Weights Module: 5-year rolling-window PCA re-estimation

用途 / Purpose:
    07 中的权重基于全样本（2000-2023）一次性估计，隐含"权重恒定"假设。
    本模块以 5 年滚动窗口（如 2000-2004, 2001-2005, ...）重新估计各维度
    PCA 权重，用于检验权重随时间的稳定性，并为稳健性检验（12 号脚本）
    提供权重漂移证据。

依赖 / Dependencies: pandas, numpy, scikit-learn
"""

import logging
from pathlib import Path

import numpy as np
import pandas as pd
from sklearn.decomposition import PCA

logging.basicConfig(level=logging.INFO, format="%(asctime)s  %(message)s")
log = logging.getLogger("09_time_varying")

DATA_DIR = Path("./data")
WINDOW = 5
RANDOM_SEED = 42

DIMENSION_COLUMNS = {
    "V": [
        "gov_expenditure_pct_gdp__minmax",
        "social_protection_expenditure__minmax",
        "fiscal_balance_pct_gdp__minmax",
        "public_debt_pct_gdp__minmax",
        "tax_revenue_pct_gdp__minmax",
        "unemployment_benefit_coverage__minmax",
    ],
    "C": [
        "gdp_growth_pct__minmax",
        "gross_capital_formation_pct_gdp__minmax",
        "patents_residents__minmax",
        "rd_expenditure_pct_gdp__minmax",
        "labor_productivity_growth__minmax",
        "high_tech_exports_pct__minmax",
    ],
    "F": [
        "regulatory_quality__minmax",
        "government_effectiveness__minmax",
        "rule_of_law__minmax",
        "corruption_control__minmax",
        "bureaucracy_delay_index__minmax",
    ],
}


def rolling_pca_weight(window_df: pd.DataFrame, cols: list) -> np.ndarray:
    sub = window_df[cols].dropna()
    if len(sub) < len(cols) + 1:
        return np.full(len(cols), np.nan)
    pca = PCA(n_components=1, random_state=RANDOM_SEED)
    pca.fit(sub.to_numpy())
    loadings = np.abs(pca.components_[0])
    return loadings / loadings.sum()


def main():
    panel = pd.read_csv(DATA_DIR / "panel_normalized_minmax.csv")
    years = sorted(panel["year"].unique())

    records = []
    for start in range(years[0], years[-1] - WINDOW + 2):
        end = start + WINDOW - 1
        window_df = panel[(panel["year"] >= start) & (panel["year"] <= end)]
        for dim, cols in DIMENSION_COLUMNS.items():
            available = [c for c in cols if c in window_df.columns]
            weights = rolling_pca_weight(window_df, available)
            for col, w in zip(available, weights):
                records.append(
                    {"window_start": start, "window_end": end, "dimension": dim,
                     "indicator": col, "weight": round(float(w), 4) if not np.isnan(w) else None}
                )
        log.info("Rolling window %d-%d processed", start, end)

    out = pd.DataFrame(records)
    out.to_csv(DATA_DIR / "rolling_weights_by_window.csv", index=False)
    log.info(
        "Time-varying weight estimation complete -> data/rolling_weights_by_window.csv (%d windows)",
        len(years) - WINDOW + 1,
    )


if __name__ == "__main__":
    main()