"""
Panel unit root tests.

Provides LLC (Levin-Lin-Chu), IPS (Im-Pesaran-Shin), Fisher-type,
and Hadri stationarity tests for panel data.

Equivalent to Stata's ``xtunitroot`` and R's ``plm::purtest()``.

References
----------
Levin, A., Lin, C.F. & Chu, C.S.J. (2002).
"Unit Root Tests in Panel Data: Asymptotic and Finite-Sample Properties."
*Journal of Econometrics*, 108(1), 1-24. [@levin2002unit]

Im, K.S., Pesaran, M.H. & Shin, Y. (2003).
"Testing for Unit Roots in Heterogeneous Panels."
*Journal of Econometrics*, 115(1), 53-74.
"""

import warnings
from typing import Optional, List, Dict, Any
import numpy as np
import pandas as pd
from scipy import stats


class PanelUnitRootResult:
    """Results from panel unit root test."""

    def __init__(self, test_type, statistic, p_value, n_units, n_periods,
                 individual_stats, lags):
        self.test_type = test_type
        self.statistic = statistic
        self.p_value = p_value
        self.n_units = n_units
        self.n_periods = n_periods
        self.individual_stats = individual_stats
        self.lags = lags

    def summary(self) -> str:
        lines = [
            f"Panel Unit Root Test: {self.test_type}",
            "=" * 55,
            f"H0: Panels contain unit roots",
            f"Ha: {'Panels are stationary' if self.test_type != 'Hadri' else 'Some panels contain unit roots'}",
            "",
            f"Statistic: {self.statistic:.4f}",
            f"P-value:   {self.p_value:.4f}",
            f"N units:   {self.n_units}",
            f"T periods: {self.n_periods}",
            "",
            f"Conclusion: {'Reject H0' if self.p_value < 0.05 else 'Fail to reject H0'} at 5%",
            "=" * 55,
        ]
        return "\n".join(lines)


