import numpy as np
import pandas as pd
from arch import arch_model
from sklearn.preprocessing import StandardScaler

class LSGarchModel:
    """
    LS-GARCH(1,1) model with Student-t distribution for tail risk estimation.
    Uses OLS mean equation with exogenous features scaled for optimization stability.
    """
    def __init__(self, p: int = 1, q: int = 1, dist: str = 'studentst'):
        self.p = p
        self.q = q
        self.dist = dist
        self.scaler = StandardScaler()
        self.model = None
        self.results = None
        self.params = None
        self.omega = None
        self.alpha = None
        self.beta = None
        self.nu = None

    def fit(self, X: pd.DataFrame, y: pd.Series):
        """Fits the GARCH model. Scales X to prevent ML optimizer overflows."""
        X_scaled = pd.DataFrame(
            self.scaler.fit_transform(X), 
            columns=X.columns, 
            index=X.index
        )
        self.model = arch_model(
            y, 
            x=X_scaled, 
            mean='LS', 
            vol='Garch', 
            p=self.p, 
            q=self.q, 
            dist=self.dist
        )
        self.results = self.model.fit(disp='off')
        
        # Save key params
        self.params = self.results.params
        self.omega = self.params.get('omega', 0.0)
        self.alpha = self.params.get('alpha[1]', 0.0)
        self.beta = self.params.get('beta[1]', 0.0)
        self.nu = self.params.get('nu', None)
        return self

    def predict(self, X: pd.DataFrame) -> pd.Series:
        """Predicts the mean equation (mu_t) out-of-sample using exogenous features."""
        if self.results is None:
            raise ValueError("Model is not fitted yet. Call fit() first.")
        X_scaled = pd.DataFrame(
            self.scaler.transform(X), 
            columns=X.columns, 
            index=X.index
        )
        coeffs = self.params[X_scaled.columns]
        const = self.params.get('Const', 0.0)
        return X_scaled.dot(coeffs) + const

    def summary(self) -> str:
        """Returns the summary text of the GARCH model."""
        if self.results is None:
            return "Model not fitted yet."
        return self.results.summary().as_text()