"""Configuration dataclasses and their YAML representation.

Every physical and numerical parameter of a run lives in one of the dataclasses below, grouped
by what it describes: the material, the cymbal profile, the seeded imperfections, the mesh, the
list of hits (each with its own exciter), the listener, the numerics and the output. Defaults
reproduce the 18-inch medium-thin crash of the original prototype. A SimulationConfig can be
built directly in Python or loaded from a YAML file whose keys mirror the field names; a hit's
exciter is chosen by its ``kind`` key ("stick", "mallet" or "force").
"""
from __future__ import annotations

import dataclasses
import math
import types
import typing
from dataclasses import dataclass, field
from pathlib import Path

import yaml


# ------------------------------------------------------------------------------- material
@dataclass
class Material:
    """Isotropic material of the shell (default: B20 bronze) and its two loss parameters."""
    E: float = 1.1e11           # Young's modulus [Pa]
    nu: float = 0.34            # Poisson ratio
    rho: float = 8700.0         # density [kg/m^3]
    sigma0: float = 0.9         # frequency-independent loss [1/s]; T60 of the lows ~ 6.9/sigma0
    sigma1: float = 5e-4        # loss proportional to omega [m^2/s]; sets the ceiling of the cascade


# -------------------------------------------------------------------------------- profile
@dataclass
class Profile:
    """Axisymmetric part of the cymbal: bow, bell and thickness taper. Mean rim height is z0 = 0."""
    radius: float = 0.2286              # outer radius [m] (18 inch)
    bell_height: float = 0.013          # bell height above the bow at the shoulder [m]
    bell_top_radius: float = 0.025      # radius of the flat bell plateau [m]
    bell_shoulder_radius: float = 0.080 # mean radius of the bell/bow junction [m]
    bow_height: float = 0.030           # bow rise from the rim to the shoulder [m]
    bow_power: float = 1.4              # bow shape exponent: 1 = cone, > 1 = curved, flat at the shoulder
    h_rim: float = 1.0e-3               # thickness at the rim [m]
    h_shoulder: float = 1.7e-3          # thickness at the shoulder [m]
    h_bell: float = 2.6e-3              # thickness in the bell [m]
    max_slope: float = 0.4              # warn when max |grad z0| exceeds this (shallow-shell bound)


@dataclass
class Imperfections:
    """Seeded symmetry breaking: hammer dents, rim warp, shoulder ovality and a bell offset."""
    seed: int = 1
    n_dents: int = 400                  # Gaussian hammer marks on the bow, random sign
    dent_depth: float = 0.25e-3         # typical dent depth [m] (x 0.5 to 1.5 jitter)
    dent_sigma: float = 7e-3            # dent Gaussian sigma [m] (x 0.7 to 1.3 jitter); keep >= element size
    dent_margin: float = 0.01           # dents stay this far from the shoulder and the rim [m]
    rim_warp: float = 0.6e-3            # m = 2 rim-warp amplitude [m]; mode m gets 2 * rim_warp / m
    rim_warp_modes: tuple[int, ...] = (2, 3, 4, 5)
    shoulder_ovality: float = 0.03      # relative m = 2 modulation of the shoulder radius
    bell_offset: float = 1.5e-3         # bell centre offset from the geometric centre [m]


@dataclass
class MeshConfig:
    element_size: float = 6e-3          # target edge length of the triangles [m]
    smoothing: int = 1                  # Laplacian smoothing passes; more improves shape, shrinks the smallest elements


# ------------------------------------------------------------------------------- exciters
@dataclass
class StickExciter:
    """Hard stick: Hertz contact (alpha = 1.5) along a short line, as in a shoulder crash."""
    kind: str = "stick"
    mass: float = 0.030                 # effective mass [kg]
    k: float = 5e8                      # total contact stiffness [N/m^alpha]
    alpha: float = 1.5                  # contact exponent
    patch_length: float = 0.03          # length of the contact line [m]
    patch_points: int = 7               # springs along the line
    patch_angle: float = 0.6            # direction of the line relative to the radial direction [rad]
    gap: float = 1e-4                   # initial clearance above the surface at release [m]


