"""Hansen-Jagannathan (1991) distance and SDF volatility bounds.

Implements the HJ distance for evaluating SDF proxies and the HJ volatility
bound (minimum variance frontier for SDFs).

References
----------
Hansen, L. P., & Jagannathan, R. (1991). Implications of security market data
    for models of dynamic economies. Journal of Political Economy, 99(2), 225-262.
Kan, R., & Robotti, C. (2009). Model comparison using the Hansen-Jagannathan
    distance. Review of Financial Studies, 22(12), 5069-5111.
"""

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


@dataclass
class HJResult:
    """Result container for Hansen-Jagannathan distance and bounds."""

    distance: float  # HJ distance (sqrt of HJD^2)
    distance_se: float  # standard error of distance
    p_value: float  # p-value for the test that HJD = 0
    df: int  # degrees of freedom used in the chi-squared test (N - n_params)
    lower_bound: np.ndarray[Any, np.dtype[Any]]  # HJ lower bound per mean_range grid
    upper_bound: np.ndarray[Any, np.dtype[Any]]  # not applicable for linear SDF; NaN
    mean_range: np.ndarray[Any, np.dtype[Any]]  # grid of E[m] values


def _newey_west_var_scalar(
    x: np.ndarray[Any, np.dtype[Any]],
    n_lags: int,
) -> float:
    """Newey-West long-run variance for a scalar time series.

    Parameters
    ----------
    x : np.ndarray
        Demeaned time series, shape (T,).
    n_lags : int
        Number of lags for Bartlett kernel.

    Returns
    -------
    float
        Long-run variance estimate.
    """
    T = len(x)

    # Lag-0 autocovariance
    var = float(np.sum(x ** 2) / T)

    # Add lag-j terms with Bartlett weights
    for j in range(1, n_lags + 1):
        w_j = 1.0 - j / (n_lags + 1.0)
        # Autocovariance at lag j
        gamma_j = float(np.sum(x[j:] * x[:-j]) / T)
        var += 2.0 * w_j * gamma_j

    return var


def hj_distance(
    returns: pd.DataFrame,
    sdf_proxy: pd.Series[Any],
    n_lags_nw: int | None = None,
    n_params: int = 0,
) -> HJResult:
    """Hansen-Jagannathan distance for an SDF proxy.

    HJD^2 = g' (E[R R'])^{-1} g
    where g_i = (1/T) sum_t m_t * R_{it} - 1

    Parameters
    ----------
    returns : pd.DataFrame
        T x N test asset returns. Must be **gross** returns (1 + r), not excess
        returns, because the pricing condition is E[m * R] = 1_N. If you have
        excess returns r_e, use E[m * r_e] = 0_N instead (pass benchmark=0_N).
    sdf_proxy : pd.Series
        SDF proxy values, length T.
    n_lags_nw : int | None
        Lags for Newey-West SE. None uses automatic bandwidth.
    n_params : int
        Number of free SDF parameters estimated from the data.
        For a non-parametric/fixed SDF proxy, use 0 (default).
        For m = b'f with K factors (b estimated), use n_params=K.
        Degrees of freedom for the chi-squared test: N - n_params.

    Returns
    -------
    HJResult
        HJ distance results.

    References
    ----------
    Hansen, L. P., & Jagannathan, R. (1991). JPE, 99(2), 225-262.
    Kan, R., & Robotti, C. (2009). RFS, 22(12), 5069-5111.
    """
    R = np.asarray(returns.values, dtype=np.float64)  # (T, N)
    m = np.asarray(sdf_proxy.values, dtype=np.float64)  # (T,)
    T, N = R.shape

    # Automatic lag selection
    if n_lags_nw is None:
        n_lags_nw = int(math.floor(4.0 * (T / 100.0) ** (2.0 / 9.0)))

    # Pricing errors: g = (1/T) sum_t m_t * R_t - 1_N
    g = (m[:, np.newaxis] * R).mean(axis=0) - np.ones(N)  # (N,)

    # Second moment matrix: E[R R'] = (1/T) sum_t R_t R_t'
    E_RR = (R.T @ R) / T  # (N, N)

    # HJD^2 = g' (E[RR'])^{-1} g
    E_RR_inv = np.linalg.inv(E_RR)
    hjd_sq = float(g @ E_RR_inv @ g)

    # HJ distance (take sqrt, ensuring non-negative)
    hjd = math.sqrt(max(hjd_sq, 0.0))

    # Standard error via delta method on HJD^2
    # phi_t = m_t * R_t - 1_N  are the per-observation moment conditions
    phi: np.ndarray[Any, np.dtype[Any]] = m[:, np.newaxis] * R - np.ones((T, N))

    # Derivative of HJD^2 w.r.t. g: d(HJD^2)/dg = 2 * E_RR_inv @ g
    d_g = 2.0 * E_RR_inv @ g  # (N,)

    # Influence function for HJD^2: psi_t = d_g' * (phi_t - g)
    # phi_t centered: phi_t - E[phi_t] = phi_t - g (since E[phi_t] = g by definition)
    psi = (phi - g) @ d_g  # (T,)

    # Long-run variance of psi_t using Newey-West
    lr_var = _newey_west_var_scalar(psi - psi.mean(), n_lags_nw)

    # SE of HJD^2, then apply delta method for SE of HJD = sqrt(HJD^2)
    se_hjd_sq = math.sqrt(lr_var / T)
    # delta method: SE(sqrt(x)) = SE(x) / (2*sqrt(x))
    distance_se = se_hjd_sq / (2.0 * hjd) if hjd > 1e-12 else 0.0

    # p-value: under H0 that HJD=0, T * HJD^2 ~ chi-squared(N - n_params)
    # (Kan & Robotti 2009 specification test); df adjusted for estimated parameters
    chi2_stat = T * hjd_sq
    df_hj = max(N - n_params, 1)
    p_value = float(1.0 - stats.chi2.cdf(chi2_stat, df=df_hj))

    # Empty bounds (not computed in distance function)
    empty_grid = np.array([0.0])

    return HJResult(
        distance=hjd,
        distance_se=distance_se,
        p_value=p_value,
        df=df_hj,
        lower_bound=np.full(1, np.nan),
        upper_bound=np.full(1, np.nan),
        mean_range=empty_grid,
    )


