"""
Kalman Pairs Backtesting Framework.

Full walk-forward backtest with:
  - Realistic execution (transaction costs, capital allocation)
  - Per-trade P&L tracking
  - Side-by-side comparison vs rolling OLS (60-day, 120-day windows)
  - Performance report: CAGR, Sharpe, Sortino, max DD, win rate, turnover
  - Regime-shift stress tests
"""

from __future__ import annotations

import numpy as np
from typing import Optional
from .kalman_filter import KalmanPairsTrader
from .pair_discovery import compute_rolling_ols_hedge

# Import project metrics
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent))
from common.metrics import sharpe, sortino, max_drawdown, win_rate


def backtest_kalman_pairs(
    X: np.ndarray,
    Y: np.ndarray,
    trader: KalmanPairsTrader,
    trade_size_usd: float = 100.0,
    transaction_cost_bps: float = 2.5,
    initial_capital: float = 10000.0,
) -> dict:
    """
    Run a walk-forward backtest for a single pair using Kalman filter.

    Args:
        X, Y: Price series (must be same length).
        trader: Pre-configured KalmanPairsTrader (already initialized).
        trade_size_usd: Notional per leg in USD.
        transaction_cost_bps: Fee per leg in basis points.
        initial_capital: Starting capital.

    Returns:
        dict with: trades list, equity_curve, metrics, final_equity.
    """
    n = min(len(X), len(Y))
    trader.reset()

    capital = initial_capital
    peak_capital = initial_capital
    equity_curve: list[dict] = []
    trades: list[dict] = []
    open_trade: Optional[dict] = None

    fee_rate = transaction_cost_bps / 10000.0  # bps → decimal

    for t in range(n):
        x_t = float(X[t])
        y_t = float(Y[t])
        result = trader.step(x_t, y_t)

        signal = result["signal"]
        beta = result["beta"]

        if signal != 0:
            if open_trade is None:
                # Open position
                entry_x = x_t
                entry_y = y_t
                size_x = trade_size_usd / entry_x if entry_x > 0 else 0
                size_y = trade_size_usd / entry_y if entry_y > 0 else 0

                # Hedge: use current beta
                # If signal = +1: LONG Y (size_y), SHORT X (size_x * beta)
                # If signal = -1: SHORT Y (size_y), LONG X (size_x * beta)
                hedge_notional = size_x * entry_x * abs(beta) if beta else 0
                fee = (trade_size_usd + hedge_notional) * fee_rate

                capital -= fee

                open_trade = {
                    "entry_time": t,
                    "signal": signal,
                    "entry_x": entry_x,
                    "entry_y": entry_y,
                    "beta_at_entry": beta,
                    "size_x": size_x,
                    "size_y": size_y,
                    "fee_paid": fee,
                }
            elif open_trade is not None and signal == -open_trade["signal"]:
                # Close position
                # PnL: (Y exit - Y entry) * size_y * sign + (X entry - X exit) * size_x * beta * sign
                exit_sign = open_trade["signal"]
                pnl_y = (y_t - open_trade["entry_y"]) * open_trade["size_y"] * exit_sign
                pnl_x = (open_trade["entry_x"] - x_t) * open_trade["size_x"] * abs(open_trade["beta_at_entry"]) * exit_sign
                gross_pnl = pnl_y + pnl_x

                exit_notional = abs(y_t * open_trade["size_y"]) + abs(x_t * open_trade["size_x"] * open_trade["beta_at_entry"])
                fee = exit_notional * fee_rate
                net_pnl = gross_pnl - fee

                capital += net_pnl

                trades.append({
                    "entry_time": open_trade["entry_time"],
                    "exit_time": t,
                    "signal": open_trade["signal"],
                    "entry_x": open_trade["entry_x"],
                    "exit_x": x_t,
                    "entry_y": open_trade["entry_y"],
                    "exit_y": y_t,
                    "beta": open_trade["beta_at_entry"],
                    "gross_pnl": round(gross_pnl, 4),
                    "net_pnl": round(net_pnl, 4),
                    "fee": round(open_trade["fee_paid"] + fee, 6),
                    "duration_bars": t - open_trade["entry_time"],
                })
                open_trade = None

        # Track equity
        unrealized = 0.0
        if open_trade is not None:
            exit_sign = open_trade["signal"]
            ur_y = (y_t - open_trade["entry_y"]) * open_trade["size_y"] * exit_sign
            ur_x = (open_trade["entry_x"] - x_t) * open_trade["size_x"] * abs(open_trade["beta_at_entry"]) * exit_sign
            unrealized = ur_y + ur_x

        peak_capital = max(peak_capital, capital + unrealized)
        equity_curve.append({
            "t": t,
            "equity": round(capital + unrealized, 4),
            "alpha": round(result["alpha"], 6),
            "beta": round(result["beta"], 6),
            "spread": round(result["spread"], 6),
            "z_score": round(result["z_score"], 4),
        })

    # Force close open trade at end
    if open_trade is not None:
        exit_sign = open_trade["signal"]
        y_t = float(Y[-1])
        x_t = float(X[-1])
        pnl_y = (y_t - open_trade["entry_y"]) * open_trade["size_y"] * exit_sign
        pnl_x = (open_trade["entry_x"] - x_t) * open_trade["size_x"] * abs(open_trade["beta_at_entry"]) * exit_sign
        gross_pnl = pnl_y + pnl_x
        exit_notional = abs(y_t * open_trade["size_y"]) + abs(x_t * open_trade["size_x"] * open_trade["beta_at_entry"])
        fee = exit_notional * fee_rate
        capital += gross_pnl - fee
        trades.append({
            "entry_time": open_trade["entry_time"],
            "exit_time": n - 1,
            "signal": open_trade["signal"],
            "entry_x": open_trade["entry_x"],
            "exit_x": x_t,
            "entry_y": open_trade["entry_y"],
            "exit_y": y_t,
            "beta": open_trade["beta_at_entry"],
            "gross_pnl": round(gross_pnl, 4),
            "net_pnl": round(gross_pnl - fee, 4),
            "fee": round(open_trade["fee_paid"] + fee, 6),
            "duration_bars": n - 1 - open_trade["entry_time"],
        })

    # ── Metrics ──
    eq = np.array([e["equity"] for e in equity_curve])
    returns = np.diff(eq) / eq[:-1] if len(eq) > 1 else np.array([0.0])

    total_pnl = capital - initial_capital
    pnl_pct = total_pnl / initial_capital * 100
    dd = max_drawdown(eq.tolist())
    sh = sharpe(returns.tolist())
    so = sortino(returns.tolist())
    wr = win_rate(trades)
    cagr = ((capital / initial_capital) ** (1 / max(n / (365 * 24), 0.01)) - 1) * 100 if n > 0 and capital > 0 else 0.0

    return {
        "total_pnl": round(total_pnl, 4),
        "pnl_pct": round(pnl_pct, 2),
        "cagr": round(cagr, 2),
        "sharpe": round(sh, 4),
        "sortino": round(so, 4),
        "max_drawdown": round(dd, 4),
        "win_rate": round(wr, 4),
        "total_trades": len(trades),
        "final_equity": round(capital, 4),
        "transaction_costs": round(sum(t["fee"] for t in trades), 4),
        "avg_trade_duration": round(np.mean([t["duration_bars"] for t in trades]), 1) if trades else 0,
        "trades": trades[-200:],
        "equity_curve": equity_curve,
        "alpha_history": [e["alpha"] for e in equity_curve],
        "beta_history": [e["beta"] for e in equity_curve],
        "spread_history": [e["spread"] for e in equity_curve],
        "z_score_history": [e["z_score"] for e in equity_curve],
    }


