"""Hansen (1982) GMM estimation of linear SDF models.

Implements one-step and two-step GMM for the linear stochastic discount factor:
    m_t = 1 - b'(f_t - E[f])

References
----------
Hansen, L. P. (1982). Large sample properties of generalized method of moments
    estimators. Econometrica, 50(4), 1029-1054.
Ferson, W. E. (2019). Empirical Asset Pricing: Models and Methods. MIT Press, Ch. 7.
"""

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 GMMResult:
    """Result container for GMM SDF estimation."""

    b: np.ndarray[Any, np.dtype[Any]]  # SDF loadings (K,)
    b_se: np.ndarray[Any, np.dtype[Any]]  # standard errors (K,)
    t_stats: np.ndarray[Any, np.dtype[Any]]  # t-statistics b / b_se (K,)
    j_statistic: float  # J-test statistic
    j_p_value: float  # p-value from chi-squared(N-K)
    pricing_errors: pd.Series[Any]  # per-asset pricing errors
    weighting_matrix: np.ndarray[Any, np.dtype[Any]]  # optimal W


def _newey_west_cov(
    h: np.ndarray[Any, np.dtype[Any]],
    n_lags: int,
) -> np.ndarray[Any, np.dtype[Any]]:
    """Newey-West long-run covariance of moment conditions h (T x N).

    Uses Bartlett kernel weights: w(j) = 1 - j/(q+1).

    Parameters
    ----------
    h : np.ndarray
        Demeaned moment conditions, shape (T, N).
    n_lags : int
        Number of lags for HAC estimator.

    Returns
    -------
    np.ndarray
        Long-run covariance matrix S, shape (N, N).
    """
    T, N = h.shape

    # Lag-0 autocovariance: Gamma_0 = (1/T) h'h
    S = (h.T @ h) / T  # (N, N)

    # Add lag-j terms with Bartlett kernel weights
    for j in range(1, n_lags + 1):
        w_j = 1.0 - j / (n_lags + 1.0)
        # Gamma_j = (1/T) sum_t h_t * h_{t-j}'
        gamma_j = (h[j:].T @ h[:-j]) / T  # (N, N)
        # Symmetric: add Gamma_j + Gamma_j'
        S += w_j * (gamma_j + gamma_j.T)

    result: np.ndarray[Any, np.dtype[Any]] = S
    return result


