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

# ------------------------------------------------------------
# Generate series (same)
# ------------------------------------------------------------
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 = noise
#series = signal + noise

# ------------------------------------------------------------
# B‑spline basis (4th degree) – fixed over entire time range
# ------------------------------------------------------------
degree = 4
n_knots = 25
knots = np.linspace(t.min(), t.max(), n_knots)
knots_aug = np.r_[np.tile(knots[0], degree), knots, np.tile(knots[-1], degree)]

def basis_at_time(x):
    """Return array of basis values at scalar x (or array)."""
    n_basis = len(knots) + degree - 1
    basis = np.zeros((len(np.atleast_1d(x)), n_basis))
    for i in range(n_basis):
        coef = np.zeros(n_basis)
        coef[i] = 1.0
        spl = BSpline(knots_aug, coef, degree, extrapolate=False)
        basis[:, i] = spl(x)
    return np.nan_to_num(basis)

# ------------------------------------------------------------
# Recursive ridge regression with forgetting factor
# ------------------------------------------------------------
class RecursiveRidge:
    def __init__(self, n_features, alpha=1.0, forget=0.99):
        self.alpha = alpha          # ridge penalty
        self.forget = forget        # forgetting factor (exponential)
        self.P = np.eye(n_features) / alpha   # inverse covariance
        self.w = np.zeros(n_features)
    
    def update(self, x, y):
        """Online update with one new sample (x: feature vector)."""
        x = x.reshape(-1, 1)
        # forgetting: rescale P
        self.P = self.P / self.forget
        # Kalman gain
        gain = self.P @ x / (x.T @ self.P @ x + 1.0)
        # update weights
        self.w = self.w + gain.flatten() * (y - x.T @ self.w)
        # update P (Woodbury)
        self.P = self.P - np.outer(gain, gain) * (x.T @ self.P @ x + 1.0)
        return self.w
    
    def predict(self, x):
        return x @ self.w

# ------------------------------------------------------------
# Online prediction – one step ahead
# ------------------------------------------------------------
n_basis = len(knots) + degree - 1
model = RecursiveRidge(n_features=n_basis, alpha=0.1, forget=0.995)

predictions = []
true_vals = []

# Start after we have at least one basis evaluation
for idx in range(1, len(series) - 1):
    # Feature = basis evaluated at current time t[idx]
    feat = basis_at_time(t[idx])
    # Target = next value in series
    target = series[idx + 1]   # predict ahead
    
    # Store true value for later comparison
    true_vals.append(target)
    
    # Predict using current model
    if idx > 1:
        pred = model.predict(feat)
        predictions.append(pred)
    else:
        predictions.append(0.0)  # initial dummy
    
    # Update model with the (current time, next value) pair
    # But note: we use current basis to predict next value
    # So we train on (feat, target)
    model.update(feat, target)

# Align lengths
predictions = np.array(predictions[1:])   # remove first dummy
true_vals = np.array(true_vals[1:])

# ------------------------------------------------------------
# Evaluate
# ------------------------------------------------------------
mse = np.mean((predictions - true_vals)**2)
print(f"Recursive ridge spline MSE: {mse:.6f}")

plt.figure(figsize=(12, 5))
plt.plot(true_vals[:300], label='True', alpha=0.7)
plt.plot(predictions[:300], label='Predicted (online spline)', alpha=0.7)
plt.title("Stable online prediction using spline basis + recursive ridge")
plt.legend()
plt.show()

# Plot coefficient evolution (optional)
plt.figure(figsize=(12, 3))
plt.plot(np.array(model.w_history).T if hasattr(model, 'w_history') else [])
plt.title("Spline coefficient evolution (exponential forgetting)")
plt.show()