"""
Econometric Modeling using Statsmodels

Author: RSK World
Website: https://rskworld.in
Email: help@rskworld.in
Phone: +91 93305 39277
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from statsmodels.tsa.vector_ar.var_model import VAR
from statsmodels.tsa.statespace.varmax import VARMAX
from statsmodels.tsa.regime_switching import MarkovRegression
from statsmodels.stats.stattools import coint
from statsmodels.tsa.vector_ar.irf import IRAnalysis
from statsmodels.tsa.vector_ar.vecm import coint_johansen
import warnings
warnings.filterwarnings('ignore')


class EconometricModel:
    """
    Econometric Modeling Tools
    
    Author: RSK World
    Website: https://rskworld.in
    Email: help@rskworld.in
    Phone: +91 93305 39277
    """
    
    def __init__(self):
        self.model = None
        self.results = None
        self.data = None
    
    def fit_var(self, data, maxlags=None, ic='aic'):
        """
        Fit Vector Autoregression (VAR) Model
        
        Parameters:
        -----------
        data : DataFrame
            Multivariate time series data
        maxlags : int
            Maximum number of lags
        ic : str
            Information criterion ('aic', 'bic', 'fpe', 'hqic')
        """
        self.data = data
        model = VAR(data)
        
        if maxlags is None:
            maxlags = min(12, int(len(data) / 3))
        
        # Select optimal lag order
        lag_order = model.select_order(maxlags=maxlags)
        optimal_lag = getattr(lag_order, ic)
        
        print(f"Optimal lag order ({ic.upper()}): {optimal_lag}")
        print("\nLag Selection Results:")
        print(lag_order.summary())
        
        # Fit model with optimal lag
        self.model = model
        self.results = model.fit(optimal_lag)
        
        return self.results
    
    def var_summary(self):
        """Print VAR model summary"""
        if self.results is not None:
            print(self.results.summary())
        else:
            print("Model not fitted yet. Call fit_var() first.")
    
    def var_forecast(self, steps=10):
        """Generate VAR forecasts"""
        if self.results is None:
            raise ValueError("Model not fitted yet. Call fit_var() first.")
        
        forecast = self.results.forecast(self.data.values[-self.results.k_ar:], 
                                        steps=steps)
        forecast_index = pd.date_range(start=self.data.index[-1], 
                                      periods=steps+1, freq='D')[1:]
        forecast_df = pd.DataFrame(forecast, index=forecast_index, 
                                 columns=self.data.columns)
        
        return forecast_df
    
    def var_irf(self, periods=10, orth=False):
        """
        Impulse Response Functions
        
        Parameters:
        -----------
        periods : int
            Number of periods
        orth : bool
            Use orthogonalized IRF
        """
        if self.results is None:
            raise ValueError("Model not fitted yet. Call fit_var() first.")
        
        irf = self.results.irf(periods)
        
        if orth:
            irf = irf.orth_irfs
        
        # Plot IRF
        n_vars = len(self.data.columns)
        fig, axes = plt.subplots(n_vars, n_vars, figsize=(15, 12))
        
        for i in range(n_vars):
            for j in range(n_vars):
                axes[i, j].plot(irf[:, i, j])
                axes[i, j].axhline(y=0, color='r', linestyle='--', alpha=0.5)
                axes[i, j].set_title(f'{self.data.columns[j]} -> {self.data.columns[i]}')
                axes[i, j].grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.show()
        
        return irf
    
    def var_granger_causality(self):
        """Test for Granger Causality"""
        if self.results is None:
            raise ValueError("Model not fitted yet. Call fit_var() first.")
        
        print("Granger Causality Tests:")
        print("=" * 70)
        
        gc = self.results.test_causality(self.data.columns[0], 
                                        self.data.columns[1:], 
                                        kind='f')
        print(gc.summary())
        
        return gc
    
    def cointegration_test(self, data, det_order=-1, k_ar_diff=1):
        """
        Test for Cointegration using Johansen Test
        
        Parameters:
        -----------
        data : DataFrame
            Multivariate time series
        det_order : int
            Deterministic order (-1: no constant, 0: constant, 1: linear trend)
        k_ar_diff : int
            Number of lags
        """
        print("Johansen Cointegration Test:")
        print("=" * 70)
        
        result = coint_johansen(data, det_order, k_ar_diff)
        
        print(f"\nEigenvalues: {result.eig}")
        print(f"Trace Statistics: {result.lr1}")
        print(f"Critical Values (Trace): {result.cvt[:, 0]}")
        print(f"Max Eigenvalue Statistics: {result.lr2}")
        print(f"Critical Values (Max Eigenvalue): {result.cvm[:, 0]}")
        
        # Test results
        trace_pvalues = result.lr1 < result.cvt[:, 0]
        max_eig_pvalues = result.lr2 < result.cvm[:, 0]
        
        print("\nCointegration Results:")
        for i in range(len(trace_pvalues)):
            print(f"r <= {i}: Trace={trace_pvalues[i]}, Max Eigenvalue={max_eig_pvalues[i]}")
        
        return result
    
    def engle_granger_cointegration(self, y, x):
        """
        Engle-Granger Cointegration Test
        
        Parameters:
        -----------
        y : array-like
            First time series
        x : array-like
            Second time series
        """
        from statsmodels.regression.linear_model import OLS
        from statsmodels.tsa.stattools import adfuller
        
        # Step 1: Estimate cointegrating regression
        model = OLS(y, x).fit()
        residuals = model.resid
        
        # Step 2: Test residuals for stationarity
        adf_result = adfuller(residuals)
        
        print("Engle-Granger Cointegration Test:")
        print("=" * 70)
        print(f"ADF Statistic: {adf_result[0]:.6f}")
        print(f"p-value: {adf_result[1]:.6f}")
        print(f"Critical Values:")
        for key, value in adf_result[4].items():
            print(f"  {key}: {value:.6f}")
        
        if adf_result[1] < 0.05:
            print("\nResult: Series are cointegrated (reject null hypothesis)")
        else:
            print("\nResult: Series are not cointegrated (fail to reject null hypothesis)")
        
        return {
            'adf_statistic': adf_result[0],
            'p_value': adf_result[1],
            'critical_values': adf_result[4],
            'is_cointegrated': adf_result[1] < 0.05,
            'cointegrating_vector': model.params
        }
    
    def fit_varmax(self, data, order=(1, 1), trend='c'):
        """
        Fit VARMAX Model
        
        Parameters:
        -----------
        data : DataFrame
            Multivariate time series
        order : tuple
            (p, q) order for VARMAX
        trend : str
            'c' for constant, 'nc' for no constant
        """
        self.data = data
        self.model = VARMAX(data, order=order, trend=trend)
        self.results = self.model.fit(maxiter=1000, disp=False)
        
        return self.results
    
    def varmax_summary(self):
        """Print VARMAX summary"""
        if self.results is not None:
            print(self.results.summary())
        else:
            print("Model not fitted yet. Call fit_varmax() first.")
    
    def forecast_variance_decomposition(self, periods=10):
        """Forecast Error Variance Decomposition"""
        if self.results is None:
            raise ValueError("Model not fitted yet. Call fit_var() first.")
        
        fevd = self.results.fevd(periods)
        
        # Plot FEVD
        n_vars = len(self.data.columns)
        fig, axes = plt.subplots(n_vars, 1, figsize=(10, 4*n_vars))
        
        for i, var in enumerate(self.data.columns):
            fevd_data = fevd[var]
            fevd_data.plot(kind='bar', stacked=True, ax=axes[i])
            axes[i].set_title(f'Forecast Error Variance Decomposition: {var}')
            axes[i].set_ylabel('Proportion')
            axes[i].legend(bbox_to_anchor=(1.05, 1), loc='upper left')
            axes[i].grid(True, alpha=0.3, axis='y')
        
        plt.tight_layout()
        plt.show()
        
        return fevd


if __name__ == "__main__":
    # Example usage
    print("Econometric Modeling Example")
    print("=" * 50)
    
    # Generate sample multivariate time series
    np.random.seed(42)
    dates = pd.date_range('2020-01-01', periods=200, freq='D')
    
    # Create correlated time series
    y1 = np.cumsum(np.random.randn(200) * 0.1) + 100
    y2 = y1 + np.random.randn(200) * 0.5  # Cointegrated with y1
    y3 = np.cumsum(np.random.randn(200) * 0.1) + 50
    
    data = pd.DataFrame({
        'Series1': y1,
        'Series2': y2,
        'Series3': y3
    }, index=dates)
    
    # Create econometric model
    econ_model = EconometricModel()
    
    # Fit VAR model
    print("\n" + "="*70)
    econ_model.fit_var(data, maxlags=5)
    econ_model.var_summary()
    
    # Test for cointegration
    print("\n" + "="*70)
    econ_model.cointegration_test(data[['Series1', 'Series2']])
    
    # Engle-Granger test
    print("\n" + "="*70)
    econ_model.engle_granger_cointegration(data['Series1'], data['Series2'])
    
    # Impulse Response Functions
    print("\n" + "="*70)
    econ_model.var_irf(periods=10)