reports/analysis.py

"""Professional research reports: IC, backtest, factor diagnostics, GP evolution."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns

from config.settings import load_settings

plt.style.use("seaborn-v0_8-whitegrid")
sns.set_palette("deep")


def _ensure_dir(path: Path) -> Path:
    path.mkdir(parents=True, exist_ok=True)
    return path


def _save_fig(fig: plt.Figure, path: Path, dpi: int = 150):
    fig.tight_layout()
    fig.savefig(path, dpi=dpi, bbox_inches="tight")
    plt.close(fig)


def plot_ic_series(pred_label: pd.DataFrame, out_path: Path, title: str = "Rank IC Time Series"):
    """Plot daily rank IC and rolling ICIR."""
    df = pred_label.copy()
    if "datetime" in df.columns:
        df = df.set_index("datetime")
    ic = df.groupby(level=0).apply(lambda x: x["label"].rank().corr(x["score"].rank(), method="spearman"))
    ic = ic.dropna()
    roll = ic.rolling(20, min_periods=5).mean()
    roll_std = ic.rolling(20, min_periods=5).std()
    icir = roll / roll_std.replace(0, np.nan)

    fig, axes = plt.subplots(3, 1, figsize=(14, 10), sharex=True)
    axes[0].bar(ic.index, ic.values, width=1.0, color="#4C72B0", alpha=0.7)
    axes[0].axhline(ic.mean(), color="red", ls="--", label=f"Mean IC={ic.mean():.4f}")
    axes[0].set_ylabel("Daily Rank IC")
    axes[0].set_title(title)
    axes[0].legend()

    axes[1].plot(roll.index, roll.values, color="#DD8452", lw=1.5)
    axes[1].set_ylabel("Rolling Mean IC (20d)")

    axes[2].plot(icir.index, icir.values, color="#55A868", lw=1.5)
    axes[2].set_ylabel("Rolling ICIR (20d)")
    axes[2].set_xlabel("Date")

    _save_fig(fig, out_path)


def plot_cumulative_returns(report_df: pd.DataFrame, out_path: Path, title: str = "Portfolio vs Benchmark"):
    """Plot cumulative return and drawdown from qlib PortAna report."""
    df = report_df.copy()
    if "return" not in df.columns:
        raise KeyError(f"report_df missing 'return' column: {df.columns.tolist()}")

    cum = (1 + df["return"]).cumprod()
    bench = (1 + df.get("bench", pd.Series(0, index=df.index))).cumprod() if "bench" in df.columns else None
    dd = cum / cum.cummax() - 1

    fig, axes = plt.subplots(2, 1, figsize=(14, 8), sharex=True, gridspec_kw={"height_ratios": [2, 1]})
    axes[0].plot(cum.index, cum.values, label="Strategy", lw=2)
    if bench is not None:
        axes[0].plot(bench.index, bench.values, label="Benchmark", lw=1.5, alpha=0.8)
    axes[0].set_title(title)
    axes[0].set_ylabel("Cumulative Return")
    axes[0].legend()

    axes[1].fill_between(dd.index, dd.values, 0, color="#C44E52", alpha=0.4)
    axes[1].plot(dd.index, dd.values, color="#C44E52", lw=1)
    axes[1].set_ylabel("Drawdown")
    axes[1].set_xlabel("Date")

    _save_fig(fig, out_path)


def plot_factor_correlation_heatmap(feature_df: pd.DataFrame, out_path: Path, max_factors: int = 30):
    factor_cols = [c for c in feature_df.columns if c.startswith("factor_")][:max_factors]
    if len(factor_cols) < 2:
        return

    sample = feature_df[factor_cols].dropna()
    if len(sample) > 50_000:
        sample = sample.sample(50_000, random_state=42)

    corr = sample.corr(method="spearman")
    fig, ax = plt.subplots(figsize=(12, 10))
    sns.heatmap(corr, cmap="RdBu_r", center=0, vmin=-1, vmax=1, ax=ax, square=True)
    ax.set_title("GP Factor Spearman Correlation")
    _save_fig(fig, out_path)


def plot_gp_evolution(factor_zoo_path: Path, out_path: Path):
    if not factor_zoo_path.exists():
        return
    zoo = pd.read_csv(factor_zoo_path)
    if zoo.empty or "generation" not in zoo.columns:
        return

    agg = zoo.groupby("generation").agg(
        best_fitness=("fitness", "max"),
        mean_fitness=("fitness", "mean"),
        best_is_icir=("is_icir", "max"),
        best_oos_icir=("oos_icir", "max"),
    )

    fig, axes = plt.subplots(2, 1, figsize=(12, 8), sharex=True)
    axes[0].plot(agg.index, agg["best_fitness"], marker="o", label="Best fitness")
    axes[0].plot(agg.index, agg["mean_fitness"], marker="s", label="Mean fitness", alpha=0.7)
    axes[0].set_ylabel("Fitness")
    axes[0].set_title("GP Evolution")
    axes[0].legend()

    axes[1].plot(agg.index, agg["best_is_icir"], marker="o", label="Best IS ICIR")
    axes[1].plot(agg.index, agg["best_oos_icir"], marker="s", label="Best OOS ICIR")
    axes[1].set_ylabel("ICIR")
    axes[1].set_xlabel("Generation")
    axes[1].legend()

    _save_fig(fig, out_path)


def plot_quantile_returns(pred_label: pd.DataFrame, out_path: Path, n_groups: int = 5):
    df = pred_label.copy()
    if isinstance(df.index, pd.MultiIndex):
        df = df.reset_index()
    df["group"] = df.groupby("datetime")["score"].transform(
        lambda x: pd.qcut(x.rank(method="first"), n_groups, labels=False, duplicates="drop")
    )
    grp = df.groupby(["datetime", "group"])["label"].mean().unstack()
    cum = (1 + grp).cumprod()

    fig, ax = plt.subplots(figsize=(14, 6))
    for col in cum.columns:
        ax.plot(cum.index, cum[col], label=f"Q{int(col)+1}", lw=1.5)
    ax.set_title(f"Quantile Portfolio Cumulative Return (Top=Q{n_groups})")
    ax.set_ylabel("Cumulative Return")
    ax.set_xlabel("Date")
    ax.legend()
    _save_fig(fig, out_path)


def plot_ic_distribution(pred_label: pd.DataFrame, out_path: Path):
    df = pred_label.copy()
    if "datetime" in df.columns:
        df = df.set_index("datetime")
    ic = df.groupby(level=0).apply(lambda x: x["label"].rank().corr(x["score"].rank(), method="spearman"))
    ic = ic.dropna()

    fig, ax = plt.subplots(figsize=(8, 5))
    ax.hist(ic.values, bins=40, color="#4C72B0", alpha=0.8, edgecolor="white")
    ax.axvline(ic.mean(), color="red", ls="--", label=f"Mean={ic.mean():.4f}")
    ax.axvline(ic.median(), color="green", ls="--", label=f"Median={ic.median():.4f}")
    ax.set_title("Daily Rank IC Distribution")
    ax.set_xlabel("Rank IC")
    ax.legend()
    _save_fig(fig, out_path)


def _load_recorder_artifacts(recorder_id: str, experiment_name: str, mlruns_root: Path) -> dict[str, Any]:
    from qlib.workflow import R

    with R.start(experiment_name=experiment_name, recorder_id=recorder_id, resume=True):
        artifacts = {}
        for name in ["pred.pkl", "label.pkl", "report_normal_1day.pkl", "positions_normal_1day.pkl"]:
            try:
                artifacts[name] = R.load_object(name)
            except Exception:
                pass
        try:
            artifacts["sig_analysis"] = R.load_object("sig_analysis.pkl")
        except Exception:
            pass
    return artifacts


def generate_experiment_report(
    recorder_id: str,
    experiment_name: str,
    output_dir: Path | None = None,
    run_id: str | None = None,
) -> Path:
    settings = load_settings()
    out_dir = _ensure_dir(output_dir or settings.path(settings.raw["output"]["reports_dir"], experiment_name))
    mlruns_root = Path(settings.mlruns_uri.replace("file://", ""))

    artifacts = _load_recorder_artifacts(recorder_id, experiment_name, mlruns_root)

    summary = {"recorder_id": recorder_id, "experiment_name": experiment_name, "artifacts": list(artifacts.keys())}

    if "pred.pkl" in artifacts and "label.pkl" in artifacts:
        pred = artifacts["pred.pkl"]
        label = artifacts["label.pkl"]
        pred_label = pred.join(label, how="inner")
        pred_label.columns = ["score", "label"]
        plot_ic_series(pred_label, out_dir / "ic_series.png", title=f"{experiment_name} Rank IC")
        plot_ic_distribution(pred_label, out_dir / "ic_distribution.png")
        plot_quantile_returns(pred_label.reset_index(), out_dir / "quantile_returns.png")

        ic = pred_label.groupby(level=0).apply(
            lambda x: x["label"].rank().corr(x["score"].rank(), method="spearman")
        )
        summary["ic_mean"] = float(ic.mean())
        summary["ic_std"] = float(ic.std())
        summary["icir"] = float(ic.mean() / ic.std()) if ic.std() else None

    if "report_normal_1day.pkl" in artifacts:
        report_df = artifacts["report_normal_1day.pkl"]
        plot_cumulative_returns(report_df, out_dir / "cumulative_return_drawdown.png", title=experiment_name)
        if "return" in report_df.columns:
            summary["ann_return"] = float(report_df["return"].mean() * 252)
            summary["ann_vol"] = float(report_df["return"].std() * np.sqrt(252))
            cum = (1 + report_df["return"]).cumprod()
            summary["max_drawdown"] = float((cum / cum.cummax() - 1).min())

    if run_id:
        gp_dir = settings.gp_output_dir(run_id)
        feature_path = gp_dir / "ML_Features_qlib.parquet"
        if feature_path.exists():
            feat_df = pd.read_parquet(feature_path)
            plot_factor_correlation_heatmap(feat_df, out_dir / "factor_correlation.png")
        plot_gp_evolution(gp_dir / "factor_zoo.csv", out_dir / "gp_evolution.png")

    with open(out_dir / "summary.json", "w", encoding="utf-8") as f:
        json.dump(summary, f, indent=2, ensure_ascii=False)

    print(f"Report saved to {out_dir}")
    return out_dir


def generate_gp_only_report(run_id: str | None = None, output_dir: Path | None = None) -> Path:
    settings = load_settings()
    run_id = run_id or settings.raw.get("experiment", {}).get("run_id", "qlib_gp_run_0")
    out_dir = _ensure_dir(output_dir or settings.path(settings.raw["output"]["reports_dir"], f"gp_{run_id}"))
    gp_dir = settings.gp_output_dir(run_id)

    plot_gp_evolution(gp_dir / "factor_zoo.csv", out_dir / "gp_evolution.png")

    feature_path = gp_dir / "ML_Features_qlib.parquet"
    if not feature_path.exists():
        feature_path = gp_dir / "ML_Features_qlib.csv"
    if feature_path.exists():
        feat_df = pd.read_parquet(feature_path) if feature_path.suffix == ".parquet" else pd.read_csv(feature_path)
        plot_factor_correlation_heatmap(feat_df, out_dir / "factor_correlation.png")

        factor_cols = [c for c in feat_df.columns if c.startswith("factor_")]
        if factor_cols and "target_return" in feat_df.columns:
            ic_rows = []
            for col in factor_cols:
                tmp = feat_df[["date", col, "target_return"]].dropna()
                daily_ic = tmp.groupby("date", group_keys=False).apply(
                    lambda x: x[col].rank().corr(x["target_return"].rank(), method="spearman"),
                    include_groups=False,
                )
                ic_rows.append({"factor": col, "ic_mean": daily_ic.mean(), "icir": daily_ic.mean() / daily_ic.std()})
            ic_df = pd.DataFrame(ic_rows).sort_values("icir", ascending=False)
            ic_df.to_csv(out_dir / "factor_ic_summary.csv", index=False)

            fig, ax = plt.subplots(figsize=(10, max(4, len(factor_cols) * 0.25)))
            ax.barh(ic_df["factor"], ic_df["icir"], color="#4C72B0")
            ax.set_xlabel("ICIR")
            ax.set_title("GP Factor ICIR Ranking")
            ax.invert_yaxis()
            _save_fig(fig, out_dir / "factor_icir_ranking.png")

    print(f"GP report saved to {out_dir}")
    return out_dir