#!/usr/bin/env python3
"""Warm-start expansion of the enhancer-small DAC decoder (188M -> >=300M).

Takes the released `Scicom-intl/CallEnhancer` decoder checkpoint (3072 ch, 188M,
500k steps) and DEEPENS it by appending extra ResidualUnits inside the widest
DecoderBlocks, so all 500k steps of training are kept (the neucodec-44k
`extend_decoder.py` trick, adapted to a CNN pyramid).

Why not "duplicate the last layer" literally: the DAC decoder is not a
homogeneous transformer stack — widths halve per block (3072->1536->768->384->
192->96), so the literal last layers are tiny (a 96-ch ResidualUnit is ~0.07M).
The parameter mass lives in the wide early blocks: one ResidualUnit at 1536 ch
is ~18.9M (@400 Hz frame rate — cheap compute) and one at 768 ch is ~4.7M
(@2 kHz). The default `1536:9,768:5` takes 188M -> ~381M
(9x18.9 + 5x4.7 = +193M), i.e. slightly over 2x the released decoder.

Why duplicate + ZERO-GATE (not random init, not plain copy):
  * ResidualUnits compute `x + f(x)`. Zeroing the new unit's output 1x1 conv
    (weight_norm g + bias) makes it an exact identity, so the expanded model is
    FUNCTION-IDENTICAL to the 500k checkpoint at step 0 (`--verify` proves it).
    The 500k-strength discriminator (also carried over) therefore stays balanced
    instead of hammering a suddenly-degraded generator.
  * The rest of the new unit is copied from the trained same-dilation unit in
    its block, so when the gate opens, gradients start from trained filter
    directions rather than noise.

Usage (on the pod, dac importable):
    hf download Scicom-intl/CallEnhancer decoder_callcentre/decoder_train_state.pt --local-dir ck
    python expand_decoder.py --ckpt ck/decoder_callcentre/decoder_train_state.pt \
        --extra-res 1536:6 --out /Sidon/decoder_callcentre/last.pt --verify

    # then continue training with the SAME graph:
    python train_decoder_callcentre.py --dec-channels 3072 --extra-res 1536:6 \
        --out /Sidon/decoder_callcentre --steps 800000 ...

The output is a trainer-resumable `last.pt`: expanded decoder + the 500k
discriminator and its optimizer state; the GENERATOR optimizer is dropped (its
param groups changed) — the trainer starts a fresh opt_g, so use a warmup/lower
LR for the first ~5-10k steps.
"""
from __future__ import annotations

import argparse

import torch

import dac

RATES = [8, 5, 4, 3, 2]
DILATIONS = (1, 3, 9)   # DAC's per-block dilation stack; appended units cycle it
N_BASE_UNITS = 3        # ResidualUnits per DecoderBlock in stock DAC
FIRST_UNIT_IDX = 2      # DecoderBlock.block = [Snake1d, WNConvTranspose1d, RU, RU, RU]


def parse_extra_res(spec: str) -> dict[int, int]:
    """'1536:6,768:4' -> {1536: 6, 768: 4} (block output channels -> extra units)."""
    out: dict[int, int] = {}
    for part in (spec or "").split(","):
        part = part.strip()
        if not part:
            continue
        ch, n = part.split(":")
        out[int(ch)] = int(n)
    return out


def build_decoder(input_channel: int, channels: int, extra_res: dict[int, int] | None = None,
                  rates=RATES):
    """Stock DAC decoder graph + `extra_res[out_ch]` extra ResidualUnits appended at
    the END of each matching DecoderBlock's Sequential — existing state-dict keys are
    untouched, new units land at `model.<i>.block.{5,6,...}`. Random init (callers
    either load a full expanded checkpoint, or warm-start via `expand_state`)."""
    dec = dac.model.dac.Decoder(input_channel=input_channel, channels=channels, rates=list(rates))
    extra_res = extra_res or {}
    used = set()
    for i in range(len(rates)):
        out_ch = channels // 2 ** (i + 1)
        n_extra = extra_res.get(out_ch, 0)
        if n_extra <= 0:
            continue
        used.add(out_ch)
        blk = dec.model[1 + i]          # model[0] is the input conv
        for k in range(n_extra):
            unit = dac.model.dac.ResidualUnit(out_ch, dilation=DILATIONS[k % len(DILATIONS)])
            blk.block.add_module(str(len(blk.block)), unit)
    unknown = set(extra_res) - used
    if unknown:
        raise ValueError(f"extra_res channels {sorted(unknown)} do not match any "
                         f"DecoderBlock output width of a {channels}-ch decoder")
    return dec


