"""Pooled cross-sectional regression and multiple testing correction.

Provides pooled OLS with clustered standard errors and multiple hypothesis
testing adjustments relevant to cross-sectional asset pricing studies.

References
----------
Petersen, M. A. (2009). Estimating standard errors in finance panel data
    sets: Comparing approaches. Review of Financial Studies, 22(1), 435-480.
Harvey, C. R., Liu, Y., & Zhu, H. (2016). ... and the cross-section of
    expected returns. Review of Financial Studies, 29(1), 5-68.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import numpy as np
import pandas as pd
from scipy import stats as sp_stats

from eapctf.timeseries.alpha import newey_west_se


@dataclass
class CSRegressionResult:
    """Result container for a pooled cross-sectional regression."""

    coefs: pd.Series[Any]
    se: pd.Series[Any]
    t_stats: pd.Series[Any]
    r_squared: float
    n_obs: int


def _clustered_se(
    X: np.ndarray[Any, np.dtype[Any]],
    residuals: np.ndarray[Any, np.dtype[Any]],
    cluster_ids: np.ndarray[Any, np.dtype[Any]],
) -> np.ndarray[Any, np.dtype[Any]]:
    """Compute one-way clustered standard errors.

    Uses the sandwich estimator with clustering:
        V = (X'X)^{-1} B (X'X)^{-1}
    where B = sum_g (X_g' e_g)(X_g' e_g)' scaled by G/(G-1) * N/(N-K).

    Parameters
    ----------
    X : np.ndarray
        Regressor matrix (N, K) including constant.
    residuals : np.ndarray
        OLS residuals (N,).
    cluster_ids : np.ndarray
        Cluster group identifiers (N,).

    Returns
    -------
    np.ndarray
        Standard errors (K,).
    """
    N, K = X.shape
    XtX_inv = np.linalg.inv(X.T @ X)

    # Unique clusters
    unique_clusters = np.unique(cluster_ids)
    G = len(unique_clusters)

    # Meat: sum of outer products of cluster-level score vectors
    B = np.zeros((K, K))
    for g in unique_clusters:
        mask = cluster_ids == g
        Xg = X[mask]
        eg = residuals[mask].reshape(-1, 1)
        # Score vector for cluster g: X_g' e_g
        score_g = (Xg * eg).sum(axis=0).reshape(-1, 1)  # (K, 1)
        B += score_g @ score_g.T

    # Small-sample correction factor: G/(G-1) * (N-1)/(N-K)
    correction = (G / (G - 1.0)) * ((N - 1.0) / (N - K))
    B *= correction

    # Sandwich: V = (X'X)^{-1} B (X'X)^{-1}
    V = XtX_inv @ B @ XtX_inv
    se: np.ndarray[Any, np.dtype[Any]] = np.sqrt(np.diag(V))
    return se


def _double_clustered_se(
    X: np.ndarray[Any, np.dtype[Any]],
    residuals: np.ndarray[Any, np.dtype[Any]],
    cluster1: np.ndarray[Any, np.dtype[Any]],
    cluster2: np.ndarray[Any, np.dtype[Any]],
) -> np.ndarray[Any, np.dtype[Any]]:
    """Two-way clustered SE using Cameron-Gelbach-Miller (2011) formula.

    V_double = V_cluster1 + V_cluster2 - V_intersection

    Parameters
    ----------
    X : np.ndarray
        Regressor matrix (N, K) including constant.
    residuals : np.ndarray
        OLS residuals (N,).
    cluster1 : np.ndarray
        First dimension cluster ids (e.g., entity).
    cluster2 : np.ndarray
        Second dimension cluster ids (e.g., date).

    Returns
    -------
    np.ndarray
        Standard errors (K,).
    """
    N, K = X.shape
    XtX_inv = np.linalg.inv(X.T @ X)

    # Compute each one-way clustering meat
    def _meat(cids: np.ndarray[Any, np.dtype[Any]]) -> np.ndarray[Any, np.dtype[Any]]:
        unique_c = np.unique(cids)
        G = len(unique_c)
        B = np.zeros((K, K))
        for g in unique_c:
            mask = cids == g
            Xg = X[mask]
            eg = residuals[mask].reshape(-1, 1)
            score = (Xg * eg).sum(axis=0).reshape(-1, 1)
            B += score @ score.T
        correction = (G / (G - 1.0)) * ((N - 1.0) / (N - K))
        meat_val: np.ndarray[Any, np.dtype[Any]] = B * correction
        return meat_val

    B1 = _meat(cluster1)
    B2 = _meat(cluster2)

    # Intersection cluster: unique (c1, c2) pairs
    intersection = np.array([f"{a}_{b}" for a, b in zip(cluster1, cluster2)])
    B12 = _meat(intersection)

    # Cameron-Gelbach-Miller: V = V1 + V2 - V12
    V = XtX_inv @ (B1 + B2 - B12) @ XtX_inv

    # Ensure positive diagonal (numerical guard)
    diag_v = np.diag(V)
    diag_v = np.maximum(diag_v, 0.0)
    se: np.ndarray[Any, np.dtype[Any]] = np.sqrt(diag_v)
    return se


def cs_regression(
    data: pd.DataFrame,
    ret_col: str = "ret",
    char_cols: list[str] | None = None,
    date_col: str = "date",
    entity_col: str = "permno",
    controls: list[str] | None = None,
    n_lags_nw: int | None = None,
    double_cluster: bool = False,
) -> CSRegressionResult:
    """Pooled cross-sectional OLS with Newey-West or clustered SE.

    Runs a pooled OLS regression of returns on characteristics (and controls),
    with robust standard errors.

    Parameters
    ----------
    data : pd.DataFrame
        Long-format panel with date, entity, returns, and characteristics.
    ret_col : str
        Name of the return column.
    char_cols : list[str] | None
        Characteristic columns to include as regressors.
    date_col : str
        Name of the date column.
    entity_col : str
        Name of the entity identifier column.
    controls : list[str] | None
        Additional control variables.
    n_lags_nw : int | None
        Lags for Newey-West SE. If None and double_cluster is False, uses
        automatic bandwidth. Ignored when double_cluster is True.
    double_cluster : bool
        If True, use two-way (entity, date) clustered SE instead of NW.

    Returns
    -------
    CSRegressionResult
        Coefficients with robust SE, R-squared, and observation count.

    References
    ----------
    Petersen, M. A. (2009). Estimating standard errors in finance panel data
        sets: Comparing approaches. RFS, 22(1), 435-480.
    """
    # Build regressor list
    rhs_cols: list[str] = []
    if char_cols is not None:
        rhs_cols.extend(char_cols)
    if controls is not None:
        rhs_cols.extend(controls)

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

    # Drop rows with missing values in relevant columns
    df = data.dropna(subset=[ret_col, *rhs_cols]).copy()
    N = len(df)

    if N < len(rhs_cols) + 2:
        msg = f"Insufficient observations: {N} < {len(rhs_cols) + 2}."
        raise ValueError(msg)

    y = np.asarray(df[ret_col].values, dtype=np.float64)
    X_raw = np.asarray(df[rhs_cols].values, dtype=np.float64)

    # Add intercept
    X = np.column_stack([np.ones(N), X_raw])

    # OLS coefficients
    coefs = np.linalg.lstsq(X, y, rcond=None)[0]

    # Residuals and R-squared
    residuals = y - X @ coefs
    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

    # Standard errors
    if double_cluster:
        # Two-way clustering on entity and date
        cluster_entity = np.asarray(df[entity_col].values)
        cluster_date = np.asarray(df[date_col].values)
        se = _double_clustered_se(X, residuals, cluster_entity, cluster_date)
    else:
        # Newey-West HAC standard errors
        se = newey_west_se(residuals, X, n_lags=n_lags_nw)

    # Extract slope coefficients (skip intercept at index 0)
    coef_vals = coefs[1:]
    se_vals = se[1:]
    t_vals = np.where(se_vals > 0.0, coef_vals / se_vals, 0.0)

    coef_series: pd.Series[Any] = pd.Series(coef_vals, index=rhs_cols, name="coef")
    se_series: pd.Series[Any] = pd.Series(se_vals, index=rhs_cols, name="se")
    t_series: pd.Series[Any] = pd.Series(t_vals, index=rhs_cols, name="t_stat")

    return CSRegressionResult(
        coefs=coef_series,
        se=se_series,
        t_stats=t_series,
        r_squared=r_squared,
        n_obs=N,
    )


def multiple_testing_correction(
    t_stats: pd.Series[Any],
    method: str = "bhy",
    alpha: float = 0.05,
) -> pd.DataFrame:
    """Multiple testing correction for cross-sectional t-statistics.

    Adjusts for the multiple testing problem when evaluating many
    candidate anomalies or factors simultaneously.

    Methods
    -------
    BHY (Benjamini-Hochberg-Yekutieli):
        Stepup procedure on sorted |t|.
        p_i = 2*(1 - Phi(|t_i|)) (two-sided)
        Sort p ascending, find largest k s.t. p_(k) <= (k/n) * alpha / c_n
        where c_n = sum(1/j, j=1..n).
        Reject all H_(j) for j <= k.

    Bonferroni:
        Reject if |t_i| > Phi^{-1}(1 - alpha/(2n)).

    Holm:
        Stepdown procedure. Sort p ascending.
        Reject H_(j) while p_(j) <= alpha / (n - j + 1).

    Parameters
    ----------
    t_stats : pd.Series
        Named series of t-statistics for each hypothesis.
    method : str
        One of 'bhy', 'bonferroni', 'holm'.
    alpha : float
        Family-wise or FDR significance level.

    Returns
    -------
    pd.DataFrame
        Columns: {t_stat, p_value, adjusted_threshold, significant}.

    References
    ----------
    Harvey, C. R., Liu, Y., & Zhu, H. (2016). ... and the cross-section of
        expected returns. Review of Financial Studies, 29(1), 5-68.
    """
    n = len(t_stats)
    names = t_stats.index.tolist()

    # Two-sided p-values from |t|
    abs_t = np.abs(t_stats.values.astype(np.float64))
    p_values = 2.0 * (1.0 - sp_stats.norm.cdf(abs_t))

    if method == "bhy":
        # BHY: Benjamini-Hochberg-Yekutieli step-up procedure
        # c_n = sum(1/j for j=1..n) -- accounts for arbitrary dependence
        c_n = float(np.sum(1.0 / np.arange(1, n + 1)))

        # Sort p-values ascending
        sorted_idx = np.argsort(p_values)
        sorted_p = p_values[sorted_idx]

        # BHY thresholds: (k/n) * alpha / c_n for k=1..n
        thresholds = np.array([(k / n) * alpha / c_n for k in range(1, n + 1)])

        # Find largest k where p_(k) <= threshold_(k)
        reject_up_to = -1
        for k in range(n):
            if sorted_p[k] <= thresholds[k]:
                reject_up_to = k

        # Reject all hypotheses with rank <= reject_up_to
        significant = np.zeros(n, dtype=bool)
        adjusted_thresh = np.full(n, np.nan)

        for k in range(n):
            orig_idx = sorted_idx[k]
            adjusted_thresh[orig_idx] = thresholds[k]
            if k <= reject_up_to:
                significant[orig_idx] = True

    elif method == "bonferroni":
        # Bonferroni: reject if p_i <= alpha / n
        bonf_threshold = alpha / n
        significant = p_values <= bonf_threshold
        adjusted_thresh = np.full(n, bonf_threshold)

    elif method == "holm":
        # Holm step-down procedure
        sorted_idx = np.argsort(p_values)
        sorted_p = p_values[sorted_idx]

        # Holm thresholds: alpha / (n - k + 1) for the k-th smallest p
        significant = np.zeros(n, dtype=bool)
        adjusted_thresh = np.full(n, np.nan)

        for k in range(n):
            holm_thresh = alpha / (n - k)
            orig_idx = sorted_idx[k]
            adjusted_thresh[orig_idx] = holm_thresh
            if sorted_p[k] <= holm_thresh:
                significant[orig_idx] = True
            else:
                # Holm is step-down: stop rejecting after first non-rejection
                break

    else:
        msg = f"Unknown method: {method}. Use 'bhy', 'bonferroni', or 'holm'."
        raise ValueError(msg)

    result = pd.DataFrame(
        {
            "t_stat": t_stats.values,
            "p_value": p_values,
            "adjusted_threshold": adjusted_thresh,
            "significant": significant,
        },
        index=names,
    )

    return result
