"""
Pre-compute backtest results for all strategies, for multiple capital scenarios.

Each strategy is backtested ONCE per market with that market's live starting
capital (MONITOR_CAPITAL: US $10k, CN ¥70k); the
resulting equity series is then replayed with different cash-injection
schedules so we get multiple scenarios without re-running each strategy.

Usage:
    uv run python precompute_backtest.py [--output PATH] [--years N]

Output JSON shape:
    {
        "computedAt": "...",
        "scenarios": {
            "lump_sum_10k": {label, capital, injectionSchedule, markets: {us, cn}},
            "dca_feb_aug":  {label, capital, injectionSchedule, markets: {us, cn}}
        }
    }
"""

import argparse
import json
import os
import sys
from datetime import datetime

import numpy as np
import pandas as pd

import data_manager
import metrics
import universe_history as uh
from bridge import (
    MIN_TRADE_VALUE,
    STRATEGY_META,
    _build_strategy_result,
    _sanitize_for_json,
    benchmark_curve,
    compute_rebalance_trades,
)
from main import backtest
from trader import (
    strategy_extra_tickers,
    strategy_supports_market,
    ETF_STRATEGY_UNIVERSES,
    MONITOR_CAPITAL,
    STRATEGY_REGISTRY,
    convert_quote_currency,
    filter_tradable_tickers,
    price_request,
)
from universe import UNIVERSES


# Initial capital per market, in local currency — the same amounts the live
# monitor starts with (US $10k, CN ¥70k ≈ $10k). Scenario amounts below are
# multiples of it, so both markets model the same real money.
MARKET_CAPITAL = dict(MONITOR_CAPITAL)


SCENARIOS = [
    {
        "id": "lump_sum_10k",
        "label": "起始资金一次性投入",
        "description": "US $10,000 / CN ¥70,000 一次性投入，无后续追加",
        "injection": None,
    },
    {
        "id": "dca_feb_aug",
        "label": "起始资金 + 每年 2/8 月追加一半",
        "description": "US $10,000 + 每年 2/8 月各 $5,000；CN ¥70,000 + 每年 2/8 月各 ¥35,000",
        "injection": {
            "amount_ratio": 0.5,
            "frequency": "yearly_months",
            "months": [2, 8],
            "startYear": 1,
        },
    },
]


def _injection_for_market(injection: dict | None, capital: float) -> dict | None:
    """Materialize a scenario's injection schedule in a market's currency."""
    if not injection:
        return None
    out = {k: v for k, v in injection.items() if k != "amount_ratio"}
    out["amount"] = round(capital * injection["amount_ratio"], 2)
    return out


def log(msg: str) -> None:
    print(f"[precompute] {msg}", file=sys.stderr)


def _equity_series_from_curve(curve: list[dict]) -> pd.Series:
    """Reconstruct a pandas Series from a precomputed equity curve list."""
    if not curve:
        return pd.Series(dtype=float)
    dates = pd.to_datetime([p["date"] for p in curve])
    values = [float(p["value"]) for p in curve]
    return pd.Series(values, index=dates)


