"""Train/export the fixed ldcc ONNX contract; no-argument invocation runs self-test.

Python requires only torch, numpy, scipy and onnx. No onnxruntime is needed.
Parity: --verify-dump tensor.bin --data SELFTEST_DIR --rate 48000 (or 44100).
The fixture is SELFTEST_DIR/fixture.wav; --rate reinterprets samples, NOT resamples.
Train and export real data with --data DATA_DIR --output-dir bin.
"""
from __future__ import annotations

import argparse
import copy
import json
import math
import random
import tempfile
from functools import lru_cache
from pathlib import Path

import numpy as np
from scipy.io import wavfile
from scipy.signal import lfilter
import torch
from torch import nn
import onnx
from onnx.reference import ReferenceEvaluator

LABELS = ["latin", "electronic", "hiphop", "rock", "classical"]
N_FFT, N_MELS, FRAMES = 2048, 128, 300
FMIN, FMAX, LOG_FLOOR = 30.0, 11000.0, 1e-10
VERSION = 1


def read_wav(path: Path) -> tuple[np.ndarray, int]:
    sr, audio = wavfile.read(path)
    if sr not in (44100, 48000):
        raise ValueError(f"{path}: only 44100/48000 Hz supported, got {sr}")
    if audio.dtype == np.uint8:
        audio = (audio.astype(np.float64) - 128.0) / 128.0
    elif np.issubdtype(audio.dtype, np.signedinteger):
        # scipy returns 24-bit WAV left-justified in int32, matching this scaling.
        audio = audio.astype(np.float64) / float(2 ** (audio.dtype.itemsize * 8 - 1))
    elif np.issubdtype(audio.dtype, np.floating):
        audio = audio.astype(np.float64)
    else:
        raise ValueError(f"Unsupported WAV dtype: {audio.dtype}")
    if audio.ndim == 2:
        audio = audio.mean(axis=1)
    if audio.ndim != 1 or not np.all(np.isfinite(audio)):
        raise ValueError(f"{path}: invalid/nonfinite PCM")
    # C# ingestion is float32 mono. Promote only AFTER this rounding for FFT parity.
    return audio.astype(np.float32), sr


