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

class MultiScaleSplineForecaster:
    """
    Forecast using a sum of B‑splines with different knot densities.
    """
    def __init__(self, t_domain, degree=4, n_knots_list=[10, 25, 50], alpha=1.0):
        """
        t_domain : (t_min, t_max) tuple covering training and forecast horizon
        degree   : spline degree (same for all scales)
        n_knots_list : list of number of interior knots per scale
        alpha    : ridge regularisation strength
        """
        self.degree = degree
        self.alpha = alpha
        self.t_min, self.t_max = t_domain
        self.n_knots_list = n_knots_list
        
        # Pre‑compute basis functions for each scale
        self.bases = []  # each element: (knots_aug, n_basis)
        for nk in n_knots_list:
            knots = np.linspace(self.t_min, self.t_max, nk)
            knots_aug = np.r_[np.tile(knots[0], degree),
                              knots,
                              np.tile(knots[-1], degree)]
            n_basis = len(knots) + degree - 1
            self.bases.append((knots_aug, n_basis))
        
        self.total_basis = sum(nb for _, nb in self.bases)
        self.coef_ = None         # will be set during fit/refit
        
    def _basis_matrix(self, t, scale_idx):
        """Evaluate basis of a given scale at times t."""
        knots_aug, n_basis = self.bases[scale_idx]
        t = np.atleast_1d(t)
        basis = np.zeros((len(t), n_basis))
        for i in range(n_basis):
            coef = np.zeros(n_basis)
            coef[i] = 1.0
            spl = BSpline(knots_aug, coef, self.degree, extrapolate=True)
            basis[:, i] = spl(t)
        return basis
    
    def _full_basis(self, t):
        """Concatenate all scales' basis functions."""
        t = np.atleast_1d(t)
        X_list = []
        for idx in range(len(self.n_knots_list)):
            X_list.append(self._basis_matrix(t, idx))
        return np.hstack(X_list)
    
    def fit(self, t_train, y_train):
        """Batch fit on historical data."""
        X = self._full_basis(t_train)
        model = Ridge(alpha=self.alpha, fit_intercept=False)
        model.fit(X, y_train)
        self.coef_ = model.coef_
        # Compute residual std for intervals
        y_pred = model.predict(X)
        self.residual_std = np.std(y_train - y_pred)
        return self
    
    def predict(self, t_pred):
        """Predict at times t_pred."""
        X = self._full_basis(t_pred)
        return X @ self.coef_
    
    def forecast_with_interval(self, t_future):
        """Return mean, lower, upper (95% CI)."""
        y_mean = self.predict(t_future)
        margin = 1.96 * self.residual_std
        return y_mean, y_mean - margin, y_mean + margin


# ------------------------------------------------------------
# Example: fix the previous failure
# ------------------------------------------------------------
if __name__ == "__main__":
    np.random.seed(42)
    t = np.linspace(0, 20, 500)
    signal = np.sin(10*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
    split = int(0.8 * len(t))
    t_train, t_test = t[:split], t[split:]
    y_train, y_test = series[:split], series[split:]
    
    # Create forecaster with three scales: coarse, medium, fine
    forecaster = MultiScaleSplineForecaster(
        t_domain=(t.min(), t.max()),
        degree=4,
        n_knots_list=[8, 20, 40],
        alpha=0.5
    )
    
    # Train on all training data
    forecaster.fit(t_train, y_train)
    
    # Forecast next 50 steps (starting after training)
    dt = t[1] - t[0]
    t_future = np.arange(t_train[-1] + dt, t_train[-1] + 51*dt, dt)
    y_mean, y_low, y_high = forecaster.forecast_with_interval(t_future)
    
    # 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_future, y_mean, 'r--', label='Multi‑scale forecast')
    plt.fill_between(t_future, y_low, y_high, color='red', alpha=0.2, label='95% CI')
    plt.title("Multi‑scale spline forecasting (coarse + fine)")
    plt.legend()
    plt.show()
    
    print("First 5 forecasted values (with 95% CI):")
    for i in range(5):
        print(f"t={t_future[i]:.2f}, value={y_mean[i]:.4f} ± {1.96*forecaster.residual_std:.4f}")