def run_market_raw(
    market: str,
    strategy_names: list[str],
    years: int,
    capital: float,
    tx_cost: float,
    fixed_fee: float,
    fee_base: float = 0.0,
    fee_per_share: float = 0.0,
) -> dict | None:
    """Run all strategies once; return raw equity series + metadata per strategy."""
    stock_strategies = []
    etf_strategies = []
    for name in strategy_names:
        base_name = name.removeprefix("sim_")
        if base_name in ETF_STRATEGY_UNIVERSES:
            etf_strategies.append(name)
        else:
            stock_strategies.append(name)

    raw_strategies: list[dict] = []
    benchmark_label = "SPY"
    bench_curve = []
    start_date = None
    end_date = None

    # --- Stock strategies ---
    if stock_strategies:
        universe = UNIVERSES[market]
        tickers = universe["fetch"]()
        benchmark_ticker = universe["benchmark"]
        benchmark_label = universe["benchmark_label"]

        pit_intervals = None
        if market == "us":
            pit_intervals = uh.load_sp500_history()
            historical_tickers = uh.all_tickers_ever(pit_intervals)
            all_tickers = sorted(set(tickers + historical_tickers + [benchmark_ticker])
                                 | {t for n in stock_strategies for t in strategy_extra_tickers(n)})
            log(f"PIT universe: {len(all_tickers)} tickers")
        else:
            all_tickers = sorted(set(tickers + [benchmark_ticker])
                                 | {t for n in stock_strategies for t in strategy_extra_tickers(n)})
            log(f"WARNING: {market.upper()} lacks PIT history; using current constituents")

        log(f"Fetching price data for {len(all_tickers)} tickers...")
        data = data_manager.update(market, all_tickers, with_open=False)
        if isinstance(data, tuple):
            data = data[0]

        if years:
            cutoff = data.index[-1] - pd.DateOffset(years=years)
            data = data[data.index >= cutoff]

        if pit_intervals is not None:
            log("Applying PIT membership mask...")
            data = uh.mask_prices(data, pit_intervals)
            tickers = [
                t for t in data.columns
                if t != benchmark_ticker and data[t].notna().any()
            ]
        else:
            tickers = filter_tradable_tickers(data, tickers)

        log(f"Universe: {len(tickers)} stocks, {data.index[0].date()} to {data.index[-1].date()}")

        for name in stock_strategies:
            log(f"  Running: {name}")
            try:
                strategy = STRATEGY_REGISTRY[name]()
                # Mixed stock+ETF strategies also read their fixed ETFs.
                cols = tickers + [t for t in strategy_extra_tickers(name)
                                  if t in data.columns and t not in tickers]

                equity = backtest(
                    strategy, data[cols],
                    initial_capital=capital,
                    transaction_cost=tx_cost,
                    fixed_fee=fixed_fee,
                    fee_base=fee_base,
                    fee_per_share=fee_per_share,
                    min_trade_value=MIN_TRADE_VALUE,
                )

                weights = strategy.generate_signals(data[cols])
                weights = weights.reindex(data[cols].index).fillna(0.0)
                trade_events = compute_rebalance_trades(weights)

                raw_strategies.append({"name": name, "equity": equity, "trades": trade_events})
                log(f"    OK: {name} (final: ${equity.iloc[-1]:,.0f})")
            except Exception as e:
                log(f"    FAILED: {name}: {e}")

        # Benchmark on exactly the strategies' dates
        bench_curve = benchmark_curve(data[benchmark_ticker], data.index, capital)
        start_date = data.index[0].strftime("%Y-%m-%d")
        end_date = data.index[-1].strftime("%Y-%m-%d")

    # --- ETF strategies ---
    if etf_strategies:
        for name in etf_strategies:
            log(f"  Running ETF strategy: {name}")
            try:
                base_name = name.removeprefix("sim_")
                etf_tickers = ETF_STRATEGY_UNIVERSES[base_name]
                all_etf = sorted(set(etf_tickers + ["SPY"]))

                etf_data = data_manager.update("etfs", price_request(market, all_etf), with_open=False)
                if isinstance(etf_data, tuple):
                    etf_data = etf_data[0]
                etf_data = convert_quote_currency(etf_data)
                if years:
                    cutoff = etf_data.index[-1] - pd.DateOffset(years=years)
                    etf_data = etf_data[etf_data.index >= cutoff]

                strategy = STRATEGY_REGISTRY[name]()

                tradable = filter_tradable_tickers(etf_data, etf_tickers)
                equity = backtest(
                    strategy, etf_data[tradable],
                    initial_capital=capital,
                    transaction_cost=tx_cost,
                    fixed_fee=fixed_fee,
                    fee_base=fee_base,
                    fee_per_share=fee_per_share,
                    min_trade_value=MIN_TRADE_VALUE,
                )

                weights = strategy.generate_signals(etf_data[tradable])
                weights = weights.reindex(etf_data[tradable].index).fillna(0.0)
                trade_events = compute_rebalance_trades(weights)

                raw_strategies.append({"name": name, "equity": equity, "trades": trade_events})
                log(f"    OK: {name} (final: ${equity.iloc[-1]:,.0f})")
            except Exception as e:
                log(f"    FAILED: {name}: {e}")

        if not stock_strategies:
            benchmark_label = "SPY (Benchmark)"
            spy_data = data_manager.update("etfs", ["SPY"], with_open=False)
            if isinstance(spy_data, tuple):
                spy_data = spy_data[0]
            if years:
                cutoff = spy_data.index[-1] - pd.DateOffset(years=years)
                spy_data = spy_data[spy_data.index >= cutoff]
            if "SPY" in spy_data.columns:
                bench_curve = benchmark_curve(spy_data["SPY"], spy_data.index, capital)
                start_date = spy_data.index[0].strftime("%Y-%m-%d")
                end_date = spy_data.index[-1].strftime("%Y-%m-%d")

    if not raw_strategies:
        return None

    return {
        "market": market,
        "strategy_runs": raw_strategies,
        "benchmark_label": benchmark_label,
        "benchmark_curve": bench_curve,
        "start_date": start_date,
        "end_date": end_date,
        "tx_cost": tx_cost,
        "fixed_fee": fixed_fee,
    }


