from __future__ import annotations

import argparse
import json
from dataclasses import replace
from pathlib import Path

import pandas as pd

from qfr.strategy.etf_trend import Constraints, TrendParams, UniverseAsset, run_backtest


def load_universe(config_path: Path):
    conf = json.loads(config_path.read_text(encoding="utf-8"))
    universe = [UniverseAsset(**a) for a in conf["assets"]]
    cons = conf.get("constraints", {})
    constraints = Constraints(
        max_positions=int(cons.get("max_positions", 4)),
        must_commodity=int(cons.get("must_include", {}).get("commodity", 0)),
        must_rates=int(cons.get("must_include", {}).get("rates", 0)),
        must_equity=int(cons.get("must_include", {}).get("equity", 0)),
    )
    return universe, constraints, cons.get("risk_proxy", "510300.SH"), cons.get("rates_fallback", "511010.SH")


def load_prices(raw_dir: Path, universe: list[UniverseAsset], start: str, end: str):
    out = {}
    for a in universe:
        fn = raw_dir / f"{a.ts_code.replace('.', '')}.parquet"
        df = pd.read_parquet(fn)
        df = df.copy()
        df["trade_date"] = df["trade_date"].astype(str)
        df = df[(df["trade_date"] >= start) & (df["trade_date"] <= end)]
        out[a.ts_code] = df
    return out


def perf_stats(equity: pd.Series):
    r = equity.pct_change().dropna()
    ann_ret = float((equity.iloc[-1] / equity.iloc[0]) ** (252 / len(r)) - 1)
    ann_vol = float(r.std(ddof=1) * (252 ** 0.5))
    dd = float((equity / equity.cummax() - 1.0).min())
    return ann_ret, ann_vol, dd


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default="configs/etf_universe.json")
    ap.add_argument("--rawdir", default="data/raw")
    ap.add_argument("--start", default="20200101")
    ap.add_argument("--end", default="20251231")
    args = ap.parse_args()

    universe, constraints, risk_proxy, rates_fallback = load_universe(Path(args.config))
    prices = load_prices(Path(args.rawdir), universe, args.start, args.end)

    base = TrendParams(rebalance_every=1, max_positions=4)

    # A very small candidate set (fast to run)
    candidates = [
        (5, 20, 3.0),
        (5, 20, 2.5),
        (3, 15, 2.5),
        (8, 30, 3.0),
        (10, 40, 3.0),
        (5, 30, 3.0),
    ]

    rows = []
    for sma_fast, sma_slow, atr_mult in candidates:
        params = replace(base, sma_fast=sma_fast, sma_slow=sma_slow, atr_mult=atr_mult)
        equity, _w = run_backtest(
            prices,
            universe,
            constraints,
            params,
            rates_fallback=rates_fallback,
            risk_proxy=risk_proxy,
        )
        ann_ret, ann_vol, dd = perf_stats(equity["equity"])
        rows.append({
            "ann_return": ann_ret,
            "ann_vol": ann_vol,
            "max_drawdown": dd,
            "sma_fast": sma_fast,
            "sma_slow": sma_slow,
            "atr_mult": atr_mult,
        })

    df = pd.DataFrame(rows).sort_values(["ann_return"], ascending=False)
    print(df.to_string(index=False))


if __name__ == "__main__":
    main()