"""A small DDSP-style generative model of the struck-sheet clips.

    python3 ml/ddsp.py train --data DATASET_DIR --out RUN_DIR [--minutes 55] [--latent 8]
    python3 ml/ddsp.py sample --model RUN_DIR/model.pt --out DIR [--n 8] [--seed 1]

Encoder: a convolutional network on the cached log-frequency spectrogram
(384 x 133) gives the mean and log-variance of an 8-dimensional latent z
(a VAE, so that draws from the prior and interpolations are valid sounds).

Decoder: z -> the controls of a parametric synthesizer:
  - K partials with frequencies f_k (sorted, 30 Hz to 8 kHz), initial
    amplitudes a_k, decay rates d_k, and a pitch-glide amount g_k, so that
    partial k is a_k exp(-d_k t) sin(phase of f_k (1 + g_k exp(-t / tau)));
  - a noise burst with a 48-band spectral shape, an amplitude, and a decay,
    for the attack transient and the broadband part of the ring.
Synthesis is plain arithmetic at 24 kHz (the simulations carry little energy
above 6 kHz), so the same model can be run in JavaScript.

Loss: multi-scale log-spectral distance between the synthesized and the
original clip (FFT sizes 256 to 4096), plus a small KL term.
"""

import argparse
import json
import math
import os
import sys
import time
import wave

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as Fn
from scipy.signal import resample_poly

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from features import build_cache  # noqa: E402

torch.backends.cudnn.enabled = False
SR = 24000
DUR = 3.0
N = int(SR * DUR)
K = 128
NOISE_BANDS = 48
F_LO, F_HI = 30.0, 8000.0


def read_wav24(path):
    with wave.open(path) as w:
        x = np.frombuffer(w.readframes(w.getnframes()), dtype="<i2").astype(np.float32) / 32768
    y = resample_poly(x, 1, 2).astype(np.float32)
    out = np.zeros(N, np.float32)
    out[: min(N, len(y))] = y[:N]
    return out


def write_wav(path, y, sr=SR):
    y = np.asarray(y, np.float32)
    y = y / (np.abs(y).max() + 1e-9) * 0.95
    with wave.open(path, "wb") as w:
        w.setnchannels(1); w.setsampwidth(2); w.setframerate(sr)
        w.writeframes((y * 32767).astype("<i2").tobytes())


class Encoder(nn.Module):
    def __init__(self, latent):
        super().__init__()
        ch, layers, c_in = (16, 32, 64, 96, 128), [], 1
        for c in ch:
            layers += [nn.Conv2d(c_in, c, 3, padding=1), nn.BatchNorm2d(c), nn.ReLU(inplace=True), nn.MaxPool2d(2)]
            c_in = c
        self.conv = nn.Sequential(*layers)
        self.fc = nn.Sequential(nn.Linear(c_in * 12 * 4, 256), nn.ReLU(inplace=True), nn.Linear(256, 2 * latent))

    def forward(self, x):  # (B, 1, 384, 133)
        h = self.conv(x)  # (B, 128, 12, 4)
        return self.fc(h.flatten(1)).chunk(2, -1)


