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

# ------------------------------------------------------------
# 1. Generate a synthetic non‑linear time series
# ------------------------------------------------------------
np.random.seed(42)
t = np.linspace(0, 20, 500)
# True signal: combination of sine, chirp, and slow drift
signal = (np.sin(t) + 0.5 * np.sin(1.5 * t**2 / 20) + 0.1 * t)
# Add small noise
noise = 0.05 * np.random.randn(len(t))
series = signal + noise

plt.figure(figsize=(12, 4))
plt.plot(t, series, label='observed series')
plt.title('Synthetic time series')
plt.legend()
plt.show()

# ------------------------------------------------------------
# 2. Build 4th‑degree B‑spline basis functions
# ------------------------------------------------------------
degree = 4
# Choose knots – here uniformly spaced across the whole time axis
n_knots = 20
knots = np.linspace(t.min(), t.max(), n_knots)
# For a degree d B‑spline we need to augment knots to avoid matrix singularities
knots_aug = np.r_[np.tile(knots[0], degree), knots, np.tile(knots[-1], degree)]

def spline_basis_matrix(x, knots, degree):
    """
    Evaluate all B‑spline basis functions of given degree at points x.
    Returns a matrix of shape (len(x), n_basis) where n_basis = len(knots) + degree - 1.
    """
    n_basis = len(knots) + degree - 1
    basis = np.zeros((len(x), n_basis))
    for i in range(n_basis):
        # Build a BSpline object that returns only the i‑th basis function
        # We do this by setting all coefficients to 0 except the i‑th
        coef = np.zeros(n_basis)
        coef[i] = 1.0
        spl = BSpline(knots_aug, coef, degree, extrapolate=False)
        # Evaluate only where the spline is defined (use 0 elsewhere)
        basis[:, i] = spl(x)
    # Replace NaN (from extrapolation) with 0
    basis = np.nan_to_num(basis)
    return basis

# ------------------------------------------------------------
# 3. Prepare training data: sliding window to predict next value
# ------------------------------------------------------------
window_size = 10   # use the last 10 time points to predict the next one
X_list = []
y_list = []

for i in range(window_size, len(series) - 1):
    # Input: the last 'window_size' actual values
    window = series[i - window_size : i]
    # Target: the very next value
    target = series[i]
    # For each window, we evaluate the spline basis at the **time index** of the target?
    # In a pure spline‑as‑predictor, we can use the current window to form a feature vector.
    # A simple approach: treat the window as a set of evaluation points and take the mean
    # of the basis functions across that window as features.
    # (More sophisticated: learn a mapping from the window to the basis weights.)
    # Here we use: features = average of basis functions over the window.
    # This makes the linear model learn a weighted combination of "activity" in the window.
    time_points_in_window = t[i - window_size : i]
    basis_vals = spline_basis_matrix(time_points_in_window, knots, degree)
    features = np.mean(basis_vals, axis=0)
    X_list.append(features)
    y_list.append(target)

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

print(f"Feature matrix shape: {X.shape} (samples × basis functions)")
print(f"Number of basis functions: {X.shape[1]}")

# ------------------------------------------------------------
# 4. Train initial linear model (spline weight learning)
# ------------------------------------------------------------
model = LinearRegression()
model.fit(X, y)

# ------------------------------------------------------------
# 5. ODE‑CCT inspired adaptive retraining
# ------------------------------------------------------------
def prediction_entropy(errors):
    """Estimate entropy from recent errors (normalised histogram)."""
    hist, _ = np.histogram(errors, bins='auto', density=True)
    prob = hist / hist.sum()
    prob = prob[prob > 0]
    entropy = -np.sum(prob * np.log(prob + 1e-12))
    return entropy

# We'll simulate online prediction: walk through the test set (hold‑out part)
train_split = int(0.8 * len(X))
X_train, X_test = X[:train_split], X[train_split:]
y_train, y_test = y[:train_split], y[train_split:]

# Re‑train on training set (initial understanding)
model.fit(X_train, y_train)

# Online prediction with adaptive retraining
window_errors = []      # store recent errors to compute entropy
error_history = []      # all errors for final evaluation
predictions = []
entropy_threshold = 1.5  # tune based on your data – try 1.0…2.0
retrain_window = 50      # number of recent samples used to recompute entropy

for idx in range(len(X_test)):
    # Predict next value
    x_current = X_test[idx].reshape(1, -1)
    y_pred = model.predict(x_current)[0]
    predictions.append(y_pred)
    true_val = y_test[idx]
    error = true_val - y_pred
    error_history.append(error)
    
    # Keep a sliding window of the last 'retrain_window' errors
    window_errors.append(error)
    if len(window_errors) > retrain_window:
        window_errors.pop(0)
    
    # Every 10 steps, compute entropy and decide to retrain
    if idx % 10 == 0 and len(window_errors) >= 10:
        ent = prediction_entropy(window_errors)
        print(f"Step {idx}, error entropy = {ent:.3f}")
        if ent > entropy_threshold:
            print(f"  -> Entropy high! Retraining on most recent {retrain_window} samples.")
            # Collapse: retrain model using the last 'retrain_window' (X, y) pairs
            # We need the actual features & targets from recent history.
            # Here we simply use the latest 'retrain_window' samples from the test set (simulate online)
            start = max(0, idx - retrain_window)
            X_recent = X_test[start:idx+1]   # include current if needed
            y_recent = y_test[start:idx+1]
            if len(X_recent) > 10:  # enough to retrain
                model.fit(X_recent, y_recent)
                # Clear error window because the system is re‑understood
                window_errors = []
                print("   Retraining finished, error window reset.")

# ------------------------------------------------------------
# 6. Evaluate final performance
# ------------------------------------------------------------
y_true = y_test[:len(predictions)]
mse = mean_squared_error(y_true, predictions)
print(f"\nFinal MSE on test set: {mse:.6f}")

# Plot results
plt.figure(figsize=(12, 5))
plt.plot(y_true, label='True values', alpha=0.7)
plt.plot(predictions, label='Predictions (with adaptive retraining)', alpha=0.7)
plt.title("Time series prediction using a 4th‑degree B‑spline")
plt.legend()
plt.show()

# Plot error distribution
plt.figure(figsize=(10, 4))
plt.hist(error_history, bins=30, alpha=0.7, label='prediction errors')
plt.xlabel('Error')
plt.ylabel('Frequency')
plt.title('Histogram of prediction errors')
plt.legend()
plt.show()