"""
Utility Functions for MixVECM
=============================

This module provides helper functions for matrix operations, lag construction,
and other computations needed for VECM estimation.

References
----------
[1] Johansen, S. (1995). Likelihood-Based Inference in Cointegrated Vector
    Autoregressive Models. Oxford University Press.
[2] Lütkepohl, H. (2005). New Introduction to Multiple Time Series Analysis.
    Springer.
"""

import numpy as np
from numpy.linalg import inv, eig, qr, svd, matrix_rank
from scipy.linalg import null_space
from typing import Tuple, Optional, Union, List
import warnings


def vec(A: np.ndarray) -> np.ndarray:
    """
    Vectorize a matrix by stacking its columns.
    
    Parameters
    ----------
    A : np.ndarray
        Input matrix of shape (m, n)
        
    Returns
    -------
    np.ndarray
        Column vector of shape (m*n, 1)
        
    Notes
    -----
    This follows the standard vec operator definition where:
    vec(A) = [a_1', a_2', ..., a_n']' where a_i is the i-th column of A.
    
    Examples
    --------
    >>> A = np.array([[1, 2], [3, 4]])
    >>> vec(A)
    array([[1], [3], [2], [4]])
    """
    return A.flatten('F').reshape(-1, 1)


def unvec(v: np.ndarray, m: int, n: int) -> np.ndarray:
    """
    Unvectorize a vector back to matrix form.
    
    Parameters
    ----------
    v : np.ndarray
        Input vector of shape (m*n,) or (m*n, 1)
    m : int
        Number of rows in output matrix
    n : int
        Number of columns in output matrix
        
    Returns
    -------
    np.ndarray
        Matrix of shape (m, n)
        
    Examples
    --------
    >>> v = np.array([1, 3, 2, 4])
    >>> unvec(v, 2, 2)
    array([[1, 2], [3, 4]])
    """
    v = np.asarray(v).flatten()
    return v.reshape((m, n), order='F')


def lag_matrix(Y: np.ndarray, p: int, include_contemporaneous: bool = False
               ) -> Tuple[np.ndarray, np.ndarray]:
    """
    Construct lagged variable matrices for VAR/VECM estimation.
    
    Parameters
    ----------
    Y : np.ndarray
        Time series data of shape (T, n) where T is the number of observations
        and n is the number of variables
    p : int
        Number of lags
    include_contemporaneous : bool, default False
        If True, include Y_t in the output
        
    Returns
    -------
    Y_current : np.ndarray
        Current values Y_t of shape (T-p, n)
    Y_lagged : np.ndarray
        Lagged values [Y_{t-1}, Y_{t-2}, ..., Y_{t-p}] of shape (T-p, n*p)
        
    Examples
    --------
    >>> Y = np.array([[1, 2], [3, 4], [5, 6], [7, 8]])
    >>> Y_curr, Y_lag = lag_matrix(Y, 2)
    >>> Y_curr.shape
    (2, 2)
    >>> Y_lag.shape
    (2, 4)
    """
    Y = np.asarray(Y)
    T, n = Y.shape
    
    if p >= T:
        raise ValueError(f"Number of lags p={p} must be less than T={T}")
    
    Y_current = Y[p:]
    Y_lagged = np.zeros((T - p, n * p))
    
    for l in range(1, p + 1):
        Y_lagged[:, (l-1)*n:l*n] = Y[p-l:T-l]
    
    return Y_current, Y_lagged


def diff_matrix(Y: np.ndarray) -> np.ndarray:
    """
    Compute first differences of a time series matrix.
    
    Parameters
    ----------
    Y : np.ndarray
        Time series data of shape (T, n)
        
    Returns
    -------
    np.ndarray
        First differences ΔY_t = Y_t - Y_{t-1} of shape (T-1, n)
        
    Examples
    --------
    >>> Y = np.array([[1, 2], [3, 5], [6, 9]])
    >>> diff_matrix(Y)
    array([[2, 3], [3, 4]])
    """
    Y = np.asarray(Y)
    return np.diff(Y, axis=0)


def compute_residuals(Y: np.ndarray, Y_fitted: np.ndarray) -> np.ndarray:
    """
    Compute residuals from fitted values.
    
    Parameters
    ----------
    Y : np.ndarray
        Actual values of shape (T, n)
    Y_fitted : np.ndarray
        Fitted values of shape (T, n)
        
    Returns
    -------
    np.ndarray
        Residuals U = Y - Y_fitted of shape (T, n)
    """
    return Y - Y_fitted