def hj_bounds(
    returns: pd.DataFrame,
    mean_range: tuple[float, float] = (0.95, 1.05),
    n_grid: int = 50,
) -> HJResult:
    """Hansen-Jagannathan (1991) SDF volatility bounds.

    For each candidate mean E[m] = mu in mean_range:
        min std(m) s.t. E[m * R_i] = 1 for all i
        lower_bound(mu) = sqrt( (1 - mu*mu_R)' (E[RR'])^{-1} (1 - mu*mu_R) )

    The derivation uses:
        E[m*R] = 1  =>  E[m]*E[R] + Cov(m, R) = 1
    The minimum-variance SDF for a given E[m] = mu satisfies:
        m* = mu + (1 - mu*mu_R)' (E[RR'])^{-1} R
    with std(m*) = sqrt( (1 - mu*mu_R)' (E[RR'])^{-1} (1 - mu*mu_R) )

    Parameters
    ----------
    returns : pd.DataFrame
        T x N test asset returns.
    mean_range : tuple[float, float]
        Range of E[m] values to evaluate.
    n_grid : int
        Number of grid points.

    Returns
    -------
    HJResult
        HJ bounds results with lower_bound tracing the volatility frontier.

    References
    ----------
    Hansen, L. P., & Jagannathan, R. (1991). JPE, 99(2), 225-262.
    """
    R = returns.values.astype(np.float64)  # (T, N)
    T, N = R.shape

    # Sample moments
    # mu_R = E[R]  -- mean returns (N,)
    mu_R = R.mean(axis=0)

    # Second moment matrix: E[R R'] = (1/T) R'R
    E_RR = (R.T @ R) / T  # (N, N)
    E_RR_inv = np.linalg.inv(E_RR)

    # Grid of candidate E[m] values
    mu_grid = np.linspace(mean_range[0], mean_range[1], n_grid)

    # Compute lower bound for each mu
    lower_bound = np.empty(n_grid)
    for i, mu in enumerate(mu_grid):
        # Vector: 1_N - mu * mu_R
        # This comes from the pricing condition E[m*R] = 1
        # => E[m]*E[R] + Cov(m,R) = 1
        delta = np.ones(N) - mu * mu_R  # (N,)

        # sigma^2(m*) = delta' (E[RR'])^{-1} delta
        var_m = float(delta @ E_RR_inv @ delta)

        # std(m*) = sqrt(max(var_m, 0))  -- numerical guard
        lower_bound[i] = math.sqrt(max(var_m, 0.0))

    # Upper bound not applicable for linear SDF; fill with NaN
    upper_bound = np.full(n_grid, np.nan)

    return HJResult(
        distance=0.0,
        distance_se=0.0,
        p_value=1.0,
        df=N,
        lower_bound=lower_bound,
        upper_bound=upper_bound,
        mean_range=mu_grid,
    )
