"""
Data Generation for Mixed VECM
==============================

This module implements data generating processes for mixed Vector Error Correction
Models containing both I(0) and I(1) variables.

The module implements two approaches from Chen (2022):
1. Generation from underlying stationary VAR (Example 3.1)
2. Generation from VECM with restricted β (Example 3.2)

References
----------
[1] Chen, P. (2022). Vector Error Correction Models with Stationary and 
    Nonstationary Variables. SSRN Working Paper 4239531.
"""

import numpy as np
from numpy.linalg import inv
from scipy.linalg import cholesky
from typing import Tuple, Optional, Dict, Any, List
from dataclasses import dataclass


@dataclass
class DataGenerationResult:
    """
    Results from data generation.
    
    Attributes
    ----------
    Y : np.ndarray
        Generated time series of shape (T, n)
    U : np.ndarray
        Generated error terms of shape (T, n)
    underlying : np.ndarray, optional
        Underlying process (ΔZ, X) if applicable
    params : dict
        Parameters used for generation
    i1_indices : list of int
        Indices of I(1) variables
    i0_indices : list of int
        Indices of I(0) variables
    """
    Y: np.ndarray
    U: np.ndarray
    underlying: Optional[np.ndarray]
    params: Dict[str, Any]
    i1_indices: List[int]
    i0_indices: List[int]


def generate_mixed_vecm_data(
    T: int = 200,
    n: int = 6,
    n_i1: int = 4,
    n_coint: int = 2,
    method: str = 'underlying_var',
    seed: Optional[int] = None,
    burn_in: int = 100,
    **kwargs
) -> DataGenerationResult:
    """
    Generate data from a mixed VECM with I(0) and I(1) variables.
    
    Parameters
    ----------
    T : int, default 200
        Number of observations
    n : int, default 6
        Total number of variables
    n_i1 : int, default 4
        Number of I(1) variables
    n_coint : int, default 2
        Number of cointegrating relations among I(1) variables
    method : str, default 'underlying_var'
        Generation method:
        - 'underlying_var': From underlying stationary VAR (Example 3.1)
        - 'restricted_beta': From VECM with restricted β (Example 3.2)
    seed : int, optional
        Random seed for reproducibility
    burn_in : int, default 100
        Number of burn-in observations
    **kwargs : dict
        Additional parameters passed to generation functions
        
    Returns
    -------
    DataGenerationResult
        Generated data and parameters
        
    Examples
    --------
    >>> result = generate_mixed_vecm_data(T=200, n=6, n_i1=4, n_coint=2, seed=42)
    >>> print(result.Y.shape)
    (200, 6)
    >>> print(f"I(1) variables: {result.i1_indices}")
    >>> print(f"I(0) variables: {result.i0_indices}")
    """
    if seed is not None:
        np.random.seed(seed)
    
    n_i0 = n - n_i1
    
    if n_i0 < 0:
        raise ValueError(f"n_i1={n_i1} cannot exceed n={n}")
    if n_coint > n_i1:
        raise ValueError(f"n_coint={n_coint} cannot exceed n_i1={n_i1}")
    if n_coint + n_i0 > n:
        raise ValueError("n_coint + n_i0 cannot exceed n")
    
    if method == 'underlying_var':
        return generate_from_underlying_var(T, n, n_i1, n_coint, seed, burn_in, **kwargs)
    elif method == 'restricted_beta':
        return generate_from_restricted_beta(T, n, n_i1, n_coint, seed, burn_in, **kwargs)
    else:
        raise ValueError(f"Unknown method: {method}")


