# =============================================================================
# Author: Simon-Pierre Boucher
# Contact: contact@spboucher.ai
# =============================================================================
"""Econometric primitives shared by every analysis script.

All estimators are deliberately implemented with :func:`numpy.linalg.lstsq`
and explicit sandwich formulas — exactly as in the original study — so that
re-running the pipeline reproduces the published numbers bit-for-bit. Do not
"modernize" the numerics (e.g. switch to statsmodels) without re-validating
every result table against ``results/``.
"""

import numpy as np
import pandas as pd
from numpy.linalg import lstsq


# --------------------------------------------------------------------------
# Transformations
# --------------------------------------------------------------------------
def winsorize(s: pd.Series, q: float = 0.01) -> pd.Series:
    """Clip a series at its ``q`` and ``1-q`` empirical quantiles."""
    lo, hi = s.quantile(q), s.quantile(1 - q)
    return s.clip(lo, hi)


def standardize(X: np.ndarray) -> np.ndarray:
    """Column-standardize (population std, zero-variance columns left at 1)."""
    means = X.mean(axis=0)
    stds = X.std(axis=0)
    stds[stds == 0] = 1
    return (X - means) / stds


def add_constant(X: np.ndarray) -> np.ndarray:
    """Prepend an intercept column of ones."""
    return np.column_stack([np.ones(len(X)), X])


def design_matrix(data: pd.DataFrame, features: list) -> np.ndarray:
    """Standardized feature matrix with intercept, ready for ``lstsq``."""
    return add_constant(standardize(data[features].values))


# --------------------------------------------------------------------------
# OLS with robust inference
# --------------------------------------------------------------------------
def ols(X: np.ndarray, y: np.ndarray) -> np.ndarray:
    """Least-squares coefficients."""
    coefs, _, _, _ = lstsq(X, y, rcond=None)
    return coefs


def r_squared(y: np.ndarray, y_pred: np.ndarray) -> float:
    ss_res = np.sum((y - y_pred) ** 2)
    ss_tot = np.sum((y - y.mean()) ** 2)
    return 1 - ss_res / ss_tot if ss_tot > 0 else 0


def adjusted_r_squared(r2: float, n: int, k: int) -> float:
    return 1 - (1 - r2) * (n - 1) / (n - k - 1)


def _xtx_inverse(X: np.ndarray) -> np.ndarray:
    try:
        return np.linalg.inv(X.T @ X)
    except np.linalg.LinAlgError:
        return np.linalg.pinv(X.T @ X)


def hc1_tstats(X: np.ndarray, y: np.ndarray, coefs: np.ndarray,
               n_features: int) -> tuple:
    """Heteroskedasticity-robust (HC1) standard errors and t-statistics.

    ``n_features`` excludes the intercept; the small-sample factor is
    ``n / (n - k - 1)`` as in the original study.
    """
    n = len(y)
    e = y - X @ coefs
    XtX_inv = _xtx_inverse(X)
    # X' diag(e²) X computed by broadcasting (identical to the dense-diagonal
    # formula, but O(nk) memory instead of O(n²)).
    S = (X * (e ** 2)[:, None]).T @ X * n / (n - n_features - 1)
    se = np.sqrt(np.abs(np.diag(XtX_inv @ S @ XtX_inv)))
    t_stats = coefs / np.where(se > 0, se, 1)
    return se, t_stats


def newey_west_tstats(X: np.ndarray, y: np.ndarray, coefs: np.ndarray,
                      n_lags: int = 5) -> tuple:
    """Newey-West HAC standard errors with Bartlett kernel weights."""
    n, k = X.shape
    e = y - X @ coefs
    XtX_inv = _xtx_inverse(X)

    S = np.zeros((k, k))
    for j in range(n):
        S += e[j] ** 2 * np.outer(X[j], X[j])
    for lag in range(1, n_lags + 1):
        w = 1 - lag / (n_lags + 1)
        for j in range(lag, n):
            cross = e[j] * e[j - lag] * np.outer(X[j], X[j - lag])
            S += w * (cross + cross.T)

    V = XtX_inv @ S @ XtX_inv
    se = np.sqrt(np.abs(np.diag(V)))
    t_stats = coefs / np.where(se > 0, se, 1)
    return se, t_stats


