"""
Statistical Testing for Mixed VECM
===================================

This module implements statistical tests for mixed Vector Error Correction Models
containing both I(0) and I(1) variables.

Main Features
-------------
1. Johansen cointegration test wrapper with enhanced output
2. Likelihood ratio tests for I(0) components (Chen, 2022)
3. Restricted beta estimation using Boswijk and Doornik (2004)
4. Cointegration rank selection procedures

References
----------
[1] Chen, P. (2022). Vector Error Correction Models with Stationary and 
    Nonstationary Variables. SSRN Working Paper 4239531.
[2] Johansen, S. (1995). Likelihood-Based Inference in Cointegrated Vector
    Autoregressive Models. Oxford University Press.
[3] Boswijk, H.P. and Doornik, J.A. (2004). Identifying, estimating and testing
    restricted cointegrated systems: An overview. Statistica Neerlandica.
"""

import numpy as np
from numpy.linalg import inv, eig, det, lstsq
from scipy import stats
from scipy.linalg import cholesky, solve
from typing import Tuple, Optional, Union, List, NamedTuple
from dataclasses import dataclass
import warnings

from .utils import (
    vec, unvec, lag_matrix, diff_matrix,
    orthogonal_complement, normalize_beta,
    compute_var_covariance, kronecker
)


@dataclass
class JohansenTestResults:
    """
    Results from Johansen cointegration test.
    
    Attributes
    ----------
    trace_stat : np.ndarray
        Trace statistics for each null hypothesis r ≤ k
    trace_crit : np.ndarray
        Critical values (90%, 95%, 99%) for trace statistics
    max_eig_stat : np.ndarray
        Maximum eigenvalue statistics
    max_eig_crit : np.ndarray
        Critical values (90%, 95%, 99%) for max eigenvalue statistics
    eigenvalues : np.ndarray
        Eigenvalues from the reduced rank regression
    eigenvectors : np.ndarray
        Eigenvectors (cointegrating vectors)
    r0t : np.ndarray
        Residuals from regressing ΔY on lagged differences
    rkt : np.ndarray
        Residuals from regressing Y_{t-1} on lagged differences
    S00 : np.ndarray
        Cross-product matrix S00
    S01 : np.ndarray
        Cross-product matrix S01
    S10 : np.ndarray
        Cross-product matrix S10
    S11 : np.ndarray
        Cross-product matrix S11
    n_obs : int
        Number of observations used
    n_vars : int
        Number of variables
    det_order : int
        Deterministic order used
    k_ar_diff : int
        Number of lagged differences
    """
    trace_stat: np.ndarray
    trace_crit: np.ndarray
    max_eig_stat: np.ndarray
    max_eig_crit: np.ndarray
    eigenvalues: np.ndarray
    eigenvectors: np.ndarray
    r0t: np.ndarray
    rkt: np.ndarray
    S00: np.ndarray
    S01: np.ndarray
    S10: np.ndarray
    S11: np.ndarray
    n_obs: int
    n_vars: int
    det_order: int
    k_ar_diff: int
    
    @property
    def trace_crit_90(self) -> np.ndarray:
        """Critical values at 90% significance level."""
        return self.trace_crit[:, 0]
    
    @property
    def trace_crit_95(self) -> np.ndarray:
        """Critical values at 95% significance level."""
        return self.trace_crit[:, 1]
    
    @property
    def trace_crit_99(self) -> np.ndarray:
        """Critical values at 99% significance level."""
        return self.trace_crit[:, 2]
    
    @property
    def max_eig_crit_90(self) -> np.ndarray:
        """Maximum eigenvalue critical values at 90% significance level."""
        return self.max_eig_crit[:, 0]
    
    @property
    def max_eig_crit_95(self) -> np.ndarray:
        """Maximum eigenvalue critical values at 95% significance level."""
        return self.max_eig_crit[:, 1]
    
    @property
    def max_eig_crit_99(self) -> np.ndarray:
        """Maximum eigenvalue critical values at 99% significance level."""
        return self.max_eig_crit[:, 2]
    
    def summary(self, significance: float = 0.05) -> str:
        """
        Generate a summary of the Johansen test results.
        
        Parameters
        ----------
        significance : float, default 0.05
            Significance level for hypothesis testing
            
        Returns
        -------
        str
            Formatted summary string
        """
        # Map significance to column index
        sig_map = {0.10: 0, 0.05: 1, 0.01: 2}
        sig_idx = sig_map.get(significance, 1)
        
        lines = []
        lines.append("=" * 70)
        lines.append("Johansen Cointegration Test Results")
        lines.append("=" * 70)
        lines.append(f"Number of observations: {self.n_obs}")
        lines.append(f"Number of variables: {self.n_vars}")
        lines.append(f"Lag order (differences): {self.k_ar_diff}")
        lines.append(f"Deterministic order: {self.det_order}")
        lines.append("")
        
        # Trace test
        lines.append("-" * 70)
        lines.append("Trace Test:")
        lines.append("-" * 70)
        lines.append(f"{'H0: rank ≤':<15}{'Test Stat':>15}{'Crit Val (' + str(int((1-significance)*100)) + '%)':>15}{'Reject H0':>15}")
        lines.append("-" * 70)
        
        for i in range(self.n_vars):
            reject = "Yes" if self.trace_stat[i] > self.trace_crit[i, sig_idx] else "No"
            lines.append(f"{i:<15}{self.trace_stat[i]:>15.4f}{self.trace_crit[i, sig_idx]:>15.4f}{reject:>15}")
        
        lines.append("")
        
        # Maximum eigenvalue test
        lines.append("-" * 70)
        lines.append("Maximum Eigenvalue Test:")
        lines.append("-" * 70)
        lines.append(f"{'H0: rank = ':<15}{'Test Stat':>15}{'Crit Val (' + str(int((1-significance)*100)) + '%)':>15}{'Reject H0':>15}")
        lines.append("-" * 70)
        
        for i in range(self.n_vars):
            reject = "Yes" if self.max_eig_stat[i] > self.max_eig_crit[i, sig_idx] else "No"
            lines.append(f"{i:<15}{self.max_eig_stat[i]:>15.4f}{self.max_eig_crit[i, sig_idx]:>15.4f}{reject:>15}")
        
        lines.append("")
        lines.append("=" * 70)
        
        return "\n".join(lines)
    
    def get_cointegration_rank(self, test: str = 'trace', 
                                significance: float = 0.05) -> int:
        """
        Determine cointegration rank from test results.
        
        Parameters
        ----------
        test : str, default 'trace'
            Test to use: 'trace' or 'max_eig'
        significance : float, default 0.05
            Significance level
            
        Returns
        -------
        int
            Estimated cointegration rank
        """
        sig_map = {0.10: 0, 0.05: 1, 0.01: 2}
        sig_idx = sig_map.get(significance, 1)
        
        if test == 'trace':
            stat = self.trace_stat
            crit = self.trace_crit[:, sig_idx]
        else:
            stat = self.max_eig_stat
            crit = self.max_eig_crit[:, sig_idx]
        
        # Find first non-rejection
        for i in range(len(stat)):
            if stat[i] <= crit[i]:
                return i
        return len(stat)


