"""
PIT-compliant Sharpe 1.5+ blend: V5 ETF timing + PIT stock-picking + cross-asset momentum.

Combines three uncorrelated alpha sources with a vol-target overlay.
All components are PIT-safe (ETF-only or membership-masked).

Run:
    uv run python -m research.sharpe_blend
"""
from __future__ import annotations

import os
import sys

import numpy as np
import pandas as pd

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from research.permanent_yearly import load_etfs
from research.pit_backtest import load_pit_prices, pit_universe
from research.pit_optimization import PITEnsemble, compute_metrics
from research.trend_rider_robustness import portfolio_returns, evaluate_weights
from research.trend_rider_v6_eval import load_combined_panel
from strategies.cross_asset_momentum import CrossAssetMomentum
from strategies.trend_rider_v5 import TrendRiderV5


# ---------------------------------------------------------------------------
# Data loading
# ---------------------------------------------------------------------------

def load_all_data() -> tuple[pd.DataFrame, pd.DataFrame]:
    """Return (etf_panel, pit_stock_prices) aligned to common dates."""
    # ETF panel for V5 and cross-asset
    etf_panel = load_combined_panel()

    # Ensure cross-asset ETFs are present (TLT, IEF)
    extra_etfs = ["TLT", "IEF"]
    missing = [t for t in extra_etfs if t not in etf_panel.columns]
    if missing:
        extra = load_etfs(missing, start="2013-06-01")
        extra = extra.reindex(etf_panel.index).ffill()
        etf_panel = etf_panel.join(extra, how="left")

    # PIT-masked stock prices
    pit_prices = load_pit_prices()
    pit_masked = pit_universe(pit_prices)

    return etf_panel, pit_masked


# ---------------------------------------------------------------------------
# Strategy runners — produce daily returns series
# ---------------------------------------------------------------------------

def run_v5(panel: pd.DataFrame, start: str = "2017-06-01") -> pd.Series:
    """TrendRiderV5 daily returns."""
    v5 = TrendRiderV5()
    weights = v5.generate_signals(panel)
    rets = portfolio_returns(weights, panel, transaction_cost=0.001)
    return rets.loc[start:]


def run_pit_stock(pit_prices: pd.DataFrame, start: str = "2017-06-01") -> pd.Series:
    """PIT stock-picking (cross-sectional momentum) daily returns."""
    strat = PITEnsemble(
        top_n=12, rebal_freq=42, mom_blend=1.0,
        asym_vol=True, asym_vol_floor=0.50,
        dd_dampen=False,
    )
    weights = strat.generate_signals(pit_prices)
    daily_rets = (weights * pit_prices.pct_change().fillna(0.0)).sum(axis=1)
    return daily_rets.loc[start:]


def run_cross_asset(panel: pd.DataFrame, start: str = "2017-06-01") -> pd.Series:
    """Cross-asset time-series momentum daily returns."""
    strat = CrossAssetMomentum(lookback=252, top_k=3, rebal_freq=21, vol_scale=True)
    weights = strat.generate_signals(panel)
    rets = portfolio_returns(weights, panel, transaction_cost=0.001)
    return rets.loc[start:]


# ---------------------------------------------------------------------------
# Vol-target overlay (standalone, operates on combined returns)
# ---------------------------------------------------------------------------

def vol_target_returns(
    combined_rets: pd.Series,
    target_vol: float = 0.18,
    vol_window: int = 20,
) -> pd.Series:
    """Scale combined returns by min(1, target_vol / realized_vol)."""
    realized = combined_rets.rolling(vol_window).std(ddof=1) * np.sqrt(252)
    realized = realized.shift(1).fillna(target_vol)
    scale = (target_vol / realized.replace(0.0, np.nan)).clip(upper=1.0).fillna(1.0)
    return combined_rets * scale


# ---------------------------------------------------------------------------
# Blend engine
# ---------------------------------------------------------------------------

def blend_returns(
    rets_v5: pd.Series,
    rets_stock: pd.Series,
    rets_xasset: pd.Series,
    w_v5: float = 0.50,
    w_stock: float = 0.30,
    w_xasset: float = 0.20,
) -> pd.Series:
    """Weighted blend of three strategy return streams."""
    # Align to common dates
    idx = rets_v5.index.intersection(rets_stock.index).intersection(rets_xasset.index)
    return (w_v5 * rets_v5.loc[idx]
            + w_stock * rets_stock.loc[idx]
            + w_xasset * rets_xasset.loc[idx])


