#!/usr/bin/env python3
"""Research-only dynamic ensemble overlay from strategy equity curves.

This does not route orders through the main backtester. It combines already
realized daily return series from multiple strategy runs using lagged rolling
Sharpe weights.
"""

from __future__ import annotations

import argparse
import json
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import pandas as pd


TRADING_DAYS_PER_YEAR = 252.0


@dataclass(frozen=True)
class CurveInput:
    label: str
    path: Path


def _parse_curve_arg(value: str) -> CurveInput:
    if "=" not in value:
        raise argparse.ArgumentTypeError("curve must be LABEL=/abs/path/to/daily_equity_curve.parquet")
    label, raw_path = value.split("=", 1)
    label = label.strip()
    path = Path(raw_path.strip())
    if not label:
        raise argparse.ArgumentTypeError("curve label must be non-empty")
    if not path.exists():
        raise argparse.ArgumentTypeError(f"curve file not found: {path}")
    return CurveInput(label=label, path=path)


def _load_curve(curve: CurveInput) -> pd.DataFrame:
    df = pd.read_parquet(curve.path)
    if "date" not in df.columns or "equity" not in df.columns:
        raise ValueError(f"{curve.path} must contain date/equity columns")
    out = df[["date", "equity"]].copy()
    out["date"] = pd.to_datetime(out["date"])
    out = out.sort_values("date")
    out[curve.label] = out["equity"].pct_change().fillna(0.0)
    return out[["date", curve.label]]


def _rolling_sharpe(series: pd.Series, window: int) -> pd.Series:
    mean = series.rolling(window, min_periods=window).mean()
    std = series.rolling(window, min_periods=window).std(ddof=0)
    sharpe = mean / std.replace(0.0, np.nan)
    return sharpe * np.sqrt(TRADING_DAYS_PER_YEAR)


def build_ensemble_returns(
    daily_returns: pd.DataFrame,
    window: int = 63,
    fallback: str = "equal",
) -> tuple[pd.DataFrame, pd.DataFrame]:
    strategy_cols = [c for c in daily_returns.columns if c != "date"]
    if len(strategy_cols) < 2:
        raise ValueError("need at least two strategy return series")

    weights = pd.DataFrame({"date": daily_returns["date"]})
    rolling = {}
    for col in strategy_cols:
        rolling[col] = _rolling_sharpe(daily_returns[col], window).shift(1)
    sharpes = pd.DataFrame({"date": daily_returns["date"], **rolling})

    positive = sharpes[strategy_cols].clip(lower=0.0)
    positive_sum = positive.sum(axis=1)
    if fallback == "equal":
        fallback_weights = pd.DataFrame(
            np.full((len(daily_returns), len(strategy_cols)), 1.0 / len(strategy_cols)),
            columns=strategy_cols,
        )
    else:
        raise ValueError(f"unsupported fallback: {fallback}")

    normalized = positive.div(positive_sum.replace(0.0, np.nan), axis=0)
    resolved = normalized.where(positive_sum.gt(0), fallback_weights)
    weights = pd.concat([weights, resolved], axis=1)

    ensemble = daily_returns.copy()
    ensemble["ensemble_return"] = 0.0
    for col in strategy_cols:
        ensemble["ensemble_return"] += ensemble[col] * weights[col]
    return ensemble, weights


def summarize_ensemble(ensemble_returns: pd.DataFrame, initial_equity: float) -> dict:
    out = ensemble_returns[["date", "ensemble_return"]].copy()
    out["equity"] = initial_equity * (1.0 + out["ensemble_return"]).cumprod()
    out["peak"] = out["equity"].cummax()
    out["drawdown_pct"] = (out["equity"] / out["peak"] - 1.0) * 100.0

    total_return_pct = (out["equity"].iloc[-1] / out["equity"].iloc[0] - 1.0) * 100.0
    years = max(len(out) / TRADING_DAYS_PER_YEAR, 1.0 / TRADING_DAYS_PER_YEAR)
    annualized_return_pct = ((out["equity"].iloc[-1] / out["equity"].iloc[0]) ** (1.0 / years) - 1.0) * 100.0
    daily_mean = out["ensemble_return"].mean()
    daily_std = out["ensemble_return"].std(ddof=0)
    sharpe = float(daily_mean / daily_std * np.sqrt(TRADING_DAYS_PER_YEAR)) if daily_std > 0 else 0.0

    out["year"] = out["date"].dt.year
    yearly_returns = {}
    for year, year_df in out.groupby("year"):
        yearly_returns[str(year)] = (
            (1.0 + year_df["ensemble_return"]).prod() - 1.0
        ) * 100.0

    return {
        "total_return_pct": total_return_pct,
        "annualized_return_pct": annualized_return_pct,
        "max_drawdown_pct": float(-out["drawdown_pct"].min()),
        "sharpe_ratio": sharpe,
        "days": int(len(out)),
        "yearly_returns_pct": yearly_returns,
    }


def main() -> int:
    parser = argparse.ArgumentParser(description="Research-only rolling-Sharpe ensemble overlay")
    parser.add_argument("--curve", action="append", type=_parse_curve_arg, required=True, help="LABEL=/abs/path/to/daily_equity_curve.parquet")
    parser.add_argument("--window", type=int, default=63, help="Rolling Sharpe lookback in trading days")
    parser.add_argument("--initial-equity", type=float, default=100000.0)
    parser.add_argument("--fallback", choices=["equal"], default="equal")
    parser.add_argument("--output-dir", required=True, help="Directory to write ensemble artifacts")
    args = parser.parse_args()

    curves = args.curve
    merged = None
    for curve in curves:
        loaded = _load_curve(curve)
        merged = loaded if merged is None else merged.merge(loaded, on="date", how="inner")

    if merged is None or merged.empty:
        raise SystemExit("no overlapping dates across curves")

    ensemble, weights = build_ensemble_returns(merged, window=args.window, fallback=args.fallback)
    summary = summarize_ensemble(ensemble, args.initial_equity)

    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    ensemble_out = ensemble[["date", "ensemble_return"]].copy()
    ensemble_out["equity"] = args.initial_equity * (1.0 + ensemble_out["ensemble_return"]).cumprod()
    ensemble_out.to_parquet(out_dir / "daily_equity_curve.parquet", index=False)
    weights.to_parquet(out_dir / "weights.parquet", index=False)
    (out_dir / "summary.json").write_text(json.dumps(summary, indent=2))

    print(json.dumps(summary, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())