def johansen_test(Y: np.ndarray, det_order: int = 0, 
                  k_ar_diff: int = 1) -> JohansenTestResults:
    """
    Perform Johansen cointegration test.
    
    This function wraps the statsmodels implementation and provides
    additional outputs needed for mixed VECM analysis.
    
    Parameters
    ----------
    Y : np.ndarray
        Time series data of shape (T, n)
    det_order : int, default 0
        Deterministic terms:
        -1: no deterministic terms
         0: constant term only
         1: constant and linear trend
    k_ar_diff : int, default 1
        Number of lagged differences in the VECM
        
    Returns
    -------
    JohansenTestResults
        Object containing test results
        
    Notes
    -----
    The Johansen test determines the cointegration rank by testing:
    H0: rank(Π) ≤ r  vs  H1: rank(Π) > r
    
    for r = 0, 1, ..., n-1, where Π = αβ' is the long-run matrix.
    
    For mixed VECM, the I(0) space dimension = number of cointegrating relations
    + number of I(0) variables.
    
    Examples
    --------
    >>> import numpy as np
    >>> np.random.seed(42)
    >>> Y = np.random.randn(200, 4).cumsum(axis=0)  # Random walk
    >>> results = johansen_test(Y, det_order=0, k_ar_diff=2)
    >>> print(results.summary())
    
    References
    ----------
    [1] Johansen, S. (1995). Likelihood-Based Inference in Cointegrated Vector
        Autoregressive Models. Oxford University Press.
    """
    Y = np.asarray(Y)
    T, n = Y.shape
    
    try:
        from statsmodels.tsa.vector_ar.vecm import coint_johansen
        
        # Call statsmodels implementation
        result = coint_johansen(Y, det_order, k_ar_diff)
        
        # Extract results
        return JohansenTestResults(
            trace_stat=result.lr1,
            trace_crit=result.cvt,
            max_eig_stat=result.lr2,
            max_eig_crit=result.cvm,
            eigenvalues=result.eig,
            eigenvectors=result.evec,
            r0t=result.r0t,
            rkt=result.rkt,
            S00=result.r0t.T @ result.r0t / (T - k_ar_diff - 1),
            S01=result.r0t.T @ result.rkt / (T - k_ar_diff - 1),
            S10=result.rkt.T @ result.r0t / (T - k_ar_diff - 1),
            S11=result.rkt.T @ result.rkt / (T - k_ar_diff - 1),
            n_obs=T - k_ar_diff - 1,
            n_vars=n,
            det_order=det_order,
            k_ar_diff=k_ar_diff
        )
    
    except ImportError:
        # Fallback to custom implementation
        return _johansen_test_native(Y, det_order, k_ar_diff)


