"""Second generative model: an oscillator bank on a fixed log-frequency grid.

    python3 ml/ddsp2.py train --data DATASET_DIR --out RUN_DIR [--minutes 20] [--latent 8]
    python3 ml/ddsp2.py sample --model RUN_DIR/model.pt --out DIR [--n 8] [--seed 1] [--z v1,...]

Instead of letting the decoder choose partial frequencies (which the spectral
loss cannot steer, see ddsp.py), the synthesizer has one oscillator at each of
M = 96 bins per octave from 30 Hz to 8 kHz (about 770), and the decoder emits
for each an amplitude, a decay rate, and a frequency offset within its bin.
A partial is then an active bin, its exact frequency the bin plus the offset,
and the amplitudes have well-behaved gradients. The decoder is a 1D
transposed-convolutional network over the frequency axis, so that patterns of
lines can shift along it easily (a change of the sheet's frequency scale is a
translation on this axis). A shaped noise burst covers the attack. The loss
adds the smoothed log-frequency "line picture" of ml/linesim.py, computed in
torch, to the multi-scale spectral loss.
"""

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

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as Fn

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from features import build_cache  # noqa: E402
from ddsp import SR, DUR, N, NOISE_BANDS, Encoder, magnitude, spectral_loss, load_audio, write_wav, write_index  # noqa: E402

torch.backends.cudnn.enabled = False
F_LO, F_HI = 30.0, 8000.0
BPO = 96
M = int(math.log2(F_HI / F_LO) * BPO)  # oscillators
GRID = F_LO * 2 ** (np.arange(M) / BPO)
HALF_BIN = 2 ** (0.5 / BPO) - 1  # relative half width of a bin


class Decoder(nn.Module):
    def __init__(self, latent, base=25, ch=64):
        super().__init__()
        self.base = base
        self.fc = nn.Sequential(nn.Linear(latent, 256), nn.ReLU(inplace=True), nn.Linear(256, ch * base), nn.ReLU(inplace=True))
        self.ch = ch
        ups = []
        c = ch
        for k in range(5):  # base * 32 = 800 >= M
            ups += [nn.ConvTranspose1d(c + 1, ch, 4, stride=2, padding=1), nn.GroupNorm(8, ch), nn.ReLU(inplace=True)]
            c = ch
        self.ups = nn.ModuleList(ups)
        self.head = nn.Conv1d(ch + 1, 3, 5, padding=2)
        self.noise = nn.Sequential(nn.Linear(latent, 128), nn.ReLU(inplace=True), nn.Linear(128, NOISE_BANDS + 2))

    def forward(self, z):
        h = self.fc(z).view(-1, self.ch, self.base)
        i = 0
        while i < len(self.ups):
            coord = torch.linspace(0, 1, h.shape[-1], device=h.device).expand(h.shape[0], 1, -1)
            h = self.ups[i](torch.cat([h, coord], 1)); h = self.ups[i + 1](h); h = self.ups[i + 2](h)
            i += 3
        coord = torch.linspace(0, 1, h.shape[-1], device=h.device).expand(h.shape[0], 1, -1)
        out = self.head(torch.cat([h, coord], 1))[:, :, :M]  # (B, 3, M)
        amp = torch.sigmoid(out[:, 0] - 3.0)
        decay = 0.3 + 40 * torch.sigmoid(out[:, 1] - 1.0)
        offset = torch.tanh(out[:, 2]) * HALF_BIN
        nz = self.noise(z)
        return dict(amp=amp, decay=decay, offset=offset, noise_shape=torch.sigmoid(nz[:, :NOISE_BANDS] - 2.0),
                    noise_amp=torch.sigmoid(nz[:, NOISE_BANDS] - 2.0), noise_decay=2 + 60 * torch.sigmoid(nz[:, NOISE_BANDS + 1]))


_GRID_T = None