def compute_aic_bic(residuals: np.ndarray, k: int, T: Optional[int] = None
                    ) -> Tuple[float, float]:
    """
    Compute AIC and BIC information criteria.
    
    Parameters
    ----------
    residuals : np.ndarray
        Residual matrix of shape (T, n)
    k : int
        Number of estimated parameters
    T : int, optional
        Number of observations. If None, inferred from residuals.
        
    Returns
    -------
    aic : float
        Akaike Information Criterion
    bic : float
        Bayesian Information Criterion
        
    Notes
    -----
    Following Lütkepohl (2005):
    - AIC = log|Σ| + 2k/T
    - BIC = log|Σ| + k*log(T)/T
    
    where Σ = U'U/T is the residual covariance matrix.
    """
    residuals = np.asarray(residuals)
    if T is None:
        T = residuals.shape[0]
    n = residuals.shape[1]
    
    # Residual covariance matrix
    Sigma = residuals.T @ residuals / T
    log_det_sigma = np.log(np.linalg.det(Sigma))
    
    aic = log_det_sigma + 2 * k / T
    bic = log_det_sigma + k * np.log(T) / T
    
    return aic, bic


def select_lag_order(Y: np.ndarray, max_lag: int = 12, 
                     criterion: str = 'bic',
                     det_order: int = 0) -> int:
    """
    Select optimal lag order for VAR model using information criteria.
    
    Parameters
    ----------
    Y : np.ndarray
        Time series data of shape (T, n)
    max_lag : int, default 12
        Maximum lag order to consider
    criterion : str, default 'bic'
        Information criterion to use: 'aic' or 'bic'
    det_order : int, default 0
        Deterministic order: -1 (none), 0 (constant), 1 (trend)
        
    Returns
    -------
    int
        Optimal lag order
        
    Notes
    -----
    The function estimates VAR(p) for p = 1, ..., max_lag and selects
    the model minimizing the chosen information criterion.
    """
    Y = np.asarray(Y)
    T, n = Y.shape
    
    best_criterion = np.inf
    best_lag = 1
    
    for p in range(1, max_lag + 1):
        if p >= T - 10:  # Need sufficient observations
            break
            
        # Estimate VAR(p)
        Y_curr, Y_lag = lag_matrix(Y, p)
        T_eff = Y_curr.shape[0]
        
        # Add deterministic components
        if det_order >= 0:
            const = np.ones((T_eff, 1))
            X = np.hstack([Y_lag, const])
        else:
            X = Y_lag
            
        if det_order == 1:
            trend = np.arange(1, T_eff + 1).reshape(-1, 1)
            X = np.hstack([X, trend])
        
        # OLS estimation
        try:
            beta_hat = np.linalg.lstsq(X, Y_curr, rcond=None)[0]
            residuals = Y_curr - X @ beta_hat
            
            # Number of parameters: n*n*p + deterministic
            k = n * n * p
            if det_order >= 0:
                k += n  # constant
            if det_order == 1:
                k += n  # trend
                
            aic, bic = compute_aic_bic(residuals, k, T_eff)
            
            ic = aic if criterion.lower() == 'aic' else bic
            
            if ic < best_criterion:
                best_criterion = ic
                best_lag = p
        except np.linalg.LinAlgError:
            continue
    
    return best_lag


def orthogonal_complement(A: np.ndarray) -> np.ndarray:
    """
    Compute the orthogonal complement of matrix A.
    
    Parameters
    ----------
    A : np.ndarray
        Input matrix of shape (n, r) with rank r < n
        
    Returns
    -------
    np.ndarray
        Orthogonal complement A_⊥ of shape (n, n-r) such that A'A_⊥ = 0
        
    Notes
    -----
    Following Johansen (1995), the orthogonal complement A_⊥ is an n × (n-r)
    matrix of full column rank such that A'A_⊥ = 0.
    
    The implementation uses the null space of A'.
    
    Examples
    --------
    >>> A = np.array([[1, 0], [0, 1], [1, 1]])
    >>> A_perp = orthogonal_complement(A)
    >>> np.allclose(A.T @ A_perp, 0)
    True
    """
    A = np.asarray(A)
    
    # Compute null space of A'
    A_perp = null_space(A.T)
    
    if A_perp.size == 0:
        # A has full column rank equal to number of rows
        return np.zeros((A.shape[0], 0))
    
    return A_perp


