import numpy as np
import pandas as pd
from scipy.stats import t

def calculate_garch_volatility(residuals: pd.Series, omega: float, alpha: float, beta: float, initial_vol: float) -> pd.Series:
    """
    Recursively calculates GARCH conditional volatility on the full series of residuals.
    Handles leading NaNs by starting the recursion at the first valid index.
    
    Formula:
        sigma^2_t = omega + alpha * e^2_{t-1} + beta * sigma^2_{t-1}
    """
    first_valid = residuals.first_valid_index()
    if first_valid is None:
        return pd.Series(np.nan, index=residuals.index)
        
    valid_res = residuals.loc[first_valid:]
    cond_var = np.zeros(len(valid_res))
    cond_var[0] = initial_vol ** 2
    
    for t_idx in range(1, len(valid_res)):
        prev_res = valid_res.iloc[t_idx - 1]
        if np.isnan(prev_res):
            cond_var[t_idx] = cond_var[t_idx - 1]
        else:
            cond_var[t_idx] = omega + alpha * (prev_res ** 2) + beta * cond_var[t_idx - 1]
        
    vol_series = pd.Series(np.sqrt(cond_var), index=valid_res.index)
    return vol_series.reindex(residuals.index)


def simulate_pvar(day_mean_prices: pd.Series, day_vols: pd.Series, nu: float, 
                  n_simulations: int = 10000, confidence_level: float = 0.95):
    """
    Runs a Monte Carlo path simulation using the Student-t distribution 
    parameterized by GARCH volatility and SARIMA expected price mean.
    
    Returns:
        simulated_paths: (24, N_SIMULATIONS) array
        var_vals: (24,) array containing Value-at-Risk limits for each hour
    """
    n_hours = len(day_mean_prices)
    simulated_paths = np.zeros((n_hours, n_simulations))
    var_vals = np.zeros(n_hours)
    
    # Calculate the target percentile (e.g. 95% confidence -> 5th percentile)
    percentile = 100 * (1.0 - confidence_level)
    
    np.random.seed(42)
    for i in range(n_hours):
        mu_t = day_mean_prices.iloc[i]
        sigma_t = day_vols.iloc[i]
        
        # Draw from Student-t
        draws = t.rvs(df=nu, loc=mu_t, scale=sigma_t, size=n_simulations)
        simulated_paths[i, :] = draws
        var_vals[i] = np.percentile(draws, percentile)
        
    return simulated_paths, var_vals


def get_prosumer_action(var_value: float, mean_value: float) -> int:
    """
    Prosumer Traffic Light control logic:
        2 (Green): Inject (Value-at-Risk is positive, risk of loss is 0)
        1 (Yellow): Buffer (Expected price positive, but Value-at-Risk is negative; buffer energy)
        0 (Red): Curtail (Both expected price and Value-at-Risk are negative; halt injects)
    """
    if var_value > 0:
        return 2
    elif mean_value > 0 and var_value <= 0:
        return 1
    else:
        return 0