class Decoder(nn.Module):
    def __init__(self, latent):
        super().__init__()
        self.mlp = nn.Sequential(nn.Linear(latent, 256), nn.LayerNorm(256), nn.ReLU(inplace=True), nn.Linear(256, 512), nn.LayerNorm(512), nn.ReLU(inplace=True))
        self.freq = nn.Linear(512, K)
        self.amp = nn.Linear(512, K)
        self.decay = nn.Linear(512, K)
        self.glide = nn.Linear(512, K)
        self.noise_shape = nn.Linear(512, NOISE_BANDS)
        self.noise_amp = nn.Linear(512, 1)
        self.noise_decay = nn.Linear(512, 1)
        self.glide_tau = nn.Linear(512, 1)

    def forward(self, z):
        h = self.mlp(z)
        # Sorted frequencies: cumulative positive increments spanning the log range.
        inc = Fn.softplus(self.freq(h)) + 1e-3
        pos = torch.cumsum(inc, -1) / inc.sum(-1, keepdim=True)  # (B, K) in (0, 1]
        f = F_LO * (F_HI / F_LO) ** (pos * 0.98 + 0.005)
        a = torch.sigmoid(self.amp(h) - 2.0)  # amplitudes in (0, 1)
        d = 0.3 + 40 * torch.sigmoid(self.decay(h) - 1.0)  # 1/s
        g = 0.15 * torch.tanh(self.glide(h))  # relative initial detuning
        tau = 0.02 + 0.5 * torch.sigmoid(self.glide_tau(h))  # s
        ns = torch.sigmoid(self.noise_shape(h) - 2.0)
        na = torch.sigmoid(self.noise_amp(h) - 2.0)
        nd = 2 + 60 * torch.sigmoid(self.noise_decay(h))
        return dict(f=f, a=a, d=d, g=g, tau=tau, noise_shape=ns, noise_amp=na, noise_decay=nd)


def synthesize(c, n=N, noise_seed=None):
    """Renders a batch of control dicts to audio (B, n) at SR."""
    B = c["f"].shape[0]
    dev = c["f"].device
    t = torch.arange(n, device=dev, dtype=torch.float32) / SR
    # Partials: instantaneous frequency with a decaying glide; phase by cumulative sum.
    f_inst = c["f"][:, :, None] * (1 + c["g"][:, :, None] * torch.exp(-t[None, None, :] / c["tau"][:, :, None]))  # (B, K, n)
    phase = 2 * math.pi * torch.cumsum(f_inst, -1) / SR
    env = c["a"][:, :, None] * torch.exp(-c["d"][:, :, None] * t[None, None, :])
    nyq = (c["f"] < SR / 2 - 200).float()[:, :, None]
    y = (env * nyq * torch.sin(phase)).sum(1)
    # Noise burst: white noise shaped in the STFT domain by a fixed spectral shape and a decaying envelope.
    gen = torch.Generator(device=dev)
    if noise_seed is not None:
        gen.manual_seed(noise_seed)
    noise = torch.randn(B, n, device=dev, generator=gen)
    n_fft, hop = 512, 128
    win = torch.hann_window(n_fft, device=dev)
    S = torch.view_as_real(torch.stft(noise, n_fft, hop, window=win, return_complex=True))  # (B, 257, frames, 2)
    shape = Fn.interpolate(c["noise_shape"][:, None, :], size=S.shape[1], mode="linear", align_corners=True)[:, 0, :]  # (B, 257)
    tf = torch.arange(S.shape[2], device=dev, dtype=torch.float32) * hop / SR
    envn = c["noise_amp"] * torch.exp(-c["noise_decay"] * tf[None, :])  # (B, frames)
    S = S * (shape[:, :, None] * envn[:, None, :])[..., None]
    yn = torch.istft(torch.view_as_complex(S.contiguous()), n_fft, hop, window=win, length=n)
    return y + 0.05 * yn


