import numpy as np
import matplotlib.pyplot as plt
from scipy.interpolate import BSpline
from sklearn.linear_model import Ridge
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import GridSearchCV

# ------------------------------------------------------------
# 1. Generate synthetic time series (same as before)
# ------------------------------------------------------------
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

# ------------------------------------------------------------
# 2. B‑spline basis (4th degree)
# ------------------------------------------------------------
degree = 4
n_knots = 20
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 spline_basis_matrix(x, knots, degree):
    n_basis = len(knots) + degree - 1
    basis = np.zeros((len(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)

# ------------------------------------------------------------
# 3. Feature construction (average basis over sliding window)
# ------------------------------------------------------------
window_size = 10
X_list, y_list = [], []
for i in range(window_size, len(series) - 1):
    window_t = t[i - window_size : i]
    basis_vals = spline_basis_matrix(window_t, knots, degree)
    features = np.mean(basis_vals, axis=0)
    X_list.append(features)
    y_list.append(series[i])

X = np.array(X_list)
y = np.array(y_list)

# Train / test split
split = int(0.8 * len(X))
X_train, X_test = X[:split], X[split:]
y_train, y_test = y[:split], y[split:]

# ------------------------------------------------------------
# 4. Ridge regression with cross‑validated alpha (relaxed spline)
# ------------------------------------------------------------
ridge = Ridge()
param_grid = {'alpha': np.logspace(-3, 3, 20)}
grid = GridSearchCV(ridge, param_grid, cv=5, scoring='neg_mean_squared_error')
grid.fit(X_train, y_train)
best_ridge = grid.best_estimator_
print(f"Best alpha (regularisation strength): {grid.best_params_['alpha']:.4f}")

# ------------------------------------------------------------
# 5. Online prediction with confidence‑based relaxation
# ------------------------------------------------------------
def rolling_error_variance(errors, window=20):
    """Compute variance of recent errors (local accuracy)."""
    if len(errors) < window:
        return np.var(errors) if errors else 1.0
    return np.var(errors[-window:])

predictions = []
confidences = []      # 1 = high confidence, 0 = relaxed (fallback)
error_history = []
window_errors = []

variance_threshold = 0.01   # tune: if error variance > this, relax prediction
fallback_value = np.mean(y_train)   # simple baseline

best_ridge.fit(X_train, y_train)   # initial fit

for idx in range(len(X_test)):
    x_cur = X_test[idx].reshape(1, -1)
    y_pred = best_ridge.predict(x_cur)[0]
    true_val = y_test[idx]
    error = true_val - y_pred
    error_history.append(error)
    window_errors.append(error)
    if len(window_errors) > 50:
        window_errors.pop(0)
    
    # Compute local accuracy (rolling variance)
    local_var = rolling_error_variance(window_errors, window=20)
    
    # Decide confidence
    if local_var < variance_threshold:
        # High accuracy region – trust spline prediction
        predictions.append(y_pred)
        confidences.append(1.0)
    else:
        # Relaxed mode – use fallback (moving average of recent true values)
        # or simply the global mean
        relaxed_pred = np.mean(y_test[max(0, idx-20):idx+1]) if idx > 0 else fallback_value
        predictions.append(relaxed_pred)
        confidences.append(0.0)
        print(f"Step {idx}: local error variance = {local_var:.5f} -> relaxing prediction")

# ------------------------------------------------------------
# 6. Evaluation
# ------------------------------------------------------------
y_true = y_test[:len(predictions)]
mse = mean_squared_error(y_true, predictions)
print(f"\nFinal MSE (with relaxed spline): {mse:.6f}")

# Plot predictions vs true
plt.figure(figsize=(12, 5))
plt.plot(y_true, label='True', alpha=0.7)
plt.plot(predictions, label='Predicted (relaxed spline)', alpha=0.7)
plt.title("Time series prediction – ridge regularised spline + confidence relaxation")
plt.legend()
plt.show()

# Plot confidence over time
plt.figure(figsize=(12, 3))
plt.fill_between(range(len(confidences)), 0, confidences, color='green', alpha=0.3, label='high confidence')
plt.plot(confidences, 'k-', lw=0.5)
plt.ylim(-0.1, 1.1)
plt.ylabel('Confidence')
plt.xlabel('Test step')
plt.title('Where the model trusts its spline prediction (1 = high accuracy)')
plt.legend()
plt.show()