"""
Impulse Response Functions for Mixed VECM
==========================================

This module implements impulse response function (IRF) analysis and forecast
error variance decomposition (FEVD) for mixed Vector Error Correction Models.

Main Features
-------------
1. Orthogonalized IRF using Cholesky decomposition
2. Generalized IRF (order invariant)
3. Forecast error variance decomposition
4. Confidence intervals via bootstrap

References
----------
[1] Chen, P. (2022). Vector Error Correction Models with Stationary and 
    Nonstationary Variables. SSRN Working Paper 4239531.
[2] Lütkepohl, H. (2005). New Introduction to Multiple Time Series Analysis.
    Springer.
"""

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

from .utils import companion_matrix


@dataclass
class IRFResults:
    """
    Results from Impulse Response Function computation.
    
    Attributes
    ----------
    irf : np.ndarray
        IRF coefficients of shape (horizons+1, n, n)
        irf[h, i, j] = response of variable i to shock j at horizon h
    horizons : int
        Number of horizons
    n_vars : int
        Number of variables
    orthogonalized : bool
        Whether IRF is orthogonalized
    method : str
        Method used ('cholesky' or 'generalized')
    irf_lower : np.ndarray, optional
        Lower confidence bound
    irf_upper : np.ndarray, optional
        Upper confidence bound
    confidence : float, optional
        Confidence level
    """
    irf: np.ndarray
    horizons: int
    n_vars: int
    orthogonalized: bool
    method: str
    irf_lower: Optional[np.ndarray] = None
    irf_upper: Optional[np.ndarray] = None
    confidence: Optional[float] = None
    
    def plot(self, response_var: Optional[int] = None, 
             impulse_var: Optional[int] = None,
             figsize: Tuple[int, int] = (12, 8)) -> None:
        """
        Plot impulse response functions.
        
        Parameters
        ----------
        response_var : int, optional
            Index of response variable (if None, plot all)
        impulse_var : int, optional
            Index of impulse variable (if None, plot all)
        figsize : tuple, default (12, 8)
            Figure size
        """
        try:
            import matplotlib.pyplot as plt
        except ImportError:
            raise ImportError("matplotlib required for plotting")
        
        n = self.n_vars
        horizons = np.arange(self.horizons + 1)
        
        if response_var is not None and impulse_var is not None:
            # Single plot
            fig, ax = plt.subplots(figsize=(8, 5))
            ax.plot(horizons, self.irf[:, response_var, impulse_var], 'b-', linewidth=2)
            if self.irf_lower is not None and self.irf_upper is not None:
                ax.fill_between(horizons, 
                               self.irf_lower[:, response_var, impulse_var],
                               self.irf_upper[:, response_var, impulse_var],
                               alpha=0.3)
            ax.axhline(y=0, color='k', linestyle='--', linewidth=0.5)
            ax.set_xlabel('Horizon')
            ax.set_ylabel('Response')
            ax.set_title(f'Response of Var {response_var+1} to Var {impulse_var+1} shock')
            ax.grid(True, alpha=0.3)
        else:
            # Multiple plots
            fig, axes = plt.subplots(n, n, figsize=figsize)
            for i in range(n):
                for j in range(n):
                    ax = axes[i, j] if n > 1 else axes
                    ax.plot(horizons, self.irf[:, i, j], 'b-', linewidth=1.5)
                    if self.irf_lower is not None and self.irf_upper is not None:
                        ax.fill_between(horizons,
                                       self.irf_lower[:, i, j],
                                       self.irf_upper[:, i, j],
                                       alpha=0.3)
                    ax.axhline(y=0, color='k', linestyle='--', linewidth=0.5)
                    if i == n - 1:
                        ax.set_xlabel('Horizon')
                    if j == 0:
                        ax.set_ylabel(f'Var {i+1}')
                    if i == 0:
                        ax.set_title(f'Var {j+1} shock')
                    ax.grid(True, alpha=0.3)
            
            plt.tight_layout()
        
        plt.show()
    
    def to_dataframe(self, response_var: int, impulse_var: int):
        """
        Convert IRF to pandas DataFrame.
        
        Parameters
        ----------
        response_var : int
            Response variable index
        impulse_var : int
            Impulse variable index
            
        Returns
        -------
        pd.DataFrame
            DataFrame with IRF values
        """
        try:
            import pandas as pd
        except ImportError:
            raise ImportError("pandas required for DataFrame conversion")
        
        data = {'horizon': np.arange(self.horizons + 1),
                'irf': self.irf[:, response_var, impulse_var]}
        
        if self.irf_lower is not None:
            data['lower'] = self.irf_lower[:, response_var, impulse_var]
            data['upper'] = self.irf_upper[:, response_var, impulse_var]
        
        return pd.DataFrame(data)


