"""
Core Mixed VECM Estimation
==========================

This module implements the estimation of Vector Error Correction Models with
mixed I(0) and I(1) variables.

The methodology follows Chen (2022), which extends the standard Johansen
procedure to handle systems containing both stationary and nonstationary
variables.

Main Classes
------------
MixedVECM : Main estimation class
MixedVECMResults : Results container

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.
"""

import numpy as np
from numpy.linalg import inv, eig, det, lstsq
from scipy import stats
from typing import Tuple, Optional, Union, List, Dict, Any
from dataclasses import dataclass, field
import warnings

from .utils import (
    vec, unvec, lag_matrix, diff_matrix, orthogonal_complement,
    normalize_beta, compute_var_covariance, kronecker,
    compute_aic_bic, select_lag_order, companion_matrix
)
from .testing import (
    johansen_test, JohansenTestResults,
    test_i0_components, I0ComponentTestResults,
    determine_i0_count, test_multiple_i0_components,
    _construct_beta_restriction_matrix, _estimate_restricted_vecm
)


@dataclass
class MixedVECMResults:
    """
    Results from Mixed VECM estimation.
    
    Attributes
    ----------
    alpha : np.ndarray
        Adjustment matrix (n × r)
    beta : np.ndarray
        Cointegrating matrix (n × r)
    Pi : np.ndarray
        Long-run matrix Π = αβ'
    Gamma : list of np.ndarray
        Short-run coefficient matrices [Γ_1, ..., Γ_{p-1}]
    C : np.ndarray
        Constant vector
    D : np.ndarray, optional
        Trend coefficient vector
    Omega : np.ndarray
        Residual covariance matrix (n × n)
    residuals : np.ndarray
        Residuals (T × n)
    fitted_values : np.ndarray
        Fitted values (T × n)
    n_obs : int
        Number of observations used
    n_vars : int
        Number of variables
    n_lags : int
        Number of lags in VECM
    coint_rank : int
        Cointegration rank
    n_i0 : int
        Number of I(0) variables
    n_i1 : int
        Number of I(1) variables
    n_coint : int
        Number of cointegrating relations among I(1) variables
    i0_indices : list of int
        Indices of I(0) variables
    i1_indices : list of int
        Indices of I(1) variables
    det_order : int
        Deterministic order used
    log_likelihood : float
        Log-likelihood value
    aic : float
        Akaike Information Criterion
    bic : float
        Bayesian Information Criterion
    johansen_results : JohansenTestResults
        Results from Johansen test
    i0_test_results : list of I0ComponentTestResults
        Results from I(0) component tests
    alpha_se : np.ndarray, optional
        Standard errors for alpha
    beta_se : np.ndarray, optional
        Standard errors for beta
    alpha_tstat : np.ndarray, optional
        t-statistics for alpha
    beta_tstat : np.ndarray, optional
        t-statistics for beta
    """
    alpha: np.ndarray
    beta: np.ndarray
    Pi: np.ndarray
    Gamma: List[np.ndarray]
    C: np.ndarray
    D: Optional[np.ndarray]
    Omega: np.ndarray
    residuals: np.ndarray
    fitted_values: np.ndarray
    n_obs: int
    n_vars: int
    n_lags: int
    coint_rank: int
    n_i0: int
    n_i1: int
    n_coint: int
    i0_indices: List[int]
    i1_indices: List[int]
    det_order: int
    log_likelihood: float
    aic: float
    bic: float
    johansen_results: JohansenTestResults
    i0_test_results: List[I0ComponentTestResults]
    alpha_se: Optional[np.ndarray] = None
    beta_se: Optional[np.ndarray] = None
    alpha_tstat: Optional[np.ndarray] = None
    beta_tstat: Optional[np.ndarray] = None
    
    def summary(self, significance: float = 0.05) -> str:
        """
        Generate comprehensive summary of Mixed VECM estimation.
        
        Parameters
        ----------
        significance : float, default 0.05
            Significance level for tests
            
        Returns
        -------
        str
            Formatted summary string
        """
        lines = []
        lines.append("=" * 80)
        lines.append("Mixed Vector Error Correction Model (VECM) Estimation Results")
        lines.append("=" * 80)
        lines.append("")
        
        # Model specification
        lines.append("Model Specification:")
        lines.append("-" * 40)
        lines.append(f"  Number of observations: {self.n_obs}")
        lines.append(f"  Number of variables: {self.n_vars}")
        lines.append(f"  Number of lags: {self.n_lags}")
        lines.append(f"  Deterministic order: {self.det_order}")
        lines.append("")
        
        # Variable classification
        lines.append("Variable Classification:")
        lines.append("-" * 40)
        lines.append(f"  I(1) variables: {len(self.i1_indices)} (indices: {self.i1_indices})")
        lines.append(f"  I(0) variables: {len(self.i0_indices)} (indices: {self.i0_indices})")
        lines.append(f"  Cointegrating relations: {self.n_coint}")
        lines.append(f"  Total I(0) space dimension: {self.coint_rank}")
        lines.append("")
        
        # Model fit
        lines.append("Model Fit:")
        lines.append("-" * 40)
        lines.append(f"  Log-likelihood: {self.log_likelihood:.4f}")
        lines.append(f"  AIC: {self.aic:.4f}")
        lines.append(f"  BIC: {self.bic:.4f}")
        lines.append("")
        
        # Cointegrating vectors (β)
        lines.append("Cointegrating Matrix (β):")
        lines.append("-" * 40)
        lines.append(self._format_matrix(self.beta, 
                                          row_labels=[f"Var {i+1}" for i in range(self.n_vars)],
                                          col_labels=[f"EC{i+1}" for i in range(self.coint_rank)]))
        lines.append("")
        
        # Adjustment coefficients (α)
        lines.append("Adjustment Matrix (α):")
        lines.append("-" * 40)
        lines.append(self._format_matrix(self.alpha,
                                          row_labels=[f"ΔVar {i+1}" for i in range(self.n_vars)],
                                          col_labels=[f"EC{i+1}" for i in range(self.coint_rank)]))
        lines.append("")
        
        # Long-run matrix (Π)
        lines.append("Long-run Matrix (Π = αβ'):")
        lines.append("-" * 40)
        lines.append(self._format_matrix(self.Pi,
                                          row_labels=[f"ΔVar {i+1}" for i in range(self.n_vars)],
                                          col_labels=[f"Var {i+1}" for i in range(self.n_vars)]))
        lines.append("")
        
        # Residual covariance
        lines.append("Residual Covariance Matrix (Ω):")
        lines.append("-" * 40)
        lines.append(self._format_matrix(self.Omega,
                                          row_labels=[f"Var {i+1}" for i in range(self.n_vars)],
                                          col_labels=[f"Var {i+1}" for i in range(self.n_vars)]))
        lines.append("")
        
        lines.append("=" * 80)
        
        return "\n".join(lines)
    
    def _format_matrix(self, M: np.ndarray, 
                       row_labels: Optional[List[str]] = None,
                       col_labels: Optional[List[str]] = None,
                       precision: int = 4) -> str:
        """Format matrix for display."""
        rows, cols = M.shape
        
        # Column width
        width = max(precision + 6, 10)
        
        lines = []
        
        # Header row
        if col_labels is not None:
            header = " " * 12 + "".join([f"{label:>{width}}" for label in col_labels])
            lines.append(header)
        
        # Data rows
        for i in range(rows):
            if row_labels is not None:
                row_str = f"{row_labels[i]:<12}"
            else:
                row_str = " " * 12
            row_str += "".join([f"{M[i, j]:>{width}.{precision}f}" for j in range(cols)])
            lines.append(row_str)
        
        return "\n".join(lines)
    
    def get_error_correction_terms(self, Y: np.ndarray) -> np.ndarray:
        """
        Compute error correction terms β'Y_{t-1}.
        
        Parameters
        ----------
        Y : np.ndarray
            Time series data (T × n)
            
        Returns
        -------
        np.ndarray
            Error correction terms (T-1 × r)
        """
        Y = np.asarray(Y)
        return Y[:-1] @ self.beta
    
    def get_cointegrating_relations(self) -> List[str]:
        """
        Get human-readable cointegrating relations.
        
        Returns
        -------
        list of str
            Cointegrating relations as equations
        """
        relations = []
        for j in range(self.coint_rank):
            terms = []
            for i in range(self.n_vars):
                coef = self.beta[i, j]
                if abs(coef) > 1e-10:
                    if coef > 0:
                        terms.append(f"+ {coef:.4f}*Y{i+1}")
                    else:
                        terms.append(f"- {abs(coef):.4f}*Y{i+1}")
            
            relation = " ".join(terms).lstrip("+ ")
            if j < self.n_coint:
                relations.append(f"Cointegrating Relation {j+1}: {relation} = 0")
            else:
                relations.append(f"I(0) Component {j - self.n_coint + 1}: {relation} = I(0)")
        
        return relations


