"""
Final check of the recommended US portfolio: three sleeves with different
return sources, equal capital, rebalanced monthly between sleeves.

  momentum  : overnight-return momentum + distance from 126d low, top 10, 10-day
              rebalance, hold while rank <= 30 (the live overnight_recovery_top10)
  defensive : sector-neutral quality (GP/A, CFO/A, ROE, accruals, leverage) +
              low risk (idiosyncratic vol, beta), top 20, monthly, buffer 2
              (price-only alternative: low risk alone)
  trend     : TrendRider V7 VT24 on leveraged ETFs (live trend_rider_v7_vt24)

    PYTHONPATH=. uv run python -m research.lab.trend_rider_long   # once: data/research/tr_long.pkl
    PYTHONPATH=. uv run --with numba python -m research.lab.us_final

The live version is ``us_three_sleeve`` (strategies/sleeve_portfolio.py).
"""

from __future__ import annotations

import json
import os
import pickle
import warnings

import numpy as np
import pandas as pd

from research.lab import fundamentals as fu
from research.lab.engine import Costs
from research.lab.factor_scan import RESULTS, load
from research.lab.report import deflated_sharpe, stats
from research.lab.rotation import factor_ranks
from research.lab.us_sweep import _blend, run, sector_rank

warnings.filterwarnings("ignore")

SPAN = ("2011-01-03", "2026-09-24")
LONG = ("2006-09-27", "2026-09-24")
BLOCKS = [("2006-09-27", "2010-12-31"), ("2011-01-01", "2014-12-31"), ("2015-01-01", "2018-12-31"),
          ("2019-01-01", "2022-12-31"), ("2023-01-01", "2026-09-24")]


def sleeve_scores(p):
    pr = factor_ranks(p, "us")
    F = fu.load(p)
    Q = _blend(*(sector_rank(F[k], p) for k in ("gp_a", "cfo_a", "roe", "accruals", "leverage")))
    low = _blend(pr["ivol"], pr["low_beta"])
    return {
        "momentum": (_blend(pr["overnight_mom"], pr["rec_126"]), dict(top_n=10, buffer=2.0, every=10)),
        "defensive": (_blend(Q, low), dict(top_n=20, buffer=2.0, every=21)),
        "lowrisk": (low, dict(top_n=20, buffer=2.0, every=21)),
    }


def monthly_blend(R: pd.DataFrame, weights: dict[str, float]) -> pd.Series:
    """Sleeves drift inside each month and are reset to ``weights`` at month
    ends (the sleeves are separate accounts; only capital moves)."""
    cols = list(weights)
    w0 = np.array([weights[c] for c in cols])
    out = []
    for _, g in R[cols].groupby(R.index.to_period("M")):
        grow = (1 + g).cumprod()
        val = grow.mul(w0, axis=1).sum(axis=1)
        prev = np.r_[1.0, val.values[:-1]]
        out.append(pd.Series(val.values / prev - 1, index=g.index))
    return pd.concat(out)


def describe(r: pd.Series, blocks=BLOCKS) -> dict:
    s = stats((1 + r).cumprod())
    bl = {f"{a[:4]}-{b[:4]}": stats((1 + r.loc[a:b]).cumprod())["sharpe"]
          for a, b in blocks if len(r.loc[a:b]) > 200}
    yrs = ((1 + r).groupby(r.index.year).prod() - 1)
    return {**{k: float(v) for k, v in s.items()}, "blocks": bl, "yearly": {int(k): float(v) for k, v in yrs.items()}}


def line(name, d) -> str:
    return (f"{name:44s} cagr {d['cagr']*100:5.1f}% sh {d['sharpe']:.2f} vol {d['vol']*100:4.1f}% "
            f"mdd {d['mdd']*100:5.1f}% calmar {d['calmar']:.2f} worst yr {min(d['yearly'].values())*100:5.1f}% "
            f"blocks [{' '.join(f'{v:.2f}' for v in d['blocks'].values())}]")