@dataclass
class MalletExciter:
    """Soft felt mallet: power-law contact on a small disk; k is set from a target compression."""
    kind: str = "mallet"
    mass: float = 0.30                  # head plus effective arm mass [kg]
    alpha: float = 2.5                  # felt exponent (1.5 = Hertz, 2 to 3 = felt)
    compression: float = 12e-3          # felt compression at full stop on a rigid wall [m]
    radius: float = 0.012               # head contact radius [m]; 7 springs (centre + hexagon)
    gap: float = 1e-4


@dataclass
class ForceExciter:
    """Prescribed force at one point, independent of the motion: a raised cosine
    F(t) = peak (1 - cos(2 pi t / duration)) / 2 for 0 <= t <= duration after the hit time,
    pushing the surface down. There is no contact and no rebound. This is the excitation of
    modal plate models such as Ducceschi and Touze (2015), and is meant for comparison with them;
    ``Hit.velocity`` is ignored."""
    kind: str = "force"
    peak: float = 100.0                 # peak force [N]
    duration: float = 2e-3              # total length of the pulse [s]


Exciter = typing.Union[StickExciter, MalletExciter, ForceExciter]
EXCITER_KINDS: dict[str, type] = {"stick": StickExciter, "mallet": MalletExciter, "force": ForceExciter}


@dataclass
class Hit:
    """One strike: when, where, how fast, and with what."""
    time: float = 0.0                   # release time of the exciter [s]
    r: float = 0.137                    # radial position of the contact centre [m]
    phi: float = 0.3                    # azimuth of the contact centre [rad]
    velocity: float = 8.0               # approach speed [m/s]
    exciter: Exciter = field(default_factory=MalletExciter)


# ------------------------------------------------------------------------------- listener
@dataclass
class Listener:
    """Listener point and air; the mean rim plane is z = 0 with the cymbal centre at the origin."""
    position: tuple[float, float, float] = (0.40, 0.30, 0.90)   # [m]
    c_air: float = 343.0
    rho_air: float = 1.2
    back_diffraction_gain: float = 1.0  # weight of the diffracted back-side term (0 = baffled only)


# ------------------------------------------------------------------------------- numerics
@dataclass
class Numerics:
    fs: int = 44100                     # audio sample rate [Hz]; the sub-step rate is an integer multiple
    duration: float = 2.0               # simulated time [s]
    cfl_safety: float = 0.6             # fraction of the linear stability limit actually used
    stability_bound: str = "global"     # "global": lambda_max(M^-1 K) by power iteration; "element": the Irons element bound (larger, used before 2026-10)
    rot_inertia_factor: float = 1 / 64  # lumped rotational inertia I = m (h^2/12 + f A_e); see docs
    inplane_mass_scale: float = 1.0     # > 1 slows in-plane waves and relaxes the time step; see docs
    steps_per_block: int = 20000        # sub-steps per compiled block
    dtype: str = "float32"              # simulation precision: "float32" or "float64"
    nonlinear: bool = True              # False drops the quadratic strain terms (linear reference run)
    tangent: bool = False               # also integrate the tangent: d/d delta of the run when every prescribed peak force is x (1 + delta)
    n_progress: int = 20                # number of progress/energy reports during the run


@dataclass
class Output:
    prefix: str = "crash"               # output files are <prefix>_<name>.wav / .png / .npz
    wav: bool = True
    npz: bool = True
    cascade_plot: bool = True           # needs matplotlib
    record_acceleration: bool = False   # keep the normal acceleration of every node at the audio rate
    displacement_rate: float = 60.0     # frames per second of transverse displacement kept for display; 0 = none
    probes: list[tuple[float, float]] = field(default_factory=list)   # points (x, y) [m] where the normal acceleration is kept at the audio rate