class MixedVECM:
    """
    Mixed Vector Error Correction Model with I(0) and I(1) variables.
    
    This class implements the methodology from Chen (2022) for estimating
    VECMs that contain both stationary I(0) and nonstationary I(1) variables.
    
    Parameters
    ----------
    Y : np.ndarray
        Time series data of shape (T, n)
    k_ar_diff : int, default 1
        Number of lagged differences in the VECM
    det_order : int, default 0
        Deterministic terms:
        -1: no deterministic terms
         0: constant only
         1: constant and trend
    auto_lag : bool, default False
        If True, automatically select lag order using BIC
    max_lag : int, default 12
        Maximum lag order for automatic selection
        
    Attributes
    ----------
    Y : np.ndarray
        Input data
    n_obs : int
        Number of observations
    n_vars : int
        Number of variables
    k_ar_diff : int
        Number of lagged differences
    det_order : int
        Deterministic order
        
    Examples
    --------
    >>> import numpy as np
    >>> from mixvecm import MixedVECM
    >>> 
    >>> # Generate or load data
    >>> Y = np.random.randn(200, 6).cumsum(axis=0)
    >>> Y[:, -2:] = np.random.randn(200, 2)  # Last 2 variables are I(0)
    >>> 
    >>> # Fit model
    >>> model = MixedVECM(Y, k_ar_diff=2, det_order=0)
    >>> results = model.fit()
    >>> print(results.summary())
    
    References
    ----------
    [1] Chen, P. (2022). Vector Error Correction Models with Stationary and 
        Nonstationary Variables. SSRN Working Paper 4239531.
    """
    
    def __init__(self, Y: np.ndarray, k_ar_diff: int = 1, 
                 det_order: int = 0, auto_lag: bool = False,
                 max_lag: int = 12):
        """Initialize Mixed VECM."""
        self.Y = np.asarray(Y)
        self.n_obs, self.n_vars = self.Y.shape
        self.det_order = det_order
        
        if auto_lag:
            self.k_ar_diff = select_lag_order(Y, max_lag, 'bic', det_order)
        else:
            self.k_ar_diff = k_ar_diff
        
        self._results = None
    
    def fit(self, coint_rank: Optional[int] = None,
            n_i0: Optional[int] = None,
            i0_indices: Optional[List[int]] = None,
            significance: float = 0.05,
            normalize: bool = True,
            compute_se: bool = True) -> MixedVECMResults:
        """
        Estimate the Mixed VECM.
        
        Parameters
        ----------
        coint_rank : int, optional
            Cointegration rank (I(0) space dimension). 
            If None, determined by Johansen test.
        n_i0 : int, optional
            Number of I(0) variables. 
            If None, determined by sequential testing.
        i0_indices : list of int, optional
            Indices of I(0) variables (assumed to be last n_i0 if None)
        significance : float, default 0.05
            Significance level for hypothesis tests
        normalize : bool, default True
            Whether to normalize cointegrating vectors
        compute_se : bool, default True
            Whether to compute standard errors
            
        Returns
        -------
        MixedVECMResults
            Estimation results
            
        Notes
        -----
        The estimation procedure follows these steps:
        
        1. Run Johansen test to determine I(0) space dimension
        2. Test for I(0) components using likelihood ratio tests
        3. Estimate restricted VECM with orthogonal I(0) and cointegration spaces
        4. Compute standard errors if requested
        
        The key insight from Chen (2022) is that a mixed VECM can be viewed as
        a standard cointegrated VECM with restrictions on β such that:
        
            β' = [[B21, B22, 0], [0, 0, B33]]
            
        where the I(0) components and cointegrating relations are orthogonal.
        """
        # Step 1: Johansen test
        joh_results = johansen_test(self.Y, self.det_order, self.k_ar_diff)
        
        # Determine cointegration rank if not specified
        if coint_rank is None:
            coint_rank = joh_results.get_cointegration_rank('trace', significance)
        
        if coint_rank == 0:
            warnings.warn("No cointegration found. Consider using standard VAR.")
        
        # Step 2: Test for I(0) components
        i0_test_results = []
        if n_i0 is None and coint_rank > 0:
            # Sequential testing
            i0_test_results = test_multiple_i0_components(
                self.Y, coint_rank, self.det_order, self.k_ar_diff,
                significance=significance
            )
            
            if i0_test_results:
                # Find the highest n_i0 where H0 is not rejected
                for result in reversed(i0_test_results):
                    if result.p_value > significance:
                        n_i0 = result.k_i0
                        break
                else:
                    n_i0 = 0
            else:
                n_i0 = 0
        elif n_i0 is None:
            n_i0 = 0
        
        # Determine I(0) indices
        if i0_indices is None:
            i0_indices = list(range(self.n_vars - n_i0, self.n_vars))
        
        i1_indices = [i for i in range(self.n_vars) if i not in i0_indices]
        
        n_coint = coint_rank - n_i0
        
        # Step 3: Estimate the model
        if n_i0 > 0 and coint_rank > 0:
            # Mixed VECM with restrictions
            results = self._estimate_restricted(
                coint_rank, n_i0, n_coint, i0_indices, i1_indices,
                joh_results, normalize
            )
        else:
            # Standard VECM
            results = self._estimate_unrestricted(
                coint_rank, i0_indices, i1_indices, joh_results, normalize
            )
        
        # Add test results
        results.johansen_results = joh_results
        results.i0_test_results = i0_test_results
        results.n_i0 = n_i0
        results.n_i1 = len(i1_indices)
        results.n_coint = n_coint
        results.i0_indices = i0_indices
        results.i1_indices = i1_indices
        
        # Step 4: Compute standard errors
        if compute_se and coint_rank > 0:
            self._compute_standard_errors(results)
        
        self._results = results
        return results
    
    def _estimate_unrestricted(self, coint_rank: int,
                                i0_indices: List[int],
                                i1_indices: List[int],
                                joh_results: JohansenTestResults,
                                normalize: bool) -> MixedVECMResults:
        """
        Estimate unrestricted VECM (standard Johansen procedure).
        
        Parameters
        ----------
        coint_rank : int
            Cointegration rank
        i0_indices : list of int
            I(0) variable indices
        i1_indices : list of int
            I(1) variable indices
        joh_results : JohansenTestResults
            Johansen test results
        normalize : bool
            Whether to normalize beta
            
        Returns
        -------
        MixedVECMResults
            Estimation results
        """
        n = self.n_vars
        r = coint_rank
        T_eff = joh_results.n_obs
        
        S00 = joh_results.S00
        S01 = joh_results.S01
        S10 = joh_results.S10
        S11 = joh_results.S11
        
        if r == 0:
            # No cointegration - estimate VAR in differences
            alpha = np.zeros((n, 0))
            beta = np.zeros((n, 0))
        else:
            # Get eigenvectors (cointegrating vectors)
            beta = joh_results.eigenvectors[:, :r]
            
            if normalize:
                beta = normalize_beta(beta, 'first')
            
            # Compute alpha
            beta_S11_beta = beta.T @ S11 @ beta
            try:
                alpha = S01 @ beta @ inv(beta_S11_beta)
            except np.linalg.LinAlgError:
                alpha = S01 @ beta @ np.linalg.pinv(beta_S11_beta)
        
        # Long-run matrix
        Pi = alpha @ beta.T
        
        # Estimate short-run dynamics (Gamma matrices)
        Gamma = self._estimate_short_run(Pi)
        
        # Constant and trend
        C, D = self._estimate_deterministic(Pi, Gamma)
        
        # Residuals and fitted values
        residuals, fitted = self._compute_residuals_fitted(Pi, Gamma, C, D)
        
        # Covariance matrix
        Omega = compute_var_covariance(residuals)
        
        # Log-likelihood and information criteria
        log_lik = self._compute_log_likelihood(residuals, Omega)
        n_params = self._count_parameters(r, len(Gamma))
        aic, bic = compute_aic_bic(residuals, n_params)
        
        return MixedVECMResults(
            alpha=alpha,
            beta=beta,
            Pi=Pi,
            Gamma=Gamma,
            C=C,
            D=D,
            Omega=Omega,
            residuals=residuals,
            fitted_values=fitted,
            n_obs=T_eff,
            n_vars=n,
            n_lags=self.k_ar_diff + 1,
            coint_rank=r,
            n_i0=len(i0_indices),
            n_i1=len(i1_indices),
            n_coint=r - len(i0_indices),
            i0_indices=i0_indices,
            i1_indices=i1_indices,
            det_order=self.det_order,
            log_likelihood=log_lik,
            aic=aic,
            bic=bic,
            johansen_results=joh_results,
            i0_test_results=[]
        )
    
    def _estimate_restricted(self, coint_rank: int, n_i0: int, n_coint: int,
                              i0_indices: List[int], i1_indices: List[int],
                              joh_results: JohansenTestResults,
                              normalize: bool) -> MixedVECMResults:
        """
        Estimate restricted Mixed VECM with orthogonal I(0) and coint spaces.
        
        Parameters
        ----------
        coint_rank : int
            Total cointegration rank (n_coint + n_i0)
        n_i0 : int
            Number of I(0) variables
        n_coint : int
            Number of cointegrating relations
        i0_indices : list of int
            I(0) variable indices
        i1_indices : list of int
            I(1) variable indices
        joh_results : JohansenTestResults
            Johansen test results
        normalize : bool
            Whether to normalize beta
            
        Returns
        -------
        MixedVECMResults
            Estimation results
        """
        n = self.n_vars
        r = coint_rank
        T_eff = joh_results.n_obs
        
        S00 = joh_results.S00
        S01 = joh_results.S01
        S10 = joh_results.S10
        S11 = joh_results.S11
        
        # Construct restriction matrices
        H, h0 = _construct_beta_restriction_matrix(n, n_coint, n_i0, i0_indices)
        
        # Estimate restricted model
        alpha, beta, Omega, converged, n_iter = _estimate_restricted_vecm(
            S00, S01, S10, S11, H, h0, n, r, n_coint, n_i0, i0_indices,
            max_iter=100, tol=1e-8
        )
        
        if not converged:
            warnings.warn("Restricted VECM estimation did not converge.")
        
        if normalize:
            beta = normalize_beta(beta, 'first')
            # Recompute alpha with normalized beta
            beta_S11_beta = beta.T @ S11 @ beta
            try:
                alpha = S01 @ beta @ inv(beta_S11_beta)
            except np.linalg.LinAlgError:
                pass
        
        # Long-run matrix
        Pi = alpha @ beta.T
        
        # Estimate short-run dynamics
        Gamma = self._estimate_short_run(Pi)
        
        # Constant and trend
        C, D = self._estimate_deterministic(Pi, Gamma)
        
        # Residuals and fitted values
        residuals, fitted = self._compute_residuals_fitted(Pi, Gamma, C, D)
        
        # Update covariance matrix from residuals
        Omega = compute_var_covariance(residuals)
        
        # Log-likelihood and information criteria
        log_lik = self._compute_log_likelihood(residuals, Omega)
        n_params = self._count_parameters(r, len(Gamma), restricted=True, 
                                           n_coint=n_coint, n_i0=n_i0)
        aic, bic = compute_aic_bic(residuals, n_params)
        
        return MixedVECMResults(
            alpha=alpha,
            beta=beta,
            Pi=Pi,
            Gamma=Gamma,
            C=C,
            D=D,
            Omega=Omega,
            residuals=residuals,
            fitted_values=fitted,
            n_obs=T_eff,
            n_vars=n,
            n_lags=self.k_ar_diff + 1,
            coint_rank=r,
            n_i0=n_i0,
            n_i1=len(i1_indices),
            n_coint=n_coint,
            i0_indices=i0_indices,
            i1_indices=i1_indices,
            det_order=self.det_order,
            log_likelihood=log_lik,
            aic=aic,
            bic=bic,
            johansen_results=joh_results,
            i0_test_results=[]
        )
    
    def _estimate_short_run(self, Pi: np.ndarray) -> List[np.ndarray]:
        """
        Estimate short-run dynamics (Gamma matrices).
        
        Parameters
        ----------
        Pi : np.ndarray
            Long-run matrix
            
        Returns
        -------
        list of np.ndarray
            Short-run coefficient matrices [Γ_1, ..., Γ_{p-1}]
        """
        Y = self.Y
        n = self.n_vars
        p = self.k_ar_diff
        
        # Compute differences
        dY = diff_matrix(Y)
        T = dY.shape[0]
        
        # Dependent variable: ΔY_t - Π Y_{t-1}
        Y_lag = Y[:-1]  # Y_{t-1}
        pi_effect = Y_lag @ Pi.T  # Π Y_{t-1}
        
        dY_adj = dY[p:] - pi_effect[p:]
        T_eff = dY_adj.shape[0]
        
        if p <= 1:
            return []
        
        # Construct lagged differences
        X = np.zeros((T_eff, n * (p - 1)))
        for l in range(1, p):
            X[:, (l-1)*n:l*n] = dY[p-l:-l if l < p else None]
        
        # Add constant if det_order >= 0
        if self.det_order >= 0:
            X = np.hstack([X, np.ones((T_eff, 1))])
        if self.det_order == 1:
            X = np.hstack([X, np.arange(1, T_eff + 1).reshape(-1, 1)])
        
        # OLS estimation
        coeffs = lstsq(X, dY_adj, rcond=None)[0]
        
        # Extract Gamma matrices
        Gamma = []
        for l in range(p - 1):
            Gamma.append(coeffs[l*n:(l+1)*n].T)
        
        return Gamma
    
    def _estimate_deterministic(self, Pi: np.ndarray,
                                 Gamma: List[np.ndarray]
                                 ) -> Tuple[np.ndarray, Optional[np.ndarray]]:
        """
        Estimate constant and trend coefficients.
        
        Parameters
        ----------
        Pi : np.ndarray
            Long-run matrix
        Gamma : list of np.ndarray
            Short-run matrices
            
        Returns
        -------
        C : np.ndarray
            Constant vector
        D : np.ndarray or None
            Trend coefficient (if det_order == 1)
        """
        Y = self.Y
        n = self.n_vars
        p = self.k_ar_diff
        
        dY = diff_matrix(Y)
        T = dY.shape[0]
        
        # Compute residuals without deterministic terms
        Y_lag = Y[:-1]
        residuals = dY[p:] - Y_lag[p:] @ Pi.T
        
        for l, Gamma_l in enumerate(Gamma):
            residuals -= dY[p-l-1:-l-1] @ Gamma_l.T
        
        # Estimate C (and D if trend)
        T_eff = residuals.shape[0]
        
        if self.det_order == -1:
            C = np.zeros(n)
            D = None
        elif self.det_order == 0:
            C = residuals.mean(axis=0)
            D = None
        else:
            # Constant and trend
            X = np.hstack([np.ones((T_eff, 1)), 
                          np.arange(1, T_eff + 1).reshape(-1, 1)])
            coeffs = lstsq(X, residuals, rcond=None)[0]
            C = coeffs[0]
            D = coeffs[1]
        
        return C, D
    
    def _compute_residuals_fitted(self, Pi: np.ndarray,
                                   Gamma: List[np.ndarray],
                                   C: np.ndarray,
                                   D: Optional[np.ndarray]
                                   ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Compute residuals and fitted values.
        
        Parameters
        ----------
        Pi : np.ndarray
            Long-run matrix
        Gamma : list of np.ndarray
            Short-run matrices
        C : np.ndarray
            Constant
        D : np.ndarray or None
            Trend coefficient
            
        Returns
        -------
        residuals : np.ndarray
            Residuals
        fitted : np.ndarray
            Fitted values
        """
        Y = self.Y
        n = self.n_vars
        p = self.k_ar_diff
        
        dY = diff_matrix(Y)
        T = dY.shape[0]
        T_eff = T - p
        
        Y_lag = Y[:-1]
        
        # Fitted values
        fitted = np.zeros((T_eff, n))
        
        # Error correction term
        fitted += Y_lag[p:] @ Pi.T
        
        # Short-run dynamics
        for l, Gamma_l in enumerate(Gamma):
            fitted += dY[p-l-1:-l-1] @ Gamma_l.T
        
        # Deterministic terms
        fitted += C
        if D is not None:
            fitted += np.outer(np.arange(1, T_eff + 1), D)
        
        # Residuals
        residuals = dY[p:] - fitted
        
        return residuals, fitted
    
    def _compute_log_likelihood(self, residuals: np.ndarray,
                                 Omega: np.ndarray) -> float:
        """
        Compute log-likelihood.
        
        Parameters
        ----------
        residuals : np.ndarray
            Residuals
        Omega : np.ndarray
            Covariance matrix
            
        Returns
        -------
        float
            Log-likelihood value
        """
        T, n = residuals.shape
        
        # Log-likelihood for multivariate normal
        log_lik = -T * n / 2 * np.log(2 * np.pi)
        log_lik -= T / 2 * np.log(det(Omega))
        
        # Quadratic form
        Omega_inv = inv(Omega)
        for t in range(T):
            log_lik -= 0.5 * residuals[t] @ Omega_inv @ residuals[t]
        
        return log_lik
    
    def _count_parameters(self, r: int, n_gamma: int,
                          restricted: bool = False,
                          n_coint: int = 0, n_i0: int = 0) -> int:
        """
        Count number of free parameters.
        
        Parameters
        ----------
        r : int
            Cointegration rank
        n_gamma : int
            Number of Gamma matrices
        restricted : bool
            Whether model is restricted
        n_coint : int
            Number of cointegrating relations
        n_i0 : int
            Number of I(0) variables
            
        Returns
        -------
        int
            Number of free parameters
        """
        n = self.n_vars
        
        # Alpha: n * r parameters (unrestricted)
        n_alpha = n * r
        
        # Beta: n * r - r^2 parameters (identified)
        if restricted:
            # Restricted beta has fewer parameters
            n_beta = (n - n_i0) * n_coint + n_i0 * n_i0 - r * r
        else:
            n_beta = n * r - r * r
        
        # Gamma matrices
        n_gamma_params = n_gamma * n * n
        
        # Deterministic
        n_det = 0
        if self.det_order >= 0:
            n_det += n  # constant
        if self.det_order == 1:
            n_det += n  # trend
        
        # Covariance (lower triangle)
        n_omega = n * (n + 1) // 2
        
        return n_alpha + n_beta + n_gamma_params + n_det + n_omega
    
    def _compute_standard_errors(self, results: MixedVECMResults) -> None:
        """
        Compute standard errors for alpha and beta.
        
        Parameters
        ----------
        results : MixedVECMResults
            Results object to update with standard errors
        """
        # This is a simplified implementation
        # Full implementation would use the asymptotic formulas from Johansen (1995)
        
        T = results.n_obs
        n = results.n_vars
        r = results.coint_rank
        
        if r == 0:
            return
        
        # Approximate standard errors using asymptotic formulas
        S11 = results.johansen_results.S11
        Omega = results.Omega
        
        # For alpha: var(α) ≈ Ω ⊗ (β'S11β)^{-1} / T
        beta = results.beta
        beta_S11_beta = beta.T @ S11 @ beta
        
        try:
            beta_S11_beta_inv = inv(beta_S11_beta)
            var_alpha = kronecker(Omega, beta_S11_beta_inv) / T
            
            # Standard errors
            se_alpha = np.sqrt(np.diag(var_alpha)).reshape(n, r)
            results.alpha_se = se_alpha
            
            # t-statistics
            results.alpha_tstat = results.alpha / se_alpha
        except np.linalg.LinAlgError:
            pass


def estimate_mixed_vecm(Y: np.ndarray, k_ar_diff: int = 1,
                        det_order: int = 0, **kwargs) -> MixedVECMResults:
    """
    Convenience function to estimate Mixed VECM.
    
    Parameters
    ----------
    Y : np.ndarray
        Time series data (T × n)
    k_ar_diff : int, default 1
        Number of lagged differences
    det_order : int, default 0
        Deterministic order
    **kwargs : dict
        Additional arguments passed to MixedVECM.fit()
        
    Returns
    -------
    MixedVECMResults
        Estimation results
        
    Examples
    --------
    >>> import numpy as np
    >>> from mixvecm import estimate_mixed_vecm
    >>> 
    >>> Y = np.random.randn(200, 6)
    >>> results = estimate_mixed_vecm(Y, k_ar_diff=2)
    >>> print(results.summary())
    """
    model = MixedVECM(Y, k_ar_diff, det_order)
    return model.fit(**kwargs)


def fit_vecm(Y: np.ndarray, coint_rank: int, k_ar_diff: int = 1,
             det_order: int = 0, **kwargs) -> MixedVECMResults:
    """
    Fit standard VECM with specified cointegration rank.
    
    Parameters
    ----------
    Y : np.ndarray
        Time series data (T × n)
    coint_rank : int
        Cointegration rank
    k_ar_diff : int, default 1
        Number of lagged differences
    det_order : int, default 0
        Deterministic order
    **kwargs : dict
        Additional arguments
        
    Returns
    -------
    MixedVECMResults
        Estimation results
    """
    model = MixedVECM(Y, k_ar_diff, det_order)
    return model.fit(coint_rank=coint_rank, n_i0=0, **kwargs)