def generate_from_underlying_var(
    T: int = 200,
    n: int = 6,
    n_i1: int = 4,
    n_coint: int = 2,
    seed: Optional[int] = None,
    burn_in: int = 100,
    theta1: Optional[np.ndarray] = None,
    theta2: Optional[np.ndarray] = None,
    B: Optional[np.ndarray] = None,
    Sigma: Optional[np.ndarray] = None
) -> DataGenerationResult:
    """
    Generate mixed VECM data from underlying stationary VAR.
    
    This implements the data generating process from Example 3.1 in Chen (2022).
    
    The underlying VAR is:
        [ΔZ_t, X1_t, X2_t]' = θ^(1) [ΔZ_{t-1}, X1_{t-1}, X2_{t-1}]' + 
                               θ^(2) [ΔZ_{t-2}, X1_{t-2}, X2_{t-2}]' + e_t
    
    The observed Y_t is:
        Y_t = B [Z_t, X1_t, X2_t]'
    
    The condition B31 = 0 ensures the last n_i0 components of Y are I(0).
    
    Parameters
    ----------
    T : int, default 200
        Number of observations
    n : int, default 6
        Total number of variables
    n_i1 : int, default 4
        Number of I(1) variables
    n_coint : int, default 2
        Number of cointegrating relations
    seed : int, optional
        Random seed
    burn_in : int, default 100
        Burn-in period
    theta1 : np.ndarray, optional
        First lag coefficient matrix (n × n)
    theta2 : np.ndarray, optional
        Second lag coefficient matrix (n × n)
    B : np.ndarray, optional
        Transformation matrix (n × n)
    Sigma : np.ndarray, optional
        Error covariance matrix (n × n)
        
    Returns
    -------
    DataGenerationResult
        Generated data and parameters
        
    Notes
    -----
    The key insight from Chen (2022) is that B31 = 0 is necessary and sufficient
    for the last n_i0 components of Y to be I(0). This means:
    
        Y = B [Z, X1, X2]'
        
    where B = [[B11, B12, B13],
               [B21, B22, B23],
               [0,   B32, B33]]
    
    ensures that the last block only depends on X (stationary), not on Z (I(1)).
    """
    if seed is not None:
        np.random.seed(seed)
    
    n_i0 = n - n_i1
    r = n_i1  # Number of unit roots (I(1) common trends)
    h = n_coint + n_i0  # Dimension of I(0) space
    
    # Dimensions in underlying process
    # ΔZ is r-dimensional (r = n_i1)
    # X is (n - r)-dimensional = n_coint + n_i0 dimensional
    
    # Generate coefficient matrices for underlying VAR if not provided
    if theta1 is None or theta2 is None:
        theta1, theta2 = _generate_stable_var_coefficients(n, seed)
    
    # Generate transformation matrix B with B31 = 0 if not provided
    if B is None:
        B = _generate_transformation_matrix(n, r, n_i0, seed)
    
    # Generate error covariance matrix
    if Sigma is None:
        Sigma = _generate_covariance_matrix(n, seed)
    
    # Generate underlying process
    T_total = T + burn_in
    
    # Underlying process: [ΔZ_t, X_t]'
    underlying = np.zeros((T_total, n))
    e = np.random.multivariate_normal(np.zeros(n), Sigma, T_total)
    
    for t in range(2, T_total):
        underlying[t] = theta1 @ underlying[t-1] + theta2 @ underlying[t-2] + e[t]
    
    # Cumulate the first r components to get Z from ΔZ
    Z_X = underlying.copy()
    Z_X[:, :r] = np.cumsum(underlying[:, :r], axis=0)
    
    # Transform to Y
    Y = (B @ Z_X.T).T
    
    # Remove burn-in
    Y = Y[burn_in:]
    underlying_out = underlying[burn_in:]
    e_out = e[burn_in:]
    
    params = {
        'theta1': theta1,
        'theta2': theta2,
        'B': B,
        'Sigma': Sigma,
        'n_i1': n_i1,
        'n_i0': n_i0,
        'n_coint': n_coint,
        'method': 'underlying_var'
    }
    
    return DataGenerationResult(
        Y=Y,
        U=e_out,
        underlying=underlying_out,
        params=params,
        i1_indices=list(range(n_i1)),
        i0_indices=list(range(n_i1, n))
    )