def backtest_rolling_ols(
    X: np.ndarray,
    Y: np.ndarray,
    window: int = 60,
    z_entry: float = 2.0,
    z_exit: float = 0.5,
    trade_size_usd: float = 100.0,
    transaction_cost_bps: float = 2.5,
    initial_capital: float = 10000.0,
) -> dict:
    """
    Baseline: classic rolling OLS pairs trading.

    Uses a fixed-lookback rolling beta instead of Kalman adaptation.
    """
    n = len(X)
    betas = compute_rolling_ols_hedge(X, Y, window)
    fee_rate = transaction_cost_bps / 10000.0

    capital = initial_capital
    equity_curve: list[dict] = []
    trades: list[dict] = []
    open_trade: Optional[dict] = None

    spreads: list[float] = []
    z_lookback = 100

    for t in range(window, n):
        x_t = float(X[t])
        y_t = float(Y[t])
        beta = betas[t] if not np.isnan(betas[t]) else 1.0

        spread = y_t - beta * x_t
        spreads.append(spread)

        # Z-score
        lb = min(z_lookback, len(spreads))
        rec = spreads[-lb:]
        mu = np.mean(rec)
        sigma = np.std(rec, ddof=1)
        z = (spread - mu) / sigma if sigma > 1e-12 else 0.0

        signal = 0
        if open_trade is None:
            if z > z_entry:
                signal = -1  # short Y, long X
            elif z < -z_entry:
                signal = +1  # long Y, short X
        else:
            if abs(z) < z_exit:
                signal = -open_trade["signal"]

        if signal != 0:
            if open_trade is None:
                size_x = trade_size_usd / x_t if x_t > 0 else 0
                size_y = trade_size_usd / y_t if y_t > 0 else 0
                hedge_notional = size_x * x_t * abs(beta)
                fee = (trade_size_usd + hedge_notional) * fee_rate
                capital -= fee
                open_trade = {
                    "entry_time": t, "signal": signal,
                    "entry_x": x_t, "entry_y": y_t,
                    "beta": beta, "size_x": size_x, "size_y": size_y,
                    "fee_paid": fee,
                }
            elif signal == -open_trade["signal"]:
                es = open_trade["signal"]
                pnl_y = (y_t - open_trade["entry_y"]) * open_trade["size_y"] * es
                pnl_x = (open_trade["entry_x"] - x_t) * open_trade["size_x"] * abs(open_trade["beta"]) * es
                gross_pnl = pnl_y + pnl_x
                exit_notional = abs(y_t * open_trade["size_y"]) + abs(x_t * open_trade["size_x"] * open_trade["beta"])
                fee = exit_notional * fee_rate
                capital += gross_pnl - fee
                trades.append({
                    "entry_time": open_trade["entry_time"], "exit_time": t,
                    "signal": open_trade["signal"], "gross_pnl": round(gross_pnl, 4),
                    "net_pnl": round(gross_pnl - fee, 4),
                    "duration_bars": t - open_trade["entry_time"],
                })
                open_trade = None

        equity_curve.append({"t": t, "equity": round(capital, 4)})

    eq = np.array([e["equity"] for e in equity_curve])
    returns = np.diff(eq) / eq[:-1] if len(eq) > 1 else np.zeros(1)
    total_pnl = capital - initial_capital
    dd = max_drawdown(eq.tolist())

    return {
        "total_pnl": round(total_pnl, 4),
        "pnl_pct": round(total_pnl / initial_capital * 100, 2),
        "sharpe": round(sharpe(returns.tolist()), 4),
        "sortino": round(sortino(returns.tolist()), 4),
        "max_drawdown": round(dd, 4),
        "win_rate": round(win_rate(trades), 4),
        "total_trades": len(trades),
        "final_equity": round(capital, 4),
        "trades": trades[-200:],
        "equity_curve": equity_curve,
    }


