"""Train the decoder student on the deployment distribution.

The pack holds, per row: the acoustic student's mels (the exact inference
input), the teacher's phoneme ids/durations, and, as audio targets, full
vocoder renders of those exact mels. The objective is therefore "reproduce the
vocoder's output on the mels you will actually receive", and there is no
acoustic->decoder distribution gap to fix later.

This is the full supervision stack used for the released weights:

  waveform L1                       (wav-weight 1.0 for this model)
  multi-resolution spectral loss
  sig-temporal (cosine-gram)        (0.4)  temporal texture
  feature-hint (channel moments)    (0.1)  distillation from the vocoder
  quiet-ceiling                     (0.010) stops quiet-frame hiss
  high-band-excess                  (0.1)  penalizes HF overshoot
  MRD adversarial + feature match   (0.075 / 0.75 from step 1500)

Requires `vocos` (pip install vocos) for the feature hooks; the pack itself
needs no vocoder at train time.

Stages used for the release (see train/README.md for wall-clock):
    recovery  --mix-prob 0.0 --const-lr --noise-zeros
    ... then short continuation arms and an averaging pass over the arms
        (the released weights are a weight average of five same-recipe runs;
         averaging paid +0.09 SCOREQ on this line, see the README).
"""
import argparse
import glob
import os
import time

import numpy as np
import torch

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import torch.nn as nn
import torch.nn.functional as F

from microtts.models import DecoderStudent


# --------------------------------------------------------------------------
# Teacher feature hooks (vocoder internals, for the hint + sig-temporal terms)
# --------------------------------------------------------------------------
class TeacherFeats:
    """Capture the vocoder's internal activations at chosen stages."""

    def __init__(self, vocos):
        self.feats = {}
        self.handles = []
        b = vocos.backbone
        self.handles.append(b.embed.register_forward_hook(self._h("embed")))
        self.handles.append(b.convnext[3].register_forward_hook(self._h("cn3")))
        self.handles.append(b.final_layer_norm.register_forward_hook(self._h("final")))

    def _h(self, key):
        def fn(mod, inp, out):
            self.feats[key] = out
        return fn

    def clear(self):
        self.feats = {}


def standardized_time_series(v):
    mean = v.mean(dim=1, keepdim=True)
    std = v.std(dim=1, keepdim=True).clamp_min(1e-4)
    return (v - mean) / std


def temporal_cosine_gram(value):
    """Per-timestep channel-vector self-similarity [B,T,T]."""
    if value.ndim != 3:
        raise RuntimeError(f"expected [B, C, T], got {value.shape}")
    series = value.transpose(1, 2)
    series = series - series.mean(dim=-1, keepdim=True)
    series = F.normalize(series, dim=-1, eps=1e-6)
    return torch.matmul(series, series.transpose(1, 2))


def signature_temporal_loss(student_feats, teacher_feats, keys=("embed", "cn3")):
    losses = []
    for key in keys:
        s = student_feats.get(key)
        t = teacher_feats.get(key)
        if s is None or t is None:
            continue
        if s.ndim == 3:
            s = s[0]
        if t.ndim == 3:
            t = t[0]
        T = min(s.shape[-1], t.shape[-1])
        if T < 4:
            continue
        s = s[:, :T].unsqueeze(0).float()
        t = t[:, :T].unsqueeze(0).float()
        losses.append(F.l1_loss(temporal_cosine_gram(s), temporal_cosine_gram(t)))
    if not losses:
        return torch.zeros((), device=next(iter(student_feats.values())).device)
    return sum(losses) / len(losses)


def channel_moment_hint_loss(student_feats, teacher_feats):
    """Standardized channel-mean + log-RMS envelopes."""
    losses = []
    for key in ("embed", "cn3"):
        s = student_feats.get(key)
        t = teacher_feats.get(key)
        if s is None or t is None:
            continue
        if s.ndim == 3:
            s = s[0]
        if t.ndim == 3:
            t = t[0]
        if s.shape[-1] != t.shape[-1]:
            T = min(s.shape[-1], t.shape[-1])
            s, t = s[:, :T], t[:, :T]
        s_mean = standardized_time_series(s.mean(dim=0, keepdim=True))
        t_mean = standardized_time_series(t.mean(dim=0, keepdim=True))
        s_rms = standardized_time_series(torch.log1p(torch.sqrt(s.square().mean(dim=0, keepdim=True) + 1e-6)))
        t_rms = standardized_time_series(torch.log1p(torch.sqrt(t.square().mean(dim=0, keepdim=True) + 1e-6)))
        losses.append(F.l1_loss(s_mean, t_mean))
        losses.append(F.l1_loss(s_rms, t_rms))
    if not losses:
        return torch.zeros((), device=next(iter(student_feats.values())).device)
    return sum(losses) / len(losses)