def _johansen_test_native(Y: np.ndarray, det_order: int, 
                          k_ar_diff: int) -> JohansenTestResults:
    """
    Native implementation of Johansen test when statsmodels unavailable.
    
    Parameters
    ----------
    Y : np.ndarray
        Time series data of shape (T, n)
    det_order : int
        Deterministic order
    k_ar_diff : int
        Number of lagged differences
        
    Returns
    -------
    JohansenTestResults
        Test results
    """
    Y = np.asarray(Y)
    T, n = Y.shape
    
    # Compute differences
    dY = diff_matrix(Y)
    
    # Create lagged levels and differences
    p = k_ar_diff
    
    # Y_{t-1}
    Y_lag = Y[p:-1] if p > 0 else Y[:-1]
    
    # ΔY_t
    dY_current = dY[p:]
    
    # Lagged differences ΔY_{t-1}, ..., ΔY_{t-p}
    T_eff = dY_current.shape[0]
    
    X_lag = np.zeros((T_eff, n * p)) if p > 0 else None
    if p > 0:
        for l in range(1, p + 1):
            X_lag[:, (l-1)*n:l*n] = dY[p-l:T-1-l]
    
    # Add deterministic components
    if det_order >= 0:
        const = np.ones((T_eff, 1))
        X_det = const
        if det_order == 1:
            trend = np.arange(1, T_eff + 1).reshape(-1, 1)
            X_det = np.hstack([const, trend])
    else:
        X_det = None
    
    # Combine regressors for auxiliary regressions
    if X_lag is not None and X_det is not None:
        X = np.hstack([X_lag, X_det])
    elif X_lag is not None:
        X = X_lag
    elif X_det is not None:
        X = X_det
    else:
        X = None
    
    # Auxiliary regressions
    if X is not None:
        # Regress ΔY on X
        beta0 = lstsq(X, dY_current, rcond=None)[0]
        r0t = dY_current - X @ beta0
        
        # Regress Y_{t-1} on X
        beta1 = lstsq(X, Y_lag, rcond=None)[0]
        rkt = Y_lag - X @ beta1
    else:
        r0t = dY_current
        rkt = Y_lag
    
    # Compute moment matrices
    S00 = r0t.T @ r0t / T_eff
    S01 = r0t.T @ rkt / T_eff
    S10 = rkt.T @ r0t / T_eff
    S11 = rkt.T @ rkt / T_eff
    
    # Solve generalized eigenvalue problem
    S11_inv = inv(S11)
    S00_inv = inv(S00)
    
    # Compute S00^{-1/2} using Cholesky
    S00_chol = cholesky(S00, lower=True)
    S00_inv_sqrt = inv(S00_chol)
    
    # Matrix for eigenvalue problem
    M = S00_inv_sqrt @ S01 @ S11_inv @ S10 @ S00_inv_sqrt.T
    
    # Eigenvalue decomposition
    eigenvalues, eigenvectors_M = eig(M)
    
    # Sort eigenvalues in descending order
    idx = np.argsort(eigenvalues)[::-1]
    eigenvalues = np.real(eigenvalues[idx])
    eigenvectors_M = np.real(eigenvectors_M[:, idx])
    
    # Transform back to get beta
    eigenvectors = S00_inv_sqrt.T @ eigenvectors_M
    
    # Compute test statistics
    trace_stat = np.zeros(n)
    max_eig_stat = np.zeros(n)
    
    for r in range(n):
        # Trace statistic: -T * sum_{i=r+1}^n log(1 - λ_i)
        trace_stat[r] = -T_eff * np.sum(np.log(1 - eigenvalues[r:]))
        # Max eigenvalue statistic: -T * log(1 - λ_{r+1})
        if r < n:
            max_eig_stat[r] = -T_eff * np.log(1 - eigenvalues[r])
    
    # Critical values (from Johansen, 1995 Table 15.1 for case II: constant, no trend)
    # These are approximate - for exact values use statsmodels
    trace_crit = _get_critical_values(n, 'trace', det_order)
    max_eig_crit = _get_critical_values(n, 'max_eig', det_order)
    
    return JohansenTestResults(
        trace_stat=trace_stat,
        trace_crit=trace_crit,
        max_eig_stat=max_eig_stat,
        max_eig_crit=max_eig_crit,
        eigenvalues=eigenvalues,
        eigenvectors=eigenvectors,
        r0t=r0t,
        rkt=rkt,
        S00=S00,
        S01=S01,
        S10=S10,
        S11=S11,
        n_obs=T_eff,
        n_vars=n,
        det_order=det_order,
        k_ar_diff=k_ar_diff
    )