def inverse_vol_weights(
    rets_v5: pd.Series,
    rets_stock: pd.Series,
    rets_xasset: pd.Series,
    window: int = 63,
) -> tuple[float, float, float]:
    """Compute inverse-vol weights from trailing realized vol."""
    vols = pd.DataFrame({
        "v5": rets_v5.rolling(window).std() * np.sqrt(252),
        "stock": rets_stock.rolling(window).std() * np.sqrt(252),
        "xasset": rets_xasset.rolling(window).std() * np.sqrt(252),
    }).iloc[-1]
    inv = 1.0 / vols.replace(0, np.nan)
    w = inv / inv.sum()
    return w["v5"], w["stock"], w["xasset"]


# ---------------------------------------------------------------------------
# Sweep
# ---------------------------------------------------------------------------

BLEND_CONFIGS = [
    ("V5=50/Stock=30/XA=20", 0.50, 0.30, 0.20),
    ("V5=40/Stock=40/XA=20", 0.40, 0.40, 0.20),
    ("V5=60/Stock=20/XA=20", 0.60, 0.20, 0.20),
    ("V5=50/Stock=25/XA=25", 0.50, 0.25, 0.25),
    ("V5=45/Stock=35/XA=20", 0.45, 0.35, 0.20),
    ("V5=55/Stock=25/XA=20", 0.55, 0.25, 0.20),
]

VOL_TARGETS = [None, 0.15, 0.18, 0.20, 0.22, 0.25]


def run_sweep(rets_v5, rets_stock, rets_xasset) -> pd.DataFrame:
    """Sweep blend configs × vol targets, return summary DataFrame."""
    rows = []

    # Add inverse-vol config
    iv_w = inverse_vol_weights(rets_v5, rets_stock, rets_xasset)
    configs = list(BLEND_CONFIGS) + [
        (f"InvVol({iv_w[0]:.0%}/{iv_w[1]:.0%}/{iv_w[2]:.0%})", *iv_w)
    ]

    for name, wv, ws, wx in configs:
        combined = blend_returns(rets_v5, rets_stock, rets_xasset, wv, ws, wx)
        for tgt in VOL_TARGETS:
            if tgt is not None:
                final = vol_target_returns(combined, target_vol=tgt)
                label = f"{name} | VT={tgt}"
            else:
                final = combined
                label = f"{name} | no-VT"
            m = compute_metrics(final)
            m["label"] = label
            m["w_v5"] = wv
            m["w_stock"] = ws
            m["w_xasset"] = wx
            m["vol_target"] = tgt
            rows.append(m)

    df = pd.DataFrame(rows)
    df = df.sort_values("sharpe", ascending=False).reset_index(drop=True)
    return df


# ---------------------------------------------------------------------------
# Validation helpers
# ---------------------------------------------------------------------------

def is_oos_split(rets: pd.Series, split_date="2023-01-01"):
    """Split returns into IS and OOS."""
    is_rets = rets[rets.index < split_date]
    oos_rets = rets[rets.index >= split_date]
    return is_rets, oos_rets