# --------------------------------------------------------------------------
# Spectral-resolution discriminator (UnivNet-style MRD)
# --------------------------------------------------------------------------
class SpectralResolutionDiscriminator(nn.Module):
    def __init__(self, n_fft, hop_length, win_length, channels=(32, 32, 32, 32)):
        super().__init__()
        self.n_fft, self.hop_length, self.win_length = n_fft, hop_length, win_length
        ks, st, pd = (3, 9), (1, 2), (1, 4)
        layers = []
        in_ch = 1
        for out_ch in channels:
            layers.append(nn.Conv2d(in_ch, out_ch, kernel_size=ks, stride=st, padding=pd))
            in_ch = out_ch
        self.layers = nn.ModuleList(layers)
        self.post = nn.Conv2d(in_ch, 1, kernel_size=ks, padding=pd)

    def spec(self, x):  # x: [1, T]
        if x.shape[-1] < self.n_fft:
            return None
        w = torch.hann_window(self.n_fft, device=x.device, dtype=x.dtype)
        s = torch.stft(x[0], self.n_fft, self.hop_length, win_length=self.win_length,
                       window=w, return_complex=True)
        return torch.log1p(s.abs()).unsqueeze(0).unsqueeze(0)  # [1,1,F,T]

    def forward(self, x):
        h = self.spec(x)
        if h is None:
            return None, None
        feats = []
        for layer in self.layers:
            h = F.leaky_relu(layer(h), 0.1)
            feats.append(h)
        score = self.post(h)
        feats.append(score)
        return score, feats


class MRD(nn.Module):
    def __init__(self, configs=((512, 128, 512), (1024, 256, 1024), (2048, 512, 2048))):
        super().__init__()
        self.discs = nn.ModuleList([SpectralResolutionDiscriminator(*c) for c in configs])

    def forward(self, x):
        scores, feats = [], []
        for d in self.discs:
            s, f = d(x)
            if s is None:
                continue
            scores.append(s)
            feats.append(f)
        return scores, feats


def generator_lsgan_loss(fake_scores):
    if not fake_scores:
        return torch.zeros(())
    total = sum((s - 1.0).square().mean() for s in fake_scores)
    return total / len(fake_scores)


def discriminator_lsgan_loss(real_scores, fake_scores):
    if not real_scores:
        return torch.zeros(())
    total = sum((r - 1.0).square().mean() + f.square().mean()
                for r, f in zip(real_scores, fake_scores))
    return total / len(real_scores)


def feature_matching_loss(real_feats, fake_feats):
    losses = []
    for rf, ff in zip(real_feats, fake_feats):
        for a, b in zip(rf, ff):
            losses.append(F.l1_loss(b, a.detach()))
    if not losses:
        return torch.zeros(())
    return sum(losses) / len(losses)


# --------------------------------------------------------------------------
# Spectral losses
# --------------------------------------------------------------------------
def stft_loss(y, yh, device):
    loss = 0.0
    n = min(y.shape[-1], yh.shape[-1])
    for n_fft, hop in [(512, 128), (1024, 256), (2048, 512)]:
        if n < n_fft:
            continue  # skip FFTs larger than the input
        Y = torch.stft(y[..., :n], n_fft, hop, window=torch.hann_window(n_fft, device=device),
                       return_complex=True, center=False)
        Yh = torch.stft(yh[..., :n], n_fft, hop, window=torch.hann_window(n_fft, device=device),
                        return_complex=True, center=False)
        T = min(Y.shape[-1], Yh.shape[-1])
        mag, mag_r = Y[..., :T].abs(), Yh[..., :T].abs()
        sc = (mag - mag_r).norm(dim=-2) / (mag.norm(dim=-2) + mag_r.norm(dim=-2) + 1e-6)
        lm = (mag.clamp_min(1e-6).log10() - mag_r.clamp_min(1e-6).log10()).abs().mean()
        loss = loss + sc.mean() + lm
    return loss / 3.0


