"""Fama-MacBeth (1973) two-pass cross-sectional regression.

Implements the classic two-pass procedure for estimating risk premia,
with Shanken (1992) errors-in-variables correction and optional
conditional FM via instrument interactions.

References
----------
Fama, E. F., & MacBeth, J. D. (1973). Risk, return, and equilibrium:
    Empirical tests. Journal of Political Economy, 81(3), 607-636.
Shanken, J. (1992). On the estimation of beta-pricing models.
    Review of Financial Studies, 5(1), 1-55. Eq. 8.
Ferson, W. E. (2019). Empirical asset pricing: Models and methods.
    MIT Press. Ch. 12 (conditional FM with instruments).
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Any

import numpy as np
import pandas as pd

from eapctf.timeseries.alpha import newey_west_se


@dataclass
class FMBResult:
    """Result container for Fama-MacBeth regression."""

    lambdas: pd.DataFrame  # K rows x {coef, se_nw, t_nw, se_shanken, t_shanken}
    # se_nw / t_nw : Newey-West HAC standard errors and t-statistics on the
    #     time series of per-period lambda estimates.  These are NOT classical
    #     OLS SEs; they account for serial correlation in lambda_t via the
    #     Bartlett kernel.  Report these in papers as "Newey-West standard errors".
    # t_shanken : pd.Series
    #     Shanken-corrected t-statistics. Contains NaN when shanken_applied=False.
    #     Note: Shanken correction applied only to beta regressors, not char regressors.
    #     Always check shanken_applied before interpreting these values.
    r_squared: float  # time-series average cross-sectional R-squared
    avg_r_squared_cs: float  # alias for r_squared
    n_periods: int
    avg_n_assets: float
    lambda_ts: pd.DataFrame  # T x K time-series of cross-sectional coefficients
    intercept_avg: float = 0.0  # time-series mean of the cross-sectional intercepts
    shanken_applied: bool = True  # False when shanken_correction=False; t_shanken is NaN


def _nw_se_timeseries(
    lambda_ts: np.ndarray[Any, np.dtype[Any]],
    n_lags: int | None = None,
) -> np.ndarray[Any, np.dtype[Any]]:
    """Newey-West standard errors on the time series of FM lambda estimates.

    For each regressor k, compute HAC SE on the T-length series lambda_{k,t}.
    This is equivalent to regressing lambda_{k,t} on a constant with NW SE.

    Parameters
    ----------
    lambda_ts : np.ndarray
        Shape (T, K) array of per-period cross-sectional coefficient estimates.
    n_lags : int | None
        Number of lags. If None, uses floor(4*(T/100)^(2/9)).

    Returns
    -------
    np.ndarray
        Shape (K,) standard errors.
    """
    T, K = lambda_ts.shape

    # Regressor is just a constant (intercept-only regression)
    X = np.ones((T, 1))

    # Compute NW SE for each lambda series using the existing utility
    se_out = np.zeros(K)
    for k in range(K):
        # Residuals from mean: lambda_{k,t} - mean(lambda_k)
        mean_k = np.mean(lambda_ts[:, k])
        resid_k = lambda_ts[:, k] - mean_k
        se_k = newey_west_se(resid_k, X, n_lags=n_lags)
        se_out[k] = float(se_k[0])

    return se_out


def fama_macbeth(
    data: pd.DataFrame,
    ret_col: str = "ret",
    factor_cols: list[str] | None = None,
    char_cols: list[str] | None = None,
    date_col: str = "date",
    entity_col: str = "permno",
    shanken_correction: bool = True,
    instruments: list[str] | None = None,
    n_lags_nw: int | None = None,
    cluster: str | None = None,
) -> FMBResult:
    """Two-pass Fama-MacBeth cross-sectional regression.

    Pass 1 (if factor_cols): For each entity, regress excess returns on factors
        over the full time series to obtain estimated betas per entity.
    Pass 2: Each date, regress cross-section of returns on betas/chars to get lambda_t.
    Average lambda_t over T periods for point estimates.
    Newey-West SE on the lambda_t series.

    Shanken (1992) EIV correction (when shanken_correction=True):
        c = 1 + lambda' Omega_f^{-1} lambda
        SE_shanken = SE_ols * sqrt(c)
        where Omega_f = sample covariance of factor returns.

    Conditional FM (when instruments is not None):
        Interact each char_col with lagged instrument columns before Pass 2.

    Parameters
    ----------
    data : pd.DataFrame
        Long-format panel with date, entity, returns, and characteristics/factors.
    ret_col : str
        Name of the return column.
    factor_cols : list[str] | None
        Factor return columns for Pass 1 time-series regressions.
    char_cols : list[str] | None
        Characteristic columns for Pass 2 cross-sectional regressions.
    date_col : str
        Name of the date column.
    entity_col : str
        Name of the entity identifier column.
    shanken_correction : bool
        Whether to apply Shanken (1992) EIV correction.
    instruments : list[str] | None
        Instrument columns for conditional FM.
    n_lags_nw : int | None
        Lags for Newey-West SE. None uses automatic bandwidth.
    cluster : str | None
        Clustering for standard errors (currently unused, reserved).

    Returns
    -------
    FMBResult
        Dataclass with risk premia estimates, t-statistics, and diagnostics.

    References
    ----------
    Fama, E. F., & MacBeth, J. D. (1973). JPE 81(3), 607-636.
    Shanken, J. (1992). RFS 5(1), 1-33. Eq. 8.
    Ferson, W. (2019). Ch. 12 (conditional FM with instruments).
    """
    df = data.copy()

    # Determine which columns to use as regressors in the cross-sectional pass
    rhs_cols: list[str] = []

    # --- Pass 1: Time-series regressions to estimate betas (if factor_cols given) ---
    if factor_cols is not None and len(factor_cols) > 0:
        # For each entity, run a full-sample TS regression: r_{i,t} = a_i + beta_i' f_t + e_{i,t}
        beta_records: list[dict[str, Any]] = []

        for entity, entity_df in df.groupby(entity_col):
            entity_df = entity_df.sort_values(date_col)
            y = np.asarray(entity_df[ret_col].values, dtype=np.float64)
            factor_mat = np.asarray(entity_df[factor_cols].values, dtype=np.float64)

            # Skip entities with too few observations
            if len(y) < len(factor_cols) + 2:
                continue

            # OLS with constant: y = alpha + beta' factors + e
            T_i = len(y)
            X = np.column_stack([np.ones(T_i), factor_mat])
            try:
                coefs = np.linalg.lstsq(X, y, rcond=None)[0]
            except np.linalg.LinAlgError:
                continue

            # Store betas (skip intercept at index 0)
            record: dict[str, Any] = {entity_col: entity}
            for j, fc in enumerate(factor_cols):
                record[f"beta_{fc}"] = float(coefs[j + 1])
            beta_records.append(record)

        # Merge betas back onto the panel (betas are entity-level, constant over time)
        beta_df = pd.DataFrame(beta_records)
        beta_col_names = [f"beta_{fc}" for fc in factor_cols]
        df = df.merge(beta_df, on=entity_col, how="inner")
        rhs_cols.extend(beta_col_names)

    # Add characteristic columns
    if char_cols is not None:
        rhs_cols.extend(char_cols)

    # --- Conditional FM: interact chars with instruments ---
    if instruments is not None and len(instruments) > 0:
        interaction_cols: list[str] = []
        base_cols = list(rhs_cols)  # copy current RHS before adding interactions
        for inst_col in instruments:
            for base_col in base_cols:
                inter_name = f"{base_col}_x_{inst_col}"
                df[inter_name] = df[base_col] * df[inst_col]
                interaction_cols.append(inter_name)
        rhs_cols.extend(interaction_cols)

    if len(rhs_cols) == 0:
        msg = "Must provide at least one of factor_cols or char_cols."
        raise ValueError(msg)

    # --- Pass 2: Cross-sectional regressions date by date ---
    dates = sorted(df[date_col].unique())
    T = len(dates)
    K = len(rhs_cols)

    # Storage for per-period estimates
    lambda_array = np.full((T, K), np.nan)
    intercept_array = np.full(T, np.nan)
    r2_array = np.full(T, np.nan)
    n_assets_array = np.full(T, np.nan)

    for t_idx, dt in enumerate(dates):
        cs = df.loc[df[date_col] == dt].dropna(subset=[ret_col, *rhs_cols])
        n_cs = len(cs)

        # Need at least K+1 observations for OLS with intercept
        if n_cs < K + 2:
            continue

        y_cs = np.asarray(cs[ret_col].values, dtype=np.float64)
        X_cs = np.column_stack([np.ones(n_cs), np.asarray(cs[rhs_cols].values, dtype=np.float64)])

        # OLS: y = intercept + lambda' X + epsilon
        try:
            coefs_cs = np.linalg.lstsq(X_cs, y_cs, rcond=None)[0]
        except np.linalg.LinAlgError:
            continue

        intercept_array[t_idx] = coefs_cs[0]
        lambda_array[t_idx, :] = coefs_cs[1:]

        # Cross-sectional R-squared
        fitted = X_cs @ coefs_cs
        ss_res = float(np.sum((y_cs - fitted) ** 2))
        ss_tot = float(np.sum((y_cs - np.mean(y_cs)) ** 2))
        r2_array[t_idx] = 1.0 - ss_res / ss_tot if ss_tot > 0.0 else 0.0
        n_assets_array[t_idx] = float(n_cs)

    # Filter out periods with NaN (insufficient obs)
    valid_mask = ~np.isnan(lambda_array[:, 0])
    lambda_valid = lambda_array[valid_mask]
    r2_valid = r2_array[valid_mask]
    n_assets_valid = n_assets_array[valid_mask]
    valid_dates = [dates[i] for i in range(T) if valid_mask[i]]
    T_valid = len(valid_dates)

    if T_valid == 0:
        msg = "No valid cross-sectional periods found."
        raise ValueError(msg)

    # --- Point estimates: time-series average of lambda_t ---
    lambda_mean = np.mean(lambda_valid, axis=0)

    # --- Standard errors: HAC (Newey-West) on the time series of lambda_t ---
    # This is a defensible upgrade over the textbook std(lambda_t)/sqrt(T):
    # the OLS SE from a regression of lambda_t on a constant with Newey-West
    # weights is equivalent to std(lambda_t)/sqrt(T) when lags=0, and
    # additionally accounts for serial correlation in lambda_t for lags>0.
    # The ``se_ols`` label is kept for backward-compatible naming; these SEs
    # are Newey-West HAC, not classical OLS SEs.
    se_nw = _nw_se_timeseries(lambda_valid, n_lags=n_lags_nw)
    t_nw = np.where(se_nw > 0.0, lambda_mean / se_nw, 0.0)

    # --- Shanken (1992) correction ---
    # Only applicable when factor_cols is provided (betas are estimated)
    se_shanken = se_nw.copy()
    t_shanken = t_nw.copy()
    shanken_applied_actual = False  # set True only if correction block executes

    if shanken_correction and factor_cols is not None and len(factor_cols) > 0:
        # Compute sample covariance matrix of factors.
        # Factors are aggregate time-series variables (one value per date);
        # use .first() to extract a clean date-indexed series rather than
        # averaging across entities (which would bias Sigma_f if factor values
        # vary across entities in the panel).
        factor_data = df.groupby(date_col)[factor_cols].first()
        omega_f = np.cov(factor_data.values.T, ddof=1)

        # Ensure omega_f is at least 2-D for single factor case
        if omega_f.ndim == 0:
            omega_f = np.array([[float(omega_f)]])

        # Extract the lambda values corresponding to beta regressors
        # The first len(factor_cols) entries in rhs_cols are beta_{factor} columns
        n_beta = len(factor_cols)
        lambda_beta = lambda_mean[:n_beta]

        # Shanken correction factor: c = 1 + lambda' Omega_f^{-1} lambda
        try:
            omega_f_inv = np.linalg.inv(omega_f)
            c = 1.0 + float(lambda_beta @ omega_f_inv @ lambda_beta)
        except np.linalg.LinAlgError:
            c = 1.0

        # Inflate SE by sqrt(c) — ONLY for beta regressors (first n_beta columns).
        # Shanken EIV correction accounts for estimation error in betas; it is
        # theoretically unjustified for characteristic regressors (Shanken 1992, Eq. 8).
        se_shanken = se_nw.copy()
        se_shanken[:n_beta] = se_nw[:n_beta] * math.sqrt(c)
        t_shanken = np.where(se_shanken > 0.0, lambda_mean / se_shanken, 0.0)
        shanken_applied_actual = True
    elif not shanken_correction:
        # Shanken SE reported as NaN when correction is disabled
        se_shanken = np.full(K, np.nan)
        t_shanken = np.full(K, np.nan)

    # --- Build output DataFrames ---
    lambdas_df = pd.DataFrame(
        {
            "coef": lambda_mean,
            "se_nw": se_nw,
            "t_nw": t_nw,
            "se_shanken": se_shanken,
            "t_shanken": t_shanken,
        },
        index=rhs_cols,
    )

    # Time series of lambda estimates
    lambda_ts_df = pd.DataFrame(lambda_valid, index=valid_dates, columns=rhs_cols)

    avg_r2 = float(np.mean(r2_valid))
    avg_n = float(np.mean(n_assets_valid))

    # Time-series mean of cross-sectional intercepts (zero-beta rate in CAPM tests)
    valid_intercepts = intercept_array[valid_mask]
    avg_intercept = float(np.nanmean(valid_intercepts)) if len(valid_intercepts) > 0 else 0.0

    return FMBResult(
        lambdas=lambdas_df,
        r_squared=avg_r2,
        avg_r_squared_cs=avg_r2,
        n_periods=T_valid,
        avg_n_assets=avg_n,
        lambda_ts=lambda_ts_df,
        intercept_avg=avg_intercept,
        shanken_applied=shanken_applied_actual,
    )