def block_bootstrap(rets: pd.Series, n_boot: int = 5000, block_size: int = 63) -> np.ndarray:
    """Block bootstrap of annualized Sharpe ratio."""
    n = len(rets)
    arr = rets.values
    sharpes = np.empty(n_boot)
    rng = np.random.default_rng(42)
    n_blocks = int(np.ceil(n / block_size))

    for i in range(n_boot):
        starts = rng.integers(0, n - block_size, size=n_blocks)
        sample = np.concatenate([arr[s:s + block_size] for s in starts])[:n]
        mu = sample.mean()
        sigma = sample.std(ddof=1)
        sharpes[i] = mu / sigma * np.sqrt(252) if sigma > 0 else 0.0
    return sharpes


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    print("=" * 80)
    print("PIT-Compliant Multi-Strategy Blend — Sharpe 1.5+ Target")
    print("=" * 80)

    # Load data
    print("\n[1] Loading data...")
    etf_panel, pit_masked = load_all_data()

    # Run individual strategies
    print("\n[2] Running individual strategies...")
    rets_v5 = run_v5(etf_panel)
    rets_stock = run_pit_stock(pit_masked)
    rets_xasset = run_cross_asset(etf_panel)

    # Individual metrics
    print("\n--- Individual Strategy Metrics ---")
    for name, r in [("V5 ETF Timing", rets_v5),
                    ("PIT Stock Momentum", rets_stock),
                    ("Cross-Asset Momentum", rets_xasset)]:
        m = compute_metrics(r)
        print(f"  {name:<25s}  Sharpe={m['sharpe']:5.2f}  CAGR={m['cagr']*100:5.1f}%  "
              f"Vol={m['vol']*100:5.1f}%  MaxDD={m['max_dd']*100:5.1f}%")

    # Correlation diagnostic
    print("\n--- Correlation Matrix (daily returns) ---")
    corr_df = pd.DataFrame({
        "V5": rets_v5, "Stock": rets_stock, "XAsset": rets_xasset
    }).dropna()
    corr = corr_df.corr()
    print(corr.to_string(float_format=lambda x: f"{x:.3f}"))

    # Rolling correlation
    print("\n--- Rolling 63d Correlations (mean / max) ---")
    for pair in [("V5", "Stock"), ("V5", "XAsset"), ("Stock", "XAsset")]:
        roll = corr_df[pair[0]].rolling(63).corr(corr_df[pair[1]])
        print(f"  {pair[0]:>8s} vs {pair[1]:<8s}: mean={roll.mean():.3f}  max={roll.max():.3f}")

    # Sweep
    print("\n[3] Running blend sweep...")
    results = run_sweep(rets_v5, rets_stock, rets_xasset)

    print("\n--- Top 15 Configurations ---")
    print(f"  {'Label':<50s} {'Sharpe':>7s} {'CAGR':>7s} {'Vol':>7s} {'MaxDD':>7s} {'Calmar':>7s}")
    for _, row in results.head(15).iterrows():
        print(f"  {row['label']:<50s} {row['sharpe']:7.2f} "
              f"{row['cagr']*100:6.1f}% {row['vol']*100:6.1f}% "
              f"{row['max_dd']*100:6.1f}% {row['calmar']:6.2f}")

    # Best config validation
    best = results.iloc[0]
    print(f"\n--- Best Config: {best['label']} ---")
    best_rets = blend_returns(rets_v5, rets_stock, rets_xasset,
                              best["w_v5"], best["w_stock"], best["w_xasset"])
    if best["vol_target"] is not None:
        best_rets = vol_target_returns(best_rets, target_vol=best["vol_target"])

    # IS/OOS
    print("\n[4] IS/OOS Validation (split: 2023-01-01)...")
    is_rets, oos_rets = is_oos_split(best_rets)
    is_m = compute_metrics(is_rets)
    oos_m = compute_metrics(oos_rets)
    print(f"  IS  (2017-2022): Sharpe={is_m['sharpe']:5.2f}  CAGR={is_m['cagr']*100:5.1f}%  MaxDD={is_m['max_dd']*100:5.1f}%")
    print(f"  OOS (2023-2026): Sharpe={oos_m['sharpe']:5.2f}  CAGR={oos_m['cagr']*100:5.1f}%  MaxDD={oos_m['max_dd']*100:5.1f}%")
    print(f"  OOS/IS ratio:    {oos_m['sharpe']/is_m['sharpe']:.2f}" if is_m['sharpe'] > 0 else "")

    # Bootstrap
    print("\n[5] Block Bootstrap (5000 resamples, block=63d)...")
    boot = block_bootstrap(best_rets, n_boot=5000)
    print(f"  Median Sharpe: {np.median(boot):.2f}")
    print(f"  5th pctile:    {np.percentile(boot, 5):.2f}")
    print(f"  95th pctile:   {np.percentile(boot, 95):.2f}")
    print(f"  P(Sharpe>1.0): {(boot > 1.0).mean()*100:.1f}%")
    print(f"  P(Sharpe>1.3): {(boot > 1.3).mean()*100:.1f}%")
    print(f"  P(Sharpe>1.5): {(boot > 1.5).mean()*100:.1f}%")

    # Parameter sensitivity
    print("\n[6] Parameter Sensitivity (±perturbation on blend weights)...")
    base_w = (best["w_v5"], best["w_stock"], best["w_xasset"])
    perturbations = [
        ("base", 0, 0, 0),
        ("+10% V5", 0.10, -0.05, -0.05),
        ("-10% V5", -0.10, 0.05, 0.05),
        ("+10% Stock", -0.05, 0.10, -0.05),
        ("-10% Stock", 0.05, -0.10, 0.05),
    ]
    for pname, dv, ds, dx in perturbations:
        wv = max(0.05, base_w[0] + dv)
        ws = max(0.05, base_w[1] + ds)
        wx = max(0.05, base_w[2] + dx)
        total = wv + ws + wx
        wv, ws, wx = wv/total, ws/total, wx/total
        r = blend_returns(rets_v5, rets_stock, rets_xasset, wv, ws, wx)
        if best["vol_target"] is not None:
            r = vol_target_returns(r, target_vol=best["vol_target"])
        m = compute_metrics(r)
        print(f"  {pname:<15s}: Sharpe={m['sharpe']:5.2f}  CAGR={m['cagr']*100:5.1f}%")

    print("\n" + "=" * 80)
    print("Done.")


if __name__ == "__main__":
    main()