from __future__ import annotations

from pathlib import Path

import pandas as pd

from dragon_branch_configs import (
    alpha_first_glued_refined_hot_cap_config,
    alpha_first_selective_veto_config,
    workbook_preserving_config,
)
from dragon_shared import END_DATE, START_DATE, evaluation_years
from dragon_strategy import DragonRuleEngine


def _load_indicator_snapshot(base_dir: Path) -> pd.DataFrame:
    df = pd.read_csv(base_dir / "dragon_indicator_snapshot.csv", encoding="utf-8-sig")
    df["date"] = pd.to_datetime(df["date"])
    return df.set_index("date", drop=False)


def _max_drawdown(equity: pd.Series) -> float:
    running_max = equity.cummax()
    drawdown = equity / running_max - 1.0
    return float(drawdown.min())


def _max_drawdown_duration(equity: pd.Series) -> int:
    running_max = equity.cummax()
    underwater = equity < running_max
    best = cur = 0
    for flag in underwater.tolist():
        cur = cur + 1 if flag else 0
        best = max(best, cur)
    return int(best)


def _run_branch(indicator_df: pd.DataFrame, config) -> pd.DataFrame:
    engine = DragonRuleEngine(config=config)
    _, trades = engine.run(indicator_df)
    trades = trades[
        (trades["buy_date"] >= START_DATE)
        & (trades["buy_date"] <= END_DATE)
        & (trades["sell_date"] >= START_DATE)
        & (trades["sell_date"] <= END_DATE)
    ].copy()
    trades["sell_dt"] = pd.to_datetime(trades["sell_date"])
    trades = trades.sort_values("sell_dt").reset_index(drop=True)
    trades["equity"] = (1.0 + trades["return_pct"].astype(float)).cumprod()
    trades["drawdown"] = trades["equity"] / trades["equity"].cummax() - 1.0
    trades["sell_year"] = trades["sell_dt"].dt.year.astype(int)
    trades["sell_month"] = trades["sell_dt"].dt.to_period("M").astype(str)
    return trades


def _summarize(branch: str, trades: pd.DataFrame) -> dict[str, object]:
    years = evaluation_years(START_DATE, END_DATE)
    compounded = float(trades["equity"].iloc[-1] - 1.0) if not trades.empty else float("nan")
    cagr = float((1.0 + compounded) ** (1.0 / years) - 1.0) if not trades.empty else float("nan")
    max_dd = _max_drawdown(trades["equity"]) if not trades.empty else float("nan")
    calmar = float(cagr / abs(max_dd)) if trades.shape[0] and max_dd < 0 else float("inf")
    return {
        "branch": branch,
        "trades": int(len(trades)),
        "compounded_return": compounded,
        "cagr": cagr,
        "max_drawdown": max_dd,
        "drawdown_duration_trades": _max_drawdown_duration(trades["equity"]) if not trades.empty else 0,
        "calmar": calmar,
        "best_trade": float(trades["return_pct"].max()) if not trades.empty else float("nan"),
        "worst_trade": float(trades["return_pct"].min()) if not trades.empty else float("nan"),
    }


def main() -> None:
    base_dir = Path(__file__).resolve().parent
    indicator_df = _load_indicator_snapshot(base_dir)

    branches = {
        "workbook_preserving": workbook_preserving_config(),
        "alpha_first_selective_veto": alpha_first_selective_veto_config(),
        "alpha_first_glued_refined_hot_cap": alpha_first_glued_refined_hot_cap_config(),
    }

    equity_frames: list[pd.DataFrame] = []
    summary_rows: list[dict[str, object]] = []
    monthly_rows: list[dict[str, object]] = []
    yearly_rows: list[dict[str, object]] = []

    for name, cfg in branches.items():
        trades = _run_branch(indicator_df, cfg)
        trades["branch"] = name
        equity_frames.append(trades[["branch", "buy_date", "sell_date", "return_pct", "equity", "drawdown", "sell_year", "sell_month"]].copy())
        summary_rows.append(_summarize(name, trades))

        month = trades.groupby("sell_month", dropna=False)["return_pct"].apply(lambda s: float((1.0 + s).prod() - 1.0)).reset_index()
        month["branch"] = name
        monthly_rows.append(month)

        year = trades.groupby("sell_year", dropna=False)["return_pct"].apply(lambda s: float((1.0 + s).prod() - 1.0)).reset_index()
        year["branch"] = name
        yearly_rows.append(year)

    equity = pd.concat(equity_frames, ignore_index=True)
    summary = pd.DataFrame(summary_rows)
    monthly = pd.concat(monthly_rows, ignore_index=True)
    yearly = pd.concat(yearly_rows, ignore_index=True)

    equity.to_csv(base_dir / "dragon_equity_curve_review.csv", index=False, encoding="utf-8-sig")
    summary.to_csv(base_dir / "dragon_drawdown_review.csv", index=False, encoding="utf-8-sig")
    monthly.to_csv(base_dir / "dragon_monthly_return_review.csv", index=False, encoding="utf-8-sig")
    yearly.to_csv(base_dir / "dragon_yearly_return_review.csv", index=False, encoding="utf-8-sig")

    lines = [
        "# Dragon Equity Curve Review",
        "",
        f"- Evaluation window: `{START_DATE}` to `{END_DATE}`.",
        "- Equity is compounded in trade sequence order.",
        "",
        "## Summary",
    ]

    for _, row in summary.iterrows():
        lines.append(
            f"- `{row['branch']}`: compounded `{row['compounded_return']:.2%}`, CAGR `{row['cagr']:.2%}`, "
            f"max_drawdown `{row['max_drawdown']:.2%}`, drawdown_duration `{int(row['drawdown_duration_trades'])}` trades, "
            f"Calmar `{row['calmar']:.2f}`"
        )

    lines.extend(["", "## Quant Judgment"])
    best = summary.sort_values(["cagr", "calmar"], ascending=[False, False]).iloc[0]
    lines.append(
        f"- Best growth profile in this pack: `{best['branch']}` with CAGR `{best['cagr']:.2%}`, max_drawdown `{best['max_drawdown']:.2%}`, Calmar `{best['calmar']:.2f}`."
    )
    lines.append("- Use this review to judge whether higher alpha is being purchased with unacceptable drawdown concentration.")

    (base_dir / "dragon_equity_curve_review.md").write_text("\n".join(lines) + "\n", encoding="utf-8")


if __name__ == "__main__":
    main()