def normalize_beta(beta: np.ndarray, 
                   normalization: str = 'first',
                   indices: Optional[List[int]] = None) -> np.ndarray:
    """
    Normalize cointegrating vectors.
    
    Parameters
    ----------
    beta : np.ndarray
        Cointegrating matrix of shape (n, r)
    normalization : str, default 'first'
        Normalization scheme:
        - 'first': Normalize first r rows to identity
        - 'specified': Use specified indices for identity rows
        - 'none': No normalization
    indices : list of int, optional
        Row indices for identity normalization when normalization='specified'
        
    Returns
    -------
    np.ndarray
        Normalized cointegrating matrix
        
    Notes
    -----
    Following Johansen (1995), β is identified up to a full rank transformation.
    The normalization β' = [I_r, β*'] is commonly used.
    """
    beta = np.asarray(beta)
    n, r = beta.shape
    
    if normalization == 'none':
        return beta
    
    if normalization == 'first':
        indices = list(range(r))
    elif normalization == 'specified':
        if indices is None:
            raise ValueError("indices must be provided for 'specified' normalization")
    else:
        raise ValueError(f"Unknown normalization: {normalization}")
    
    # Extract the r rows for normalization
    B = beta[indices, :]
    
    # Check if B is invertible
    if np.linalg.matrix_rank(B) < r:
        warnings.warn("Normalization matrix is rank deficient. "
                      "Returning unnormalized beta.")
        return beta
    
    # Normalize: beta_new = beta @ inv(B)
    beta_normalized = beta @ np.linalg.inv(B)
    
    return beta_normalized


def compute_var_covariance(residuals: np.ndarray) -> np.ndarray:
    """
    Compute variance-covariance matrix of residuals.
    
    Parameters
    ----------
    residuals : np.ndarray
        Residual matrix of shape (T, n)
        
    Returns
    -------
    np.ndarray
        Covariance matrix Ω of shape (n, n)
        
    Notes
    -----
    Computes Ω = U'U/T (ML estimator)
    """
    residuals = np.asarray(residuals)
    T = residuals.shape[0]
    return residuals.T @ residuals / T


def compute_long_run_covariance(residuals: np.ndarray, 
                                 kernel: str = 'bartlett',
                                 bandwidth: Optional[int] = None) -> np.ndarray:
    """
    Compute long-run covariance matrix using kernel estimation.
    
    Parameters
    ----------
    residuals : np.ndarray
        Residual matrix of shape (T, n)
    kernel : str, default 'bartlett'
        Kernel function: 'bartlett' or 'parzen'
    bandwidth : int, optional
        Bandwidth parameter. If None, uses Newey-West automatic selection.
        
    Returns
    -------
    np.ndarray
        Long-run covariance matrix of shape (n, n)
    """
    residuals = np.asarray(residuals)
    T, n = residuals.shape
    
    if bandwidth is None:
        bandwidth = int(np.floor(4 * (T / 100) ** (2/9)))
    
    # Center residuals
    residuals = residuals - residuals.mean(axis=0)
    
    # Compute autocovariances
    Sigma = np.zeros((n, n))
    for j in range(-bandwidth, bandwidth + 1):
        if kernel == 'bartlett':
            w = 1 - abs(j) / (bandwidth + 1)
        elif kernel == 'parzen':
            z = abs(j) / (bandwidth + 1)
            if z <= 0.5:
                w = 1 - 6 * z**2 + 6 * z**3
            else:
                w = 2 * (1 - z)**3
        else:
            raise ValueError(f"Unknown kernel: {kernel}")
        
        if j >= 0:
            Gamma_j = residuals[j:].T @ residuals[:T-j] / T if j > 0 else residuals.T @ residuals / T
        else:
            Gamma_j = residuals[:T+j].T @ residuals[-j:] / T
        
        Sigma += w * Gamma_j
    
    return Sigma


