"""Shared utilities for the CNNA Laola HuggingFace app.

This file keeps the didactic modules comparable:
- same base defaults,
- same response profile,
- same reference schedule,
- same diagnostic plot contract.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Tuple

import matplotlib.pyplot as plt
import numpy as np


@dataclass(frozen=True)
class BaseDefaults:
    n_people: int = 80
    tau: float = 0.20
    activation_time: float = 1.20
    refractory_time: float = 2.00
    seat_spacing: float = 1.00
    launch_count: int = 5
    launch_interval: float = 4.00
    dt: float = 0.03
    t_max: float = 22.0
    snapshot_time: float = 18.0


BASE = BaseDefaults()


def response_profile(s: np.ndarray, activation_time: float) -> np.ndarray:
    """Compact-support stand-up/plateau/sit-down profile p(s)."""
    y = np.zeros_like(s, dtype=float)
    if activation_time <= 0:
        return y

    rise_end = 0.25 * activation_time
    fall_start = 0.75 * activation_time

    rising = (0.0 <= s) & (s < rise_end)
    plateau = (rise_end <= s) & (s < fall_start)
    falling = (fall_start <= s) & (s <= activation_time)

    y[rising] = s[rising] / rise_end
    y[plateau] = 1.0
    y[falling] = 1.0 - (s[falling] - fall_start) / (activation_time - fall_start)
    return np.clip(y, 0.0, 1.0)


def exact_regular_schedule(n_people: int, tau: float) -> np.ndarray:
    """Reference schedule t_i = i τ.

    This deliberately avoids recursive cumsum. Recursive accumulation of 0.2 in
    binary float creates harmless ~1e-14 round-off drift. For the exact regular
    reference this would be a diagnostic artefact, not a model effect.
    """
    n = int(n_people)
    return float(tau) * np.arange(n, dtype=float)


def launch_offsets(launch_count: int, launch_interval: float) -> np.ndarray:
    count = max(1, int(launch_count))
    interval = max(float(launch_interval), 1e-12)
    return interval * np.arange(count, dtype=float)


def time_grid(t_max: float, dt: float) -> np.ndarray:
    # Include the right endpoint robustly without relying on exact binary divisibility.
    step = max(float(dt), 1e-9)
    return np.arange(0.0, float(t_max) + 0.5 * step, step, dtype=float)


def activity_from_schedule(
    times: np.ndarray,
    activation_times: np.ndarray,
    launches: np.ndarray,
    activation_time: float,
    *,
    ring_period: float | None = None,
) -> np.ndarray:
    """Build H[t, i] from activation times and launch offsets."""
    heights = np.zeros((len(times), len(activation_times)), dtype=float)
    for offset in launches:
        local_times = times[:, None] - activation_times[None, :] - float(offset)
        if ring_period is not None:
            local_times = local_times % ring_period
        heights = np.maximum(heights, response_profile(local_times, activation_time))
    return heights


def clean_near_zero(values: np.ndarray, scale: float = 1.0) -> np.ndarray:
    """Set numerical round-off at the float-noise scale to exact zero."""
    arr = np.array(values, dtype=float, copy=True)
    tol = max(1e-12, 1000.0 * np.finfo(float).eps * max(1.0, float(scale), float(arr.size)))
    arr[np.abs(arr) <= tol] = 0.0
    return arr


def timing_error(actual_activation: np.ndarray, reference_activation: np.ndarray) -> np.ndarray:
    scale = max(
        1.0,
        float(np.max(np.abs(actual_activation))) if actual_activation.size else 1.0,
        float(np.max(np.abs(reference_activation))) if reference_activation.size else 1.0,
    )
    return clean_near_zero(actual_activation - reference_activation, scale=scale)


def finite_difference_residual(times: np.ndarray, positions: np.ndarray, heights: np.ndarray, c: float) -> float:
    """Normalized residual for ∂t h + c∂x h = 0 as coarse-grained diagnostic."""
    if len(times) < 3 or len(positions) < 3:
        return float("nan")
    dt = float(times[1] - times[0])
    dx = float(positions[1] - positions[0])
    dh_dt = (heights[2:, 1:-1] - heights[:-2, 1:-1]) / (2.0 * dt)
    dh_dx = (heights[1:-1, 2:] - heights[1:-1, :-2]) / (2.0 * dx)
    residual = dh_dt + c * dh_dx
    scale = np.mean(np.abs(dh_dt)) + np.mean(np.abs(c * dh_dx)) + 1e-12
    return float(np.mean(np.abs(residual)) / scale)


def transport_diffusion_residuals(
    times: np.ndarray,
    positions: np.ndarray,
    heights: np.ndarray,
    c: float,
    d_eff: float,
) -> Tuple[float, float]:
    """Normalized residuals for transport and transport-diffusion diagnostics."""
    if len(times) < 3 or len(positions) < 3:
        return float("nan"), float("nan")

    dt = float(times[1] - times[0])
    dx = float(positions[1] - positions[0])

    h_center = heights[1:-1, 1:-1]
    dh_dt = (heights[2:, 1:-1] - heights[:-2, 1:-1]) / (2.0 * dt)
    dh_dx = (heights[1:-1, 2:] - heights[1:-1, :-2]) / (2.0 * dx)
    d2h_dx2 = (heights[1:-1, 2:] - 2.0 * h_center + heights[1:-1, :-2]) / (dx**2)

    r_transport = dh_dt + c * dh_dx
    r_diffusion = dh_dt + c * dh_dx - float(d_eff) * d2h_dx2

    scale = np.mean(np.abs(dh_dt)) + np.mean(np.abs(c * dh_dx)) + 1e-12
    return float(np.mean(np.abs(r_transport)) / scale), float(np.mean(np.abs(r_diffusion)) / scale)


def count_visible_packets(values: np.ndarray, threshold: float = 0.05) -> int:
    active = np.asarray(values) > threshold
    if active.size == 0:
        return 0
    starts = active & np.concatenate(([True], ~active[:-1]))
    return int(np.sum(starts))


def make_heatmap(
    times: np.ndarray,
    positions: np.ndarray,
    heights: np.ndarray,
    title: str,
):
    fig, ax = plt.subplots(figsize=(10, 4.8))
    im = ax.imshow(
        heights,
        aspect="auto",
        origin="lower",
        extent=[positions[0], positions[-1], times[0], times[-1]],
    )
    fig.colorbar(im, ax=ax, label="lokale Höhe / Aktivität hᵢ(t)")
    ax.set_xlabel("Knoten / Sitzposition xᵢ")
    ax.set_ylabel("Zeit t [s]")
    ax.set_title(title)
    fig.tight_layout()
    return fig


def make_snapshot_comparison(
    times: np.ndarray,
    positions: np.ndarray,
    simulated: np.ndarray,
    reference: np.ndarray,
    snapshot_time: float,
    sim_label: str,
):
    idx = int(np.argmin(np.abs(times - float(snapshot_time))))
    fig, ax = plt.subplots(figsize=(10, 3.8))
    ax.plot(positions, reference[idx, :], linestyle="--", linewidth=1.4, label="idealer Script-1-Wellenzug")
    ax.plot(positions, simulated[idx, :], marker="o", linewidth=1.5, markersize=3, label=sim_label)
    ax.set_ylim(-0.05, 1.08)
    ax.set_xlabel("Knoten / Sitzposition xᵢ")
    ax.set_ylabel("hᵢ(t)")
    ax.set_title(f"Momentaufnahme bei t = {times[idx]:.2f} s")
    ax.grid(True, alpha=0.25)
    ax.legend()
    fig.tight_layout()
    return fig


def make_timing_error_plot(positions: np.ndarray, errors: np.ndarray, title: str):
    fig, ax = plt.subplots(figsize=(10, 3.8))
    ax.plot(positions, errors, marker="o", linewidth=1.4, markersize=3)
    ax.axhline(0.0, linewidth=1.0)
    ax.set_xlabel("Knoten / Sitzposition xᵢ")
    ax.set_ylabel("kumulativer Timing-Fehler [s]")
    ax.set_title(title)
    ax.grid(True, alpha=0.25)
    fig.tight_layout()
    return fig