def gmm_estimate(
    returns: pd.DataFrame,
    factors: pd.DataFrame,
    date_col: str | None = None,
    n_lags_nw: int | None = None,
    two_step: bool = True,
    instruments: pd.DataFrame | None = None,
    weighting: str = "efficient",
) -> GMMResult:
    """Hansen (1982) GMM estimation of linear SDF.

    SDF: m_t = 1 - b'(f_t - E[f])
    Moments: g_T(b) = (1/T) sum_t m_t(b) * R_t - 1_N

    Step 1 (W = I): b = (Sigma_Rf' Sigma_Rf)^{-1} Sigma_Rf' (mu_R - 1_N)
    Step 2 options:
      weighting="efficient": W = S^{-1} (optimal GMM; minimizes asymptotic variance)
      weighting="hj":        W = (E[RR'])^{-1} (Hansen-Jagannathan weighting;
                             directly minimizes HJ distance)

    Parameters
    ----------
    returns : pd.DataFrame
        T x N test asset returns.
    factors : pd.DataFrame
        T x K factor returns.
    date_col : str | None
        If provided, drop this column before estimation.
    n_lags_nw : int | None
        Lags for Newey-West. None uses automatic bandwidth.
    two_step : bool
        If True, run two-step GMM with optimal weighting matrix.
    instruments : pd.DataFrame | None
        Optional instruments (not used in basic linear SDF estimation).
    weighting : str
        Weighting matrix type: "efficient" (W = S^{-1}, default) or "hj"
        (W = (E[RR'])^{-1}, Hansen-Jagannathan distance weighting).

    Returns
    -------
    GMMResult
        GMM estimation results including SDF loadings, SEs, t-statistics, and J-test.

    References
    ----------
    Hansen, L. P. (1982). Econometrica, 50(4), 1029-1054.
    Hansen, L. P., & Jagannathan, R. (1997). JF, 52(2), 557-590.
    """
    if weighting not in ("efficient", "hj"):
        msg = f"weighting must be 'efficient' or 'hj', got {weighting!r}"
        raise ValueError(msg)
    # Drop date column if specified
    if date_col is not None:
        returns = returns.drop(columns=[date_col], errors="ignore")
        factors = factors.drop(columns=[date_col], errors="ignore")

    # Align returns and factors on their common index (date-based) and sort
    # into chronological order.  Two failure modes are prevented here:
    #   1. Mismatched row order: .concat aligns by index label, so periods are
    #      correctly paired even if the two DataFrames arrive in different orders.
    #   2. Wrong temporal order: Newey-West autocorrelation computation depends on
    #      the physical row sequence; sort_index() enforces chronological order.
    _ret_cols = list(returns.columns)
    _fac_cols = list(factors.columns)
    _aligned = pd.concat([returns, factors], axis=1).dropna().sort_index()
    returns = _aligned[_ret_cols]
    factors = _aligned[_fac_cols]

    # Convert to numpy arrays
    R = returns.values.astype(np.float64)  # (T, N)
    F = factors.values.astype(np.float64)  # (T, K)
    T, N = R.shape
    _T_f, K = F.shape

    # Automatic Newey-West lag selection: Andrews (1991) rule-of-thumb
    if n_lags_nw is None:
        n_lags_nw = int(math.floor(4.0 * (T / 100.0) ** (2.0 / 9.0)))

    # Demean factors: f_tilde_t = f_t - mean(f)
    mu_f = F.mean(axis=0)  # (K,)
    F_dm = F - mu_f  # (T, K) demeaned factors

    # Sample moments:
    # mu_R = (1/T) sum_t R_t  -- mean returns (N,)
    mu_R = R.mean(axis=0)

    # Sigma_Rf = (1/T) sum_t R_t * f_tilde_t'  -- second moment matrix (N, K)
    # This is the covariance between returns and demeaned factors
    Sigma_Rf = (R.T @ F_dm) / T  # (N, K)

    # Target vector for GMM: mu_R - 1_N (pricing condition E[mR]=1 => E[R]-1 = Sigma_Rf b)
    target = mu_R - np.ones(N)  # (N,)

    # ---- Step 1: Identity weighting matrix ----
    # b1 = (Sigma_Rf' I Sigma_Rf)^{-1} Sigma_Rf' I target
    # = (Sigma_Rf' Sigma_Rf)^{-1} Sigma_Rf' target  (OLS-like)
    bread_1 = Sigma_Rf.T @ Sigma_Rf  # (K, K)
    b1 = np.linalg.solve(bread_1, Sigma_Rf.T @ target)  # (K,)

    # Weighting matrix and final b
    if two_step:
        # Compute moment conditions at first-step estimates
        # m_t(b1) = 1 - b1'(f_t - mu_f)
        m_t = 1.0 - F_dm @ b1  # (T,)

        # h_t = m_t * R_t - 1_N  (per-observation moment conditions)
        h = (m_t[:, np.newaxis] * R) - np.ones((T, N))  # (T, N)

        # Newey-West long-run covariance of h_t (needed for SEs regardless of weighting)
        S = _newey_west_cov(h, n_lags_nw)  # (N, N)

        if weighting == "hj":
            # HJ weighting: W = (E[RR'])^{-1}
            # This directly minimizes the Hansen-Jagannathan distance
            # E[RR'] = (1/T) R'R
            ERR = (R.T @ R) / T  # (N, N)
            W = np.linalg.inv(ERR)  # (N, N)
        else:
            # Efficient GMM: W = S^{-1}
            W = np.linalg.inv(S)  # (N, N)

        # ---- Step 2: Optimal weighting ----
        # b2 = (Sigma_Rf' W Sigma_Rf)^{-1} Sigma_Rf' W target
        bread_2 = Sigma_Rf.T @ W @ Sigma_Rf  # (K, K)
        b = np.linalg.solve(bread_2, Sigma_Rf.T @ W @ target)  # (K,)
    else:
        # One-step GMM: W = I, b = b1
        W = np.eye(N)
        b = b1

        # Still need S for standard errors
        m_t = 1.0 - F_dm @ b  # (T,)
        h = (m_t[:, np.newaxis] * R) - np.ones((T, N))  # (T, N)
        S = _newey_west_cov(h, n_lags_nw)  # (N, N)

    # ---- Pricing errors ----
    # g_T(b) = (1/T) sum_t m_t(b) * R_t - 1_N
    m_final = 1.0 - F_dm @ b  # (T,)
    g_T = (m_final[:, np.newaxis] * R).mean(axis=0) - np.ones(N)  # (N,)

    # ---- J-test ----
    # J = T * g_T' W g_T ~ chi-squared(N - K) under H0
    j_stat = float(T * g_T @ W @ g_T)

    # Degrees of freedom: N moment conditions - K parameters
    df = N - K
    j_p_value = float(1.0 - stats.chi2.cdf(j_stat, df)) if df > 0 else 1.0

    # ---- Standard errors ----
    # Var(b) = (1/T) (D'WD)^{-1} D'WSWD (D'WD)^{-1}, D = Sigma_Rf
    bread = Sigma_Rf.T @ W @ Sigma_Rf  # (K, K)
    bread_inv = np.linalg.inv(bread)  # (K, K)
    meat = Sigma_Rf.T @ W @ S @ W @ Sigma_Rf  # (K, K)
    V_b = (bread_inv @ meat @ bread_inv) / T  # (K, K)
    b_se = np.sqrt(np.diag(V_b))  # (K,)

    # Build pricing errors Series
    pricing_errs: pd.Series[Any] = pd.Series(g_T, index=returns.columns, name="pricing_error")

    # t-statistics: b / b_se (element-wise; safe division)
    t_stats = np.where(b_se > 0.0, b / b_se, 0.0)

    return GMMResult(
        b=b,
        b_se=b_se,
        t_stats=t_stats,
        j_statistic=j_stat,
        j_p_value=j_p_value,
        pricing_errors=pricing_errs,
        weighting_matrix=W,
    )
