import numpy as np
import matplotlib.pyplot as plt
from scipy.interpolate import BSpline
from sklearn.linear_model import Ridge

class OnlineSplineForecaster:
    """
    Online time series forecasting with 4th degree B‑spline basis.
    - Recursive ridge for online adaptation.
    - For forecasting: refits a static ridge model on recent data.
    """
    def __init__(self, t_train, degree=4, n_knots=25, alpha=1.0, forget=0.995):
        self.degree = degree
        self.alpha = alpha
        self.forget = forget
        self.t_min = t_train.min()
        self.t_max = t_train.max()
        self.knots = np.linspace(self.t_min, self.t_max, n_knots)
        self.knots_aug = np.r_[np.tile(self.knots[0], degree),
                               self.knots,
                               np.tile(self.knots[-1], degree)]
        self.n_basis = len(self.knots) + degree - 1
        
        # Recursive state
        self.P = np.eye(self.n_basis) / alpha
        self.w = np.zeros(self.n_basis)
        
        # Store data for optional refit
        self.t_history = []
        self.y_history = []
        
    def _basis(self, t):
        t = np.atleast_1d(t)
        basis = np.zeros((len(t), self.n_basis))
        for i in range(self.n_basis):
            coef = np.zeros(self.n_basis)
            coef[i] = 1.0
            spl = BSpline(self.knots_aug, coef, self.degree, extrapolate=True)
            basis[:, i] = spl(t)
        return basis
    
    def update(self, t_current, y_next):
        x = self._basis(t_current).flatten()
        # Store for later refit
        self.t_history.append(t_current)
        self.y_history.append(y_next)
        
        # Recursive update
        self.P = self.P / self.forget
        gain = self.P @ x / (x @ self.P @ x + 1.0)
        self.w = self.w + gain * (y_next - x @ self.w)
        self.P = self.P - np.outer(gain, gain) * (x @ self.P @ x + 1.0)
        return self
    
    def predict_next(self, t_current):
        x = self._basis(t_current).flatten()
        return x @ self.w
    
    def refit_for_forecast(self, window_size=100, alpha=None):
        """
        Fit a static ridge model on the most recent `window_size` observations.
        This model will be used for forecasting.
        """
        if alpha is None:
            alpha = self.alpha
        n = len(self.t_history)
        if n < window_size:
            window_size = n
        t_recent = np.array(self.t_history[-window_size:])
        y_recent = np.array(self.y_history[-window_size:])
        X_recent = self._basis(t_recent)
        self.static_model = Ridge(alpha=alpha, fit_intercept=False)
        self.static_model.fit(X_recent, y_recent)
        # Compute residual standard deviation for prediction intervals
        y_pred = self.static_model.predict(X_recent)
        residuals = y_recent - y_pred
        self.residual_std = np.std(residuals)
        return self
    
    def forecast(self, t_future):
        """
        Forecast using the static model (must call refit_for_forecast() first).
        Returns (mean, lower, upper) for 95% interval.
        """
        if not hasattr(self, 'static_model'):
            raise RuntimeError("Call refit_for_forecast() before forecasting.")
        X_future = self._basis(t_future)
        y_mean = self.static_model.predict(X_future)
        margin = 1.96 * self.residual_std
        return y_mean, y_mean - margin, y_mean + margin


# ------------------------------------------------------------
# Example usage (same synthetic data)
# ------------------------------------------------------------
if __name__ == "__main__":
    np.random.seed(42)
    t = np.linspace(0, 20, 500)
    signal = np.sin(t) + 0.5 * np.sin(1.5 * t**2 / 20) + 0.1 * t
    noise = 0.05 * np.random.randn(len(t))
    series = signal + noise
    
    split = int(0.8 * len(t))
    t_train, t_test = t[:split], t[split:]
    y_train, y_test = series[:split], series[split:]
    
    forecaster = OnlineSplineForecaster(t_train, degree=4, n_knots=25,
                                        alpha=0.1, forget=0.995)
    
    # Online training on training set
    for i in range(1, len(y_train) - 1):
        forecaster.update(t_train[i], y_train[i+1])
    
    # Refit static model on recent data (last 100 points)
    forecaster.refit_for_forecast(window_size=100)
    
    # Forecast next 50 steps
    dt = t[1] - t[0]
    t_fore = np.arange(t_train[-1] + dt, t_train[-1] + 51*dt, dt)
    y_mean, y_low, y_high = forecaster.forecast(t_fore)
    
    # Plot
    plt.figure(figsize=(12, 6))
    plt.plot(t_train, y_train, label='Training', alpha=0.5)
    plt.plot(t_test, y_test, label='True test', alpha=0.7)
    plt.plot(t_fore, y_mean, 'r--', label='Forecast')
    plt.fill_between(t_fore, y_low, y_high, color='red', alpha=0.2, label='95% interval')
    plt.title("Stable forecasting with spline + recent ridge refit")
    plt.legend()
    plt.show()
    
    # Print first 5 forecasts
    print("First 5 forecasted values (with 95% CI):")
    for i in range(5):
        print(f"t={t_fore[i]:.2f}, forecast={y_mean[i]:.4f} ± {1.96*forecaster.residual_std:.4f}")