@dataclass
class FEVDResults:
    """
    Results from Forecast Error Variance Decomposition.
    
    Attributes
    ----------
    fevd : np.ndarray
        FEVD of shape (horizons+1, n, n)
        fevd[h, i, j] = proportion of variance of variable i explained by shock j at horizon h
    horizons : int
        Number of horizons
    n_vars : int
        Number of variables
    """
    fevd: np.ndarray
    horizons: int
    n_vars: int
    
    def plot(self, variable: Optional[int] = None,
             figsize: Tuple[int, int] = (12, 6)) -> None:
        """
        Plot forecast error variance decomposition.
        
        Parameters
        ----------
        variable : int, optional
            Index of variable to plot (if None, plot all)
        figsize : tuple, default (12, 6)
            Figure size
        """
        try:
            import matplotlib.pyplot as plt
        except ImportError:
            raise ImportError("matplotlib required for plotting")
        
        n = self.n_vars
        horizons = np.arange(self.horizons + 1)
        
        if variable is not None:
            # Single variable
            fig, ax = plt.subplots(figsize=(8, 5))
            bottom = np.zeros(self.horizons + 1)
            for j in range(n):
                ax.bar(horizons, self.fevd[:, variable, j], bottom=bottom,
                       label=f'Var {j+1}', alpha=0.8)
                bottom += self.fevd[:, variable, j]
            ax.set_xlabel('Horizon')
            ax.set_ylabel('Proportion')
            ax.set_title(f'FEVD of Variable {variable+1}')
            ax.legend()
            ax.set_ylim(0, 1)
        else:
            # All variables
            fig, axes = plt.subplots(1, n, figsize=figsize)
            for i in range(n):
                ax = axes[i] if n > 1 else axes
                bottom = np.zeros(self.horizons + 1)
                for j in range(n):
                    ax.bar(horizons, self.fevd[:, i, j], bottom=bottom,
                           label=f'Var {j+1}', alpha=0.8)
                    bottom += self.fevd[:, i, j]
                ax.set_xlabel('Horizon')
                if i == 0:
                    ax.set_ylabel('Proportion')
                ax.set_title(f'Var {i+1}')
                ax.set_ylim(0, 1)
            
            axes[-1 if n > 1 else axes].legend(bbox_to_anchor=(1.05, 1), loc='upper left')
            plt.tight_layout()
        
        plt.show()
    
    def to_dataframe(self, variable: int):
        """
        Convert FEVD to pandas DataFrame.
        
        Parameters
        ----------
        variable : int
            Variable index
            
        Returns
        -------
        pd.DataFrame
            DataFrame with FEVD values
        """
        try:
            import pandas as pd
        except ImportError:
            raise ImportError("pandas required for DataFrame conversion")
        
        data = {'horizon': np.arange(self.horizons + 1)}
        for j in range(self.n_vars):
            data[f'Var_{j+1}'] = self.fevd[:, variable, j]
        
        return pd.DataFrame(data)


def compute_irf(results, horizons: int = 20,
                orthogonalized: bool = False,
                method: str = 'cholesky',
                ordering: Optional[List[int]] = None) -> IRFResults:
    """
    Compute impulse response functions for Mixed VECM.
    
    Parameters
    ----------
    results : MixedVECMResults
        Estimation results
    horizons : int, default 20
        Number of horizons
    orthogonalized : bool, default False
        Whether to orthogonalize shocks
    method : str, default 'cholesky'
        Orthogonalization method: 'cholesky' or 'generalized'
    ordering : list of int, optional
        Variable ordering for Cholesky decomposition
        
    Returns
    -------
    IRFResults
        IRF results
        
    Notes
    -----
    For a VECM, we first convert to VAR representation and then compute IRFs.
    
    The VECM:
        ΔY_t = αβ'Y_{t-1} + Σ Γ_l ΔY_{t-l} + ε_t
        
    can be written as a VAR in levels:
        Y_t = Φ_1 Y_{t-1} + Φ_2 Y_{t-2} + ... + ε_t
        
    where Φ_1 = I + αβ' + Γ_1, Φ_l = Γ_l - Γ_{l-1} for l > 1.
    
    Examples
    --------
    >>> from mixvecm import estimate_mixed_vecm, compute_irf
    >>> results = estimate_mixed_vecm(Y, k_ar_diff=2)
    >>> irf = compute_irf(results, horizons=20, orthogonalized=True)
    >>> irf.plot()
    """
    n = results.n_vars
    Pi = results.Pi
    Gamma = results.Gamma
    Omega = results.Omega
    p = results.n_lags
    
    # Convert VECM to VAR representation
    Phi = _vecm_to_var(Pi, Gamma, n, p)
    
    # Compute IRF
    if orthogonalized:
        if method == 'cholesky':
            irf = _compute_cholesky_irf(Phi, Omega, horizons, ordering)
        else:
            irf = _compute_generalized_irf(Phi, Omega, horizons)
    else:
        irf = _compute_structural_irf(Phi, horizons)
    
    return IRFResults(
        irf=irf,
        horizons=horizons,
        n_vars=n,
        orthogonalized=orthogonalized,
        method=method if orthogonalized else 'structural'
    )