def run_comparison(
    X: np.ndarray,
    Y: np.ndarray,
    transition_covariance: float = 1e-4,
    observation_covariance: float = 1e-2,
    z_entry: float = 2.0,
    z_exit: float = 0.5,
    trade_size_usd: float = 100.0,
    transaction_cost_bps: float = 2.5,
    ols_windows: list[int] = [60, 120],
) -> dict:
    """
    Run Kalman vs rolling OLS comparison backtest.

    Returns:
        dict with kalman_results, ols_results, and comparison_summary.
    """
    trader = KalmanPairsTrader(
        transition_covariance=transition_covariance,
        observation_covariance=observation_covariance,
        z_entry=z_entry, z_exit=z_exit,
    )

    kalman = backtest_kalman_pairs(
        X, Y, trader,
        trade_size_usd=trade_size_usd,
        transaction_cost_bps=transaction_cost_bps,
    )

    ols_results = {}
    for w in ols_windows:
        ols_results[f"ols_{w}d"] = backtest_rolling_ols(
            X, Y, window=w,
            z_entry=z_entry, z_exit=z_exit,
            trade_size_usd=trade_size_usd,
            transaction_cost_bps=transaction_cost_bps,
        )

    return {
        "kalman": kalman,
        "ols": ols_results,
    }