def generate_from_restricted_beta(
    T: int = 200,
    n: int = 6,
    n_i1: int = 4,
    n_coint: int = 2,
    seed: Optional[int] = None,
    burn_in: int = 100,
    alpha: Optional[np.ndarray] = None,
    beta: Optional[np.ndarray] = None,
    Psi: Optional[np.ndarray] = None,
    C: Optional[np.ndarray] = None,
    Omega: Optional[np.ndarray] = None
) -> DataGenerationResult:
    """
    Generate mixed VECM data from VECM with restricted β.
    
    This implements the data generating process from Example 3.2 in Chen (2022).
    
    The VECM is:
        ΔY_t = C + αβ'Y_{t-1} + Ψ ΔY_{t-1} + ε_t
        
    where β is restricted such that:
        β' = [[B21, B22, 0], [0, 0, B33]]
        
    This ensures the last n_i0 components are I(0).
    
    Parameters
    ----------
    T : int, default 200
        Number of observations
    n : int, default 6
        Total number of variables
    n_i1 : int, default 4
        Number of I(1) variables
    n_coint : int, default 2
        Number of cointegrating relations
    seed : int, optional
        Random seed
    burn_in : int, default 100
        Burn-in period
    alpha : np.ndarray, optional
        Adjustment matrix (n × r) where r = n_coint + n_i0
    beta : np.ndarray, optional
        Cointegrating matrix (n × r) with restricted structure
    Psi : np.ndarray, optional
        Short-run dynamics matrix (n × n)
    C : np.ndarray, optional
        Constant vector (n,)
    Omega : np.ndarray, optional
        Error covariance matrix (n × n)
        
    Returns
    -------
    DataGenerationResult
        Generated data and parameters
        
    Notes
    -----
    According to Lemma 2.6 in Chen (2022), the restricted β structure:
    
        β' = [[B21, B22, 0], [0, 0, B33]]
        
    is necessary and sufficient for the VECM to contain n_i0 I(0) components.
    The I(0) space and cointegration space are orthogonal.
    """
    if seed is not None:
        np.random.seed(seed)
    
    n_i0 = n - n_i1
    r = n_coint + n_i0  # Total I(0) space dimension
    
    # Generate VECM parameters if not provided
    if alpha is None:
        alpha = _generate_alpha_matrix(n, r, n_i0, seed)
    
    if beta is None:
        beta = _generate_restricted_beta(n, n_i1, n_coint, n_i0, seed)
    
    if Psi is None:
        Psi = _generate_short_run_matrix(n, seed)
    
    if C is None:
        C = np.random.randn(n) * 0.1
    
    if Omega is None:
        Omega = _generate_covariance_matrix(n, seed)
    
    # Generate data from VECM
    T_total = T + burn_in + 10  # Extra for initial values
    
    Y = np.zeros((T_total, n))
    epsilon = np.random.multivariate_normal(np.zeros(n), Omega, T_total)
    
    # Initialize with small random values
    Y[:2] = np.random.randn(2, n) * 0.1
    
    for t in range(2, T_total):
        dY_lag = Y[t-1] - Y[t-2]
        ec = beta.T @ Y[t-1]  # Error correction term
        dY = C + alpha @ ec + Psi @ dY_lag + epsilon[t]
        Y[t] = Y[t-1] + dY
    
    # Remove burn-in
    Y = Y[burn_in + 10:]
    epsilon_out = epsilon[burn_in + 10:]
    
    params = {
        'alpha': alpha,
        'beta': beta,
        'Psi': Psi,
        'C': C,
        'Omega': Omega,
        'n_i1': n_i1,
        'n_i0': n_i0,
        'n_coint': n_coint,
        'method': 'restricted_beta'
    }
    
    return DataGenerationResult(
        Y=Y,
        U=epsilon_out,
        underlying=None,
        params=params,
        i1_indices=list(range(n_i1)),
        i0_indices=list(range(n_i1, n))
    )


def _generate_stable_var_coefficients(n: int, seed: Optional[int] = None
                                       ) -> Tuple[np.ndarray, np.ndarray]:
    """
    Generate coefficient matrices for a stable VAR(2).
    
    Parameters
    ----------
    n : int
        Number of variables
    seed : int, optional
        Random seed
        
    Returns
    -------
    theta1 : np.ndarray
        First lag coefficients (n × n)
    theta2 : np.ndarray
        Second lag coefficients (n × n)
    """
    if seed is not None:
        np.random.seed(seed)
    
    max_attempts = 100
    
    for _ in range(max_attempts):
        # Generate random coefficients with moderate values
        theta1 = np.random.randn(n, n) * 0.3
        theta2 = np.random.randn(n, n) * 0.1
        
        # Make diagonal elements negative for stability
        np.fill_diagonal(theta1, -np.abs(np.diag(theta1)) - 0.1)
        
        # To ensure to keep lag 2 coefficient for ΔZ small (for lag 1 VECM)
        # as in the paper: θ^(2)_*1 = 0
        # theta2[:, :r] = 0  # This would be applied if we had r information
        
        # Check stability using companion matrix
        companion = np.zeros((2*n, 2*n))
        companion[:n, :n] = theta1
        companion[:n, n:] = theta2
        companion[n:, :n] = np.eye(n)
        
        eigenvalues = np.linalg.eigvals(companion)
        if np.all(np.abs(eigenvalues) < 0.99):
            return theta1, theta2
    
    # Fallback: use the exact coefficients from the paper (Example 3.1)
    return _example_3_1_coefficients(n)