@lru_cache(maxsize=2)
def mel_bank(sr: int) -> np.ndarray:
    if sr not in (44100, 48000):
        raise ValueError("Expected 44100 or 48000 Hz")
    low = 2595.0 * math.log10(1 + FMIN / 700)
    high = 2595.0 * math.log10(1 + FMAX / 700)
    # Fractional edges: do NOT round/floor the FFT-bin locations.
    edges = np.array([
        700 * (10 ** ((low + (high - low) * i / (N_MELS + 1)) / 2595) - 1) * N_FFT / sr
        for i in range(N_MELS + 2)
    ], dtype=np.float64)
    bins = np.arange(N_FFT // 2 + 1, dtype=np.float64)
    filters = np.zeros((N_MELS, len(bins)), dtype=np.float64)
    for m in range(N_MELS):
        left, center, right = edges[m:m + 3]
        filters[m] = np.maximum(0, np.minimum((bins - left) / (center - left), (right - bins) / (right - center))) / (right - left)
    return filters


def stft_manual(x: np.ndarray, sr: int) -> np.ndarray:
    """Causal periodic-Hann complex rFFT, shape [1025, floor(len(x)/hop)]."""
    hop = sr // 100
    x = np.asarray(x, dtype=np.float32).astype(np.float64)
    if x.ndim != 1 or not np.isfinite(x).all():
        raise ValueError("Expected finite mono samples")
    if len(x) < hop:
        return np.empty((N_FFT // 2 + 1, 0), dtype=np.complex128)
    padded = np.pad(x, (N_FFT - hop, 0))
    frames = np.lib.stride_tricks.sliding_window_view(padded, N_FFT)[::hop]
    hann = 0.5 - 0.5 * np.cos(2 * np.pi * np.arange(N_FFT, dtype=np.float64) / N_FFT)
    return np.fft.rfft(frames * hann, axis=1).T


def logmel(x: np.ndarray, sr: int) -> np.ndarray:
    spectrum = stft_manual(x, sr)
    # Raw power: no /N and no doubling of positive-frequency bins.
    power = spectrum.real ** 2 + spectrum.imag ** 2
    energies = mel_bank(sr) @ power
    return np.log10(np.maximum(energies, LOG_FLOOR))


def tensor_frontend(x: np.ndarray, sr: int, mean: float, std: float) -> np.ndarray:
    needed = FRAMES * (sr // 100)
    if len(x) < needed:
        raise ValueError(f"Need {needed} samples for 300 frames, got {len(x)}")
    if not np.isfinite(mean) or not np.isfinite(std) or std <= 0:
        raise ValueError("Invalid normalization statistics")
    result = ((logmel(x[:needed], sr) - mean) / std).astype(np.float32)
    if result.shape != (N_MELS, FRAMES) or not np.isfinite(result).all():
        raise ValueError("Invalid frontend output")
    return result[None, None]


def discover(data: Path) -> list[tuple[Path, int]]:
    found = []
    for label_id, label in enumerate(LABELS):
        folder = data / label
        if folder.exists():
            found.extend((p.resolve(), label_id) for p in sorted(folder.iterdir()) if p.suffix.lower() == ".wav")
    if not found:
        raise ValueError("No WAV files in class subfolders")
    return found


def split_files(items, seed: int):
    rng = random.Random(seed)
    train, val = [], []
    for label in range(len(LABELS)):
        group = [item for item in items if item[1] == label]
        if not group:
            continue
        if len(group) < 2:
            raise ValueError(f"Class {LABELS[label]} needs at least two distinct files")
        rng.shuffle(group)
        nval = max(1, round(len(group) * 0.25))
        val.extend(group[:nval]); train.extend(group[nval:])
    if set(p for p, _ in train) & set(p for p, _ in val):
        raise AssertionError("File leakage between train and validation")
    return train, val


def random_crop(audio, sr, rng):
    count = FRAMES * (sr // 100)
    if len(audio) < count:
        raise ValueError("Training WAV shorter than 3 seconds")
    offset = int(rng.integers(0, len(audio) - count + 1))
    return audio[offset:offset + count].copy()


def training_statistics(items, seed):
    """Global scalar mean/std on TRAIN FILES ONLY; two reproducible crops/file."""
    rng = np.random.default_rng(seed)
    total = squares = 0.0
    count = 0
    for path, _ in items:
        audio, sr = read_wav(path)
        for _ in range(2):
            values = logmel(random_crop(audio, sr, rng), sr)
            total += float(values.sum()); squares += float(np.square(values).sum()); count += values.size
    mean = total / count
    std = math.sqrt(max(0, squares / count - mean * mean))
    if std < 1e-8:
        raise ValueError("Training set has zero variance")
    return mean, std


class WavDataset(torch.utils.data.Dataset):
    def __init__(self, items, mean, std, seed, augment):
        self.items, self.mean, self.std, self.seed, self.augment = items, mean, std, seed, augment
        self.rng = np.random.default_rng(seed)

    def __len__(self):
        return len(self.items)

    def __getitem__(self, index):
        path, label = self.items[index]
        audio, sr = read_wav(path)
        # Validation crop fixed across epochs; train crop randomized on each read.
        rng = self.rng if self.augment else np.random.default_rng(self.seed + index)
        x = random_crop(audio, sr, rng)
        if self.augment:
            x *= np.float32(10 ** (rng.uniform(-6, 6) / 20))
            if rng.random() < 0.5:
                x *= -1
        return torch.from_numpy(tensor_frontend(x, sr, self.mean, self.std)[0]), label


class GenreCNN(nn.Module):
    def __init__(self, active_labels):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 8, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(8, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(16, 24, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
            nn.AdaptiveAvgPool2d((1, 1)),
        )
        self.classifier = nn.Linear(24, 5)
        self.register_buffer("active", torch.tensor([i in active_labels for i in range(5)], dtype=torch.bool))

    def forward(self, x):
        logits = self.classifier(self.features(x).flatten(1))
        # Self-test keeps five outputs. Absent classes have EXACT zero probability.
        return torch.softmax(logits.masked_fill(~self.active, float("-inf")), dim=1)


def epoch(model, loader, optimizer=None):
    training = optimizer is not None
    model.train(training)
    loss_total, correct, count = 0.0, 0, 0
    with torch.set_grad_enabled(training):
        for x, labels in loader:
            if training:
                optimizer.zero_grad()
            probs = model(x)
            loss = -torch.log(probs[torch.arange(len(labels)), labels].clamp_min(1e-9)).mean()
            if training:
                loss.backward(); optimizer.step()
            loss_total += loss.item() * len(labels)
            correct += (probs.argmax(1) == labels).sum().item(); count += len(labels)
    return loss_total / count, correct / count


def validate_export(model_path: Path, model: GenreCNN, sample: np.ndarray):
    exported = onnx.load(str(model_path))
    onnx.checker.check_model(exported)
    shape = [dim.dim_value for dim in exported.graph.input[0].type.tensor_type.shape.dim]
    if shape != [1, 1, 128, 300]:
        raise AssertionError(f"Bad ONNX shape {shape}")
    output = ReferenceEvaluator(exported).run(None, {exported.graph.input[0].name: sample})[0]
    if output.shape != (1, 5) or not np.isfinite(output).all() or np.any(output < 0) or abs(float(output.sum()) - 1) > 1e-5:
        raise AssertionError(f"Invalid ONNX probabilities: {output}")
    with torch.no_grad():
        expected = model(torch.from_numpy(sample)).numpy()
    difference = float(np.max(np.abs(output - expected)))
    if difference > 1e-4:
        raise AssertionError(f"Torch/ONNX difference {difference}")
    inactive = ~model.active.cpu().numpy()
    if np.any(output[0, inactive] != 0):
        raise AssertionError("Unused classes must be exactly zero")
    print(f"ONNX OK: shape={shape}; sum={output.sum():.9f}; torch_max_abs={difference:.9g}; probs={output.tolist()}", flush=True)


def train(data: Path, output: Path, epochs: int, seed: int, batch_size: int, patience: int):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    items = discover(data)
    train_items, val_items = split_files(items, seed)
    mean, std = training_statistics(train_items, seed)
    output.mkdir(parents=True, exist_ok=True)
    manifest = {"train": [str(p) for p, _ in train_items], "val": [str(p) for p, _ in val_items]}
    (output / "split.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
    print(f"SPLIT: train_files={len(train_items)} val_files={len(val_items)} overlap=0; mean={mean:.12g} std={std:.12g}", flush=True)
    active = sorted({label for _, label in train_items})
    model = GenreCNN(active)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    train_loader = torch.utils.data.DataLoader(WavDataset(train_items, mean, std, seed, True), batch_size=batch_size, shuffle=True, num_workers=0)
    val_loader = torch.utils.data.DataLoader(WavDataset(val_items, mean, std, seed + 1, False), batch_size=batch_size, num_workers=0)
    best, stale, best_state = math.inf, 0, None
    checkpoint = output / "genre_model.pt"
    for number in range(1, epochs + 1):
        loss, accuracy = epoch(model, train_loader, optimizer)
        val_loss, val_accuracy = epoch(model, val_loader)
        print(f"EPOCH {number}/{epochs}: loss={loss:.6f} accuracy={accuracy:.3f} val_loss={val_loss:.6f} val_accuracy={val_accuracy:.3f}", flush=True)
        if val_loss < best:
            best, stale = val_loss, 0
            best_state = copy.deepcopy(model.state_dict())
            torch.save({"state_dict": best_state, "mean": mean, "std": std, "labels": LABELS, "version": VERSION, "epoch": number}, checkpoint)
        else:
            stale += 1
            if stale >= patience:
                print(f"EARLY STOP after {number} epochs", flush=True); break
    model.load_state_dict(best_state); model.eval()
    model_path, json_path = output / "genre_model.onnx", output / "genre_model.json"
    example = torch.zeros(1, 1, 128, 300, dtype=torch.float32)
    # Explicit legacy exporter: opset17, fixed shape, no dynamo/onnxscript requirement.
    torch.onnx.export(model, example, str(model_path), input_names=["logmel"], output_names=["probabilities"],
                      opset_version=17, dynamo=False, external_data=False)
    print("EXPORT: torch.onnx.export(dynamo=False), opset=17", flush=True)
    metadata = {"labels": LABELS, "mean": mean, "std": std, "version": VERSION,
                "frontend": {"fft": 2048, "hann": "periodic", "hop": "sr//100", "padding": "causal-left-N-hop",
                             "power": "abs(rfft)^2", "mel": "HTK", "bands": 128, "fmin": 30, "fmax": 11000,
                             "area": "divide by full support width in fractional bins", "log_floor": 1e-10}}
    json_path.write_text(json.dumps(metadata, indent=2), encoding="utf-8")
    audio, sr = read_wav(val_items[0][0])
    validate_export(model_path, model, tensor_frontend(audio, sr, mean, std))
    print(f"MODEL={model_path.resolve()}\nJSON={json_path.resolve()}\nCHECKPOINT={checkpoint.resolve()}", flush=True)
    return model_path, json_path


def synthetic_dataset(root: Path, seed: int):
    rng = np.random.default_rng(seed)
    sr = 48000
    time = np.arange(sr * 20, dtype=np.float64) / sr
    for label_id, label in enumerate(LABELS[:3]):
        folder = root / label; folder.mkdir(parents=True)
        for index in range(12):
            if label_id == 0:
                period = 60 / rng.uniform(85, 110)
                phase = np.mod(time, period)
                audio = np.exp(-phase * 45) * np.sin(2 * np.pi * (70 + index * 3) * time)
                audio += 0.2 * rng.standard_normal(len(time)) * np.exp(-np.mod(time + period * 0.375, period) * 80)
            elif label_id == 1:
                frequency = rng.uniform(150, 450)
                audio = sum(np.sin(2 * np.pi * frequency * harmonic * time + rng.uniform(0, 6)) / harmonic for harmonic in range(1, 6))
                audio *= 0.6 + 0.4 * np.sin(2 * np.pi * rng.uniform(0.5, 2) * time)
            else:
                audio = lfilter([1, -0.8], [1, -rng.uniform(0.1, 0.6)], rng.standard_normal(len(time)))
                audio *= 0.4 + 0.6 * np.sin(2 * np.pi * rng.uniform(1, 4) * time) ** 2
            audio = audio / max(float(np.max(np.abs(audio))), 1e-9) * rng.uniform(0.3, 0.8)
            wavfile.write(folder / f"synthetic_{index:02d}.wav", sr, np.round(audio * 32767).astype(np.int16))
    # A fixed, broadband fixture exercises low-energy bins without depending on a player.
    n = sr * 4
    t = np.arange(n) / sr
    fixture = 0.15 * np.sin(2 * np.pi * 440 * t) + 0.1 * np.sin(2 * np.pi * 4500 * t) + 0.03 * rng.standard_normal(n)
    wavfile.write(root / "fixture.wav", sr, np.round(np.clip(fixture, -1, 1) * 32767).astype(np.int16))
    print(f"SELFTEST_DIR={root.resolve()}\nFIXTURE={(root / 'fixture.wav').resolve()}", flush=True)


def verify_dump(path: Path, data: Path, rate: int | None, metadata_path: Path | None):
    if data is None:
        raise ValueError("--verify-dump needs --data pointing to the printed self-test directory")
    metadata = json.loads((metadata_path or data / "genre_model.json").read_text(encoding="utf-8"))
    if metadata["labels"] != LABELS or metadata["version"] != VERSION:
        raise ValueError("Model metadata contract mismatch")
    audio, sr = read_wav(data / "fixture.wav")
    if rate is not None:
        sr = rate
    expected = tensor_frontend(audio, sr, metadata["mean"], metadata["std"])
    if path.stat().st_size != 4 * N_MELS * FRAMES:
        raise ValueError("Dump must contain exactly 153600 bytes (38400 float32 values)")
    actual = np.fromfile(path, dtype="<f4").reshape(1, 1, N_MELS, FRAMES)
    if not np.isfinite(actual).all():
        raise ValueError("Non-finite tensor dump")
    difference = float(np.max(np.abs(actual.astype(np.float64) - expected.astype(np.float64))))
    print(f"PARITY rate={sr} shape={actual.shape} max_abs={difference:.12g} threshold=0.0001", flush=True)
    if difference >= 1e-4:
        raise AssertionError(f"Frontend parity failed: {difference}")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data", type=Path)
    parser.add_argument("--output-dir", type=Path, default=Path(__file__).resolve().parents[1] / "bin")
    parser.add_argument("--epochs", type=int, default=20)
    parser.add_argument("--batch-size", type=int, default=4)
    parser.add_argument("--patience", type=int, default=5)
    parser.add_argument("--seed", type=int, default=12345)
    parser.add_argument("--threads", type=int, default=4)
    parser.add_argument("--selftest", action="store_true")
    parser.add_argument("--verify-dump", type=Path)
    parser.add_argument("--metadata-json", type=Path)
    parser.add_argument("--rate", type=int, choices=[44100, 48000])
    args = parser.parse_args()
    if min(args.epochs, args.batch_size, args.patience, args.threads) < 1:
        parser.error("epochs/batch-size/patience/threads must be positive")
    torch.set_num_threads(args.threads)
    if args.verify_dump:
        verify_dump(args.verify_dump, args.data, args.rate, args.metadata_json)
    elif args.selftest or args.data is None:
        root = Path(tempfile.mkdtemp(prefix="ldcc-genre-selftest-"))
        synthetic_dataset(root, args.seed)
        train(root, root, 3, args.seed, args.batch_size, args.patience)
        print("SELFTEST PASS: synthetic only, not a usable musical genre model", flush=True)
    else:
        train(args.data, args.output_dir, args.epochs, args.seed, args.batch_size, args.patience)


if __name__ == "__main__":
    main()