"""End-to-end reproduction of the MAX-Effect / lottery-preference strategy.

    python run.py            # uses committed offline caches in data/
    python run.py --refresh  # re-downloads prices + factors from the network

Pipeline
--------
1. Load daily prices (Yahoo) and FF factors (Ken French), from cache or network.
2. Build the MAX(N) signal; select N *in-sample only*, evaluate *out-of-sample*.
3. Run the monthly long-short quintile backtest with transaction costs.
4. Quantify the MAX <-> idiosyncratic-vol link.
5. Run a Fama-French (+ momentum) regression to isolate alpha.
6. Write charts to results/*.png and metrics to results/metrics.json.
7. Print a concise summary to stdout.
"""
from __future__ import annotations

import argparse
import json
import random

import numpy as np
import pandas as pd

from src import config, data
from src.config import BacktestParams
from src.backtest import performance_stats, run_backtest
from src.factors import (
    daily_factors_to_monthly,
    factor_regression,
    max_ivol_correlation,
)
from src.plots import (
    plot_factor_loadings,
    plot_long_short_equity,
    plot_max_ivol_scatter,
    plot_quintile_returns,
)
from src.signals import daily_returns, idiosyncratic_vol, max_signal


def _seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)


def _slice(s: pd.Series, start=None, end=None) -> pd.Series:
    return s.loc[start:end]


def select_max_n(prices: pd.DataFrame, grid, is_end: str) -> tuple[int, dict]:
    """Select MAX(N) using only the in-sample window (net long-short Sharpe)."""
    scores = {}
    for n in grid:
        res = run_backtest(prices, BacktestParams(max_n=n))
        is_net = _slice(res.long_short_net, end=is_end)
        scores[n] = performance_stats(is_net)
    best = max(grid, key=lambda n: scores[n]["sharpe"])
    return best, scores


def main(refresh: bool = False) -> None:
    _seed_everything(config.SEED)
    config.ensure_dirs()

    # ---- 1. Data ---------------------------------------------------------
    prices = data.load_prices(refresh=refresh)
    factors_daily = data.load_ff_factors(refresh=refresh)
    prices = prices.loc[config.START:config.END]
    print(f"Loaded prices: {prices.shape[1]} tickers x {prices.shape[0]} days "
          f"({prices.index.min().date()} -> {prices.index.max().date()})")

    factors_monthly = daily_factors_to_monthly(factors_daily)

    # ---- 2. In-sample parameter selection -------------------------------
    best_n, grid_scores = select_max_n(prices, config.MAX_N_GRID, config.IS_END)
    print(f"Selected MAX(N={best_n}) on the in-sample window "
          f"(candidates {list(config.MAX_N_GRID)}).")

    # ---- 3. Backtest with the chosen N ----------------------------------
    params = BacktestParams(max_n=best_n)
    res = run_backtest(prices, params)
    qret = res.quintile_returns

    full_stats = performance_stats(res.long_short_net)
    is_stats = performance_stats(_slice(res.long_short_net, end=config.IS_END))
    oos_stats = performance_stats(_slice(res.long_short_net, start=config.OOS_START))
    gross_stats = performance_stats(res.long_short_gross)

    # ---- 4. MAX vs idiosyncratic volatility -----------------------------
    rets = daily_returns(prices)
    max_panel = max_signal(rets, n_high=best_n)
    ivol_panel = idiosyncratic_vol(rets, factors_daily)
    corr = max_ivol_correlation(max_panel, ivol_panel)
    print(f"MAX vs IVOL: pooled Pearson={corr['pooled_pearson']:.2f}, "
          f"avg cross-sectional={corr['avg_cross_sectional']:.2f}")

    # ---- 5. Factor regressions ------------------------------------------
    reg_full = factor_regression(res.long_short_net, factors_monthly)
    reg_oos = factor_regression(
        _slice(res.long_short_net, start=config.OOS_START), factors_monthly
    )

    # ---- 6. Charts -------------------------------------------------------
    plot_quintile_returns(qret, config.RESULTS_DIR / "quintile_returns.png")
    plot_long_short_equity(
        res.long_short_gross, res.long_short_net,
        config.RESULTS_DIR / "long_short_equity.png", oos_start=config.OOS_START,
    )
    plot_max_ivol_scatter(max_panel, ivol_panel, config.RESULTS_DIR / "max_vs_ivol.png")
    plot_factor_loadings(reg_full, config.RESULTS_DIR / "factor_loadings.png")

    # ---- 7. Persist metrics ---------------------------------------------
    qret.to_csv(config.RESULTS_DIR / "quintile_returns.csv")
    reg_full["table"].to_csv(config.RESULTS_DIR / "factor_regression.csv")

    metrics = {
        "config": {
            "universe_size": int(prices.shape[1]),
            "sample": [str(prices.index.min().date()), str(prices.index.max().date())],
            "in_sample_end": config.IS_END,
            "oos_start": config.OOS_START,
            "selected_max_n": int(best_n),
            "cost_bps_one_way": config.COST_BPS,
            "n_quintiles": config.N_QUINTILES,
        },
        "quintile_annualized_return": {
            c: float(qret[c].mean() * 12.0) for c in qret.columns
        },
        "long_short": {
            "gross": gross_stats,
            "net_full": full_stats,
            "net_in_sample": is_stats,
            "net_out_of_sample": oos_stats,
            "avg_monthly_turnover": float(res.turnover.mean()),
        },
        "in_sample_grid_sharpe": {
            str(n): grid_scores[n]["sharpe"] for n in grid_scores
        },
        "max_vs_ivol": corr,
        "factor_regression_full": {
            "factors": reg_full["factors"],
            "alpha_monthly": reg_full["alpha_monthly"],
            "alpha_annualized": reg_full["alpha_annualized"],
            "alpha_tstat": reg_full["alpha_tstat"],
            "alpha_pvalue": reg_full["alpha_pvalue"],
            "r_squared": reg_full["r_squared"],
            "loadings": {
                idx: {"coef": float(row["coef"]), "t_stat": float(row["t_stat"])}
                for idx, row in reg_full["table"].iterrows()
            },
        },
        "factor_regression_oos": {
            "alpha_annualized": reg_oos["alpha_annualized"],
            "alpha_tstat": reg_oos["alpha_tstat"],
            "r_squared": reg_oos["r_squared"],
            "n_obs": reg_oos["n_obs"],
        },
    }
    with open(config.RESULTS_DIR / "metrics.json", "w") as fh:
        json.dump(metrics, fh, indent=2)

    _print_summary(metrics, reg_full)