def high_band_excess_loss(prediction, target, sample_rate=24000, high_band_hz=6000.0, margin_db=1.0):
    """Penalize HF energy ratio above the target's ratio times a margin."""
    n_fft, hop = 1024, 256
    window = torch.hann_window(n_fft, device=prediction.device, dtype=prediction.dtype)
    ps = torch.stft(prediction[0], n_fft, hop, win_length=n_fft, window=window, return_complex=True)
    ts = torch.stft(target[0], n_fft, hop, win_length=n_fft, window=window, return_complex=True)
    nyq = sample_rate / 2.0
    freqs = torch.linspace(0.0, nyq, n_fft // 2 + 1, device=prediction.device, dtype=prediction.dtype)
    mask = (freqs >= high_band_hz).to(dtype=prediction.dtype).unsqueeze(-1)
    p_pow = ps.abs().square()
    t_pow = ts.abs().square()
    pr = torch.sum(p_pow * mask) / torch.sum(p_pow).clamp_min(1e-10)
    tr = torch.sum(t_pow * mask) / torch.sum(t_pow).clamp_min(1e-10)
    margin_ratio = 10.0 ** (margin_db / 10.0)
    excess = torch.log(pr.clamp_min(1e-10)) - torch.log((tr * margin_ratio).clamp_min(1e-10))
    return F.relu(excess).square()


def framed_rms(x, frame_size=1024, frame_hop=256):
    if x.ndim == 2:
        x = x[0]
    if x.ndim == 1:
        x = x.unsqueeze(0)
    if x.shape[-1] < frame_size:
        return x.square().mean(dim=-1).sqrt().unsqueeze(-1)
    frames = x.unfold(-1, frame_size, frame_hop)
    return frames.square().mean(dim=-1).sqrt().mean(dim=0, keepdim=True)


def quiet_ceiling_loss(prediction, target, quantile=0.25, frame_size=1024, frame_hop=256, margin_db=0.5):
    """Penalize predictions louder than the target on the target's quiet frames."""
    tr = framed_rms(target, frame_size, frame_hop)
    pr = framed_rms(prediction, frame_size, frame_hop)
    T = min(tr.shape[-1], pr.shape[-1])
    tr, pr = tr[..., :T], pr[..., :T]
    thr = torch.quantile(tr.flatten(), quantile)
    mask = (torch.log1p(tr) <= thr).to(prediction.dtype)
    db = 20.0 / np.log(10.0)
    excess_db = db * (torch.log(pr + 1e-8) - torch.log(tr + 1e-8)) - margin_db
    return torch.sum(F.relu(excess_db) * mask) / mask.sum().clamp_min(1.0) / 20.0


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--pack", default="packs/deploy")
    ap.add_argument("--steps", type=int, default=20000)
    ap.add_argument("--batch", type=int, default=4)
    ap.add_argument("--lr", type=float, default=2e-4)
    ap.add_argument("--device", default="cuda")
    ap.add_argument("--tag", default="dec")
    ap.add_argument("--mix-prob", type=float, default=0.0,
                    help="fraction of items that use live acoustic mels instead "
                         "of the pack's (the deployment pack needs 0.0)")
    ap.add_argument("--init", required=True,
                    help="teacher-initialized decoder checkpoint (init_decoder.py) or a "
                         "previous decoder checkpoint for continuation runs")
    ap.add_argument("--resume-step", type=int, default=0,
                    help="step number the init checkpoint was saved at (continuation runs)")
    ap.add_argument("--dec-dim", type=int, default=130)
    ap.add_argument("--dec-pw", type=int, default=337)
    ap.add_argument("--dec-blocks", type=int, default=4)
    ap.add_argument("--ac-hidden", type=int, default=72)
    ap.add_argument("--ac-depth", type=int, default=5)
    ap.add_argument("--ac-ckpt", default="acoustic.pt")
    ap.add_argument("--dur-hidden", type=int, default=12)
    ap.add_argument("--dur-ckpt", default="duration.pt")
    # loss weights (the released recipe)
    ap.add_argument("--feature-hint-weight", type=float, default=0.1)
    ap.add_argument("--quiet-ceiling-weight", type=float, default=0.010)
    ap.add_argument("--high-band-excess-weight", type=float, default=0.1)
    ap.add_argument("--adv-weight", type=float, default=0.075)
    ap.add_argument("--adv-feature-weight", type=float, default=0.75)
    ap.add_argument("--adv-lr", type=float, default=2e-4)
    ap.add_argument("--adv-start-step", type=int, default=1500)
    ap.add_argument("--const-lr", action="store_true",
                    help="constant lr (the released recipe)")
    ap.add_argument("--betas", type=str, default="0.8,0.99")
    ap.add_argument("--wav-weight", type=float, default=1.0)
    ap.add_argument("--sig-temporal-weight", type=float, default=0.4)
    ap.add_argument("--loud-weight", type=float, default=0.0)
    ap.add_argument("--seg", type=int, default=64,
                    help="training crop length in mel frames")
    ap.add_argument("--noise-zeros", action="store_true",
                    help="zeros instead of randn for the decoder noise input "
                         "(matches inference; the released recipe)")
    args = ap.parse_args()

    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
    files = sorted(glob.glob(os.path.join(args.pack, "*.npz")))
    print(f"pack: {len(files)} rows, device {device}", flush=True)

    from microtts.models import AcousticStudent, DurationStudent

    ac = AcousticStudent(hidden=args.ac_hidden, depth=args.ac_depth).to(device)
    ac.load_state_dict(torch.load(args.ac_ckpt, map_location=device, weights_only=False))
    ac.eval()
    dur = DurationStudent(hidden=args.dur_hidden).to(device)
    dur.load_state_dict(torch.load(args.dur_ckpt, map_location=device, weights_only=False))
    dur.eval()

    model = DecoderStudent(dim=args.dec_dim, pw_hidden=args.dec_pw, blocks=args.dec_blocks).to(device)
    model.load_state_dict(torch.load(args.init, map_location=device, weights_only=False))
    n = sum(p.numel() for p in model.parameters())
    print(f"decoder params: {n:,} (init {args.init})", flush=True)

    from vocos import Vocos
    teacher = Vocos.from_pretrained("charactr/vocos-mel-24khz").to(device).eval()
    for p in teacher.parameters():
        p.requires_grad_(False)
    tf = TeacherFeats(teacher)
    print("vocoder loaded with hooks", flush=True)

    mrd = MRD().to(device)
    betas = tuple(float(x) for x in args.betas.split(","))
    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=betas)
    opt_d = torch.optim.AdamW(mrd.parameters(), lr=args.adv_lr, betas=betas)
    sched = None if args.const_lr else torch.optim.lr_scheduler.ExponentialLR(opt, gamma=0.9998)
    if sched is not None and args.resume_step > 0:
        for _ in range(args.resume_step):
            sched.step()

    RUN = os.path.join("runs", args.tag)
    os.makedirs(RUN, exist_ok=True)
    t0 = time.time()
    for step in range(args.resume_step + 1, args.steps + 1):
        idx = np.random.choice(len(files), min(args.batch, len(files)), replace=False)
        cnt = 0
        content_total = None
        opt.zero_grad()
        gan_pair = None
        for ii, i in enumerate(idx):
            d = np.load(files[i])
            mel_t = torch.from_numpy(d["mel"]).to(device)
            audio = torch.from_numpy(d["audio"]).to(device)
            T = mel_t.shape[1]
            seg = args.seg
            s = 0
            if T > seg:
                s = np.random.randint(0, T - seg)
            if np.random.rand() < args.mix_prob:
                # live student mel (used only for non-deployment packs)
                ids = torch.from_numpy(d["phone_ids"]).long().to(device)
                with torch.no_grad():
                    dd = dur(ids)
                    mel_s = ac(ids, dd)
                seg_s = min(seg, mel_s.shape[1])
                if seg_s < 16:
                    continue
                s_s = np.random.randint(0, max(1, mel_s.shape[1] - seg_s + 1))
                mel = mel_s[:, s_s:s_s + seg_s]
                a0, a1 = s_s * 256, (s_s + seg_s) * 256
                if a1 > audio.shape[0]:
                    a1 = audio.shape[0]
                if a1 - a0 < 1024:
                    continue
                audio = audio[a0:a1]
            else:
                mel = mel_t[:, s:s + seg] if T > seg else mel_t
                audio = audio[s * 256:(s + mel.shape[1]) * 256]
            noise = torch.zeros(4, mel.shape[1], device=device) if args.noise_zeros \
                else torch.randn(4, mel.shape[1], device=device)
            wav = model.synthesize(mel, noise)
            n_w = min(wav.shape[-1], audio.shape[-1])
            wav_c, aud_c = wav[:n_w].unsqueeze(0), audio[:n_w].unsqueeze(0)
            if gan_pair is None:
                gan_pair = (wav_c, aud_c)

            l_wav = F.l1_loss(wav[:n_w], audio[:n_w])
            l_stft = stft_loss(wav_c, aud_c, device)
            rms_p = wav[:n_w].square().mean().clamp_min(1e-8).sqrt()
            rms_t = audio[:n_w].square().mean().clamp_min(1e-8).sqrt()
            l_loud = F.l1_loss(rms_p, rms_t)

            # feature-hint + signature-temporal: first item of the step only
            l_hint = torch.zeros((), device=device)
            l_sigt = torch.zeros((), device=device)
            need_feats = (args.feature_hint_weight > 0 or args.sig_temporal_weight > 0) and ii == 0
            if need_feats:
                with torch.no_grad():
                    tf.clear()
                    teacher.decode(mel.unsqueeze(0))
                    t_feats = {k: v.detach() for k, v in tf.feats.items()}
                s_feats = {}
                x = model.embed(mel) + model.noise_adapter(noise)
                s_feats["embed"] = x
                x2 = model.norm(x.transpose(0, 1)).transpose(0, 1)
                cn_idx = min(3, len(model.blocks) - 1)
                for bi, b in enumerate(model.blocks):
                    x2 = b(x2)
                    if bi == cn_idx:
                        s_feats["cn3"] = x2
                t_embed = t_feats.get("embed")
                if t_embed is not None:
                    t_embed = t_embed[0]
                    Tt = min(t_embed.shape[-1], s_feats["embed"].shape[-1])
                    t_feats["embed"] = t_embed[:, :Tt].unsqueeze(0)
                    s_feats["embed"] = s_feats["embed"][:, :Tt].unsqueeze(0)
                t_cn3 = t_feats.get("cn3")
                if t_cn3 is not None:
                    t_cn3 = t_cn3[0]
                    Tt = min(t_cn3.shape[-1], s_feats["cn3"].shape[-1])
                    t_feats["cn3"] = t_cn3[:, :Tt].unsqueeze(0)
                    s_feats["cn3"] = s_feats["cn3"][:, :Tt].unsqueeze(0)
                l_hint = channel_moment_hint_loss(s_feats, t_feats)
                if args.sig_temporal_weight > 0:
                    l_sigt = signature_temporal_loss(s_feats, t_feats)

            l_hf = high_band_excess_loss(wav_c, aud_c, 24000, 6000.0, 1.0) \
                if args.high_band_excess_weight > 0 else torch.zeros((), device=device)
            l_quiet = quiet_ceiling_loss(wav_c[0], aud_c[0], 0.25, 1024, 256, 0.5) \
                if args.quiet_ceiling_weight > 0 else torch.zeros((), device=device)

            item = (args.wav_weight * l_wav + 1.0 * l_stft + args.loud_weight * l_loud
                    + args.feature_hint_weight * l_hint
                    + args.sig_temporal_weight * l_sigt
                    + args.high_band_excess_weight * l_hf
                    + args.quiet_ceiling_weight * l_quiet) / max(args.batch, 1)
            content_total = item if content_total is None else content_total + item
            cnt += 1

        if cnt == 0 or content_total is None:
            continue

        content_total.backward(retain_graph=True)
        l_adv = torch.zeros((), device=device)
        l_adv_feat = torch.zeros((), device=device)
        if step >= args.adv_start_step and gan_pair is not None:
            wav_g, aud_g = gan_pair
            nw = min(wav_g.shape[-1], aud_g.shape[-1])
            wav_g, aud_g = wav_g[..., :nw], aud_g[..., :nw]
            opt_d.zero_grad()
            real_scores, _ = mrd(aud_g)
            fake_scores, _ = mrd(wav_g.detach())
            d_loss = discriminator_lsgan_loss(real_scores, fake_scores)
            d_loss.backward()
            torch.nn.utils.clip_grad_norm_(mrd.parameters(), 5.0)
            opt_d.step()
            for p in mrd.parameters():
                p.requires_grad_(False)
            fake_scores_g, fake_feats_g = mrd(wav_g)
            _, real_feats_g = mrd(aud_g)
            l_adv = generator_lsgan_loss(fake_scores_g)
            l_adv_feat = feature_matching_loss(real_feats_g, fake_feats_g)
            for p in mrd.parameters():
                p.requires_grad_(True)
            gan_loss = (args.adv_weight * l_adv + args.adv_feature_weight * l_adv_feat) / max(args.batch, 1)
            gan_loss.backward()

        torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
        opt.step()
        if sched is not None:
            sched.step()

        if step % 100 == 0 or step == 1:
            dt = time.time() - t0
            print(f"step {step}/{args.steps} | wav {l_wav.item():.4f} stft {l_stft.item():.3f} "
                  f"hint {l_hint.item():.4f} sigt {l_sigt.item():.4f} hf {l_hf.item():.4f} "
                  f"quiet {l_quiet.item():.5f} adv {l_adv.item():.3f} fm {l_adv_feat.item():.3f} "
                  f"| {step/dt:.1f} sps", flush=True)
        if step % 500 == 0:
            torch.save(model.state_dict(), os.path.join(RUN, "last.pt"))
            with open(os.path.join(RUN, "last_step.txt"), "w") as f:
                f.write(str(step))

    torch.save(model.state_dict(), os.path.join(RUN, "final.pt"))
    print("DECODER DONE")


if __name__ == "__main__":
    main()