# The fields in which a continued run may differ from the run it continues: the length of the run,
# how often energy is reported, and the naming of the output files. Two configurations that agree
# on everything else describe runs that are prefixes of one another (Simulator.run(resume=...)),
# and share a cache key (cache.py).
CONTINUATION_FIELDS = (("numerics", "duration"), ("numerics", "n_progress"), ("output", "prefix"),
                       ("output", "wav"), ("output", "npz"), ("output", "cascade_plot"))


# ----------------------------------------------------------------------------- simulation
@dataclass
class SimulationConfig:
    material: Material = field(default_factory=Material)
    profile: Profile = field(default_factory=Profile)
    imperfections: Imperfections = field(default_factory=Imperfections)
    mesh: MeshConfig = field(default_factory=MeshConfig)
    hits: list[Hit] = field(default_factory=lambda: [Hit()])
    listener: Listener = field(default_factory=Listener)
    numerics: Numerics = field(default_factory=Numerics)
    output: Output = field(default_factory=Output)

    # ---- validation
    def validate(self) -> None:
        p, n = self.profile, self.numerics
        if not self.hits:
            raise ValueError("at least one hit is required")
        if n.dtype not in ("float32", "float64"):
            raise ValueError(f"numerics.dtype must be 'float32' or 'float64', got {n.dtype!r}")
        if n.stability_bound not in ("global", "element"):
            raise ValueError(f"numerics.stability_bound must be 'global' or 'element', got {n.stability_bound!r}")
        if not (0 < n.cfl_safety <= 1):
            raise ValueError("numerics.cfl_safety must lie in (0, 1]")
        if self.mesh.element_size <= 0 or self.mesh.element_size > p.radius / 2:
            raise ValueError("mesh.element_size must be positive and well below the radius")
        if n.tangent and self.output.record_acceleration:
            raise ValueError("numerics.tangent does not support output.record_acceleration")
        if n.tangent and not any(isinstance(h.exciter, ForceExciter) for h in self.hits):
            raise ValueError("numerics.tangent differentiates the prescribed peak forces: it needs a force exciter")
        if self.output.displacement_rate < 0:
            raise ValueError("output.displacement_rate must be non-negative")
        for i, h in enumerate(self.hits):
            if h.time < 0:
                raise ValueError(f"hit {i}: time must be non-negative")
            if h.velocity <= 0:
                raise ValueError(f"hit {i}: velocity must be positive (the exciter approaches from above)")
            ex = h.exciter
            extent = (ex.patch_length / 2 if isinstance(ex, StickExciter)
                      else ex.radius if isinstance(ex, MalletExciter) else 0.0)
            if h.r + extent >= p.radius - self.mesh.element_size:
                raise ValueError(f"hit {i}: contact patch reaches the rim (r = {h.r}, extent = {extent})")
            if ex.kind not in EXCITER_KINDS:
                raise ValueError(f"hit {i}: unknown exciter kind {ex.kind!r}")
            if isinstance(ex, ForceExciter) and (ex.peak < 0 or ex.duration <= 0):
                raise ValueError(f"hit {i}: a force pulse needs peak >= 0 and duration > 0")
        for i, (x, y) in enumerate(self.output.probes):
            if math.hypot(x, y) >= p.radius:
                raise ValueError(f"output.probes[{i}] = ({x}, {y}) lies outside the cymbal")

    # ---- dict / YAML
    def to_dict(self) -> dict:
        return _plain(dataclasses.asdict(self))

    def identity_dict(self) -> dict:
        """``to_dict()`` without CONTINUATION_FIELDS: what a run and its continuation share."""
        d = self.to_dict()
        for section, name in CONTINUATION_FIELDS:
            d[section].pop(name, None)
        return d

    @classmethod
    def from_dict(cls, d: dict) -> "SimulationConfig":
        return _from_dict(cls, d)

    def to_yaml(self, path: str | Path) -> None:
        Path(path).write_text(yaml.safe_dump(self.to_dict(), sort_keys=False))

    @classmethod
    def from_yaml(cls, path: str | Path) -> "SimulationConfig":
        with open(path) as f:
            d = yaml.safe_load(f) or {}
        return cls.from_dict(d)


