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

class OnlineSplineForecaster:
    """
    Online time series forecasting using 4th degree B‑spline basis
    and recursive ridge regression with exponential forgetting.
    """
    def __init__(self, t_train, degree=4, n_knots=25, alpha=1.0, forget=0.995):
        """
        t_train : array of observed time points (used to set knots)
        degree  : spline degree (4 recommended)
        n_knots : number of interior knots
        alpha   : ridge regularisation parameter
        forget  : forgetting factor (0<forget<=1); lower = faster adaptation
        """
        self.degree = degree
        self.alpha = alpha
        self.forget = forget
        
        # Define knots covering the training time range
        self.t_min = t_train.min()
        self.t_max = t_train.max()
        self.knots = np.linspace(self.t_min, self.t_max, n_knots)
        
        # For B‑spline we need to augment knots (repeated at ends)
        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
        
        # Initialise recursive ridge
        self.P = np.eye(self.n_basis) / alpha
        self.w = np.zeros(self.n_basis)
        
        # Store history for diagnostics (optional)
        self.w_history = []
        
    def _basis(self, t):
        """Evaluate all basis functions at time(s) t (scalar or array)."""
        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  # shape (len(t), n_basis)
    
    def update(self, t_current, y_next):
        """
        Update the model with a new observation.
        t_current : time at which we make the prediction
        y_next    : actual value observed at next time step (t_current + dt)
        """
        # Feature vector at current time
        x = self._basis(t_current).flatten()
        # Recursive ridge update (with forgetting)
        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)
        self.w_history.append(self.w.copy())
        return self
    
    def predict_next(self, t_current):
        """Predict the next value at time t_current (no update)."""
        x = self._basis(t_current).flatten()
        return x @ self.w
    
    def forecast(self, t_future):
        """
        Forecast the series at future time(s) t_future using the latest coefficients.
        Returns prediction and optionally standard deviation.
        """
        X_future = self._basis(t_future)  # shape (n_future, n_basis)
        y_pred = X_future @ self.w
        # Prediction variance: x^T P x (where P is the parameter covariance)
        # Using the current P (after last update). Note: P is inflated by forgetting,
        # but this gives a heuristic uncertainty.
        var_pred = np.array([x @ self.P @ x for x in X_future])
        std_pred = np.sqrt(var_pred)
        return y_pred, std_pred


# ------------------------------------------------------------
# Example usage on your synthetic data
# ------------------------------------------------------------
if __name__ == "__main__":
    # Generate same series
    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 into training (first 80%) and test (last 20%)
    split_idx = int(0.8 * len(t))
    t_train, t_test = t[:split_idx], t[split_idx:]
    y_train, y_test = series[:split_idx], series[split_idx:]
    
    # Initialise forecaster
    forecaster = OnlineSplineForecaster(t_train, degree=4, n_knots=25,
                                        alpha=0.1, forget=0.995)
    
    # Online training: update sequentially
    for i in range(1, len(y_train) - 1):
        forecaster.update(t_train[i], y_train[i+1])  # predict i+1 from i
    
    # Now forecast the next 50 steps beyond the training set
    dt = t[1] - t[0]  # assume uniform spacing
    t_forecast = np.arange(t_train[-1] + dt,
                           t_train[-1] + 51*dt,
                           dt)
    y_forecast, y_std = forecaster.forecast(t_forecast)
    
    # Compare with actual test values (for evaluation)
    y_test_aligned = y_test[:len(t_forecast)]
    
    # 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_forecast, y_forecast, 'r--', label='Forecast')
    plt.fill_between(t_forecast, y_forecast - 2*y_std, y_forecast + 2*y_std,
                     color='red', alpha=0.2, label='±2σ uncertainty')
    plt.title("Online spline forecasting with 4th degree B‑spline")
    plt.xlabel("Time")
    plt.ylabel("Value")
    plt.legend()
    plt.show()
    
    # Print forecasted values for first few steps
    print("First 5 forecasted values:")
    for i in range(5):
        print(f"t={t_forecast[i]:.2f}, forecast={y_forecast[i]:.4f} ± {2*y_std[i]:.4f}")