import numpy as np
import matplotlib.pyplot as plt

# =============================================================================
# Fractal Iteration for Regularized Linear Regression (Ridge Regression)
# =============================================================================
#
# This script demonstrates a "fractal iteration": repeated application of an
# affine transformation on a weight vector (which can be seen as a column matrix).
# The transformation is contractive, and its fixed point is the solution of a
# ridge regression problem. The regularization parameter λ (lambda) controls
# the complexity of the model, preventing overfitting to noisy training data.
#
# The iteration is:
#
#   w_{k+1} = A w_k + b
#
# where:
#   A = I - lr * (X^T X + λ I)
#   b = lr * X^T y
#
# This is a contractive map (spectral radius of A < 1) for sufficiently small
# learning rate lr. The fixed point w* = (X^T X + λ I)^{-1} X^T y is the ridge
# regression estimator, which generalises better than ordinary least squares
# when the number of features is large relative to the number of samples.
#
# The term "fractal iteration" comes from the self‑similar nature of the process:
# each iteration applies the same affine transform, and the set of all possible
# weight vectors (the orbit) converges to a single attractor – the fixed point.
# In the context of generalisation, this attractor is the "generalisation limit"
# that does not cause overfitting because the regularisation term λ I shrinks
# the weights, effectively limiting the model's capacity.
# =============================================================================

# -----------------------------------------------------------------------------
# Generate synthetic data: a noisy sinusoid with many polynomial features
# -----------------------------------------------------------------------------
np.random.seed(42)
n_samples = 30
X = np.linspace(0, 2 * np.pi, n_samples).reshape(-1, 1)
y_true = np.sin(X).ravel()
y = y_true + 0.3 * np.random.randn(n_samples)  # add noise

# Create polynomial features up to degree 15 (many features -> overfitting)
degree = 15
X_poly = np.vstack([X.ravel() ** d for d in range(degree + 1)]).T

# Standardise features (important for numerical stability)
mean = X_poly.mean(axis=0)
std = X_poly.std(axis=0)
std[std == 0] = 1.0
X_poly = (X_poly - mean) / std

# -----------------------------------------------------------------------------
# Parameters of the fractal iteration
# -----------------------------------------------------------------------------
lambda_reg = 1.0          # Regularisation strength (higher -> simpler model)
lr = 0.01                 # Learning rate (must be small enough for contraction)
max_iter = 1000           # Number of fractal iterations
tol = 1e-8                # Stopping tolerance

# Initial guess: zero vector (no effect)
w = np.zeros(degree + 1)

# Precompute matrices for the affine map
Xtx = X_poly.T @ X_poly
Xty = X_poly.T @ y
I = np.eye(degree + 1)

A = I - lr * (Xtx + lambda_reg * I)
b = lr * Xty

# -----------------------------------------------------------------------------
# Fractal iteration: repeatedly apply w <- A w + b
# -----------------------------------------------------------------------------
history = []
for i in range(max_iter):
    w_new = A @ w + b
    history.append(w_new.copy())
    if np.linalg.norm(w_new - w) < tol:
        print(f"Converged after {i+1} iterations.")
        break
    w = w_new
else:
    print(f"Reached maximum iterations ({max_iter}).")

# Fixed point obtained by the iteration
w_fixed = w

# -----------------------------------------------------------------------------
# Verify that the fixed point matches the closed‑form ridge solution
# -----------------------------------------------------------------------------
w_ridge = np.linalg.solve(Xtx + lambda_reg * I, Xty)
print("Difference between iterated and closed‑form ridge weights:",
      np.linalg.norm(w_fixed - w_ridge))

# -----------------------------------------------------------------------------
# Visualise the generalisation behaviour: comparison with overfitted OLS
# -----------------------------------------------------------------------------
# Ordinary Least Squares (no regularisation, λ=0)
w_ols = np.linalg.lstsq(X_poly, y, rcond=None)[0]

# Generate dense test points for smooth curves
X_test = np.linspace(0, 2 * np.pi, 300).reshape(-1, 1)
X_test_poly = np.vstack([X_test.ravel() ** d for d in range(degree + 1)]).T
X_test_poly = (X_test_poly - mean) / std

y_pred_ols = X_test_poly @ w_ols
y_pred_ridge = X_test_poly @ w_ridge

plt.figure(figsize=(10, 6))
plt.scatter(X, y, color='black', label='Noisy training data')
plt.plot(X_test, np.sin(X_test), 'k--', label='True function (sin)')
plt.plot(X_test, y_pred_ols, 'r-', label='OLS (overfitted)', linewidth=2)
plt.plot(X_test, y_pred_ridge, 'g-', label=f'Ridge (λ={lambda_reg})', linewidth=2)
plt.xlabel('x')
plt.ylabel('y')
plt.title('Fractal iteration fixed point: a regularised solution that avoids overfitting')
plt.legend()
plt.grid(alpha=0.3)
plt.show()

# -----------------------------------------------------------------------------
# Additional diagnostics: norm of weights and contraction factor
# -----------------------------------------------------------------------------
print("\n--- Model complexity ---")
print(f"Norm of OLS weights        : {np.linalg.norm(w_ols):.3f}")
print(f"Norm of ridge weights (λ={lambda_reg}): {np.linalg.norm(w_ridge):.3f}")

# Spectral radius of A determines contraction speed
eigvals = np.linalg.eigvals(A)
rho = np.max(np.abs(eigvals))
print(f"\nSpectral radius of A: {rho:.6f} (<1 ensures contraction)")

# Show that the iteration is indeed a fractal (self‑similar) process
# by plotting a few intermediate weight trajectories (first two coordinates)
history = np.array(history)
plt.figure(figsize=(8, 5))
plt.plot(history[:, 0], history[:, 1], 'o-', markersize=3, linewidth=1,
         color='blue', alpha=0.6)
plt.scatter(history[0, 0], history[0, 1], color='red', s=80, label='Start')
plt.scatter(history[-1, 0], history[-1, 1], color='green', s=80, label='Fixed point')
plt.xlabel('Weight 0 (bias term)')
plt.ylabel('Weight 1 (linear term)')
plt.title('Fractal iteration in weight space (attractor is the generalisation limit)')
plt.legend()
plt.grid(alpha=0.3)
plt.show()