def _print_summary(metrics: dict, reg_full: dict) -> None:
    c = metrics["config"]
    ls = metrics["long_short"]
    print("\n" + "=" * 66)
    print("  MAX-EFFECT / LOTTERY-PREFERENCE STRATEGY  -  RESULTS SUMMARY")
    print("=" * 66)
    print(f"  Universe: {c['universe_size']} liquid large caps   "
          f"Sample: {c['sample'][0]} -> {c['sample'][1]}")
    print(f"  Signal: MAX(N={c['selected_max_n']})  |  "
          f"long Q1 (low MAX) / short Q{c['n_quintiles']} (high MAX)  |  "
          f"cost {c['cost_bps_one_way']:.0f} bps/side")
    print("-" * 66)
    print("  Annualised mean return by MAX quintile:")
    for q, v in metrics["quintile_annualized_return"].items():
        bar = "#" * max(0, int(abs(v) * 200))
        print(f"    {q}: {v*100:7.2f}%  {bar}")
    print("-" * 66)
    print("  Long-short (low MAX - high MAX), net of costs:")
    for label, key in [("Full sample", "net_full"),
                       ("In-sample  ", "net_in_sample"),
                       ("Out-of-samp", "net_out_of_sample")]:
        s = ls[key]
        print(f"    {label}: ann ret {s['ann_return']*100:6.2f}%  "
              f"vol {s['ann_vol']*100:5.2f}%  Sharpe {s['sharpe']:5.2f}  "
              f"t {s['t_stat']:5.2f}  maxDD {s['max_drawdown']*100:6.1f}%")
    print(f"    Gross full Sharpe {ls['gross']['sharpe']:.2f}  ->  "
          f"net full Sharpe {ls['net_full']['sharpe']:.2f}   "
          f"(avg turnover {ls['avg_monthly_turnover']:.2f}/mo)")
    print("-" * 66)
    mv = metrics["max_vs_ivol"]
    print(f"  MAX <-> IVOL correlation: pooled Pearson {mv['pooled_pearson']:.2f}, "
          f"avg cross-sectional {mv['avg_cross_sectional']:.2f}")
    print("-" * 66)
    fr = metrics["factor_regression_full"]
    print(f"  Fama-French + MOM regression ({', '.join(fr['factors'])}):")
    print(f"    Alpha: {fr['alpha_annualized']*100:6.2f}% / yr "
          f"(t = {fr['alpha_tstat']:.2f}, p = {fr['alpha_pvalue']:.3f}), "
          f"R^2 = {fr['r_squared']:.2f}")
    for name, ld in fr["loadings"].items():
        if name == "alpha":
            continue
        print(f"      beta[{name:<7}] = {ld['coef']:+.3f}  (t = {ld['t_stat']:+.2f})")
    oos = metrics["factor_regression_oos"]
    print(f"    OOS alpha: {oos['alpha_annualized']*100:.2f}% / yr "
          f"(t = {oos['alpha_tstat']:.2f}, n = {oos['n_obs']})")
    print("=" * 66)
    print("  Charts + metrics written to results/. Done.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--refresh", action="store_true",
                        help="re-download data from the network instead of using caches")
    args = parser.parse_args()
    main(refresh=args.refresh)