def compute_orthogonalized_irf(results, horizons: int = 20,
                                ordering: Optional[List[int]] = None,
                                bootstrap: bool = False,
                                n_bootstrap: int = 500,
                                confidence: float = 0.95,
                                seed: Optional[int] = None) -> IRFResults:
    """
    Compute orthogonalized IRF with optional bootstrap confidence intervals.
    
    Parameters
    ----------
    results : MixedVECMResults
        Estimation results
    horizons : int, default 20
        Number of horizons
    ordering : list of int, optional
        Variable ordering for Cholesky
    bootstrap : bool, default False
        Whether to compute bootstrap confidence intervals
    n_bootstrap : int, default 500
        Number of bootstrap replications
    confidence : float, default 0.95
        Confidence level
    seed : int, optional
        Random seed for bootstrap
        
    Returns
    -------
    IRFResults
        IRF results with optional confidence intervals
    """
    n = results.n_vars
    Pi = results.Pi
    Gamma = results.Gamma
    Omega = results.Omega
    p = results.n_lags
    
    # Convert VECM to VAR
    Phi = _vecm_to_var(Pi, Gamma, n, p)
    
    # Point estimate
    irf = _compute_cholesky_irf(Phi, Omega, horizons, ordering)
    
    # Bootstrap confidence intervals
    irf_lower = None
    irf_upper = None
    
    if bootstrap:
        if seed is not None:
            np.random.seed(seed)
        
        irf_bootstrap = np.zeros((n_bootstrap, horizons + 1, n, n))
        
        residuals = results.residuals
        T = residuals.shape[0]
        
        for b in range(n_bootstrap):
            # Resample residuals
            idx = np.random.choice(T, size=T, replace=True)
            resid_boot = residuals[idx]
            
            # Compute bootstrap covariance
            Omega_boot = resid_boot.T @ resid_boot / T
            
            # Perturb coefficients (simplified)
            # Full implementation would re-estimate model
            Phi_boot = [phi + np.random.randn(*phi.shape) * 0.01 for phi in Phi]
            
            irf_bootstrap[b] = _compute_cholesky_irf(Phi_boot, Omega_boot, horizons, ordering)
        
        # Compute confidence intervals
        alpha = 1 - confidence
        irf_lower = np.percentile(irf_bootstrap, 100 * alpha / 2, axis=0)
        irf_upper = np.percentile(irf_bootstrap, 100 * (1 - alpha / 2), axis=0)
    
    return IRFResults(
        irf=irf,
        horizons=horizons,
        n_vars=n,
        orthogonalized=True,
        method='cholesky',
        irf_lower=irf_lower,
        irf_upper=irf_upper,
        confidence=confidence if bootstrap else None
    )