def _example_3_1_coefficients(n: int) -> Tuple[np.ndarray, np.ndarray]:
    """
    Return coefficients similar to Example 3.1 in the paper.
    """
    # From the paper's Table for Example 3.1
    if n == 6:
        theta1 = np.array([
            [-0.06, -0.52, -0.05, 0.15, -0.05, -0.1],
            [-0.52, 0.25, 0.28, -0.07, 0.04, 0.03],
            [-0.05, 0.28, -0.33, -0.07, 0.29, -0.2],
            [0.15, -0.07, -0.07, -0.44, 0.15, -0.05],
            [-0.05, 0.04, 0.29, 0.15, -0.27, -0.44],
            [-0.1, 0.03, -0.2, -0.05, -0.44, 0.26]
        ])
        
        theta2 = np.array([
            [0, 0, 0.07, 0.09, 0, 0.04],
            [0, 0, 0.01, 0.06, -0.04, 0.01],
            [0, 0, 0.03, 0.04, 0.08, 0.1],
            [0, 0, 0.04, 0.01, 0.04, 0.06],
            [0, 0, 0.08, 0.04, 0, 0.07],
            [0, 0, 0.01, 0.06, 0.07, 0.05]
        ])
        return theta1, theta2
    
    # For other dimensions, generate stable coefficients
    theta1 = -0.3 * np.eye(n) + 0.1 * np.random.randn(n, n)
    theta2 = 0.05 * np.random.randn(n, n)
    return theta1, theta2


def _generate_transformation_matrix(n: int, r: int, n_i0: int,
                                     seed: Optional[int] = None) -> np.ndarray:
    """
    Generate transformation matrix B with B31 = 0.
    
    Parameters
    ----------
    n : int
        Total number of variables
    r : int
        Number of unit roots (I(1) common trends)
    n_i0 : int
        Number of I(0) variables
        
    Returns
    -------
    B : np.ndarray
        Transformation matrix (n × n)
    """
    if seed is not None:
        np.random.seed(seed)
    
    n_i1 = n - n_i0
    
    if n == 6 and r == 4 and n_i0 == 2:
        # Use exact B from Example 3.1
        B = np.array([
            [-0.6, 1.5, 0.42, 0.55, 0.35, -1.83],
            [0.32, 0.15, -0.23, -0.5, -0.66, -0.35],
            [1.87, -2.61, -1.6, -0.37, -0.05, 1.22],
            [-1.77, -0.11, 0.73, 0.2, -0.5, 0.34],
            [0, 0, 0.3, -0.6, 0.37, 0.38],
            [0, 0, -0.4, 0.5, 0.23, -0.79]
        ])
        return B
    
    # Generate random B with B31 = 0 structure
    B = np.random.randn(n, n)
    
    # Set B31 = 0 (last n_i0 rows, first r columns)
    # In our setup: first r columns correspond to Z (cumulated ΔZ)
    # For n_i0 I(0) components, we need B[n-n_i0:, :r] = 0
    B[n_i1:, :r] = 0
    
    # Ensure B is invertible
    while np.abs(np.linalg.det(B)) < 0.1:
        B[:n_i1, :] += 0.1 * np.random.randn(n_i1, n)
    
    return B


def _generate_alpha_matrix(n: int, r: int, n_i0: int,
                            seed: Optional[int] = None) -> np.ndarray:
    """
    Generate adjustment matrix α.
    
    Parameters
    ----------
    n : int
        Number of variables
    r : int
        Cointegration rank (= n_coint + n_i0)
    n_i0 : int
        Number of I(0) variables
        
    Returns
    -------
    alpha : np.ndarray
        Adjustment matrix (n × r)
    """
    if seed is not None:
        np.random.seed(seed)
    
    alpha = np.random.randn(n, r) * 0.2
    
    # Make some adjustments for stability
    # First n_coint columns: cointegration adjustment
    # Last n_i0 columns: I(0) component adjustment
    
    return alpha