def companion_matrix(Phi_list: List[np.ndarray]) -> np.ndarray:
    """
    Construct companion form matrix for VAR(p).
    
    Parameters
    ----------
    Phi_list : list of np.ndarray
        List of VAR coefficient matrices [Φ_1, Φ_2, ..., Φ_p]
        Each Φ_i has shape (n, n)
        
    Returns
    -------
    np.ndarray
        Companion matrix of shape (n*p, n*p)
        
    Notes
    -----
    The companion matrix is:
    
        F = | Φ_1  Φ_2  ...  Φ_{p-1}  Φ_p |
            | I    0    ...  0        0   |
            | 0    I    ...  0        0   |
            | ...  ...  ...  ...      ... |
            | 0    0    ...  I        0   |
    """
    p = len(Phi_list)
    n = Phi_list[0].shape[0]
    
    F = np.zeros((n * p, n * p))
    
    # First block row
    for i, Phi in enumerate(Phi_list):
        F[:n, i*n:(i+1)*n] = Phi
    
    # Identity blocks below diagonal
    if p > 1:
        F[n:, :-n] = np.eye(n * (p - 1))
    
    return F


def is_stable(Phi_list: List[np.ndarray], tol: float = 1.0) -> bool:
    """
    Check if VAR is stable (all eigenvalues inside unit circle).
    
    Parameters
    ----------
    Phi_list : list of np.ndarray
        List of VAR coefficient matrices
    tol : float, default 1.0
        Tolerance for eigenvalue modulus check
        
    Returns
    -------
    bool
        True if VAR is stable
    """
    F = companion_matrix(Phi_list)
    eigenvalues = np.linalg.eigvals(F)
    return np.all(np.abs(eigenvalues) < tol)


def kronecker(A: np.ndarray, B: np.ndarray) -> np.ndarray:
    """
    Compute Kronecker product A ⊗ B.
    
    Parameters
    ----------
    A : np.ndarray
        First matrix
    B : np.ndarray
        Second matrix
        
    Returns
    -------
    np.ndarray
        Kronecker product
    """
    return np.kron(A, B)


def commutation_matrix(m: int, n: int) -> np.ndarray:
    """
    Compute the commutation matrix K_{m,n}.
    
    Parameters
    ----------
    m : int
        First dimension
    n : int
        Second dimension
        
    Returns
    -------
    np.ndarray
        Commutation matrix of shape (m*n, m*n)
        
    Notes
    -----
    The commutation matrix satisfies:
    K_{m,n} vec(A) = vec(A') for any m×n matrix A
    """
    K = np.zeros((m * n, m * n))
    for i in range(m):
        for j in range(n):
            # Unit vector e_i in R^m
            e_i = np.zeros(m)
            e_i[i] = 1
            # Unit vector e_j in R^n
            e_j = np.zeros(n)
            e_j[j] = 1
            # K = sum_{i,j} (e_j ⊗ e_i)(e_i ⊗ e_j)'
            K += np.outer(np.kron(e_j, e_i), np.kron(e_i, e_j))
    return K


def duplication_matrix(n: int) -> np.ndarray:
    """
    Compute the duplication matrix D_n.
    
    Parameters
    ----------
    n : int
        Dimension
        
    Returns
    -------
    np.ndarray
        Duplication matrix of shape (n^2, n(n+1)/2)
        
    Notes
    -----
    For symmetric matrix A, vec(A) = D_n vech(A)
    """
    d = int(n * (n + 1) / 2)
    D = np.zeros((n * n, d))
    
    count = 0
    for j in range(n):
        for i in range(j, n):
            # Index in vec(A)
            vec_idx_ij = i + j * n
            vec_idx_ji = j + i * n
            D[vec_idx_ij, count] = 1
            if i != j:
                D[vec_idx_ji, count] = 1
            count += 1
    
    return D


def elimination_matrix(n: int) -> np.ndarray:
    """
    Compute the elimination matrix L_n.
    
    Parameters
    ----------
    n : int
        Dimension
        
    Returns
    -------
    np.ndarray
        Elimination matrix of shape (n(n+1)/2, n^2)
        
    Notes
    -----
    For symmetric matrix A, vech(A) = L_n vec(A)
    """
    d = int(n * (n + 1) / 2)
    L = np.zeros((d, n * n))
    
    count = 0
    for j in range(n):
        for i in range(j, n):
            vec_idx = i + j * n
            L[count, vec_idx] = 1
            count += 1
    
    return L


def block_diagonal(*matrices: np.ndarray) -> np.ndarray:
    """
    Create a block diagonal matrix from input matrices.
    
    Parameters
    ----------
    *matrices : np.ndarray
        Input matrices
        
    Returns
    -------
    np.ndarray
        Block diagonal matrix
    """
    from scipy.linalg import block_diag
    return block_diag(*matrices)