def build_scenario_for_market(
    raw: dict,
    scenario: dict,
    years: int,
    capital: float,
) -> dict:
    """Materialize a scenario (with optional injections) for one market."""
    injection_schedule = _injection_for_market(scenario.get("injection"), capital)
    benchmark_curve = raw["benchmark_curve"]
    benchmark_label = raw["benchmark_label"]

    strategy_results = []
    for run in raw["strategy_runs"]:
        equity = run["equity"]
        try:
            result = _build_strategy_result(
                run["name"], equity, run["trades"], injection_schedule, capital,
            )
            strategy_results.append(result)
        except Exception as e:
            log(f"    [scenario={scenario['id']}] FAILED {run['name']}: {e}")

    # Re-scale benchmark for DCA scenarios so the comparison stays apples-to-apples.
    final_bench_curve = benchmark_curve
    if injection_schedule and benchmark_curve:
        bench_series = _equity_series_from_curve(benchmark_curve)
        bench_returns = bench_series.pct_change().fillna(0.0)
        from bridge import build_injection_schedule
        from research.permanent_yearly import equity_with_cashflows
        contribs, _ = build_injection_schedule(bench_series, injection_schedule)
        bench_with_inj = equity_with_cashflows(bench_returns, contribs, capital)
        final_bench_curve = [
            {"date": d.strftime("%Y-%m-%d"), "value": round(float(v), 2)}
            for d, v in bench_with_inj.items()
        ]

    return {
        "meta": {
            "market": raw["market"],
            "years": years,
            "initialCapital": capital,
            "startDate": raw["start_date"],
            "endDate": raw["end_date"],
            "txCost": raw["tx_cost"],
            "fixedFee": raw["fixed_fee"],
            "injectionSchedule": injection_schedule,
        },
        "strategies": strategy_results,
        "benchmark": {
            "name": benchmark_label,
            "equityCurve": final_bench_curve,
        },
    }


def main():
    parser = argparse.ArgumentParser(description="Pre-compute backtest results")
    parser.add_argument(
        "--output", default=None,
        help="Output directory (default: ../stock-agent/data/precomputed)",
    )
    parser.add_argument(
        "--years", type=int, default=10,
        help="Backtest period in years (default: 10)",
    )
    args = parser.parse_args()

    output_dir = args.output or os.path.join(
        os.path.dirname(os.path.abspath(__file__)), "..", "stock-agent", "data", "precomputed"
    )
    os.makedirs(output_dir, exist_ok=True)

    us_strategies = []
    cn_strategies = []
    for name in STRATEGY_REGISTRY:
        meta = STRATEGY_META.get(name, {"markets": ["us", "cn"]})
        markets = [m for m in meta.get("markets", ["us", "cn"]) if strategy_supports_market(name, m)]
        if "us" in markets:
            us_strategies.append(name)
        if "cn" in markets:
            cn_strategies.append(name)

    us_capital = MARKET_CAPITAL["us"]
    cn_capital = MARKET_CAPITAL["cn"]

    # Run each market once (cached price data → fast subsequent runs across scenarios).
    log("=" * 60)
    log(f"US MARKET: {len(us_strategies)} strategies @ ${us_capital:,.0f}")
    log("=" * 60)
    # US fee schedule (IBKR-style tiered):
    #   fee = max(2.0, 1.88 + 0.009 * shares)
    us_raw = run_market_raw(
        "us", us_strategies,
        years=args.years, capital=us_capital,
        tx_cost=0.001, fixed_fee=2.0,
        fee_base=1.88, fee_per_share=0.009,
    )

    log("=" * 60)
    log(f"CN MARKET: {len(cn_strategies)} strategies @ ¥{cn_capital:,.0f}")
    log("=" * 60)
    cn_raw = run_market_raw(
        "cn", cn_strategies,
        years=args.years, capital=cn_capital,
        tx_cost=0.001, fixed_fee=5.0,
    )

    # Materialize every scenario × market by applying its injection schedule.
    scenarios_out: dict[str, dict] = {}
    for scenario in SCENARIOS:
        log(f"--- scenario: {scenario['id']} ---")
        scenarios_out[scenario["id"]] = {
            "id": scenario["id"],
            "label": scenario["label"],
            "description": scenario["description"],
            "capital": us_capital,
            "capitalByMarket": {"us": us_capital, "cn": cn_capital},
            "injectionSchedule": _injection_for_market(scenario["injection"], us_capital),
            "markets": {},
        }
        if us_raw:
            scenarios_out[scenario["id"]]["markets"]["us"] = build_scenario_for_market(
                us_raw, scenario, args.years, us_capital
            )
        if cn_raw:
            scenarios_out[scenario["id"]]["markets"]["cn"] = build_scenario_for_market(
                cn_raw, scenario, args.years, cn_capital
            )

    output = {
        "computedAt": datetime.now().isoformat(),
        "config": {"years": args.years, "initialCapital": MARKET_CAPITAL},
        "scenarios": scenarios_out,
    }

    output_path = os.path.join(output_dir, "backtest_results.json")
    with open(output_path, "w") as f:
        json.dump(_sanitize_for_json(output), f, default=str)

    file_size = os.path.getsize(output_path) / (1024 * 1024)
    log("=" * 60)
    log(f"DONE: Written to {output_path} ({file_size:.1f} MB)")
    for sid, sc in scenarios_out.items():
        for mk, mv in sc["markets"].items():
            log(f"  [{sid}] {mk.upper()}: {len(mv['strategies'])} strategies")
    log("=" * 60)


if __name__ == "__main__":
    main()