def double_clustered_tstats(X: np.ndarray, y: np.ndarray, coefs: np.ndarray,
                            tickers: np.ndarray, dates: np.ndarray) -> tuple:
    """Two-way (ticker + date) clustered SEs, Cameron-Gelbach-Miller (2011)."""
    k = X.shape[1]
    e = y - X @ coefs
    XtX_inv = _xtx_inverse(X)

    def cluster_meat(labels: np.ndarray) -> np.ndarray:
        S = np.zeros((k, k))
        for g in np.unique(labels):
            mask = labels == g
            u = (X[mask].T * e[mask]).sum(axis=1, keepdims=True)
            S += u @ u.T
        return S

    S_hc = (X * (e ** 2)[:, None]).T @ X
    V = XtX_inv @ (cluster_meat(tickers) + cluster_meat(dates) - S_hc) @ XtX_inv
    se = np.sqrt(np.abs(np.diag(V)))
    t_stats = coefs / np.where(se > 0, se, 1)
    return se, t_stats


def quantile_regression(X: np.ndarray, y: np.ndarray, tau: float,
                        max_iter: int = 50) -> np.ndarray:
    """Quantile regression via iteratively reweighted least squares."""
    coefs = ols(X, y)
    for _ in range(max_iter):
        residuals = y - X @ coefs
        weights = np.where(residuals >= 0, tau, 1 - tau)
        weights = np.maximum(weights / (np.abs(residuals) + 1e-6), 1e-6)
        Xw = X * weights[:, None]
        try:
            coefs = np.linalg.solve(Xw.T @ X, Xw.T @ y)
        except np.linalg.LinAlgError:
            break
    return coefs


def pooled_regression_summary(data: pd.DataFrame, features: list, target: str,
                              label: str = "", min_obs: int = 100):
    """Winsorize → standardize → pooled OLS with HC1; compact summary dict.

    Returns ``None`` when fewer than ``min_obs`` complete observations exist.
    """
    sub = data[features + [target]].dropna()
    if len(sub) < min_obs:
        return None
    for f in features:
        sub[f] = winsorize(sub[f])
    sub[target] = winsorize(sub[target])

    X = design_matrix(sub, features)
    y = sub[target].values
    coefs = ols(X, y)
    r2 = r_squared(y, X @ coefs)
    _, t = hc1_tstats(X, y, coefs, len(features))
    return {
        'label': label,
        'n_obs': len(sub),
        'r2': r2,
        'adj_r2': adjusted_r_squared(r2, len(y), len(features)),
        'n_significant': int(np.sum(np.abs(t[1:]) > 1.96)),
    }


# --------------------------------------------------------------------------
# Fama-MacBeth
# --------------------------------------------------------------------------
def fama_macbeth(data: pd.DataFrame, features: list, target: str,
                 min_obs_per_day: int = 10, min_periods: int = 30):
    """Daily cross-sectional regressions; time-series means and t-statistics.

    Returns ``(results_df, n_periods, n_tickers)`` or ``None`` when the panel
    is too small.
    """
    sub = data[features + [target, 'ticker', 'trade_date']].dropna()
    if len(sub) < 100:
        return None
    for f in features:
        sub[f] = winsorize(sub[f])
    sub[target] = winsorize(sub[target])

    coef_series = []
    for dt in sub['trade_date'].unique():
        day = sub[sub['trade_date'] == dt]
        if len(day) < min_obs_per_day:
            continue
        X = add_constant(standardize(day[features].values))
        try:
            coef_series.append(ols(X, day[target].values))
        except np.linalg.LinAlgError:
            continue

    if len(coef_series) < min_periods:
        return None

    coef_array = np.array(coef_series)
    avg = coef_array.mean(axis=0)
    se = coef_array.std(axis=0) / np.sqrt(len(coef_series))
    t = avg / se
    results = pd.DataFrame({
        'variable': ['const'] + features,
        'fm_coefficient': avg,
        'fm_std_error': se,
        'fm_t_stat': t,
        'fm_significant_5pct': np.abs(t) > 1.96,
        'fm_significant_1pct': np.abs(t) > 2.576,
    })
    return results, len(coef_series), sub['ticker'].nunique()