def _generate_restricted_beta(n: int, n_i1: int, n_coint: int, n_i0: int,
                               seed: Optional[int] = None) -> np.ndarray:
    """
    Generate restricted cointegrating matrix β.
    
    The restriction is:
        β' = [[B21, B22, 0], [0, 0, B33]]
        
    where:
    - First n_coint rows involve only I(1) variables
    - Last n_i0 rows are identity for I(0) variables
    
    Parameters
    ----------
    n : int
        Total number of variables
    n_i1 : int
        Number of I(1) variables
    n_coint : int
        Number of cointegrating relations
    n_i0 : int
        Number of I(0) variables
        
    Returns
    -------
    beta : np.ndarray
        Cointegrating matrix (n × r) where r = n_coint + n_i0
    """
    if seed is not None:
        np.random.seed(seed)
    
    r = n_coint + n_i0
    beta = np.zeros((n, r))
    
    # Cointegrating relations (first n_coint columns)
    # Only involve I(1) variables (first n_i1 rows)
    beta[:n_i1, :n_coint] = np.random.randn(n_i1, n_coint)
    
    # I(0) component selection (last n_i0 columns)
    # Identity for I(0) variables (last n_i0 rows)
    beta[n_i1:, n_coint:] = np.eye(n_i0)
    
    return beta


def _generate_short_run_matrix(n: int, seed: Optional[int] = None) -> np.ndarray:
    """
    Generate short-run dynamics matrix Ψ.
    
    Parameters
    ----------
    n : int
        Number of variables
    seed : int, optional
        Random seed
        
    Returns
    -------
    Psi : np.ndarray
        Short-run matrix (n × n)
    """
    if seed is not None:
        np.random.seed(seed)
    
    Psi = np.random.randn(n, n) * 0.2
    
    # Make eigenvalues small for stability
    eigenvalues = np.linalg.eigvals(Psi)
    if np.max(np.abs(eigenvalues)) > 0.9:
        Psi = Psi * 0.5
    
    return Psi


def _generate_covariance_matrix(n: int, seed: Optional[int] = None) -> np.ndarray:
    """
    Generate positive definite covariance matrix.
    
    Parameters
    ----------
    n : int
        Dimension
    seed : int, optional
        Random seed
        
    Returns
    -------
    Sigma : np.ndarray
        Covariance matrix (n × n)
    """
    if seed is not None:
        np.random.seed(seed + 1000)
    
    A = np.random.randn(n, n) * 0.5
    Sigma = A @ A.T + 0.1 * np.eye(n)
    
    return Sigma


def example_dgp_1(T: int = 200, seed: Optional[int] = None,
                  burn_in: int = 100) -> DataGenerationResult:
    """
    Generate data using Example 3.1 parameters from Chen (2022).
    
    This generates a 6-variable system with:
    - 4 I(1) variables (series 1-4)
    - 2 I(0) variables (series 5-6)
    - 2 cointegrating relations among I(1) variables
    
    Parameters
    ----------
    T : int, default 200
        Number of observations
    seed : int, optional
        Random seed
    burn_in : int, default 100
        Burn-in period
        
    Returns
    -------
    DataGenerationResult
        Generated data using Example 3.1 parameters
        
    Notes
    -----
    The parameters are taken directly from Table in Section 3.1 of Chen (2022).
    """
    # Exact parameters from Example 3.1
    theta1 = np.array([
        [-0.06, -0.52, -0.05, 0.15, -0.05, -0.1],
        [-0.52, 0.25, 0.28, -0.07, 0.04, 0.03],
        [-0.05, 0.28, -0.33, -0.07, 0.29, -0.2],
        [0.15, -0.07, -0.07, -0.44, 0.15, -0.05],
        [-0.05, 0.04, 0.29, 0.15, -0.27, -0.44],
        [-0.1, 0.03, -0.2, -0.05, -0.44, 0.26]
    ])
    
    theta2 = np.array([
        [0, 0, 0.07, 0.09, 0, 0.04],
        [0, 0, 0.01, 0.06, -0.04, 0.01],
        [0, 0, 0.03, 0.04, 0.08, 0.1],
        [0, 0, 0.04, 0.01, 0.04, 0.06],
        [0, 0, 0.08, 0.04, 0, 0.07],
        [0, 0, 0.01, 0.06, 0.07, 0.05]
    ])
    
    B = np.array([
        [-0.6, 1.5, 0.42, 0.55, 0.35, -1.83],
        [0.32, 0.15, -0.23, -0.5, -0.66, -0.35],
        [1.87, -2.61, -1.6, -0.37, -0.05, 1.22],
        [-1.77, -0.11, 0.73, 0.2, -0.5, 0.34],
        [0, 0, 0.3, -0.6, 0.37, 0.38],
        [0, 0, -0.4, 0.5, 0.23, -0.79]
    ])
    
    return generate_from_underlying_var(
        T=T, n=6, n_i1=4, n_coint=2, seed=seed, burn_in=burn_in,
        theta1=theta1, theta2=theta2, B=B
    )