def magnitude(x, n_fft):
    """STFT magnitude without complex-tensor ops (the office PyTorch cannot JIT-compile those)."""
    win = torch.hann_window(n_fft, device=x.device)
    S = torch.view_as_real(torch.stft(x, n_fft, n_fft // 4, window=win, return_complex=True))
    return torch.sqrt(S[..., 0] ** 2 + S[..., 1] ** 2 + 1e-12)


def spectral_loss(y, x):
    loss = 0.0
    for n_fft in (256, 512, 1024, 2048, 4096):
        Y = magnitude(y, n_fft)
        X = magnitude(x, n_fft)
        loss = loss + (Y - X).abs().mean() + 0.5 * (torch.log(Y + 1e-4) - torch.log(X + 1e-4)).abs().mean()
    return loss / 5


class Model(nn.Module):
    def __init__(self, latent=8):
        super().__init__()
        self.latent = latent
        self.enc = Encoder(latent)
        self.dec = Decoder(latent)

    def encode(self, spec):
        mu, logvar = self.enc(spec)
        return mu, logvar.clamp(-8, 4)


def load_audio(data, ids):
    cache = os.path.join(data, "ml", "audio24.npy")
    if os.path.exists(cache):
        a = np.load(cache, mmap_mode="r")
        if a.shape[0] == len(ids):
            return np.array(a)
    a = np.stack([read_wav24(os.path.join(data, "clips", i + ".wav")) for i in ids])
    np.save(cache, a)
    return a


def train(args):
    os.makedirs(os.path.join(args.out, "samples"), exist_ok=True)
    rows, feats = build_cache(args.data, os.path.join(args.data, "ml", "cache-logf"), kind="logf")
    ids = [r["id"] for r in rows]
    audio = load_audio(args.data, ids)
    print(f"{len(ids)} clips, spectrograms {feats.shape}, audio {audio.shape}", flush=True)
    dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    X = torch.from_numpy(feats)
    A = torch.from_numpy(audio)
    model = Model(args.latent).to(dev)
    opt = torch.optim.Adam(model.parameters(), lr=args.lr)
    rng = np.random.default_rng(0)
    fixed = [0, 1234, 2500, 4321]
    t0, step, last_sample = time.time(), 0, 0
    log = open(os.path.join(args.out, "train.log"), "a")
    while (time.time() - t0) / 60 < args.minutes:
        perm = rng.permutation(len(ids))
        for k in range(0, len(perm) - args.batch + 1, args.batch):
            b = perm[k:k + args.batch]
            spec = X[b].to(dev).float().unsqueeze(1)
            x = A[b].to(dev)
            mu, logvar = model.encode(spec)
            z = mu + torch.randn_like(mu) * torch.exp(0.5 * logvar)
            y = synthesize(model.dec(z))
            rec = spectral_loss(y, x)
            kl = -0.5 * torch.mean(1 + logvar - mu ** 2 - logvar.exp())
            beta = args.beta * min(1.0, step / 2000)
            loss = rec + beta * kl
            opt.zero_grad(); loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), 5.0)
            opt.step()
            step += 1
            if step % 50 == 0:
                msg = f"step {step} ({(time.time() - t0) / 60:.1f} min): rec {rec.item():.4f} kl {kl.item():.3f}"
                print(msg, flush=True); log.write(msg + "\n"); log.flush()
            if (time.time() - t0) / 60 - last_sample >= args.sample_every or (time.time() - t0) / 60 >= args.minutes:
                last_sample = (time.time() - t0) / 60
                torch.save({"state": model.state_dict(), "latent": args.latent, "step": step}, os.path.join(args.out, "model.pt"))
                write_samples(model, X, ids, fixed, os.path.join(args.out, "samples", f"m{int(last_sample):03d}"), dev, step)
                write_index(args.out)
            if (time.time() - t0) / 60 >= args.minutes:
                break
    print("done", step, "steps")