def _zero_output_gate(unit) -> None:
    """Zero the unit's final 1x1 conv (weight-norm g + bias) => unit is identity."""
    conv = unit.block[-1]
    with torch.no_grad():
        if hasattr(conv, "weight_g"):                      # torch old-style weight_norm
            conv.weight_g.zero_()
        elif hasattr(conv, "parametrizations"):            # new-style parametrization
            conv.parametrizations.weight.original0.zero_()
        else:
            conv.weight.zero_()
        if conv.bias is not None:
            conv.bias.zero_()


def expand_state(dec, old_sd: dict, channels: int, extra_res: dict[int, int],
                 rates=RATES) -> None:
    """Load the trained (unexpanded) `old_sd` into the expanded `dec`, then init each
    appended unit as a COPY of its block's trained same-dilation unit with the output
    gate zeroed (exact identity => function-preserving warm start)."""
    missing, unexpected = dec.load_state_dict(old_sd, strict=False)
    if unexpected:
        raise RuntimeError(f"checkpoint keys not in expanded graph: {unexpected[:5]} ...")
    base_keys = set(old_sd)
    for key in missing:
        if ".block." not in key or int(key.split(".block.")[1].split(".")[0]) < FIRST_UNIT_IDX + N_BASE_UNITS:
            raise RuntimeError(f"missing non-appended key after load: {key}")
    for i in range(len(rates)):
        out_ch = channels // 2 ** (i + 1)
        n_extra = extra_res.get(out_ch, 0)
        blk = dec.model[1 + i]
        for k in range(n_extra):
            src = blk.block[FIRST_UNIT_IDX + (k % N_BASE_UNITS)]   # dilation-matched trained unit
            new = blk.block[FIRST_UNIT_IDX + N_BASE_UNITS + k]
            new.load_state_dict(src.state_dict())
            _zero_output_gate(new)
    assert base_keys <= set(dec.state_dict())


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--ckpt", required=True,
                    help="decoder_train_state.pt / last.pt / decoder_only.pt (needs 'decoder' key)")
    ap.add_argument("--extra-res", default="1536:9,768:5",
                    help="extra ResidualUnits per block as 'out_ch:n,...' "
                         "[1536:9,768:5 => 188M -> ~381M; 1536:6 => ~301M]")
    ap.add_argument("--out", required=True, help="output trainer-resumable last.pt")
    ap.add_argument("--input-channel", type=int, default=1024)
    ap.add_argument("--verify", action="store_true",
                    help="also build the base decoder and assert identical outputs")
    ap.add_argument("--set-step", type=int, default=None,
                    help="override the stored step (default: keep, e.g. 500000)")
    a = ap.parse_args()

    ck = torch.load(a.ckpt, map_location="cpu")
    old_sd = ck["decoder"] if "decoder" in ck else ck
    channels = ck.get("dec_channels", 3072)
    if ck.get("extra_res"):
        raise SystemExit(f"{a.ckpt} is already expanded (extra_res={ck['extra_res']})")
    extra = parse_extra_res(a.extra_res)

    dec = build_decoder(a.input_channel, channels, extra)
    expand_state(dec, old_sd, channels, extra)
    n_old = sum(v.numel() for v in old_sd.values())
    n_new = sum(p.numel() for p in dec.parameters())
    print(f"[expand] {channels}ch +{a.extra_res}: {n_old/1e6:.1f}M -> {n_new/1e6:.1f}M params")

    if a.verify:
        base = build_decoder(a.input_channel, channels, None)
        base.load_state_dict(old_sd)
        base.eval(); dec.eval()
        with torch.no_grad():
            x = torch.randn(1, a.input_channel, 50)
            d = (base(x) - dec(x)).abs().max().item()
        print(f"[verify] max |base - expanded| on random input: {d:.3e}")
        assert d < 1e-5, "expanded decoder is NOT function-preserving"

    out = {"step": a.set_step if a.set_step is not None else ck.get("step", 0),
           "decoder": dec.state_dict(),
           "dec_channels": channels,
           "extra_res": a.extra_res}
    for key in ("disc", "opt_d"):      # keep the 500k discriminator (+ its optimizer);
        if key in ck:                  # opt_g is dropped — decoder param groups changed
            out[key] = ck[key]
    torch.save(out, a.out)
    print(f"[save] {a.out} (step={out['step']}, disc={'disc' in out}, opt_d={'opt_d' in out}; "
          f"opt_g intentionally dropped -> trainer starts a fresh generator optimizer)")


if __name__ == "__main__":
    main()