def example_dgp_2(T: int = 200, seed: Optional[int] = None,
                  burn_in: int = 100) -> DataGenerationResult:
    """
    Generate data using Example 3.2 parameters from Chen (2022).
    
    This generates a 6-variable system with:
    - 4 I(1) variables (series 1-4)
    - 2 I(0) variables (series 5-6)
    - 2 cointegrating relations among I(1) variables
    
    Parameters
    ----------
    T : int, default 200
        Number of observations
    seed : int, optional
        Random seed
    burn_in : int, default 100
        Burn-in period
        
    Returns
    -------
    DataGenerationResult
        Generated data using Example 3.2 parameters
        
    Notes
    -----
    The parameters are taken directly from Table in Section 3.1 of Chen (2022).
    The key feature is the restricted β structure.
    """
    # Parameters from Example 3.2
    # Note: The paper shows α and β in transposed form
    
    # α is 6 × 4 (n × r where r = 2 coint + 2 I(0))
    alpha = np.array([
        [-0.09, 0.54, 6.05, 14.74],
        [0.21, 0.55, 2.78, 7.64],
        [0.3, -0.23, -0.11, -1.32],
        [0.13, 0.32, 5.32, 12.59],
        [-0.13, -0.06, 0.12, -1.65],
        [-0.19, 0.06, -0.9, -3.39]
    ])
    
    # β is 6 × 4 with restricted structure
    # β' = [[B21, B22, 0], [0, 0, B33]]
    # First 2 columns: cointegrating relations (only I(1) vars)
    # Last 2 columns: I(0) selection (only I(0) vars)
    beta = np.array([
        [1, 1, 0, 0],
        [-1.34, -3.8, 0, 0],
        [-0.27, 0.59, 0, 0],
        [-0.29, 0.89, 0, 0],
        [0, 0, -2.13, -0.21],
        [0, 0, 1, 1]
    ])
    
    # Short-run dynamics Ψ (6 × 6)
    Psi = np.array([
        [-1.11, 0.71, -0.54, 0.04, 7.42, -9.35],
        [-0.52, 0.4, -0.26, -0.13, 4.55, -5.86],
        [0.73, -0.28, 0.41, -0.33, -4.16, 5.16],
        [-0.38, 0.38, -0.44, -0.12, 5.96, -7.33],
        [0.12, -0.13, 0.06, 0.05, -1.21, 1.49],
        [0.26, -0.07, 0.22, -0.04, -2.76, 3.26]
    ])
    
    # Constant (small values)
    C = np.zeros(6)
    
    return generate_from_restricted_beta(
        T=T, n=6, n_i1=4, n_coint=2, seed=seed, burn_in=burn_in,
        alpha=alpha, beta=beta, Psi=Psi, C=C
    )


def simulate_monte_carlo(dgp_func, n_simulations: int = 1000,
                         T: int = 200, **kwargs) -> List[DataGenerationResult]:
    """
    Run Monte Carlo simulations.
    
    Parameters
    ----------
    dgp_func : callable
        Data generating function
    n_simulations : int, default 1000
        Number of simulations
    T : int, default 200
        Sample size
    **kwargs : dict
        Additional arguments passed to dgp_func
        
    Returns
    -------
    list of DataGenerationResult
        Results from each simulation
    """
    results = []
    for i in range(n_simulations):
        result = dgp_func(T=T, seed=i, **kwargs)
        results.append(result)
    return results
