from __future__ import annotations

from dataclasses import dataclass
from functools import lru_cache
from typing import Any

import numpy as np


@dataclass(frozen=True)
class KSParams:
    n: int = 1024
    length: float = 128.0
    linear_growth: float = 1.0
    third_order_dispersion: float = 0.0
    fifth_order_dispersion: float = 0.0
    fourth_order_damping: float = 1.0
    dealias: bool = True
    max_abs_state: float = 1.0e6


def params_from_config(config: dict[str, Any]) -> KSParams:
    model = config["model"]
    family = str(model.get("family", "kuramoto_sivashinsky")).lower()
    if family not in {"kuramoto_sivashinsky", "ks"}:
        raise ValueError(f"Unsupported model family for KS workflow: {family}")
    return KSParams(
        n=int(model.get("N", model.get("n", 1024))),
        length=float(model.get("length", 128.0)),
        linear_growth=float(model.get("linear_growth", 1.0)),
        third_order_dispersion=float(model.get("third_order_dispersion", 0.0)),
        fifth_order_dispersion=float(model.get("fifth_order_dispersion", 0.0)),
        fourth_order_damping=float(model.get("fourth_order_damping", 1.0)),
        dealias=bool(model.get("dealias", True)),
        max_abs_state=float(model.get("max_abs_state", 1.0e6)),
    )


@lru_cache(maxsize=32)
def wavenumbers(params: KSParams) -> np.ndarray:
    return 2.0 * np.pi * np.fft.fftfreq(params.n, d=params.length / params.n)


@lru_cache(maxsize=32)
def mode_indices(params: KSParams) -> np.ndarray:
    return np.fft.fftfreq(params.n) * params.n


@lru_cache(maxsize=32)
def dealias_mask(params: KSParams) -> np.ndarray:
    if not params.dealias:
        return np.ones(params.n, dtype=bool)
    return np.abs(mode_indices(params)) <= params.n / 3.0


@lru_cache(maxsize=64)
def etdrk4_coefficients(params: KSParams, dt: float) -> dict[str, np.ndarray]:
    k = wavenumbers(params)
    linear = (
        float(params.linear_growth) * k**2
        + 1j * float(params.third_order_dispersion) * k**3
        - 1j * float(params.fifth_order_dispersion) * k**5
        - float(params.fourth_order_damping) * k**4
    )
    e = np.exp(dt * linear)
    e2 = np.exp(0.5 * dt * linear)

    roots = np.exp(1j * np.pi * (np.arange(1, 65, dtype=np.float64) - 0.5) / 64.0)
    lr = dt * linear[:, None] + roots[None, :]
    q = dt * np.mean((np.exp(0.5 * lr) - 1.0) / lr, axis=1)
    f1 = dt * np.mean((-4.0 - lr + np.exp(lr) * (4.0 - 3.0 * lr + lr**2)) / lr**3, axis=1)
    f2 = dt * np.mean((2.0 + lr + np.exp(lr) * (-2.0 + lr)) / lr**3, axis=1)
    f3 = dt * np.mean((-4.0 - 3.0 * lr - lr**2 + np.exp(lr) * (4.0 - lr)) / lr**3, axis=1)
    return {"e": e, "e2": e2, "q": q, "f1": f1, "f2": f2, "f3": f3}


def fft(field: np.ndarray) -> np.ndarray:
    return np.fft.fft(field, axis=-1)


def ifft(field_hat: np.ndarray) -> np.ndarray:
    return np.fft.ifft(field_hat, axis=-1).real


def spectral_filter(field: np.ndarray, params: KSParams) -> np.ndarray:
    field_hat = fft(np.asarray(field, dtype=np.float64))
    field_hat[..., ~dealias_mask(params)] = 0.0
    field_hat[..., 0] = 0.0
    return ifft(field_hat)


def project_real_hat(field_hat: np.ndarray, params: KSParams) -> np.ndarray:
    projected = fft(ifft(field_hat))
    projected[..., ~dealias_mask(params)] = 0.0
    projected[..., 0] = 0.0
    return projected


def derivative(field: np.ndarray, params: KSParams, order: int = 1) -> np.ndarray:
    k = wavenumbers(params)
    return ifft((1j * k) ** int(order) * fft(np.asarray(field, dtype=np.float64)))


def nonlinear_hat(field_hat: np.ndarray, params: KSParams) -> np.ndarray:
    k = wavenumbers(params)
    field = ifft(field_hat)
    squared_hat = fft(field**2)
    squared_hat[..., ~dealias_mask(params)] = 0.0
    out = -0.5j * k * squared_hat
    out[..., ~dealias_mask(params)] = 0.0
    out[..., 0] = 0.0
    return out