def _adf_single(y, lags=None, trend='c'):
    """ADF test for a single series. Returns (t-stat, p-value, lags)."""
    n = len(y)
    if lags is None:
        lags = int(np.floor(4 * (n / 100)**0.25))
    lags = min(lags, n // 3)

    dy = np.diff(y)
    y_lag = y[:-1]

    T = len(dy) - lags
    if T < 3:
        return np.nan, np.nan, lags

    Y = dy[lags:]
    X = y_lag[lags:].reshape(-1, 1)
    for j in range(1, lags + 1):
        X = np.column_stack([X, dy[lags-j:len(dy)-j]])
    if trend == 'c':
        X = np.column_stack([X, np.ones(T)])
    elif trend == 'ct':
        X = np.column_stack([X, np.ones(T), np.arange(1, T+1)])

    try:
        beta = np.linalg.lstsq(X, Y, rcond=None)[0]
        resid = Y - X @ beta
        se = np.sqrt(np.sum(resid**2) / (T - X.shape[1]) *
                     np.linalg.inv(X.T @ X)[0, 0])
        t_stat = beta[0] / se
    except np.linalg.LinAlgError:
        return np.nan, np.nan, lags

    # Approximate p-value using MacKinnon distribution
    # For unit root: left-tailed test
    # Use normal approximation for large T
    p_value = stats.norm.cdf(t_stat)

    return t_stat, p_value, lags


def panel_unitroot(
    data: pd.DataFrame,
    variable: str,
    id: str = 'id',
    time: str = 'time',
    test: str = 'ips',
    lags: int = None,
    trend: str = 'c',
) -> PanelUnitRootResult:
    """
    Panel unit root test.

    Equivalent to Stata's ``xtunitroot`` and R's ``plm::purtest()``.

    Parameters
    ----------
    data : pd.DataFrame
        Panel data.
    variable : str
        Variable to test.
    id : str, default 'id'
        Unit identifier.
    time : str, default 'time'
        Time identifier.
    test : str, default 'ips'
        Test type: 'llc' (Levin-Lin-Chu), 'ips' (Im-Pesaran-Shin),
        'fisher' (Fisher-type ADF), 'hadri' (stationarity test).
    lags : int, optional
        Number of ADF lags. If None, uses AIC selection.
    trend : str, default 'c'
        'n' (none), 'c' (constant), 'ct' (constant + trend).

    Returns
    -------
    PanelUnitRootResult

    Examples
    --------
    >>> import statspai as sp
    >>> result = sp.panel_unitroot(df, variable='gdp', id='country', time='year')
    >>> print(result.summary())
    """
    units = data[id].unique()
    N = len(units)
    T_avg = data.groupby(id)[variable].count().mean()

    individual_results = []
    n_short = 0
    for unit in units:
        y = data.loc[data[id] == unit].sort_values(time)[variable].dropna().values
        if len(y) < 5:
            n_short += 1
            continue
        t_stat, p_val, used_lags = _adf_single(y.astype(float), lags=lags, trend=trend)
        individual_results.append({
            'unit': unit, 't_stat': t_stat, 'p_value': p_val, 'lags': used_lags
        })

    ind_df = pd.DataFrame(individual_results)
    valid = ind_df.dropna(subset=['t_stat']) if len(ind_df) else ind_df
    n_valid = len(valid)
    n_adf_failed = len(ind_df) - n_valid

    # The panel statistic is built only from units with a finite ADF t-stat.
    # Silently shrinking the unit set hides that some series were dropped
    # (singular ADF design or too few periods) — surface it (CLAUDE.md §7).
    if n_valid == 0:
        raise ValueError(
            f"panel_unitroot('{test}'): no unit yielded a valid ADF statistic "
            f"({n_short}/{N} units had <5 periods, {n_adf_failed} had a "
            f"singular ADF design). Cannot compute a panel unit-root test."
        )
    if n_short > 0 or n_adf_failed > 0:
        warnings.warn(
            f"panel_unitroot('{test}'): computed over {n_valid}/{N} units. "
            f"Excluded {n_short} unit(s) with <5 periods and {n_adf_failed} "
            f"unit(s) whose ADF regression was singular. The reported "
            f"statistic and n_units reflect only the {n_valid} valid units.",
            RuntimeWarning, stacklevel=2,
        )

    if test == 'ips':
        # Im-Pesaran-Shin: average of individual ADF t-statistics
        t_bar = valid['t_stat'].mean()

        # IPS critical moments (approximate for large T)
        # E[t] ≈ -1.52 for trend='c', Var[t] ≈ 0.77 (from IPS tables)
        if trend == 'c':
            E_t = -1.52
            Var_t = 0.77
        elif trend == 'ct':
            E_t = -2.12
            Var_t = 0.67
        else:
            E_t = -0.41
            Var_t = 0.95

        W_stat = np.sqrt(n_valid) * (t_bar - E_t) / np.sqrt(Var_t)
        p_value = stats.norm.cdf(W_stat)

        return PanelUnitRootResult(
            test_type='Im-Pesaran-Shin (IPS)',
            statistic=W_stat,
            p_value=p_value,
            n_units=n_valid,
            n_periods=int(T_avg),
            individual_stats=ind_df,
            lags=lags,
        )

    elif test == 'llc':
        # Levin-Lin-Chu: pooled t-statistic with bias correction
        t_bar = valid['t_stat'].mean()
        # Simplified LLC: adjusted pooled t
        delta_star = np.sqrt(n_valid) * t_bar
        # Bias correction (approximate)
        mu_star = -1.0  # approximate mean under H0
        sigma_star = 1.0  # approximate std
        t_star = (delta_star - mu_star * np.sqrt(n_valid)) / sigma_star
        p_value = stats.norm.cdf(t_star)

        return PanelUnitRootResult(
            test_type='Levin-Lin-Chu (LLC)',
            statistic=t_star,
            p_value=p_value,
            n_units=n_valid,
            n_periods=int(T_avg),
            individual_stats=ind_df,
            lags=lags,
        )

    elif test == 'fisher':
        # Fisher-type: combine p-values using -2 Σ ln(p_i) ~ χ²(2N)
        p_values = valid['p_value'].values
        p_values = np.clip(p_values, 1e-10, 1 - 1e-10)
        fisher_stat = -2 * np.sum(np.log(p_values))
        df = 2 * n_valid
        p_value = 1 - stats.chi2.cdf(fisher_stat, df)

        return PanelUnitRootResult(
            test_type='Fisher-type ADF',
            statistic=fisher_stat,
            p_value=p_value,
            n_units=n_valid,
            n_periods=int(T_avg),
            individual_stats=ind_df,
            lags=lags,
        )

    elif test == 'hadri':
        # Hadri (2000) stationarity test
        # H0: stationarity vs H1: unit root in some panels
        lm_stats = []
        for unit in units:
            y = data.loc[data[id] == unit].sort_values(time)[variable].dropna().values
            if len(y) < 5:
                continue
            T_i = len(y)
            if trend == 'c':
                resid = y - y.mean()
            else:
                t_vec = np.arange(T_i)
                X = np.column_stack([np.ones(T_i), t_vec])
                b = np.linalg.lstsq(X, y, rcond=None)[0]
                resid = y - X @ b
            S = np.cumsum(resid)
            sigma2 = np.sum(resid**2) / T_i
            if sigma2 > 0:
                lm_i = np.sum(S**2) / (T_i**2 * sigma2)
                lm_stats.append(lm_i)

        if len(lm_stats) == 0:
            return PanelUnitRootResult('Hadri', np.nan, np.nan, 0, 0, None, lags)

        lm_bar = np.mean(lm_stats)
        # Hadri standardized statistic
        # Under H0: Z ~ N(0,1)
        mu = 1/6 if trend == 'c' else 1/15
        sigma2_lm = 1/45 if trend == 'c' else 11/6300
        Z = np.sqrt(n_valid) * (lm_bar - mu) / np.sqrt(sigma2_lm)
        p_value = 1 - stats.norm.cdf(Z)

        return PanelUnitRootResult(
            test_type='Hadri (stationarity)',
            statistic=Z,
            p_value=p_value,
            n_units=len(lm_stats),
            n_periods=int(T_avg),
            individual_stats=None,
            lags=lags,
        )

    else:
        raise ValueError(f"Unknown test: {test}. Use 'ips', 'llc', 'fisher', or 'hadri'.")