def main() -> None:
    p, _ = load("us")
    pf = p.slice()
    pf.costs = Costs(p.costs.buy, p.costs.sell, None, 2.0, 10_000.0)     # + $2 per order on a $10k sleeve
    S = sleeve_scores(p)
    R, Rfee = {}, {}
    for name, (score, kw) in S.items():
        span = LONG if name in ("momentum", "lowrisk") else SPAN
        R[name] = run(p, score.where(p.member), [span], **kw)["_rets"]
        Rfee[name] = run(pf, score.where(p.member), [span], **kw)["_rets"]
    ew = p.close.notna().astype(float).where(p.member)
    R["ew_sp500"] = run(p, ew, [LONG], top_n=10000, every=10)["_rets"]
    with open(os.path.join("data", "research", "tr_long.pkl"), "rb") as fh:
        tr = pickle.load(fh)["equity"]["trend_rider_v7_vt24"]
    R["trend"] = Rfee["trend"] = tr.pct_change().reindex(p.dates).loc["2010-03-01":]   # $2/order already
    R["spy"] = p.benchmark.ffill().pct_change().reindex(p.dates)
    R = pd.DataFrame(R).loc[LONG[0]:LONG[1]]
    Rfee = pd.DataFrame(Rfee).loc[LONG[0]:LONG[1]]

    out = {}
    span = R.loc[SPAN[0]:SPAN[1], ["momentum", "defensive", "lowrisk", "trend", "spy"]].fillna(0.0)
    monthly = (1 + span).resample("ME").prod() - 1
    print("monthly return correlation 2011-2026:\n", monthly.corr().round(2).to_string(), "\n")
    print("=== sleeves (10 bp costs; trend includes $2/order)")
    for c in R.columns:
        r = R[c].dropna()
        out[c] = describe(r)
        print(line(f"{c} {r.index[0].year}-", out[c]))
    three = {"momentum": 1 / 3, "defensive": 1 / 3, "trend": 1 / 3}
    three_px = {"momentum": 1 / 3, "lowrisk": 1 / 3, "trend": 1 / 3}
    print("\n=== portfolios 2011-2026, monthly rebalanced between sleeves")
    for label, w, src in [("3 sleeves (defensive = quality+lowrisk)", three, R),
                          ("3 sleeves (defensive = lowrisk, price only)", three_px, R),
                          ("3 sleeves, + $2/order on $10k sleeves", three, Rfee),
                          ("momentum + defensive (no ETF sleeve)", {"momentum": .5, "defensive": .5}, R),
                          ("momentum + trend", {"momentum": .5, "trend": .5}, R)]:
        r = monthly_blend(src.loc[SPAN[0]:SPAN[1]].fillna(0.0), w)
        out[label] = describe(r)
        print(line(label, out[label]))
    print("\n=== price-only walk-back 2006-2026 (includes 2008; trend sleeve from 2010-03, before that 50/50)")
    rr = R.loc[LONG[0]:LONG[1]]
    pre = monthly_blend(rr.loc[:"2010-02-28", ["momentum", "lowrisk"]].fillna(0.0), {"momentum": .5, "lowrisk": .5})
    post = monthly_blend(rr.loc["2010-03-01":].fillna(0.0), three_px)
    r = pd.concat([pre, post])
    out["walkback price-only"] = describe(r)
    print(line("momentum + lowrisk (+ trend from 2010)", out["walkback price-only"]))
    for c in ("momentum", "lowrisk", "ew_sp500", "spy"):
        print(line(f"  {c} 2006-", describe(rr[c].dropna())))

    # deflated Sharpe: every configuration this and the previous round tried
    from research.lab.trials import count_trials
    n = count_trials(RESULTS, "us")
    if n < 1:
        raise ValueError("A trial ledger is required for DSR")
    r = monthly_blend(R.loc[SPAN[0]:SPAN[1]].fillna(0.0), three)
    dsr = deflated_sharpe(r, n)
    out["dsr"] = {"observed_trials_lower_bound": n, "role": "development", **{k: float(v) for k, v in dsr.items()}}
    print(f"\nDSR of the 3-sleeve book: {dsr['dsr']:.3f} (Sharpe {dsr['sr_annual']:.2f} vs "
          f"expected max of {n} unskilled trials {dsr['sr0_annual']:.2f})")
    print("\nyearly returns (3 sleeves vs SPY):")
    spy_y = out["spy"]["yearly"]
    for y, v in out["3 sleeves (defensive = quality+lowrisk)"]["yearly"].items():
        print(f"  {y}: {v*100:6.1f}%  SPY {spy_y.get(y, float('nan'))*100:6.1f}%")
    with open(os.path.join(RESULTS, "us_final.json"), "w") as fh:
        json.dump(out, fh, indent=1)
    R.to_pickle(os.path.join(RESULTS, "us_final_sleeves.pkl"))


if __name__ == "__main__":
    main()