@torch.no_grad()
def write_samples(model, X, ids, fixed, out, dev, step, n_prior=6):
    model.eval()
    os.makedirs(out, exist_ok=True)
    names = []
    spec = X[fixed].to(dev).float().unsqueeze(1)
    mu, _ = model.encode(spec)
    y = synthesize(model.dec(mu), noise_seed=1).cpu().numpy()
    for i, k in enumerate(fixed):
        write_wav(os.path.join(out, f"recon-{ids[k]}.wav"), y[i]); names.append((f"reconstruction of {ids[k]}", f"recon-{ids[k]}.wav"))
    g = torch.Generator(device=dev); g.manual_seed(7)
    zp = torch.randn(n_prior, model.latent, device=dev, generator=g)
    y = synthesize(model.dec(zp), noise_seed=2).cpu().numpy()
    for i in range(n_prior):
        write_wav(os.path.join(out, f"prior-{i}.wav"), y[i]); names.append((f"prior sample {i}", f"prior-{i}.wav"))
    steps = 5
    zi = torch.stack([mu[0] + (mu[1] - mu[0]) * s / (steps - 1) for s in range(steps)])
    y = synthesize(model.dec(zi), noise_seed=3).cpu().numpy()
    for s in range(steps):
        write_wav(os.path.join(out, f"interp-{s}.wav"), y[s]); names.append((f"interpolation {ids[fixed[0]]} -> {ids[fixed[1]]}, step {s}", f"interp-{s}.wav"))
    with open(os.path.join(out, "list.json"), "w") as f:
        json.dump({"step": step, "items": names}, f)
    model.train()


def write_index(run):
    sdir = os.path.join(run, "samples")
    snaps = sorted(d for d in os.listdir(sdir) if os.path.isdir(os.path.join(sdir, d)))
    blocks = []
    for d in reversed(snaps):
        try:
            meta = json.load(open(os.path.join(sdir, d, "list.json")))
        except Exception:
            continue
        rows = "".join(f"<tr><td class='l'>{label}</td><td class='l'><audio controls preload='none' src='samples/{d}/{fn}'></audio></td></tr>" for label, fn in meta["items"])
        blocks.append(f"<h2>after {int(d[1:])} min, step {meta['step']}</h2><table>{rows}</table>")
    html = ("<!doctype html><meta charset='utf-8'><title>DDSP sheet model samples</title><style>body{font:14px system-ui,sans-serif;margin:16px} table{border-collapse:collapse} td{padding:3px 8px;border-bottom:1px solid #ddd} td.l{text-align:left} audio{height:28px;width:260px} h2{font-size:15px;margin-top:18px}</style>"
            "<h1>DDSP-style sheet model: samples during training</h1><p>Reconstructions of four fixed dataset clips (compare with clips/ID.wav in the gallery), six draws from the prior, and a five-step interpolation between two clips. Newest first.</p>" + "".join(blocks))
    with open(os.path.join(run, "samples.html"), "w") as f:
        f.write(html)


@torch.no_grad()
def sample(args):
    ck = torch.load(args.model, map_location="cpu", weights_only=False)
    model = Model(ck["latent"]); model.load_state_dict(ck["state"]); model.eval()
    os.makedirs(args.out, exist_ok=True)
    g = torch.Generator(); g.manual_seed(args.seed)
    z = torch.randn(args.n, ck["latent"], generator=g) if args.z is None else torch.tensor([[float(v) for v in args.z.split(",")]])
    y = synthesize(model.dec(z), noise_seed=args.seed).numpy()
    for i in range(y.shape[0]):
        write_wav(os.path.join(args.out, f"sample-{i}.wav"), y[i])
        print(f"sample-{i}.wav  z = " + ", ".join(f"{v:.2f}" for v in z[i].tolist()))


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="cmd", required=True)
    t = sub.add_parser("train"); t.add_argument("--data", required=True); t.add_argument("--out", required=True)
    t.add_argument("--minutes", type=float, default=55); t.add_argument("--latent", type=int, default=8); t.add_argument("--batch", type=int, default=16)
    t.add_argument("--lr", type=float, default=1e-3); t.add_argument("--beta", type=float, default=1e-3); t.add_argument("--sample-every", type=float, default=10)
    s = sub.add_parser("sample"); s.add_argument("--model", required=True); s.add_argument("--out", required=True)
    s.add_argument("--n", type=int, default=8); s.add_argument("--seed", type=int, default=1); s.add_argument("--z", default=None, help="comma-separated latent vector for one clip")
    a = ap.parse_args()
    (train if a.cmd == "train" else sample)(a)