def _get_critical_values(n: int, test: str, det_order: int) -> np.ndarray:
    """
    Get critical values for Johansen test.
    
    Returns approximate critical values for 90%, 95%, 99% significance.
    For more accurate values, use statsmodels implementation.
    """
    # Critical values from Johansen (1995) for case with constant (det_order=0)
    # Format: [90%, 95%, 99%] for each r
    
    trace_cv = {
        1: np.array([[6.50, 8.18, 11.65]]),
        2: np.array([[12.91, 14.90, 19.19],
                     [6.50, 8.18, 11.65]]),
        3: np.array([[21.07, 24.05, 29.68],
                     [12.91, 14.90, 19.19],
                     [6.50, 8.18, 11.65]]),
        4: np.array([[30.49, 34.03, 40.19],
                     [21.07, 24.05, 29.68],
                     [12.91, 14.90, 19.19],
                     [6.50, 8.18, 11.65]]),
        5: np.array([[40.06, 44.06, 51.58],
                     [30.49, 34.03, 40.19],
                     [21.07, 24.05, 29.68],
                     [12.91, 14.90, 19.19],
                     [6.50, 8.18, 11.65]]),
        6: np.array([[51.38, 55.77, 63.90],
                     [40.06, 44.06, 51.58],
                     [30.49, 34.03, 40.19],
                     [21.07, 24.05, 29.68],
                     [12.91, 14.90, 19.19],
                     [6.50, 8.18, 11.65]]),
    }
    
    max_eig_cv = {
        1: np.array([[6.50, 8.18, 11.65]]),
        2: np.array([[12.07, 14.07, 18.63],
                     [6.50, 8.18, 11.65]]),
        3: np.array([[18.06, 20.97, 25.52],
                     [12.07, 14.07, 18.63],
                     [6.50, 8.18, 11.65]]),
        4: np.array([[23.72, 27.07, 32.24],
                     [18.06, 20.97, 25.52],
                     [12.07, 14.07, 18.63],
                     [6.50, 8.18, 11.65]]),
        5: np.array([[29.41, 33.12, 38.78],
                     [23.72, 27.07, 32.24],
                     [18.06, 20.97, 25.52],
                     [12.07, 14.07, 18.63],
                     [6.50, 8.18, 11.65]]),
        6: np.array([[35.02, 39.00, 45.10],
                     [29.41, 33.12, 38.78],
                     [23.72, 27.07, 32.24],
                     [18.06, 20.97, 25.52],
                     [12.07, 14.07, 18.63],
                     [6.50, 8.18, 11.65]]),
    }
    
    if test == 'trace':
        if n in trace_cv:
            return trace_cv[n]
        else:
            # Extrapolate for larger systems
            return np.vstack([trace_cv[6]] + [trace_cv[6][-1:]] * (n - 6))
    else:
        if n in max_eig_cv:
            return max_eig_cv[n]
        else:
            return np.vstack([max_eig_cv[6]] + [max_eig_cv[6][-1:]] * (n - 6))