def compute_fevd(results, horizons: int = 20,
                 ordering: Optional[List[int]] = None) -> FEVDResults:
    """
    Compute forecast error variance decomposition.
    
    Parameters
    ----------
    results : MixedVECMResults
        Estimation results
    horizons : int, default 20
        Number of horizons
    ordering : list of int, optional
        Variable ordering for Cholesky
        
    Returns
    -------
    FEVDResults
        FEVD results
        
    Notes
    -----
    The FEVD shows the proportion of the forecast error variance of each
    variable that can be attributed to shocks in each variable.
    
    For horizon h, variable i, shock j:
        FEVD_{i,j}(h) = Σ_{s=0}^{h} θ_{i,j}(s)² / Σ_{k=1}^{n} Σ_{s=0}^{h} θ_{i,k}(s)²
        
    where θ_{i,j}(s) is the orthogonalized IRF.
    """
    n = results.n_vars
    Pi = results.Pi
    Gamma = results.Gamma
    Omega = results.Omega
    p = results.n_lags
    
    # Convert VECM to VAR
    Phi = _vecm_to_var(Pi, Gamma, n, p)
    
    # Compute orthogonalized IRF
    irf = _compute_cholesky_irf(Phi, Omega, horizons, ordering)
    
    # Compute FEVD
    fevd = np.zeros((horizons + 1, n, n))
    
    # Cumulative squared IRF
    irf_sq_cumsum = np.cumsum(irf ** 2, axis=0)
    
    for h in range(horizons + 1):
        for i in range(n):
            total_var = np.sum(irf_sq_cumsum[h, i, :])
            if total_var > 0:
                fevd[h, i, :] = irf_sq_cumsum[h, i, :] / total_var
            else:
                fevd[h, i, :] = 1 / n  # Equal contribution if zero variance
    
    return FEVDResults(
        fevd=fevd,
        horizons=horizons,
        n_vars=n
    )


def _vecm_to_var(Pi: np.ndarray, Gamma: List[np.ndarray],
                 n: int, p: int) -> List[np.ndarray]:
    """
    Convert VECM representation to VAR representation.
    
    The VECM:
        ΔY_t = Π Y_{t-1} + Σ Γ_l ΔY_{t-l} + ε_t
        
    becomes VAR:
        Y_t = Φ_1 Y_{t-1} + ... + Φ_p Y_{t-p} + ε_t
        
    where:
        Φ_1 = I + Π + Γ_1
        Φ_l = Γ_l - Γ_{l-1} for 1 < l < p
        Φ_p = -Γ_{p-1}
    
    Parameters
    ----------
    Pi : np.ndarray
        Long-run matrix (n × n)
    Gamma : list of np.ndarray
        Short-run matrices [Γ_1, ..., Γ_{p-1}]
    n : int
        Number of variables
    p : int
        VAR lag order
        
    Returns
    -------
    list of np.ndarray
        VAR coefficient matrices [Φ_1, ..., Φ_p]
    """
    Phi = []
    
    # Φ_1 = I + Π + Γ_1 (or I + Π if no Gamma)
    Phi_1 = np.eye(n) + Pi
    if len(Gamma) > 0:
        Phi_1 = Phi_1 + Gamma[0]
    Phi.append(Phi_1)
    
    # Φ_l = Γ_l - Γ_{l-1} for 1 < l < p
    for l in range(1, len(Gamma)):
        Phi_l = Gamma[l] - Gamma[l-1]
        Phi.append(Phi_l)
    
    # Φ_p = -Γ_{p-1}
    if len(Gamma) > 0:
        Phi_p = -Gamma[-1]
        Phi.append(Phi_p)
    
    return Phi


def _compute_structural_irf(Phi: List[np.ndarray], 
                            horizons: int) -> np.ndarray:
    """
    Compute non-orthogonalized (structural) IRF.
    
    Parameters
    ----------
    Phi : list of np.ndarray
        VAR coefficient matrices
    horizons : int
        Number of horizons
        
    Returns
    -------
    np.ndarray
        IRF of shape (horizons+1, n, n)
    """
    n = Phi[0].shape[0]
    p = len(Phi)
    
    # IRF at each horizon
    irf = np.zeros((horizons + 1, n, n))
    
    # Impact (h=0)
    irf[0] = np.eye(n)
    
    # Recursive computation
    # Θ_h = Σ_{l=1}^{min(h,p)} Θ_{h-l} Φ_l
    for h in range(1, horizons + 1):
        for l in range(min(h, p)):
            if l < len(Phi):
                irf[h] += irf[h - l - 1] @ Phi[l]
    
    return irf


