"""Time-series alpha regressions with Newey-West standard errors.

Provides OLS factor model estimation (r_t = alpha + beta'f_t + e_t)
with HAC-robust inference and rolling CAPM beta computation.

References
----------
Newey, W. K., & West, K. D. (1987). A simple, positive semi-definite,
    heteroskedasticity and autocorrelation consistent covariance matrix.
    Econometrica, 55(3), 703-708.
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Any

import numpy as np
import pandas as pd


@dataclass
class AlphaResult:
    """Result container for a single time-series alpha regression."""

    alpha: float
    alpha_t: float
    alpha_se: float
    betas: pd.Series[Any]
    r_squared: float
    info_ratio: float


def newey_west_se(
    residuals: np.ndarray[Any, np.dtype[Any]],
    X: np.ndarray[Any, np.dtype[Any]],
    n_lags: int | None = None,
) -> np.ndarray[Any, np.dtype[Any]]:
    """Newey-West HAC standard errors for OLS coefficients.

    Computes heteroskedasticity-and-autocorrelation-consistent (HAC)
    standard errors using the Bartlett kernel.

    Parameters
    ----------
    residuals : np.ndarray
        OLS residuals, shape (T,).
    X : np.ndarray
        Regressor matrix including constant, shape (T, K+1).
    n_lags : int | None
        Number of lags for HAC estimator. If None, uses the automatic
        bandwidth: q = floor(4 * (T/100)^(2/9)).

    Returns
    -------
    np.ndarray
        Standard errors for each coefficient, shape (K+1,).

    References
    ----------
    Newey, W. K., & West, K. D. (1987). Econometrica, 55(3), 703-708.
    """
    T, K = X.shape

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

    # Meat of the sandwich: S = sum of weighted autocovariance matrices
    # Start with lag 0 (heteroskedasticity-consistent part)
    # Gamma_j = (1/T) * sum_t e_t * e_{t-j} * x_t * x_{t-j}'
    e = residuals.reshape(-1, 1)  # (T, 1)

    # Lag-0 term: Gamma_0 = X' diag(e^2) X / T
    S = (X * (e ** 2)).T @ X  # (K, K), unscaled sum

    # Add lag j terms with Bartlett weights: w(j) = 1 - j/(q+1)
    for j in range(1, n_lags + 1):
        # Bartlett kernel weight
        w_j = 1.0 - j / (n_lags + 1.0)

        # Cross-products for lag j: sum_t e_t * e_{t-j} * x_t * x_{t-j}'
        gamma_j = (X[j:] * (e[j:] * e[:-j])).T @ X[:-j]  # (K, K)

        # Symmetric: add both Gamma_j and Gamma_j'
        S += w_j * (gamma_j + gamma_j.T)

    # Bread of the sandwich: (X'X)^{-1}
    XtX_inv = np.linalg.inv(X.T @ X)

    # HAC covariance: V = (X'X)^{-1} S (X'X)^{-1}
    V_hac = XtX_inv @ S @ XtX_inv

    # Return standard errors (square root of diagonal)
    se: np.ndarray[Any, np.dtype[Any]] = np.sqrt(np.diag(V_hac))
    return se


def _run_single_ols(
    y: np.ndarray[Any, np.dtype[Any]],
    factor_mat: np.ndarray[Any, np.dtype[Any]],
    factor_names: list[str],
    n_lags_nw: int | None,
) -> AlphaResult:
    """Run OLS regression for a single portfolio and return AlphaResult.

    Parameters
    ----------
    y : np.ndarray
        Portfolio excess returns, shape (T,).
    factor_mat : np.ndarray
        Factor return matrix, shape (T, K).
    factor_names : list[str]
        Names corresponding to columns of factor_mat.
    n_lags_nw : int | None
        Lags for Newey-West SE.
    """
    T = len(y)

    # Build regressor matrix: [constant, factors]
    ones = np.ones((T, 1))
    X = np.hstack([ones, factor_mat])  # (T, K+1)

    # OLS coefficients: beta_hat = (X'X)^{-1} X'y
    XtX_inv = np.linalg.inv(X.T @ X)
    beta_hat = XtX_inv @ (X.T @ y)  # (K+1,)

    # Extract alpha (intercept) and factor betas
    alpha = float(beta_hat[0])
    betas_arr = beta_hat[1:]

    # Residuals
    residuals = y - X @ beta_hat

    # R-squared: 1 - SS_res / SS_tot
    ss_res = float(np.sum(residuals ** 2))
    ss_tot = float(np.sum((y - np.mean(y)) ** 2))
    r_squared = 1.0 - ss_res / ss_tot if ss_tot > 0.0 else 0.0

    # Newey-West standard errors
    se = newey_west_se(residuals, X, n_lags=n_lags_nw)
    alpha_se = float(se[0])

    # t-statistic on alpha
    alpha_t = alpha / alpha_se if alpha_se > 0.0 else 0.0

    # Information ratio: alpha / std(residuals)
    resid_std = float(np.std(residuals, ddof=1))
    info_ratio = alpha / resid_std if resid_std > 0.0 else 0.0

    betas_series: pd.Series[Any] = pd.Series(betas_arr, index=factor_names, name="betas")

    return AlphaResult(
        alpha=alpha,
        alpha_t=alpha_t,
        alpha_se=alpha_se,
        betas=betas_series,
        r_squared=r_squared,
        info_ratio=info_ratio,
    )


def time_series_alpha(
    portfolio_returns: pd.Series[Any] | pd.DataFrame,
    factor_returns: pd.DataFrame,
    n_lags_nw: int | None = None,
) -> AlphaResult | list[AlphaResult]:
    """OLS factor regression with Newey-West standard errors.

    Estimates r_t = alpha + beta_1*f1_t + ... + beta_K*fK_t + e_t
    for each portfolio, using HAC standard errors on alpha.

    Parameters
    ----------
    portfolio_returns : pd.Series or pd.DataFrame
        If Series: single portfolio excess returns indexed by date.
        If DataFrame: each column is a portfolio, rows are dates.
    factor_returns : pd.DataFrame
        Factor returns, columns are factors, rows are dates.
    n_lags_nw : int | None
        Lags for Newey-West. None uses automatic bandwidth.

    Returns
    -------
    AlphaResult or list[AlphaResult]
        Single result for Series input, list for DataFrame input.
    """
    # Align indices to ensure matching time periods
    if isinstance(portfolio_returns, pd.Series):
        aligned = pd.concat([portfolio_returns.rename("__port__"), factor_returns], axis=1).dropna()
        y = np.asarray(aligned["__port__"].values)
        factor_mat = np.asarray(aligned[factor_returns.columns].values)
        factor_names = list(factor_returns.columns)
        return _run_single_ols(y, factor_mat, factor_names, n_lags_nw)

    # DataFrame case: run regression for each column
    results: list[AlphaResult] = []
    factor_names = list(factor_returns.columns)

    for col in portfolio_returns.columns:
        aligned = pd.concat(
            [portfolio_returns[col].rename("__port__"), factor_returns], axis=1
        ).dropna()
        y = np.asarray(aligned["__port__"].values)
        factor_mat = np.asarray(aligned[factor_names].values)
        results.append(_run_single_ols(y, factor_mat, factor_names, n_lags_nw))

    return results


def rolling_beta(
    returns: pd.Series[Any],
    market_returns: pd.Series[Any],
    window: int = 60,
    min_periods: int = 24,
) -> pd.Series[Any]:
    """Rolling OLS CAPM beta estimate.

    Estimates r_t = alpha + beta * rm_t + e_t over a rolling window
    and returns the time series of beta estimates.

    Parameters
    ----------
    returns : pd.Series
        Asset excess returns indexed by date.
    market_returns : pd.Series
        Market excess returns indexed by date.
    window : int
        Rolling window size in periods (default 60 months).
    min_periods : int
        Minimum number of observations required (default 24).

    Returns
    -------
    pd.Series
        Rolling beta estimates, indexed by date. NaN where insufficient data.
    """
    # Align the two series
    aligned = pd.concat(
        [returns.rename("ret"), market_returns.rename("mkt")], axis=1
    ).dropna()

    ret = aligned["ret"]
    mkt = aligned["mkt"]

    # Rolling covariance / rolling variance gives beta
    # beta_t = Cov(r, rm) / Var(rm) over the rolling window
    rolling_cov = ret.rolling(window=window, min_periods=min_periods).cov(mkt)
    rolling_var = mkt.rolling(window=window, min_periods=min_periods).var()

    # Beta = Cov(r, rm) / Var(rm)
    beta: pd.Series[Any] = rolling_cov / rolling_var
    beta.name = "beta"

    return beta