@dataclass
class I0ComponentTestResults:
    """
    Results from testing for I(0) components in mixed VECM.
    
    Attributes
    ----------
    k_i0 : int
        Number of I(0) components
    k_coint : int
        Number of cointegrating relations among I(1) variables
    lr_statistic : float
        Likelihood ratio test statistic
    p_value : float
        P-value from chi-squared distribution
    df : int
        Degrees of freedom
    omega_restricted : np.ndarray
        Restricted covariance estimate
    omega_unrestricted : np.ndarray
        Unrestricted covariance estimate
    beta_restricted : np.ndarray
        Restricted cointegrating matrix
    alpha_restricted : np.ndarray
        Restricted adjustment matrix
    converged : bool
        Whether iteration converged
    n_iter : int
        Number of iterations
    """
    k_i0: int
    k_coint: int
    lr_statistic: float
    p_value: float
    df: int
    omega_restricted: np.ndarray
    omega_unrestricted: np.ndarray
    beta_restricted: np.ndarray
    alpha_restricted: np.ndarray
    converged: bool
    n_iter: int
    
    @property
    def lr_stat(self) -> float:
        """Alias for lr_statistic for convenience."""
        return self.lr_statistic
    
    def summary(self) -> str:
        """Generate summary of test results."""
        lines = []
        lines.append("=" * 60)
        lines.append("Test for I(0) Components in Mixed VECM")
        lines.append("=" * 60)
        lines.append(f"H0: {self.k_i0} I(0) component(s) present")
        lines.append(f"Number of cointegrating relations: {self.k_coint}")
        lines.append("")
        lines.append(f"LR Statistic: {self.lr_statistic:.4f}")
        lines.append(f"Degrees of Freedom: {self.df}")
        lines.append(f"P-value: {self.p_value:.4f}")
        lines.append("")
        if self.p_value < 0.01:
            conclusion = "Reject H0 at 1% significance level"
        elif self.p_value < 0.05:
            conclusion = "Reject H0 at 5% (but not 1%) significance level"
        elif self.p_value < 0.10:
            conclusion = "Reject H0 at 10% (but not 5%) significance level"
        else:
            conclusion = "Cannot reject H0"
        lines.append(f"Conclusion: {conclusion}")
        lines.append("")
        lines.append(f"Iteration converged: {self.converged}")
        lines.append(f"Number of iterations: {self.n_iter}")
        lines.append("=" * 60)
        return "\n".join(lines)


def test_i0_components(Y: np.ndarray, k_i0: int, i0_space_dim: int,
                       det_order: int = 0, k_ar_diff: int = 1,
                       max_iter: int = 100, tol: float = 1e-8,
                       i0_indices: Optional[List[int]] = None
                       ) -> I0ComponentTestResults:
    """
    Test for the presence of k I(0) components in a mixed VECM.
    
    This implements the likelihood ratio test from Chen (2022) based on
    restricted beta estimation following Boswijk and Doornik (2004).
    
    Parameters
    ----------
    Y : np.ndarray
        Time series data of shape (T, n)
    k_i0 : int
        Number of I(0) components to test for
    i0_space_dim : int
        Total dimension of I(0) space from Johansen test
        (= number of cointegrating relations + number of I(0) variables)
    det_order : int, default 0
        Deterministic order
    k_ar_diff : int, default 1
        Number of lagged differences
    max_iter : int, default 100
        Maximum iterations for estimation
    tol : float, default 1e-8
        Convergence tolerance
    i0_indices : list of int, optional
        Indices of suspected I(0) variables (last k_i0 variables if None)
        
    Returns
    -------
    I0ComponentTestResults
        Test results
        
    Notes
    -----
    According to Lemma 2.6 in Chen (2022), the necessary and sufficient condition
    for a VECM to contain k I(0) variables is that β can be written as:
    
        β' = [[B21, B22, 0], [0, 0, B33]]
    
    The test uses the likelihood ratio:
    
        LR = T(log|Ω̃| - log|S00| + Σ log(1 - λ̂_i))
    
    which is asymptotically χ² with r*k degrees of freedom, where r is the
    cointegration rank.
    
    Examples
    --------
    >>> # First run Johansen test
    >>> joh_result = johansen_test(Y, det_order=0, k_ar_diff=2)
    >>> i0_dim = joh_result.get_cointegration_rank()
    >>> 
    >>> # Test for 2 I(0) components
    >>> result = test_i0_components(Y, k_i0=2, i0_space_dim=i0_dim)
    >>> print(result.summary())
    
    References
    ----------
    [1] Chen, P. (2022). Vector Error Correction Models with Stationary and 
        Nonstationary Variables. SSRN Working Paper 4239531.
    [2] Boswijk, H.P. and Doornik, J.A. (2004). Identifying, estimating and testing
        restricted cointegrated systems: An overview. Statistica Neerlandica.
    """
    Y = np.asarray(Y)
    T, n = Y.shape
    
    # Number of cointegrating relations among I(1) variables
    k_coint = i0_space_dim - k_i0
    
    if k_coint < 0:
        raise ValueError(f"k_i0={k_i0} exceeds i0_space_dim={i0_space_dim}")
    
    if k_i0 <= 0:
        raise ValueError("k_i0 must be positive")
    
    # Total cointegrating rank (I(0) space dimension)
    r = i0_space_dim
    
    if r > n:
        raise ValueError(f"i0_space_dim={r} cannot exceed n={n}")
    
    # Get Johansen test results for cross-product matrices
    joh = johansen_test(Y, det_order, k_ar_diff)
    
    S00 = joh.S00
    S01 = joh.S01
    S10 = joh.S10
    S11 = joh.S11
    T_eff = joh.n_obs
    
    # Unrestricted log-likelihood (from Johansen)
    eigenvalues = joh.eigenvalues[:r]
    log_det_omega_u = np.log(det(S00)) + np.sum(np.log(1 - eigenvalues))
    
    # Construct restriction matrices for beta
    # Under H0: β' = [[B21, B22, 0], [0, 0, B33]]
    # This means: (n-k) variables in cointegrating relations, k variables as I(0)
    
    # Reorder variables if i0_indices specified
    if i0_indices is None:
        i0_indices = list(range(n - k_i0, n))
    
    # Build H matrix for restrictions vec(β) = H*φ + h0
    # β is n × r where r = k_coint + k_i0
    # Restricted form: β = [[β_11, 0], [0, β_22]] where:
    #   β_11 is (n-k_i0) × k_coint for cointegration
    #   β_22 is k_i0 × k_i0 for I(0) components
    
    H, h0 = _construct_beta_restriction_matrix(n, k_coint, k_i0, i0_indices)
    
    # Estimate restricted model using Boswijk-Doornik iteration
    alpha_r, beta_r, omega_r, converged, n_iter = _estimate_restricted_vecm(
        S00, S01, S10, S11, H, h0, n, r, k_coint, k_i0, i0_indices,
        max_iter, tol
    )
    
    # Compute likelihood ratio statistic
    log_det_omega_r = np.log(det(omega_r))
    
    # LR = T(log|Ω̃| - log|Ω_u|)
    lr_stat = T_eff * (log_det_omega_r - log_det_omega_u)
    
    # Degrees of freedom: r × k
    # This is the number of restrictions imposed
    df = r * k_i0
    
    # P-value from chi-squared distribution
    p_value = 1 - stats.chi2.cdf(lr_stat, df)
    
    return I0ComponentTestResults(
        k_i0=k_i0,
        k_coint=k_coint,
        lr_statistic=lr_stat,
        p_value=p_value,
        df=df,
        omega_restricted=omega_r,
        omega_unrestricted=S00,
        beta_restricted=beta_r,
        alpha_restricted=alpha_r,
        converged=converged,
        n_iter=n_iter
    )


