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

# =============================================================================
# From a Stochastic 1D Function to a 2D Fractal (Fractal Interpolation)
# =============================================================================
#
# This program implements the reverse of the previous idea:
#   - We start with a 1D stochastic function y = f(i) for i = 0..N-1.
#   - We then construct a 2D fractal curve that interpolates these points
#     using an iterated function system (IFS) with affine transformations.
#
# The resulting curve is a fractal in the plane, whose roughness is controlled
# by the IFS scaling factors. It is the attractor of a contractive IFS.
#
# If we interpret the 1D function as a "seed", the 2D fractal is its
# "natural generalisation" to a self‑similar object. The process is contractive,
# so iterating the IFS converges to the fractal – a fixed point in function space.
# =============================================================================

def stochastic_1d_function(n_points, kind='random_walk', random_seed=42):
    """
    Generate a stochastic 1D function (values y at integer indices 0..n_points-1).
    
    Parameters:
        n_points : int
            Number of points.
        kind : str
            'white_noise' : independent Gaussian samples.
            'random_walk' : cumulative sum of Gaussian steps.
            'fbm' : fractional Brownian motion (approximated).
        random_seed : int
            Seed for reproducibility.
    """
    rng = np.random.default_rng(random_seed)
    if kind == 'white_noise':
        y = rng.normal(0, 1, n_points)
    elif kind == 'random_walk':
        steps = rng.normal(0, 1, n_points)
        y = np.cumsum(steps)
    elif kind == 'fbm':
        # Simple approximation: scaled random walk with long-range correlations
        steps = rng.normal(0, 1, n_points)
        y = np.cumsum(steps)
        # Add a trend to mimic Hurst exponent > 0.5
        y = y + 0.3 * np.linspace(-1, 1, n_points)
    else:
        raise ValueError("Unknown kind")
    # Normalise to [0,1] range for convenience
    y = (y - y.min()) / (y.max() - y.min() + 1e-8)
    return y

def fractal_interpolation(x, y, scaling=0.5, n_iter=10, n_refine=1000):
    """
    Generate a 2D fractal curve using fractal interpolation.
    
    The method constructs an IFS of N-1 affine maps, each mapping the interval
    [x_i, x_{i+1}] to the whole interval [x_0, x_{N-1}] (or the whole curve).
    The vertical scaling factor 'scaling' controls the fractal dimension.
    
    Parameters:
        x, y : ndarray
            The original interpolation points (1D stochastic function).
        scaling : float
            Vertical scaling factor (0 < scaling < 1). Lower scaling gives
            smoother curves; higher scaling gives more jagged fractals.
        n_iter : int
            Number of IFS iterations (the higher, the denser the fractal).
        n_refine : int
            Number of points in the final curve.
    
    Returns:
        x_fractal, y_fractal : ndarray
            Points on the fractal curve in 2D.
    """
    n = len(x)
    # The IFS maps: each map corresponds to an interval [x_i, x_{i+1}]
    # and maps the entire bounding box to that interval.
    # For simplicity we use the standard fractal interpolation construction.
    
    # Precompute IFS coefficients
    a = np.zeros(n-1)   # horizontal scaling
    b = np.zeros(n-1)   # horizontal translation
    c = np.zeros(n-1)   # vertical scaling (typically = scaling * (y_{i+1} - y_i) / (x_{i+1} - x_i)?)
    d = np.zeros(n-1)   # vertical translation
    e = np.zeros(n-1)   # additional vertical term (affine in x)
    f = np.zeros(n-1)   # constant vertical term
    
    for i in range(n-1):
        x0, x1 = x[i], x[i+1]
        y0, y1 = y[i], y[i+1]
        a[i] = (x1 - x0) / (x[-1] - x[0])
        b[i] = (x0 * x[-1] - x1 * x[0]) / (x[-1] - x[0])
        # Standard affine map for y: y' = c_i * y + d_i * x + e_i
        # Constraint: endpoints must match (x0, y0) and (x1, y1) when transformed.
        # We use the classic form: y' = s_i * y + (y_{i+1} - s_i * y_i)
        s = scaling  # uniform vertical scaling factor (can be made i‑dependent)
        c[i] = s
        d[i] = (y1 - s * y0) / (x1 - x0)
        e[i] = y0 - d[i] * x0
    
    # Generate the fractal by iterating the IFS on a random starting point
    # We use the "chaos game" method for speed and simplicity.
    points = []
    # Choose a random starting point in the bounding box
    cur_x = np.random.uniform(x[0], x[-1])
    cur_y = np.random.uniform(y.min(), y.max())
    
    # Pre‑select a long random sequence of map indices
    indices = np.random.choice(n-1, size=n_refine * n_iter, p=np.ones(n-1)/(n-1))
    
    for idx in indices:
        # Apply the map
        cur_x = a[idx] * cur_x + b[idx]
        cur_y = c[idx] * cur_y + d[idx] * cur_x + e[idx]
        points.append((cur_x, cur_y))
    
    # Keep only the last part (the attractor)
    points = np.array(points[-n_refine:])
    return points[:, 0], points[:, 1]

# -----------------------------------------------------------------------------
# Main: generate stochastic 1D function and plot its corresponding 2D fractal
# -----------------------------------------------------------------------------
if __name__ == "__main__":
    # 1. Create a stochastic 1D function
    n_control = 12                # number of control points
    y1d = stochastic_1d_function(n_control, kind='random_walk', random_seed=42)
    x1d = np.linspace(0, 1, n_control)
    
    # 2. Generate the 2D fractal via fractal interpolation
    scaling_factor = 0.6          # between 0 and 1 – higher = more fractal detail
    x_fractal, y_fractal = fractal_interpolation(x1d, y1d, scaling=scaling_factor,
                                                  n_iter=12, n_refine=5000)
    
    # 3. Plotting
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
    
    # Left: the original stochastic 1D function
    axes[0].plot(x1d, y1d, 'ro-', markersize=8, label='Control points')
    axes[0].fill_between(x1d, y1d, alpha=0.2)
    axes[0].set_title(f'Stochastic 1D function ({len(x1d)} points)')
    axes[0].set_xlabel('x')
    axes[0].set_ylabel('y')
    axes[0].grid(alpha=0.3)
    axes[0].legend()
    
    # Right: the corresponding 2D fractal curve
    axes[1].plot(x_fractal, y_fractal, 'b-', linewidth=0.5, alpha=0.7)
    axes[1].plot(x1d, y1d, 'ro', markersize=6, label='Original control points')
    axes[1].set_title(f'2D fractal via IFS (scaling = {scaling_factor})')
    axes[1].set_xlabel('x')
    axes[1].set_ylabel('y')
    axes[1].set_aspect('equal')
    axes[1].grid(alpha=0.3)
    axes[1].legend()
    
    plt.suptitle('From a stochastic 1D function to a 2D fractal (Fractal Interpolation)')
    plt.tight_layout()
    plt.show()
    
    # Optional: iteratively refine by applying spline smoothing to the 1D function
    # and regenerating the 2D fractal – this would be the "reverse" of the previous code.
    print("\n--- Interpretation ---")
    print("The 2D fractal is the attractor of an IFS that interpolates the 1D points.")
    print("The scaling factor controls the fractal dimension (higher → rougher).")
    print("If we repeatedly smooth the 1D function (e.g., with splines) and regenerate")
    print("the fractal, we obtain a sequence converging to a simpler attractor –")
    print("this would be the 'generalisation limit' analogue of your earlier request.")