"""
benchmarks.py
=============
Task generators for K-R Reservoir Architecture evaluation.
All outputs are normalised to zero mean, unit variance.

Tasks:
    narma(order, T, seed)   -- NARMA-n benchmark
    mackey_glass(T, seed)   -- Mackey-Glass chaotic series
    santa_fe(T, seed)       -- Lorenz-based chaotic series (Santa Fe proxy)

Usage:
    from benchmarks import narma, mackey_glass, santa_fe
    u, y = narma(order=50, T=3000, seed=0)
"""

import numpy as np


def narma(order: int, T: int = 3000, seed: int = 0):
    """
    NARMA-n benchmark (Atiya & Parlos, 2000).

    Equation:
        y(t) = 0.3*y(t-1)
               + c * y(t-1) * sum(y[t-n:t])
               + 1.5 * u(t-1) * u(t-n)
               + 0.1
        where c = 0.05 / max(1, n/10)

    Parameters
    ----------
    order : int   Memory order n (use 10, 30, 50, or 100)
    T     : int   Number of output time steps
    seed  : int   Random seed for input u

    Returns
    -------
    u_out : ndarray, shape (T,)  Normalised input
    y_out : ndarray, shape (T,)  Normalised target
    """
    np.random.seed(seed)
    WARMUP = 200
    total = T + WARMUP
    u = np.random.uniform(0, 0.5, total)
    y = np.zeros(total)
    c = 0.05 / max(1.0, order / 10.0)

    for t in range(order, total):
        y[t] = (0.3 * y[t - 1]
                + c * y[t - 1] * np.sum(y[t - order:t])
                + 1.5 * u[t - 1] * u[t - order]
                + 0.1)

    u_out = u[WARMUP: WARMUP + T]
    y_out = y[WARMUP: WARMUP + T]
    u_out = (u_out - u_out.mean()) / (u_out.std() + 1e-12)
    y_out = (y_out - y_out.mean()) / (y_out.std() + 1e-12)
    return u_out, y_out


def mackey_glass(T: int = 3000, tau: int = 17, seed: int = 0):
    """
    Mackey-Glass chaotic delay-differential equation (Mackey & Glass, 1977).
    Euler integration with dt=0.1.
    1-step-ahead prediction: u=x[t], y=x[t+1].

    Parameters
    ----------
    T    : int   Number of output time steps
    tau  : int   Delay (17 = weakly chaotic, 30 = strongly chaotic)
    seed : int   Unused (deterministic system); kept for API consistency

    Returns
    -------
    u_out : ndarray (T,)  Normalised input
    y_out : ndarray (T,)  Normalised 1-step-ahead target
    """
    WARMUP = 500
    total = T + WARMUP + tau + 10
    x = np.zeros(total)
    x[:tau] = 0.5

    dt = 0.1
    for t in range(tau, total - 1):
        x_tau = x[t - tau]
        dx = 0.2 * x_tau / (1.0 + x_tau ** 10) - 0.1 * x[t]
        x[t + 1] = x[t] + dt * dx

    series = x[WARMUP: WARMUP + T + 1]
    series = (series - series.mean()) / (series.std() + 1e-12)
    return series[:T], series[1: T + 1]


def santa_fe(T: int = 3000, seed: int = 0):
    """
    Santa Fe Laser proxy: z-component of the Lorenz attractor.
    Parameters: sigma=10, r=28, b=8/3, dt=0.01.

    Parameters
    ----------
    T    : int   Number of output time steps
    seed : int   Unused; kept for API consistency

    Returns
    -------
    u_out : ndarray (T,)  Normalised input
    y_out : ndarray (T,)  Normalised 1-step-ahead target
    """
    WARMUP = 1000
    total = T + WARMUP + 1
    sigma, r, b = 10.0, 28.0, 8.0 / 3.0
    dt = 0.01

    xyz = np.zeros((total, 3))
    xyz[0] = [0.1, 0.0, 0.0]

    for t in range(total - 1):
        x, y, z = xyz[t]
        dx = sigma * (y - x)
        dy = x * (r - z) - y
        dz = x * y - b * z
        xyz[t + 1] = xyz[t] + dt * np.array([dx, dy, dz])

    series = xyz[WARMUP: WARMUP + T + 1, 2]
    series = (series - series.mean()) / (series.std() + 1e-12)
    return series[:T], series[1: T + 1]