# --------------------------------------------------------------------------
# Time-series: Granger causality and bivariate VAR
# --------------------------------------------------------------------------
def granger_test(y: np.ndarray, x: np.ndarray, max_lag: int = 5):
    """F-test of the null that ``x`` does not Granger-cause ``y``."""
    n = len(y)
    if n < max_lag + 50:
        return None

    Y = y[max_lag:]
    X_r = np.column_stack([y[max_lag - i - 1:n - i - 1] for i in range(max_lag)])
    X_r = add_constant(X_r)
    ssr_r = np.sum((Y - X_r @ ols(X_r, Y)) ** 2)

    X_u = np.column_stack([X_r] + [x[max_lag - i - 1:n - i - 1] for i in range(max_lag)])
    ssr_u = np.sum((Y - X_u @ ols(X_u, Y)) ** 2)

    n_eff = len(Y)
    f_stat = ((ssr_r - ssr_u) / max_lag) / (ssr_u / (n_eff - 2 * max_lag - 1))

    from scipy import stats
    try:
        p_value = 1 - stats.f.cdf(f_stat, max_lag, n_eff - 2 * max_lag - 1)
    except Exception:
        p_value = np.nan
    return {'f_stat': f_stat, 'p_value': p_value, 'n_obs': n_eff}


def estimate_var(y1: np.ndarray, y2: np.ndarray, lags: int = 5,
                 irf_periods: int = 20):
    """Bivariate VAR(p) by equation-wise OLS, with companion-form IRFs."""
    n = len(y1)
    if n < lags + 50:
        return None

    Y1, Y2 = y1[lags:], y2[lags:]
    X_lags = []
    for lag in range(1, lags + 1):
        X_lags.append(y1[lags - lag:n - lag])
        X_lags.append(y2[lags - lag:n - lag])
    X = np.column_stack([np.ones(len(Y1))] + X_lags)

    coefs1 = ols(X, Y1)
    r2_1 = 1 - np.sum((Y1 - X @ coefs1) ** 2) / np.sum((Y1 - Y1.mean()) ** 2)
    coefs2 = ols(X, Y2)
    r2_2 = 1 - np.sum((Y2 - X @ coefs2) ** 2) / np.sum((Y2 - Y2.mean()) ** 2)

    # Companion matrix
    A = np.zeros((2 * lags, 2 * lags))
    for lag in range(lags):
        A[0, 2 * lag] = coefs1[1 + 2 * lag]
        A[0, 2 * lag + 1] = coefs1[1 + 2 * lag + 1]
        A[1, 2 * lag] = coefs2[1 + 2 * lag]
        A[1, 2 * lag + 1] = coefs2[1 + 2 * lag + 1]
    for i in range(2, 2 * lags):
        A[i, i - 2] = 1.0

    irf = np.zeros((irf_periods, 2, 2))
    for h in range(irf_periods):
        power = np.eye(2 * lags) if h == 0 else power @ A
        irf[h] = power[:2, :2]

    return {'coefs_eq1': coefs1, 'coefs_eq2': coefs2,
            'r2_eq1': r2_1, 'r2_eq2': r2_2,
            'irf': irf, 'n_obs': len(Y1)}