"""Evaluation metrics, paired block-bootstrap intervals, Diebold-Mariano / Giacomini-White tests, Holm correction,
calibration, and two-way clustered panel regressions."""
from __future__ import annotations

import numpy as np
import pandas as pd
import polars as pl
from scipy import stats
from sklearn.metrics import average_precision_score, brier_score_loss, log_loss, roc_auc_score


# --------------------------------------------------------------------- metrics
def binary_metrics(y: np.ndarray, p: np.ndarray) -> dict:
    p = np.clip(p, 1e-6, 1 - 1e-6)
    out = {"n": int(len(y)), "base_rate": float(y.mean()), "logloss": float(log_loss(y, p)), "brier": float(brier_score_loss(y, p))}
    if 0 < y.mean() < 1:
        out["auc"] = float(roc_auc_score(y, p))
        out["pr_auc"] = float(average_precision_score(y, p))
    # calibration intercept / slope (logistic recalibration of the logit)
    z = np.log(p / (1 - p))
    try:
        from sklearn.linear_model import LogisticRegression
        m = LogisticRegression(C=1e6, solver="lbfgs", max_iter=1000).fit(z.reshape(-1, 1), y)
        out["cal_slope"], out["cal_intercept"] = float(m.coef_[0][0]), float(m.intercept_[0])
    except Exception:
        pass
    return out


def continuous_metrics(y: np.ndarray, yhat: np.ndarray, y_train_mean: float) -> dict:
    e = y - yhat
    ss = np.sum((y - y_train_mean) ** 2)
    return {"n": int(len(y)), "mae": float(np.mean(np.abs(e))), "rmse": float(np.sqrt(np.mean(e ** 2))),
            "oos_r2": float(1 - np.sum(e ** 2) / ss) if ss > 0 else float("nan")}


def pointwise_logloss(y: np.ndarray, p: np.ndarray) -> np.ndarray:
    p = np.clip(p, 1e-6, 1 - 1e-6)
    return -(y * np.log(p) + (1 - y) * np.log(1 - p))


def reliability_table(y: np.ndarray, p: np.ndarray, bins: int = 10) -> pd.DataFrame:
    q = np.quantile(p, np.linspace(0, 1, bins + 1))
    q[0], q[-1] = -np.inf, np.inf
    b = np.digitize(p, q[1:-1])
    rows = []
    for k in range(bins):
        m = b == k
        if m.sum():
            rows.append({"decile": k + 1, "n": int(m.sum()), "mean_pred": float(p[m].mean()), "obs_rate": float(y[m].mean())})
    return pd.DataFrame(rows)


# ------------------------------------------------------------ paired inference
def quarter_blocks(dates: np.ndarray) -> np.ndarray:
    d = pd.to_datetime(dates)
    return (d.year * 4 + (d.month - 1) // 3).to_numpy()


def paired_block_bootstrap(d: np.ndarray, blocks: np.ndarray, n_boot: int = 1000, seed: int = 20261102) -> dict:
    """Mean of a per-observation loss differential with a cluster (block) bootstrap over calendar quarters.
    Returns the mean, a percentile 95 % CI and a two-sided bootstrap p-value for H0: mean = 0."""
    rng = np.random.default_rng(seed)
    ub = np.unique(blocks)
    # per-block sums and counts -> resample blocks with replacement (equal-weight by observation)
    sums = np.array([d[blocks == b].sum() for b in ub]); cnts = np.array([(blocks == b).sum() for b in ub])
    mean = float(d.mean())
    draws = np.empty(n_boot)
    for i in range(n_boot):
        idx = rng.integers(0, len(ub), len(ub))
        draws[i] = sums[idx].sum() / cnts[idx].sum()
    lo, hi = np.percentile(draws, [2.5, 97.5])
    p = float(2 * min((draws - mean <= -abs(mean)).mean(), (draws - mean >= abs(mean)).mean()))
    return {"mean": mean, "ci_lo": float(lo), "ci_hi": float(hi), "p_boot": p, "n_blocks": int(len(ub)), "n": int(len(d))}


def dm_gw_test(d: np.ndarray, blocks: np.ndarray) -> dict:
    """Unconditional Giacomini-White (= Diebold-Mariano with HAC) on the per-quarter mean loss differential.
    HAC: Newey-West with lag = floor(4 (T/100)^(2/9))."""
    ub = np.unique(blocks)
    dq = np.array([d[blocks == b].mean() for b in ub])
    T = len(dq)
    if T < 4:
        return {"t": float("nan"), "p": float("nan"), "T": T}
    m = dq.mean(); u = dq - m
    L = int(np.floor(4 * (T / 100) ** (2 / 9)))
    s2 = u @ u / T
    for l in range(1, L + 1):
        w = 1 - l / (L + 1)
        s2 += 2 * w * (u[l:] @ u[:-l]) / T
    se = np.sqrt(max(s2, 1e-18) / T)
    t = m / se
    return {"t": float(t), "p": float(2 * (1 - stats.t.cdf(abs(t), df=T - 1))), "T": int(T), "mean_quarterly": float(m)}


def holm(pvals: dict[str, float]) -> dict[str, float]:
    items = sorted(pvals.items(), key=lambda kv: kv[1])
    m = len(items); out = {}; running = 0.0
    for i, (k, p) in enumerate(items):
        adj = min(1.0, (m - i) * p); running = max(running, adj); out[k] = running
    return out


# ------------------------------------------------------------- panel regression
def demean_two_way(df: pd.DataFrame, cols: list[str], fe1: str, fe2: str, iters: int = 25) -> pd.DataFrame:
    """Alternating projections for two fixed effects (firm, sector x quarter)."""
    x = df[cols].astype(float).copy()
    for _ in range(iters):
        x = x - x.groupby(df[fe1]).transform("mean")
        x = x - x.groupby(df[fe2]).transform("mean")
    return x


def two_way_cluster_ols(df: pd.DataFrame, y: str, xs: list[str], fe1: str, fe2: str, cl1: str, cl2: str) -> pd.DataFrame:
    """OLS with two-way FE (within transformation) and Cameron-Gelbach-Miller two-way clustered SEs."""
    d = df.dropna(subset=[y] + xs + [fe1, fe2, cl1, cl2]).copy()
    z = demean_two_way(d, [y] + xs, fe1, fe2)
    X = z[xs].to_numpy(); Y = z[y].to_numpy()
    XtX_inv = np.linalg.pinv(X.T @ X)
    beta = XtX_inv @ X.T @ Y
    e = Y - X @ beta
    def meat(groups):
        S = np.zeros((X.shape[1], X.shape[1]))
        for _, idx in pd.Series(np.arange(len(groups))).groupby(groups.to_numpy()).groups.items():
            idx = np.asarray(list(idx)); s = X[idx].T @ e[idx]; S += np.outer(s, s)
        return S
    g12 = (d[cl1].astype(str) + "|" + d[cl2].astype(str))
    V = XtX_inv @ (meat(d[cl1]) + meat(d[cl2]) - meat(g12)) @ XtX_inv
    se = np.sqrt(np.clip(np.diag(V), 0, None))
    t = beta / se
    G = min(d[cl1].nunique(), d[cl2].nunique())
    p = 2 * (1 - stats.t.cdf(np.abs(t), df=max(G - 1, 1)))
    return pd.DataFrame({"var": xs, "coef": beta, "se": se, "t": t, "p": p, "n": len(d), "n_firms": d[cl1].nunique(),
                         "n_quarters": d[cl2].nunique()})