def _construct_beta_restriction_matrix(n: int, k_coint: int, k_i0: int,
                                        i0_indices: List[int]
                                        ) -> Tuple[np.ndarray, np.ndarray]:
    """
    Construct H matrix and h0 vector for restricted beta estimation.
    
    The restriction is that β can be written as:
        β' = [[B21, B22, 0], [0, 0, B33]]
    
    In vectorized form: vec(β) = H*φ + h0
    
    Parameters
    ----------
    n : int
        Number of variables
    k_coint : int
        Number of cointegrating relations
    k_i0 : int
        Number of I(0) components
    i0_indices : list of int
        Indices of I(0) variables
        
    Returns
    -------
    H : np.ndarray
        Selection matrix
    h0 : np.ndarray
        Constant vector (usually zeros)
    """
    r = k_coint + k_i0
    
    # β is n × r
    # Free parameters in β:
    # - For cointegration part: (n - k_i0) × k_coint parameters
    # - For I(0) part: k_i0 × k_i0 parameters
    
    n_free = (n - k_i0) * k_coint + k_i0 * k_i0
    
    H = np.zeros((n * r, n_free))
    h0 = np.zeros((n * r, 1))
    
    # Get indices of I(1) variables
    i1_indices = [i for i in range(n) if i not in i0_indices]
    
    col_idx = 0
    
    # Cointegrating part: first k_coint columns of β
    # Only I(1) variables have non-zero coefficients
    for j in range(k_coint):
        for i in i1_indices:
            row_idx = i + j * n  # Position in vec(β)
            H[row_idx, col_idx] = 1
            col_idx += 1
    
    # I(0) part: last k_i0 columns of β
    # Only I(0) variables have non-zero coefficients
    for j in range(k_coint, r):
        for i in i0_indices:
            row_idx = i + j * n
            H[row_idx, col_idx] = 1
            col_idx += 1
    
    return H, h0