def synthesize(c, n=N, noise_seed=None):
    global _GRID_T
    B = c["amp"].shape[0]
    dev = c["amp"].device
    if _GRID_T is None or _GRID_T.device != dev:
        _GRID_T = torch.tensor(GRID, dtype=torch.float32, device=dev)
    t = torch.arange(n, device=dev, dtype=torch.float32) / SR
    f = _GRID_T[None, :] * (1 + c["offset"])  # (B, M)
    usable = (f < SR / 2 - 200).float()
    phase = 2 * math.pi * f[:, :, None] * t[None, None, :]
    env = (c["amp"] * usable)[:, :, None] * torch.exp(-c["decay"][:, :, None] * t[None, None, :])
    y = (env * torch.sin(phase)).sum(1)
    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))
    shape = Fn.interpolate(c["noise_shape"][:, None, :], size=S.shape[1], mode="linear", align_corners=True)[:, 0, :]
    tf = torch.arange(S.shape[2], device=dev, dtype=torch.float32) * hop / SR
    envn = c["noise_amp"][:, None] * torch.exp(-c["noise_decay"][:, None] * tf[None, :])
    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


# ---------- line-picture loss (linesim.py in torch) ----------

_PIC = {}


def picture_torch(x, tol=0.02, n_fft=4096, hop=1024, slices=((0, 0.3), (0.3, 1.0), (1.0, 2.0), (2.0, 3.0))):
    dev = x.device
    key = (tol, str(dev))
    if key not in _PIC:
        f_lin = np.arange(n_fft // 2 + 1) * SR / n_fft
        centers = F_LO * 2 ** (np.arange(int(math.log2(5000 / F_LO) * 48)) / 48)
        W = np.exp(-0.5 * ((np.log(np.maximum(f_lin[None, :], 1e-3) / centers[:, None])) / tol) ** 2)
        W[:, 0] = 0
        W /= np.maximum(W.sum(1, keepdims=True), 1e-12)
        _PIC[key] = torch.tensor(W, dtype=torch.float32, device=dev)
    W = _PIC[key]
    P = magnitude(x, n_fft) ** 2  # (B, bins, frames), hop n_fft // 4
    fps = SR / (n_fft // 4)
    outs = []
    for a, b in slices:
        sl = P[:, :, int(a * fps): max(int(a * fps) + 1, int(b * fps))].mean(-1)
        outs.append(W @ sl.T)  # (n_log, B)
    db = 10 * torch.log10(torch.stack(outs, 0) + 1e-12)  # (K, n_log, B)
    db = db - db.amax((0, 1), keepdim=True)
    return torch.clamp(db, min=-70)


def picture_loss(y, x):
    return (picture_torch(y) - picture_torch(x)).abs().mean() / 10


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 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, {M} oscillators, 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)
    print(f"{sum(p.numel() for p in model.parameters()) / 1e6:.2f} M parameters", flush=True)
    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")
    done = False
    while not done:
        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)
            pic = picture_loss(y, x)
            kl = -0.5 * torch.mean(1 + logvar - mu ** 2 - logvar.exp())
            beta = args.beta * min(1.0, step / 3000)
            loss = rec + pic + beta * torch.clamp(kl, min=args.free_bits)
            opt.zero_grad(); loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), 5.0)
            opt.step()
            step += 1
            minutes = (time.time() - t0) / 60
            if step % 50 == 0:
                msg = f"step {step} ({minutes:.1f} min): rec {rec.item():.4f} pic {pic.item():.4f} ({10 * pic.item():.2f} dB) kl {kl.item():.3f}"
                print(msg, flush=True); log.write(msg + "\n"); log.flush()
            if minutes - last_sample >= args.sample_every or minutes >= args.minutes:
                last_sample = minutes
                torch.save({"state": model.state_dict(), "latent": args.latent, "step": step, "kind": "bank"}, os.path.join(args.out, "model.pt"))
                write_samples(model, X, ids, fixed, os.path.join(args.out, "samples", f"m{int(minutes):03d}"), dev, step)
                write_index(args.out)
            if minutes >= args.minutes:
                done = True
                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()


@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=20); t.add_argument("--latent", type=int, default=8); t.add_argument("--batch", type=int, default=16)
    t.add_argument("--lr", type=float, default=5e-4); t.add_argument("--beta", type=float, default=2e-3); t.add_argument("--free-bits", type=float, default=0.5)
    t.add_argument("--sample-every", type=float, default=5)
    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)
    a = ap.parse_args()
    (train if a.cmd == "train" else sample)(a)