def etdrk4_step_hat(field_hat: np.ndarray, dt: float, params: KSParams) -> np.ndarray:
    coeff = etdrk4_coefficients(params, float(dt))
    nv = nonlinear_hat(field_hat, params)
    a = coeff["e2"] * field_hat + coeff["q"] * nv
    na = nonlinear_hat(a, params)
    b = coeff["e2"] * field_hat + coeff["q"] * na
    nb = nonlinear_hat(b, params)
    c = coeff["e2"] * a + coeff["q"] * (2.0 * nb - nv)
    nc = nonlinear_hat(c, params)
    out = coeff["e"] * field_hat + coeff["f1"] * nv + 2.0 * coeff["f2"] * (na + nb) + coeff["f3"] * nc
    return project_real_hat(out, params)


def step(field: np.ndarray, dt: float, params: KSParams) -> np.ndarray:
    return ifft(etdrk4_step_hat(fft(np.asarray(field, dtype=np.float64)), dt, params))


def integrate_steps(field: np.ndarray, n_steps: int, dt: float, params: KSParams) -> np.ndarray:
    field_hat = fft(spectral_filter(field, params))
    for _ in range(int(n_steps)):
        field_hat = etdrk4_step_hat(field_hat, dt, params)
    return ifft(field_hat)


def random_initial_state(
    params: KSParams,
    rng: np.random.Generator,
    std: float = 0.5,
    spectral_slope: float = 2.0,
    max_initial_mode: int | None = None,
) -> np.ndarray:
    n = int(params.n)
    modes = np.abs(mode_indices(params))
    cutoff = int(max_initial_mode if max_initial_mode is not None else max(4, n // 16))
    active = (modes > 0) & (modes <= cutoff)
    amplitude = np.zeros_like(modes, dtype=np.float64)
    amplitude[active] = modes[active] ** (-float(spectral_slope))
    phase = rng.normal(size=n) + 1j * rng.normal(size=n)
    field = ifft(phase * amplitude)
    field -= float(np.mean(field))
    current = float(np.sqrt(np.mean(field**2)))
    if current == 0.0:
        field = rng.normal(0.0, 1.0, size=n)
        field -= float(np.mean(field))
        current = float(np.sqrt(np.mean(field**2)))
    return spectral_filter(field * (float(std) / current), params)


def saved_trajectory(
    field0: np.ndarray,
    total_time: float,
    save_interval: float,
    dt: float,
    params: KSParams,
    include_initial: bool = False,
) -> np.ndarray:
    if total_time <= 0.0:
        raise ValueError("total_time must be positive")
    if save_interval <= 0.0:
        raise ValueError("save_interval must be positive")
    steps_per_save = int(round(save_interval / dt))
    if steps_per_save <= 0:
        raise ValueError("save_interval must be at least one timestep")
    total_steps = int(round(total_time / dt))
    if total_steps % steps_per_save != 0:
        raise ValueError("total_time must contain an integer number of save intervals")

    n_save = total_steps // steps_per_save
    extra = 1 if include_initial else 0
    states = np.empty((n_save + extra, params.n), dtype=np.float32)
    field_hat = fft(spectral_filter(field0, params))
    cursor = 0
    if include_initial:
        states[cursor] = ifft(field_hat).astype(np.float32)
        cursor += 1
    for step_idx in range(1, total_steps + 1):
        field_hat = etdrk4_step_hat(field_hat, dt, params)
        if step_idx % steps_per_save == 0:
            field = ifft(field_hat)
            if not is_finite_state(field, params):
                raise FloatingPointError(f"KS trajectory became non-finite at step {step_idx}")
            states[cursor] = field.astype(np.float32)
            cursor += 1
    return states


def power_spectrum(field: np.ndarray, params: KSParams) -> dict[str, np.ndarray]:
    arr = np.asarray(field, dtype=np.float64)
    field_hat = fft(arr)
    shell = np.rint(np.abs(mode_indices(params))).astype(np.int64)
    max_shell = int(params.n // 2)
    energy_2d = np.abs(field_hat) ** 2 / float(params.n**2)
    if arr.ndim == 1:
        spectrum = np.bincount(shell, weights=energy_2d, minlength=max_shell + 1)
    else:
        rows = [np.bincount(shell, weights=row, minlength=max_shell + 1) for row in energy_2d.reshape(-1, params.n)]
        spectrum = np.mean(np.stack(rows, axis=0), axis=0)
    return {
        "wavenumber": np.arange(0, max_shell + 1, dtype=np.int64),
        "power_spectrum": spectrum[: max_shell + 1].astype(np.float64),
    }


def state_diagnostics(field: np.ndarray, params: KSParams) -> dict[str, float]:
    arr = np.asarray(field, dtype=np.float64)
    return {
        "finite": bool(np.all(np.isfinite(arr))),
        "max_abs_state": float(np.max(np.abs(arr))),
        "mean": float(np.mean(arr)),
        "std": float(np.std(arr)),
        "rms": float(np.sqrt(np.mean(arr**2))),
        "energy": float(0.5 * np.mean(arr**2)),
        "gradient_rms": float(np.sqrt(np.mean(derivative(arr, params, 1) ** 2))),
    }


def is_finite_state(field: np.ndarray, params: KSParams) -> bool:
    arr = np.asarray(field)
    return bool(np.all(np.isfinite(arr)) and np.max(np.abs(arr)) <= params.max_abs_state)
