"""GRS test for joint significance of time-series alphas.

Implements the Gibbons, Ross, and Shanken (1989) F-test for whether
a set of portfolio alphas are jointly zero in a factor model.

References
----------
Gibbons, M. R., Ross, S. A., & Shanken, J. (1989). A test of the
    efficiency of a given portfolio. Econometrica, 57(5), 1121-1152.
    Eq. 5 (GRS statistic).
"""

from __future__ import annotations

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 GRSResult:
    """Result container for the GRS test."""

    statistic: float
    p_value: float
    alphas: pd.Series[Any]
    alpha_se: pd.Series[Any]
    sharpe_factors: float


def grs_test(
    portfolio_returns: pd.DataFrame | pd.Series[Any],
    factor_returns: pd.DataFrame,
    date_col: str | None = None,
) -> GRSResult:
    """GRS F-test for joint alpha = 0 across N portfolios.

    Tests whether a set of N portfolio alphas from time-series
    regressions on K factors are jointly zero.

    GRS = ((T - N - K) / N) * (1 + mu_f' Omega^{-1} mu_f)^{-1}
          * alpha' Sigma^{-1} alpha ~ F(N, T - N - K)

    Parameters
    ----------
    portfolio_returns : pd.DataFrame or pd.Series
        Wide format (T x N): each column is a portfolio, index is dates
        (set date_col=None). Long format: DataFrame with a date column
        and portfolio columns (set date_col to the column name).
        If pd.Series, treated as a single portfolio (N=1).
    factor_returns : pd.DataFrame
        Factor returns, T x K. If date_col is set, must share the same
        date column.
    date_col : str | None
        Name of the date column for long format. If None, assumes wide
        format with DatetimeIndex or integer index as time axis.

    Returns
    -------
    GRSResult
        Contains GRS statistic, p-value, per-portfolio alphas and SEs,
        and the factor Sharpe ratio.

    References
    ----------
    Gibbons, M. R., Ross, S. A., & Shanken, J. (1989). Econometrica,
        57(5), 1121-1152. Eq. 5.
    """
    # --- Handle input formats ---

    # If Series, convert to single-column DataFrame
    if isinstance(portfolio_returns, pd.Series):
        portfolio_returns = portfolio_returns.to_frame(name=portfolio_returns.name or "port_1")

    # If long format with date column, pivot to wide
    if date_col is not None:
        if date_col in portfolio_returns.columns:
            portfolio_returns = portfolio_returns.set_index(date_col)
        if date_col in factor_returns.columns:
            factor_returns = factor_returns.set_index(date_col)

    # Align portfolio and factor returns on their common time index
    port_cols = list(portfolio_returns.columns)
    factor_cols = list(factor_returns.columns)
    aligned = pd.concat([portfolio_returns, factor_returns], axis=1).dropna()

    port_mat = np.asarray(aligned[port_cols].values)  # (T, N)
    fact_mat = np.asarray(aligned[factor_cols].values)  # (T, K)

    T, N = port_mat.shape
    K = fact_mat.shape[1]

    # Guard: GRS F-statistic requires df2 = T - N - K > 0.
    # Without this, the Sigma estimate is degenerate and the F-distribution
    # denominator is undefined, producing a silent GRS = 0.
    if T - N - K <= 0:
        raise ValueError(
            f"GRS test requires T > N + K observations. "
            f"Got T={T} periods, N={N} portfolios, K={K} factors "
            f"(need T > {N + K})."
        )

    # --- Step 1: OLS regression per portfolio ---
    # r_{i,t} = alpha_i + beta_i' f_t + e_{i,t}
    ones = np.ones((T, 1))
    X = np.hstack([ones, fact_mat])  # (T, K+1)
    XtX_inv = np.linalg.inv(X.T @ X)

    # Store alphas, residuals, and Newey-West SEs
    alphas = np.zeros(N)
    alpha_ses = np.zeros(N)
    resid_mat = np.zeros((T, N))

    for i in range(N):
        y_i = port_mat[:, i]
        beta_hat = XtX_inv @ (X.T @ y_i)  # (K+1,)
        alphas[i] = beta_hat[0]
        resid_mat[:, i] = y_i - X @ beta_hat

        # Newey-West SE for alpha (first element)
        se = newey_west_se(resid_mat[:, i], X, n_lags=None)
        alpha_ses[i] = se[0]

    # --- Step 2: Compute GRS statistic ---

    # Residual covariance matrix: Sigma = (1/T) * E' E
    # Use (T - K - 1) degrees of freedom for unbiased estimate
    Sigma = (resid_mat.T @ resid_mat) / (T - K - 1)  # (N, N)

    # Factor mean vector and covariance matrix
    mu_f = np.mean(fact_mat, axis=0)  # (K,)
    # Unbiased factor covariance: Omega = (1/(T-1)) * (F - mu_f)' (F - mu_f)
    F_demeaned = fact_mat - mu_f
    Omega = (F_demeaned.T @ F_demeaned) / (T - 1)  # (K, K)

    # Sharpe ratio of the factor portfolio: sqrt(mu_f' Omega^{-1} mu_f)
    Omega_inv = np.linalg.inv(Omega)
    sr_sq = float(mu_f @ Omega_inv @ mu_f)
    sharpe_factors = float(np.sqrt(sr_sq))

    # GRS statistic (Gibbons, Ross, Shanken 1989, Eq. 5):
    # GRS = ((T - N - K) / N) * (1 + mu_f' Omega^{-1} mu_f)^{-1}
    #       * alpha' Sigma^{-1} alpha
    Sigma_inv = np.linalg.inv(Sigma)
    alpha_quad = float(alphas @ Sigma_inv @ alphas)

    grs_stat = ((T - N - K) / N) * (1.0 / (1.0 + sr_sq)) * alpha_quad

    # p-value from F(N, T - N - K) distribution
    df1 = N
    df2 = T - N - K
    p_value = 1.0 - float(stats.f.cdf(grs_stat, df1, df2))

    # Clamp p-value to [0, 1] for numerical safety
    p_value = max(0.0, min(1.0, p_value))

    # Build output Series
    alphas_series: pd.Series[Any] = pd.Series(alphas, index=port_cols, name="alpha")
    alpha_se_series: pd.Series[Any] = pd.Series(alpha_ses, index=port_cols, name="alpha_se")

    return GRSResult(
        statistic=grs_stat,
        p_value=p_value,
        alphas=alphas_series,
        alpha_se=alpha_se_series,
        sharpe_factors=sharpe_factors,
    )