def _compute_cholesky_irf(Phi: List[np.ndarray], Omega: np.ndarray,
                          horizons: int, 
                          ordering: Optional[List[int]] = None) -> np.ndarray:
    """
    Compute orthogonalized IRF using Cholesky decomposition.
    
    Parameters
    ----------
    Phi : list of np.ndarray
        VAR coefficient matrices
    Omega : np.ndarray
        Residual covariance matrix
    horizons : int
        Number of horizons
    ordering : list of int, optional
        Variable ordering
        
    Returns
    -------
    np.ndarray
        Orthogonalized IRF of shape (horizons+1, n, n)
    """
    n = Phi[0].shape[0]
    
    # Reorder if specified
    if ordering is not None:
        order = np.array(ordering)
        Omega_reord = Omega[np.ix_(order, order)]
        Phi_reord = [phi[np.ix_(order, order)] for phi in Phi]
    else:
        order = np.arange(n)
        Omega_reord = Omega
        Phi_reord = Phi
    
    # Cholesky decomposition: Ω = PP'
    try:
        P = cholesky(Omega_reord, lower=True)
    except np.linalg.LinAlgError:
        # If not positive definite, use eigenvalue decomposition
        eigvals, eigvecs = np.linalg.eigh(Omega_reord)
        eigvals = np.maximum(eigvals, 1e-10)
        P = eigvecs @ np.diag(np.sqrt(eigvals))
    
    # Compute structural IRF first
    irf_struct = _compute_structural_irf(Phi_reord, horizons)
    
    # Orthogonalize: Θ_h P
    irf_orth = np.zeros((horizons + 1, n, n))
    for h in range(horizons + 1):
        irf_orth[h] = irf_struct[h] @ P
    
    # Reorder back if needed
    if ordering is not None:
        inv_order = np.argsort(order)
        irf_final = np.zeros_like(irf_orth)
        for h in range(horizons + 1):
            irf_final[h] = irf_orth[h][np.ix_(inv_order, inv_order)]
        return irf_final
    
    return irf_orth


def _compute_generalized_irf(Phi: List[np.ndarray], Omega: np.ndarray,
                              horizons: int) -> np.ndarray:
    """
    Compute generalized (order-invariant) IRF.
    
    The generalized IRF is defined as:
        θ^G_{i,j}(h) = σ_j^{-1} e_i' Θ_h Σ e_j
        
    where σ_j is the standard deviation of shock j.
    
    Parameters
    ----------
    Phi : list of np.ndarray
        VAR coefficient matrices
    Omega : np.ndarray
        Residual covariance matrix
    horizons : int
        Number of horizons
        
    Returns
    -------
    np.ndarray
        Generalized IRF of shape (horizons+1, n, n)
    """
    n = Phi[0].shape[0]
    
    # Compute structural IRF
    irf_struct = _compute_structural_irf(Phi, horizons)
    
    # Standard deviations
    sigma = np.sqrt(np.diag(Omega))
    
    # Generalized IRF
    irf_gen = np.zeros((horizons + 1, n, n))
    for h in range(horizons + 1):
        for j in range(n):
            irf_gen[h, :, j] = irf_struct[h] @ Omega[:, j] / sigma[j]
    
    return irf_gen


def historical_decomposition(results, Y: np.ndarray,
                              ordering: Optional[List[int]] = None
                              ) -> Tuple[np.ndarray, np.ndarray]:
    """
    Compute historical decomposition of time series.
    
    Parameters
    ----------
    results : MixedVECMResults
        Estimation results
    Y : np.ndarray
        Time series data
    ordering : list of int, optional
        Variable ordering for Cholesky
        
    Returns
    -------
    contribution : np.ndarray
        Contribution of each shock to each variable (T, n, n)
    baseline : np.ndarray
        Baseline (deterministic) component (T, n)
    """
    n = results.n_vars
    T = Y.shape[0]
    Pi = results.Pi
    Gamma = results.Gamma
    Omega = results.Omega
    p = results.n_lags
    
    # Convert to VAR
    Phi = _vecm_to_var(Pi, Gamma, n, p)
    
    # Get structural shocks
    residuals = results.residuals
    
    # Cholesky decomposition
    try:
        P = cholesky(Omega, lower=True)
    except np.linalg.LinAlgError:
        eigvals, eigvecs = np.linalg.eigh(Omega)
        eigvals = np.maximum(eigvals, 1e-10)
        P = eigvecs @ np.diag(np.sqrt(eigvals))
    
    P_inv = inv(P)
    
    # Structural shocks
    shocks = (P_inv @ residuals.T).T
    
    # Compute IRF
    irf = _compute_cholesky_irf(Phi, Omega, T, ordering)
    
    # Historical decomposition
    T_resid = residuals.shape[0]
    contribution = np.zeros((T_resid, n, n))
    
    for t in range(T_resid):
        for j in range(n):
            for s in range(t + 1):
                if t - s < irf.shape[0]:
                    contribution[t, :, j] += irf[t - s, :, j] * shocks[s, j]
    
    # Baseline (what remains)
    Y_adj = Y[p + 1:][:T_resid]
    baseline = Y_adj - np.sum(contribution, axis=2)
    
    return contribution, baseline