def _estimate_restricted_vecm(S00: np.ndarray, S01: np.ndarray,
                               S10: np.ndarray, S11: np.ndarray,
                               H: np.ndarray, h0: np.ndarray,
                               n: int, r: int, k_coint: int, k_i0: int,
                               i0_indices: List[int],
                               max_iter: int, tol: float
                               ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, bool, int]:
    """
    Estimate restricted VECM using Boswijk-Doornik iterative procedure.
    
    Parameters
    ----------
    S00, S01, S10, S11 : np.ndarray
        Cross-product matrices from Johansen procedure
    H : np.ndarray
        Restriction matrix
    h0 : np.ndarray
        Restriction constant
    n : int
        Number of variables
    r : int
        Cointegration rank
    k_coint : int
        Number of cointegrating relations
    k_i0 : int
        Number of I(0) components
    i0_indices : list of int
        Indices of I(0) variables
    max_iter : int
        Maximum iterations
    tol : float
        Convergence tolerance
        
    Returns
    -------
    alpha : np.ndarray
        Estimated adjustment matrix
    beta : np.ndarray
        Estimated cointegrating matrix
    omega : np.ndarray
        Estimated covariance matrix
    converged : bool
        Convergence indicator
    n_iter : int
        Number of iterations
    """
    # Initialize with unrestricted estimates
    S11_inv = inv(S11)
    S00_inv = inv(S00)
    
    # Initial eigenvalue decomposition
    M = S00_inv @ S01 @ S11_inv @ S10
    eigvals, eigvecs = eig(M)
    idx = np.argsort(np.real(eigvals))[::-1]
    eigvecs = np.real(eigvecs[:, idx[:r]])
    
    # Initial beta from eigenvectors (needs proper normalization)
    # Use S11^{-1} @ S10 @ S00^{-1/2} eigenvectors
    beta = S11_inv @ S10 @ eigvecs
    
    # Project onto restriction space
    phi = lstsq(H, vec(beta) - h0, rcond=None)[0]
    beta = unvec(H @ phi + h0, n, r)
    
    # Initial alpha
    alpha = S01 @ beta @ inv(beta.T @ S11 @ beta)
    
    # Initial omega
    omega = S00 - alpha @ beta.T @ S10 - S01 @ beta @ alpha.T + alpha @ beta.T @ S11 @ beta @ alpha.T
    
    converged = False
    n_iter = 0
    
    for iteration in range(max_iter):
        n_iter = iteration + 1
        
        # Store old values
        alpha_old = alpha.copy()
        beta_old = beta.copy()
        omega_old = omega.copy()
        
        # Update phi (and thus beta)
        omega_inv = inv(omega)
        
        # Following Boswijk-Doornik (2004) eq. (3.17)
        # phi_tilde = (H'(α'Ω^{-1}α ⊗ S11)H)^{-1} H'(α'Ω^{-1} ⊗ S11)(vec(Π'_LS) - (α ⊗ I)h0)
        
        Pi_LS = S01 @ S11_inv  # Unrestricted long-run matrix estimate
        
        A = kronecker(alpha.T @ omega_inv @ alpha, S11)
        HAH = H.T @ A @ H
        
        try:
            HAH_inv = inv(HAH)
        except np.linalg.LinAlgError:
            HAH_inv = np.linalg.pinv(HAH)
        
        b = kronecker(alpha.T @ omega_inv, S11) @ (vec(Pi_LS.T) - kronecker(alpha, np.eye(n)) @ h0)
        phi = HAH_inv @ H.T @ b
        
        beta = unvec(H @ phi + h0, n, r)
        
        # Update alpha (psi)
        # psi_tilde = (G'(Ω^{-1} ⊗ β'S11β)G)^{-1} G'(Ω^{-1} ⊗ β'S11)(vec(Π'_LS))
        # For unrestricted alpha, G = I
        
        beta_S11_beta = beta.T @ S11 @ beta
        
        try:
            B = kronecker(omega_inv, beta_S11_beta)
            B_inv = inv(B)
            c = kronecker(omega_inv, beta.T @ S11) @ vec(Pi_LS.T)
            alpha = unvec(B_inv @ c, n, r)
        except np.linalg.LinAlgError:
            alpha = S01 @ beta @ inv(beta.T @ S11 @ beta)
        
        # Update omega
        omega = S00 - alpha @ beta.T @ S10 - S01 @ beta @ alpha.T + alpha @ beta.T @ S11 @ beta @ alpha.T
        
        # Check convergence
        diff_alpha = np.max(np.abs(alpha - alpha_old))
        diff_beta = np.max(np.abs(beta - beta_old))
        diff_omega = np.max(np.abs(omega - omega_old))
        
        if max(diff_alpha, diff_beta, diff_omega) < tol:
            converged = True
            break
    
    return alpha, beta, omega, converged, n_iter


@dataclass
class LRTestResults:
    """
    Results from likelihood ratio test.
    
    Attributes
    ----------
    statistic : float
        LR test statistic
    p_value : float
        P-value
    df : int
        Degrees of freedom
    null_hypothesis : str
        Description of null hypothesis
    """
    statistic: float
    p_value: float
    df: int
    null_hypothesis: str
    
    def summary(self) -> str:
        """Generate summary."""
        lines = []
        lines.append("Likelihood Ratio Test")
        lines.append("=" * 40)
        lines.append(f"H0: {self.null_hypothesis}")
        lines.append(f"LR Statistic: {self.statistic:.4f}")
        lines.append(f"Degrees of Freedom: {self.df}")
        lines.append(f"P-value: {self.p_value:.4f}")
        return "\n".join(lines)