# --------------------------------------------------------------- generic dataclass <-> dict
def dict_differences(a, b, path: str = "") -> list[str]:
    """Dotted paths at which two plain nested structures differ."""
    if isinstance(a, dict) and isinstance(b, dict):
        return [d for k in sorted(set(a) | set(b), key=str)
                for d in dict_differences(a.get(k), b.get(k), f"{path}.{k}" if path else str(k))]
    return [] if a == b else [path or "(root)"]


def _plain(obj):
    """Recursively turn tuples into lists so that yaml.safe_dump accepts the structure."""
    if isinstance(obj, dict):
        return {k: _plain(v) for k, v in obj.items()}
    if isinstance(obj, (list, tuple)):
        return [_plain(v) for v in obj]
    return obj


def _build(tp, value):
    """Construct a value of annotated type ``tp`` from plain YAML data."""
    origin = typing.get_origin(tp)
    if isinstance(tp, type) and dataclasses.is_dataclass(tp):
        return _from_dict(tp, value)
    if origin is list:
        (item,) = typing.get_args(tp)
        return [_build(item, v) for v in value]
    if origin is tuple:
        return tuple(value)
    if origin in (typing.Union, getattr(types, "UnionType", typing.Union)):   # UnionType: Python >= 3.10
        # The only union in the schema is the exciter, discriminated by its "kind" key.
        if isinstance(value, dict) and "kind" in value:
            try:
                return _from_dict(EXCITER_KINDS[value["kind"]], value)
            except KeyError:
                raise ValueError(f"unknown exciter kind {value['kind']!r}; "
                                 f"expected one of {sorted(EXCITER_KINDS)}") from None
        if any(isinstance(value, t) for t in typing.get_args(tp)):
            return value
        raise ValueError("an exciter mapping needs a 'kind' key ('stick', 'mallet' or 'force')")
    if tp is float and isinstance(value, (int, float)):
        return float(value)
    return value


def _from_dict(cls, d):
    if isinstance(d, cls):
        return d
    if not isinstance(d, dict):
        raise ValueError(f"expected a mapping for {cls.__name__}, got {type(d).__name__}")
    hints = typing.get_type_hints(cls)
    names = {f.name for f in dataclasses.fields(cls)}
    unknown = set(d) - names
    if unknown:
        raise ValueError(f"unknown field(s) for {cls.__name__}: {sorted(unknown)}")
    return cls(**{k: _build(hints[k], v) for k, v in d.items()})


def default_config(exciter: str = "mallet") -> SimulationConfig:
    """The prototype's default run for the given exciter kind."""
    if exciter == "mallet":
        hit = Hit(time=0.0, r=0.60 * Profile.radius, phi=0.3, velocity=8.0, exciter=MalletExciter())
    elif exciter == "stick":
        hit = Hit(time=0.0, r=0.85 * Profile.radius, phi=0.3, velocity=7.0, exciter=StickExciter())
    else:
        raise ValueError(f"unknown exciter kind {exciter!r}")
    return SimulationConfig(hits=[hit])


def t60_estimate(material: Material, frequency: float) -> float:
    """Decay time to -60 dB implied by sigma0 and sigma1 at a given frequency, for a flat plate
    with bending wavenumber k^2 = omega / sqrt(D / rho h). Rough, but useful for tuning losses."""
    # The sigma1 term acts as 2 sigma1 k^2; for a thin plate k^2 = omega sqrt(rho h / D) which
    # depends on thickness; we quote the value for h = 1 mm.
    h = 1e-3
    D = material.E * h**3 / (12 * (1 - material.nu**2))
    omega = 2 * math.pi * frequency
    k2 = omega * math.sqrt(material.rho * h / D)
    sigma = material.sigma0 + material.sigma1 * k2
    return 6.91 / sigma