# /// script
# requires-python = ">=3.11"
# dependencies = ["numpy<2.1", "scipy"]
# ///
"""Build SoftChart's three-voice modal-resonator Taiko preview bank.

The supplied reference WAVs are used offline for RMS-envelope and level
calibration only.  Reference samples, paths, hashes, and envelope arrays are
never written to the output.  Every output sample is generated from a seeded
noise/impulse excitation, modal resonators, and a decaying pitch sweep.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import wave
from dataclasses import dataclass
from pathlib import Path

import numpy as np
from scipy.signal import butter, sosfilt


SAMPLE_RATE = 48_000
RANDOM_SEED = 20_260_713
BUILD_ORDER = ("Big Don", "Don", "Katsu")
EXPECTED_FRAMES = {"Big Don": 42_946, "Don": 39_836, "Katsu": 19_537}
SYNTH_FRAMES = EXPECTED_FRAMES
OUTPUT_NAMES = {
    "Big Don": ("don-big.wav", "don_big"),
    "Don": ("don.wav", "don"),
    "Katsu": ("katsu.wav", "katsu"),
}


@dataclass(frozen=True)
class ReferencePaths:
    big_don: Path
    don: Path
    katsu: Path

    def by_role(self) -> dict[str, Path]:
        return {"Big Don": self.big_don, "Don": self.don, "Katsu": self.katsu}


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1 << 20), b""):
            digest.update(block)
    return digest.hexdigest()


def _read_reference(path: Path, role: str) -> np.ndarray:
    if not path.is_file():
        raise FileNotFoundError(f"missing {role} reference WAV: {path}")
    with wave.open(str(path), "rb") as source:
        sample_rate = source.getframerate()
        channels = source.getnchannels()
        width = source.getsampwidth()
        frames = source.getnframes()
        payload = source.readframes(frames)
    if sample_rate != SAMPLE_RATE or channels != 1 or width != 2:
        raise ValueError(f"{role} reference must be mono PCM16 at 48 kHz")
    if frames != EXPECTED_FRAMES[role]:
        raise ValueError(
            f"{role} reference must contain {EXPECTED_FRAMES[role]} frames, got {frames}"
        )
    audio = np.frombuffer(payload, dtype="<i2").astype(np.float64) / 32768.0
    if not np.all(np.isfinite(audio)) or float(np.max(np.abs(audio))) <= 1e-9:
        raise ValueError(f"{role} reference must contain finite, non-silent audio")
    return audio


def _resonator_bank(
    excitation: np.ndarray,
    frequencies: list[float],
    gains: list[float],
    t60s: list[float],
) -> np.ndarray:
    output = np.zeros_like(excitation)
    for frequency, gain, t60 in zip(frequencies, gains, t60s, strict=True):
        radius = 10.0 ** (-3.0 / (max(t60, 0.001) * SAMPLE_RATE))
        a1 = 2.0 * radius * math.cos(2.0 * math.pi * frequency / SAMPLE_RATE)
        a2 = -(radius * radius)
        input_gain = (1.0 - radius) * 4.0
        previous = second_previous = 0.0
        mode = np.empty_like(excitation)
        for index, sample in enumerate(excitation):
            value = input_gain * sample + a1 * previous + a2 * second_previous
            mode[index] = value
            second_previous, previous = previous, value
        output += gain * mode
    return output


def _soft_normalize(audio: np.ndarray, peak: float = 0.95) -> np.ndarray:
    shaped = np.tanh(audio * 1.15)
    maximum = float(np.max(np.abs(shaped)))
    if not math.isfinite(maximum) or maximum <= 1e-12:
        raise ValueError("synthesizer produced silent or non-finite audio")
    return shaped * (peak / maximum)


def _make_synth(role: str, rng: np.random.Generator) -> np.ndarray:
    frame_count = SYNTH_FRAMES[role]
    if role == "Big Don":
        frequencies = [86, 125, 141, 176, 193, 240, 264, 291, 313, 344, 409,
                       529, 578, 700, 788, 869, 957, 1207, 1650, 2200, 2770]
        gains = [1.00, .92, .28, .22, .45, .38, .28, .22, .18, .12, .12,
                 .09, .07, .07, .08, .045, .04, .035, .018, .012, .010]
        t60s = [.72, .82, .55, .50, .46, .40, .36, .34, .31, .28, .25,
                .22, .19, .16, .14, .12, .11, .09, .07, .055, .045]
        cutoff, noise_decay = 6_500.0, .018
        sweep_start, sweep_end, sweep_decay = 302.0, 82.0, .11
        direct_amount, body_amount, sweep_amount = .13, 1.60, .23
    elif role == "Don":
        frequencies = [86, 125, 141, 174, 193, 240, 247, 276, 292, 313, 368,
                       386, 408, 487, 540, 578, 647, 694, 787, 870, 962, 2165]
        gains = [1.0, .60, .27, .15, .34, .36, .23, .18, .20, .21, .10,
                 .15, .10, .08, .08, .07, .07, .12, .12, .05, .05, .018]
        t60s = [.52, .48, .37, .34, .34, .31, .30, .27, .25, .24, .22,
                .21, .20, .18, .17, .16, .14, .14, .12, .10, .09, .052]
        cutoff, noise_decay = 7_500.0, .024
        sweep_start, sweep_end, sweep_decay = 346.0, 86.0, .075
        direct_amount, body_amount, sweep_amount = .17, 1.55, .16
    elif role == "Katsu":
        time = np.arange(frame_count) / SAMPLE_RATE
        noise = rng.standard_normal(frame_count)
        high = sosfilt(
            butter(3, 300 / (SAMPLE_RATE / 2), btype="high", output="sos"), noise,
        )
        band = sosfilt(
            butter(
                4, [450 / (SAMPLE_RATE / 2), 10_000 / (SAMPLE_RATE / 2)],
                btype="band", output="sos",
            ),
            noise,
        )
        excitation = np.zeros(frame_count)
        excitation[0] = 4.5
        excitation += 1.6 * band * np.exp(-time / .018)
        frequencies = [455, 560, 641, 735, 849, 1025, 1047, 1166, 1214, 1374,
                       2188, 2231, 2279, 2933, 2994, 3083, 3123, 3856, 4006]
        gains = [.25, .32, .38, .30, .26, .65, .52, .34, .26, .30, .85,
                 .72, .38, .38, .30, .40, .34, .30, .28]
        t60s = [.23, .18, .16, .14, .13, .20, .18, .15, .14, .13, .17,
                .16, .13, .12, .11, .12, .11, .10, .10]
        body = _resonator_bank(excitation, frequencies, gains, t60s)
        direct = .32 * high * np.exp(-time / .050) + .50 * band * np.exp(-time / .012)
        click_frames = max(1, int(.0012 * SAMPLE_RATE))
        click = np.zeros(frame_count)
        click[:click_frames] = rng.standard_normal(click_frames) * np.hanning(click_frames)
        synth = 1.2 * body + direct + .18 * click
        lead = int(.018 * SAMPLE_RATE)
        return _soft_normalize(np.concatenate([np.zeros(lead), synth])[:frame_count])
    else:
        raise ValueError(f"unsupported synth role: {role}")

    time = np.arange(frame_count) / SAMPLE_RATE
    noise = rng.standard_normal(frame_count)
    low = sosfilt(
        butter(4, cutoff / (SAMPLE_RATE / 2), btype="low", output="sos"), noise,
    )
    excitation = np.zeros(frame_count)
    excitation[0] = 3.0
    excitation += 1.2 * low * np.exp(-time / noise_decay)
    body = _resonator_bank(excitation, frequencies, gains, t60s)
    instantaneous_frequency = (
        (sweep_start - sweep_end) * np.exp(-time / .025) + sweep_end
    )
    phase = 2.0 * math.pi * np.cumsum(instantaneous_frequency) / SAMPLE_RATE
    sweep = np.sin(phase) * np.exp(-time / sweep_decay)
    synth = (
        body_amount * body
        + sweep_amount * sweep
        + direct_amount * low * np.exp(-time / noise_decay)
    )
    synth *= (
        1.0
        + .03 * np.sin(2.0 * math.pi * 7.7 * time)
        + .018 * np.sin(2.0 * math.pi * 14.9 * time + .5)
    )
    lead = int(.018 * SAMPLE_RATE)
    return _soft_normalize(np.concatenate([np.zeros(lead), synth])[:frame_count])


def _attack_burst(
    frame_count: int,
    seed: int,
    low_hz: float,
    high_hz: float,
    attack_seconds: float,
    decay_seconds: float,
) -> np.ndarray:
    time = np.arange(frame_count) / SAMPLE_RATE
    noise = np.random.default_rng(seed).standard_normal(frame_count)
    band = sosfilt(
        butter(
            3,
            [low_hz / (SAMPLE_RATE / 2), high_hz / (SAMPLE_RATE / 2)],
            btype="band",
            output="sos",
        ),
        noise,
    )
    envelope = (
        (1.0 - np.exp(-time / attack_seconds))
        * np.exp(-time / decay_seconds)
    )
    return band * envelope


def _calibrate_spectrum(role: str, synth: np.ndarray) -> np.ndarray:
    """Tune attack and modal balance while preserving the supplied method."""

    lead = int(.018 * SAMPLE_RATE)
    active_frames = len(synth) - lead
    time = np.arange(active_frames) / SAMPLE_RATE
    calibrated = synth.copy()

    if role in {"Big Don", "Don"}:
        seed = 9_201 if role == "Big Don" else 9_202
        burst = _attack_burst(
            active_frames, seed, 2_000.0, 14_000.0, .0006, .0045,
        )
        if role == "Big Don":
            calibrated[lead:] += 2.05 * burst
            calibrated[lead:] += (
                .25 * np.sin(2.0 * math.pi * 86.0 * time) * np.exp(-time / .50)
            )
        else:
            calibrated[lead:] += 2.40 * burst
            calibrated[lead:] += (
                .20 * np.sin(2.0 * math.pi * 125.0 * time) * np.exp(-time / .30)
            )
        return calibrated

    if role == "Katsu":
        low = sosfilt(
            butter(4, 3_000.0 / (SAMPLE_RATE / 2), btype="low", output="sos"),
            calibrated,
        )
        mid = sosfilt(
            butter(
                3,
                [1_400.0 / (SAMPLE_RATE / 2), 2_800.0 / (SAMPLE_RATE / 2)],
                btype="band",
                output="sos",
            ),
            calibrated,
        )
        burst = _attack_burst(
            active_frames, 9_203, 1_000.0, 10_000.0, .0003, .004,
        )
        calibrated = low
        calibrated[lead:] += 12.0 * mid[lead:] * np.exp(-time / .060)
        calibrated[lead:] += 10.0 * burst
        return calibrated

    raise ValueError(f"unsupported synth role: {role}")


def _rms_envelope(audio: np.ndarray, milliseconds: float = 6.0) -> np.ndarray:
    window = max(1, int(SAMPLE_RATE * milliseconds / 1_000.0))
    return np.sqrt(
        np.convolve(audio * audio, np.ones(window) / window, mode="same") + 1e-10
    )


def _match_reference_envelope(
    synth: np.ndarray,
    reference: np.ndarray,
) -> np.ndarray:
    frame_count = min(len(synth), len(reference))
    if frame_count <= 0:
        raise ValueError("synthesis and reference must both contain audio")
    reference = reference[:frame_count]
    matched = synth[:frame_count].copy()
    window = max(3, int(.012 * SAMPLE_RATE))
    kernel = np.hanning(window)
    kernel /= np.sum(kernel)
    for _ in range(2):
        ratio = _rms_envelope(reference) / (_rms_envelope(matched) + 1e-7)
        log_ratio = np.log(np.clip(ratio, .03, 30.0))
        matched *= np.exp(np.convolve(log_ratio, kernel, mode="same"))
    target_rms = float(np.sqrt(np.mean(reference * reference)))
    measured_rms = float(np.sqrt(np.mean(matched * matched)))
    if measured_rms <= 1e-12:
        raise ValueError("envelope calibration produced silence")
    matched *= target_rms / measured_rms
    peak_limit = min(.98, max(.90, float(np.max(np.abs(reference)))))
    measured_peak = float(np.max(np.abs(matched)))
    if measured_peak > peak_limit:
        matched *= peak_limit / measured_peak
    if not np.all(np.isfinite(matched)):
        raise ValueError("envelope calibration produced non-finite audio")
    return matched


def _anchor_frame(signal: np.ndarray) -> int:
    peak = float(np.max(np.abs(signal)))
    candidates = np.flatnonzero(np.abs(signal) >= peak * 10.0 ** (-30.0 / 20.0))
    if len(candidates) == 0:
        raise ValueError("synthesized voice has no measurable onset")
    return int(candidates[0])


def _profile(signal: np.ndarray, anchor: int) -> dict[str, float]:
    active = signal[anchor:]
    energy = active * active
    cumulative = np.cumsum(energy)
    spectrum = np.abs(np.fft.rfft(active * np.hanning(len(active)))) ** 2
    frequencies = np.fft.rfftfreq(len(active), 1.0 / SAMPLE_RATE)
    spectrum /= max(float(np.sum(spectrum)), 1e-15)
    attack = active[:int(round(.010 * SAMPLE_RATE))]
    attack_spectrum = np.abs(np.fft.rfft(attack * np.hanning(len(attack)))) ** 2
    attack_frequencies = np.fft.rfftfreq(len(attack), 1.0 / SAMPLE_RATE)
    tail = active[int(round(.120 * SAMPLE_RATE)):]
    tail_spectrum = np.abs(np.fft.rfft(tail * np.hanning(len(tail)))) ** 2
    tail_frequencies = np.fft.rfftfreq(len(tail), 1.0 / SAMPLE_RATE)
    return {
        "centroid_hz": round(float(np.sum(frequencies * spectrum)), 3),
        "attack_centroid_hz": round(float(
            np.sum(attack_frequencies * attack_spectrum)
            / max(float(np.sum(attack_spectrum)), 1e-15)
        ), 3),
        "tail_centroid_hz": round(float(
            np.sum(tail_frequencies * tail_spectrum)
            / max(float(np.sum(tail_spectrum)), 1e-15)
        ), 3),
        "peak_seconds_after_anchor": round(
            float(np.argmax(np.abs(active)) / SAMPLE_RATE), 6,
        ),
        "t90_seconds_after_anchor": round(
            float(np.searchsorted(cumulative, .90 * cumulative[-1]) / SAMPLE_RATE), 6,
        ),
        "t99_seconds_after_anchor": round(
            float(np.searchsorted(cumulative, .99 * cumulative[-1]) / SAMPLE_RATE), 6,
        ),
        "crest_db": round(
            float(20.0 * math.log10(
                np.max(np.abs(active)) / np.sqrt(np.mean(energy))
            )),
            3,
        ),
    }


def _write_pcm16(path: Path, signal: np.ndarray) -> tuple[str, str]:
    clipped = np.clip(signal, -1.0, 1.0)
    pcm = np.rint(clipped * 32767.0).astype("<i2")
    with wave.open(str(path), "wb") as output:
        output.setnchannels(1)
        output.setsampwidth(2)
        output.setframerate(SAMPLE_RATE)
        output.writeframes(pcm.tobytes())
    return _sha256(path), hashlib.sha256(pcm.tobytes()).hexdigest()


def build(reference_paths: ReferencePaths, output_dir: Path) -> None:
    references = {
        role: _read_reference(path, role)
        for role, path in reference_paths.by_role().items()
    }
    output_dir.mkdir(parents=True, exist_ok=True)
    for stale in output_dir.glob("*.wav"):
        stale.unlink()

    rng = np.random.default_rng(RANDOM_SEED)
    files: list[dict[str, object]] = []
    for role in BUILD_ORDER:
        synth = _make_synth(role, rng)
        synth = _calibrate_spectrum(role, synth)
        synth = _match_reference_envelope(synth, references[role])
        filename, kind = OUTPUT_NAMES[role]
        path = output_dir / filename
        file_hash, pcm_hash = _write_pcm16(path, synth)
        anchor = _anchor_frame(synth)
        files.append({
            "filename": filename,
            "kind": kind,
            "sha256": file_hash,
            "pcm_sha256": pcm_hash,
            "sample_rate": SAMPLE_RATE,
            "channels": 1,
            "frames": len(synth),
            "anchor_frame": anchor,
            "profile": _profile(synth, anchor),
        })

    manifest = {
        "schema_version": 2,
        "kit_id": "softchart-modal-resonator-v1",
        "license": "MIT",
        "sample_rate": SAMPLE_RATE,
        "generator": "scripts/build_taiko_synth.py",
        "method": (
            "seeded noise/impulse excitation, modal resonators, decaying pitch "
            "sweep, band-limited transient calibration, and offline RMS-envelope "
            "calibration"
        ),
        "random_seed": RANDOM_SEED,
        "reference_policy": (
            "Reference audio is used only for offline RMS-envelope and level "
            "calibration; it is not copied, mixed, embedded, or required at runtime."
        ),
        "files": sorted(files, key=lambda item: str(item["filename"])),
    }
    (output_dir / "manifest.json").write_text(
        json.dumps(manifest, indent=2, ensure_ascii=False) + "\n",
        encoding="utf-8",
    )


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--big-don-reference", type=Path, required=True)
    parser.add_argument("--don-reference", type=Path, required=True)
    parser.add_argument("--katsu-reference", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    args = parser.parse_args()
    build(
        ReferencePaths(
            big_don=args.big_don_reference,
            don=args.don_reference,
            katsu=args.katsu_reference,
        ),
        args.output_dir,
    )


if __name__ == "__main__":
    main()