"""Mean-variance spanning test.

Tests whether a new asset (or set of assets) expands the mean-variance
frontier spanned by existing benchmark assets, using the regression-based
approach of Huberman and Kandel (1987).

References
----------
Huberman, G., & Kandel, S. (1987). Mean-variance spanning.
    Journal of Finance, 42(4), 873-888.
Barillas, F., & Shanken, J. (2018). Comparing asset pricing models.
    Journal of Finance, 73(2), 715-754.
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Any

import numpy as np
import pandas as pd
from scipy import stats

from eapctf.timeseries.alpha import newey_west_se


@dataclass
class SpanningResult:
    """Result container for the mean-variance spanning test."""

    f_statistic: float
    p_value: float
    rejected: bool
    alpha: float
    alpha_t: float
    delta_sr_squared: float


def _compute_squared_sharpe(
    returns: np.ndarray[Any, np.dtype[Any]],
) -> float:
    """Compute the maximum squared Sharpe ratio from a set of returns.

    Uses the tangency portfolio: SR^2 = mu' Sigma^{-1} mu.

    Parameters
    ----------
    returns : np.ndarray
        Asset returns matrix, shape (T, N).

    Returns
    -------
    float
        Maximum squared Sharpe ratio.
    """
    mu = np.mean(returns, axis=0)
    T = returns.shape[0]
    demeaned = returns - mu
    Sigma = (demeaned.T @ demeaned) / (T - 1)
    Sigma_inv = np.linalg.inv(Sigma)
    return float(mu @ Sigma_inv @ mu)


def spanning_test(
    test_returns: pd.Series[Any] | pd.DataFrame,
    benchmark_returns: pd.DataFrame,
    n_lags_nw: int | None = None,
) -> SpanningResult:
    """Mean-variance spanning test (regression-based).

    Regresses test_returns on benchmark_returns + constant, then performs
    the Huberman-Kandel (1987) joint test: H0: alpha=0 AND sum(betas)=1.
    Rejection implies the test asset expands the mean-variance frontier.

    Uses a Wald F-statistic with Newey-West HAC covariance for robustness
    to heteroskedasticity and autocorrelation.

    Parameters
    ----------
    test_returns : pd.Series or pd.DataFrame
        Returns of the test asset(s). If DataFrame, uses the first column.
    benchmark_returns : pd.DataFrame
        Returns of benchmark assets, T x K.
    n_lags_nw : int | None
        Lags for Newey-West SE. None uses automatic bandwidth.

    Returns
    -------
    SpanningResult
        F-statistic, p-value, rejection indicator, alpha and its t-stat,
        and the improvement in squared Sharpe ratio.

    References
    ----------
    Huberman, G., & Kandel, S. (1987). Journal of Finance, 42(4), 873-888.
    """
    # Handle DataFrame input: use first column
    if isinstance(test_returns, pd.DataFrame):
        test_returns = test_returns.iloc[:, 0]

    # Align time indices
    bench_cols = list(benchmark_returns.columns)
    aligned = pd.concat(
        [test_returns.rename("__test__"), benchmark_returns], axis=1
    ).dropna()

    y = np.asarray(aligned["__test__"].values)  # (T,)
    B = np.asarray(aligned[bench_cols].values)  # (T, K)
    T, K = B.shape

    # --- OLS regression: y_t = alpha + beta' B_t + e_t ---
    ones = np.ones((T, 1))
    X = np.hstack([ones, B])  # (T, K+1)
    XtX = X.T @ X
    XtX_inv = np.linalg.inv(XtX)
    theta = XtX_inv @ (X.T @ y)  # (K+1,): [alpha, beta_1, ..., beta_K]

    alpha_hat = float(theta[0])
    residuals = y - X @ theta

    # --- Newey-West HAC covariance of theta ---
    se = newey_west_se(residuals, X, n_lags=n_lags_nw)
    alpha_se = float(se[0])
    alpha_t = alpha_hat / alpha_se if alpha_se > 0.0 else 0.0

    # --- Huberman-Kandel Wald test ---
    # H0: R * theta = r, where R = [[1, 0, ..., 0], [0, 1, 1, ..., 1]]
    #                          r = [0, 1]
    # Constraint 1: alpha = 0
    # Constraint 2: sum(betas) = 1
    q = 2  # number of restrictions
    R = np.zeros((q, K + 1))
    R[0, 0] = 1.0  # alpha = 0
    R[1, 1:] = 1.0  # sum(betas) = 1
    r = np.array([0.0, 1.0])

    # Compute full HAC covariance matrix (not just diagonal SEs)
    # Rebuild the sandwich: V = (X'X)^{-1} S (X'X)^{-1}
    if n_lags_nw is None:
        n_lags_actual = int(math.floor(4.0 * (T / 100.0) ** (2.0 / 9.0)))
    else:
        n_lags_actual = n_lags_nw

    e = residuals.reshape(-1, 1)

    # Meat: S = sum of Bartlett-weighted autocovariance matrices
    S = (X * (e ** 2)).T @ X
    for j in range(1, n_lags_actual + 1):
        w_j = 1.0 - j / (n_lags_actual + 1.0)
        gamma_j = (X[j:] * (e[j:] * e[:-j])).T @ X[:-j]
        S += w_j * (gamma_j + gamma_j.T)

    V_hac = XtX_inv @ S @ XtX_inv  # full (K+1) x (K+1) HAC covariance

    # Wald statistic: W = (R*theta - r)' [R V R']^{-1} (R*theta - r)
    diff = R @ theta - r  # (q,)
    RVR = R @ V_hac @ R.T  # (q, q)
    RVR_inv = np.linalg.inv(RVR)
    wald = float(diff @ RVR_inv @ diff)

    # Convert Wald to F: F = W / q
    f_stat = wald / q

    # p-value from F(q, T - K - 1) distribution
    df1 = q
    df2 = T - K - 1
    p_value = 1.0 - float(stats.f.cdf(f_stat, df1, df2))
    p_value = max(0.0, min(1.0, p_value))

    # Reject spanning at 5% significance
    rejected = p_value < 0.05

    # --- Delta squared Sharpe ratio ---
    # SR^2 of augmented portfolio minus SR^2 of benchmark
    augmented = np.asarray(aligned[["__test__"] + bench_cols].values)  # (T, K+1)
    sr_sq_aug = _compute_squared_sharpe(augmented)
    sr_sq_bench = _compute_squared_sharpe(B)
    delta_sr_sq = sr_sq_aug - sr_sq_bench

    return SpanningResult(
        f_statistic=f_stat,
        p_value=p_value,
        rejected=rejected,
        alpha=alpha_hat,
        alpha_t=alpha_t,
        delta_sr_squared=delta_sr_sq,
    )