def likelihood_ratio_test(ll_unrestricted: float, ll_restricted: float,
                          df: int, null_hypothesis: str = "") -> LRTestResults:
    """
    Perform likelihood ratio test.
    
    Parameters
    ----------
    ll_unrestricted : float
        Log-likelihood of unrestricted model
    ll_restricted : float
        Log-likelihood of restricted model
    df : int
        Degrees of freedom (number of restrictions)
    null_hypothesis : str, optional
        Description of null hypothesis
        
    Returns
    -------
    LRTestResults
        Test results
    """
    lr_stat = 2 * (ll_unrestricted - ll_restricted)
    lr_stat = max(0, lr_stat)  # Ensure non-negative
    p_value = 1 - stats.chi2.cdf(lr_stat, df)
    
    return LRTestResults(
        statistic=lr_stat,
        p_value=p_value,
        df=df,
        null_hypothesis=null_hypothesis
    )


def select_cointegration_rank(Y: np.ndarray, det_order: int = 0,
                               k_ar_diff: int = 1,
                               method: str = 'trace',
                               significance: float = 0.05) -> int:
    """
    Select cointegration rank using Johansen test.
    
    Parameters
    ----------
    Y : np.ndarray
        Time series data of shape (T, n)
    det_order : int, default 0
        Deterministic order
    k_ar_diff : int, default 1
        Number of lagged differences
    method : str, default 'trace'
        Test method: 'trace' or 'max_eig'
    significance : float, default 0.05
        Significance level
        
    Returns
    -------
    int
        Selected cointegration rank
    """
    result = johansen_test(Y, det_order, k_ar_diff)
    return result.get_cointegration_rank(method, significance)


def test_multiple_i0_components(Y: np.ndarray, i0_space_dim: int,
                                 det_order: int = 0, k_ar_diff: int = 1,
                                 max_i0: Optional[int] = None,
                                 significance: float = 0.05
                                 ) -> List[I0ComponentTestResults]:
    """
    Test for multiple numbers of I(0) components.
    
    Parameters
    ----------
    Y : np.ndarray
        Time series data
    i0_space_dim : int
        Total I(0) space dimension from Johansen test
    det_order : int, default 0
        Deterministic order
    k_ar_diff : int, default 1
        Number of lagged differences
    max_i0 : int, optional
        Maximum number of I(0) components to test
    significance : float, default 0.05
        Significance level
        
    Returns
    -------
    list of I0ComponentTestResults
        Results for each k_i0 value
        
    Notes
    -----
    This implements the sequential testing procedure from Chen (2022):
    Test k_i0 = 1, 2, 3, ... until the null hypothesis cannot be rejected.
    """
    n = Y.shape[1]
    
    if max_i0 is None:
        max_i0 = min(i0_space_dim, n - 1)
    
    results = []
    
    for k_i0 in range(1, max_i0 + 1):
        if k_i0 > i0_space_dim:
            break
            
        try:
            result = test_i0_components(Y, k_i0, i0_space_dim, 
                                        det_order, k_ar_diff)
            results.append(result)
            
            # If we cannot reject H0, stop
            if result.p_value > significance:
                break
        except Exception as e:
            warnings.warn(f"Test for k_i0={k_i0} failed: {e}")
            break
    
    return results


def determine_i0_count(Y: np.ndarray, i0_space_dim: int,
                       det_order: int = 0, k_ar_diff: int = 1,
                       significance: float = 0.05) -> Tuple[int, int]:
    """
    Determine the number of I(0) components and cointegrating relations.
    
    Parameters
    ----------
    Y : np.ndarray
        Time series data
    i0_space_dim : int
        Total I(0) space dimension from Johansen test
    det_order : int, default 0
        Deterministic order
    k_ar_diff : int, default 1
        Number of lagged differences
    significance : float, default 0.05
        Significance level
        
    Returns
    -------
    k_i0 : int
        Number of I(0) components
    k_coint : int
        Number of cointegrating relations
    """
    results = test_multiple_i0_components(Y, i0_space_dim, det_order,
                                           k_ar_diff, significance=significance)
    
    if not results:
        return 0, i0_space_dim
    
    # Find highest k_i0 where H0 is not rejected
    for result in reversed(results):
        if result.p_value > significance:
            return result.k_i0, result.k_coint
    
    return 0, i0_space_dim
