"""Information shares, Granger causality, and VAR-based price discovery."""

from __future__ import annotations

import numpy as np
import pandas as pd
from scipy import stats
from statsmodels.tsa.api import VAR
from statsmodels.tsa.stattools import grangercausalitytests, adfuller


def granger_causality_matrix(
    data: pd.DataFrame,
    max_lag: int = 10,
) -> tuple[pd.DataFrame, pd.DataFrame]:
    """Pairwise Granger causality F-tests.

    Returns (f_stats, p_values) DataFrames indexed by caused, columns by causing.
    """
    cols = data.columns.tolist()
    f_mat = pd.DataFrame(np.nan, index=cols, columns=cols)
    p_mat = pd.DataFrame(np.nan, index=cols, columns=cols)

    clean = data.dropna()
    if len(clean) < max_lag + 5:
        return f_mat, p_mat

    for caused in cols:
        for causing in cols:
            if caused == causing:
                continue
            try:
                result = grangercausalitytests(
                    clean[[caused, causing]].values, maxlag=max_lag, verbose=False
                )
                # Lag chosen by the AIC of the unrestricted bivariate regression
                # (revision 2026-10-17). The earlier code selected the lag with
                # the smallest p-value, a search that inflates rejection rates.
                best_lag = min(result.keys(), key=lambda k: result[k][1][1].aic)
                f_stat = result[best_lag][0]["ssr_ftest"][0]
                p_val = result[best_lag][0]["ssr_ftest"][1]
                f_mat.loc[caused, causing] = f_stat
                p_mat.loc[caused, causing] = p_val
            except Exception:
                pass

    return f_mat, p_mat


def estimate_var(
    data: pd.DataFrame,
    max_lags: int = 10,
    ic: str = "aic",
) -> dict:
    """Estimate VAR model with optimal lag selection.

    Returns dict with 'model', 'optimal_lag', 'coefficients', 'irf', 'summary'.
    """
    clean = data.dropna()
    if len(clean) < max_lags + 10:
        return {"model": None, "optimal_lag": 0, "coefficients": {}, "irf": {}}

    model = VAR(clean)
    result = model.fit(maxlags=max_lags, ic=ic)

    coefficients = {}
    params = result.params
    if isinstance(params, pd.DataFrame):
        for col in params.columns:
            coefficients[col] = params[col]
    else:
        for col_idx, col in enumerate(data.columns):
            coefficients[col] = pd.Series(params[:, col_idx])

    irf_data = {}
    try:
        irf_result = result.irf(periods=20)
        for i, shock in enumerate(data.columns):
            for j, resp in enumerate(data.columns):
                irf_data[(shock, resp)] = irf_result.irfs[:, j, i]
    except Exception:
        pass

    return {
        "model": result,
        "optimal_lag": result.k_ar,
        "coefficients": coefficients,
        "irf": irf_data,
        "summary": str(result.summary()),
    }


def compute_impulse_responses(var_result: dict, periods: int = 20) -> dict:
    """Extract impulse response functions from VAR result."""
    if var_result.get("model") is None:
        return {}
    irf = var_result["model"].irf(periods=periods)
    data = {}
    names = var_result["model"].names
    for i, shock in enumerate(names):
        for j, resp in enumerate(names):
            data[(shock, resp)] = irf.irfs[:, j, i]
    return data


def lead_lag_regression(
    y: pd.Series,
    x: pd.Series,
    leads: range = range(-5, 6),
) -> pd.DataFrame:
    """Cross-correlation and lead-lag OLS regressions.

    For each lag k: regress y_t on x_{t-k}.
    """
    results = []
    for k in leads:
        x_shifted = x.shift(k)
        valid = pd.DataFrame({"y": y, "x": x_shifted}).dropna()
        if len(valid) < 5:
            results.append({"lag": k, "coefficient": np.nan, "t_stat": np.nan, "r_squared": np.nan})
            continue

        x_vals = valid["x"].values
        y_vals = valid["y"].values
        x_dm = x_vals - x_vals.mean()
        y_dm = y_vals - y_vals.mean()
        var_x = (x_dm ** 2).sum()
        if var_x == 0:
            results.append({"lag": k, "coefficient": np.nan, "t_stat": np.nan, "r_squared": np.nan})
            continue

        beta = (x_dm * y_dm).sum() / var_x
        residuals = y_dm - beta * x_dm
        se_beta = np.sqrt((residuals ** 2).sum() / ((len(valid) - 2) * var_x))
        t_stat = beta / se_beta if se_beta > 0 else np.nan
        ss_res = (residuals ** 2).sum()
        ss_tot = (y_dm ** 2).sum()
        r_sq = 1 - ss_res / ss_tot if ss_tot > 0 else np.nan

        results.append({"lag": k, "coefficient": beta, "t_stat": t_stat, "r_squared": r_sq})

    return pd.DataFrame(results)


def hasbrouck_information_shares(
    prices: pd.DataFrame,
    max_lags: int = 5,
) -> pd.DataFrame:
    """Hasbrouck (1995) information shares via variance decomposition.

    Uses VAR on returns and computes forecast error variance decomposition
    as a proxy for information shares when series may not be cointegrated.
    """
    returns = prices.pct_change().dropna()
    clean = returns.dropna()
    if len(clean) < max_lags + 10:
        return pd.DataFrame({"market": prices.columns, "info_share": np.nan})

    try:
        model = VAR(clean)
        result = model.fit(maxlags=max_lags, ic="aic")
        fevd = result.fevd(periods=10)
        decomp = fevd.decomp
        shares = decomp[-1]
        info_shares = pd.DataFrame(
            shares,
            index=prices.columns,
            columns=prices.columns,
        )
        diag_shares = pd.Series(
            {col: info_shares.loc[col, col] for col in prices.columns},
            name="info_share",
        )
        return diag_shares.to_frame().reset_index().rename(columns={"index": "market"})
    except Exception:
        return pd.DataFrame({"market": prices.columns.tolist(), "info_share": [np.nan] * len(prices.columns)})