"""Phonon (gradium/phonon-31m-en) building blocks for recipes/phonon-31m-en.py: weights, eager reference, export programs.

The eager reference follows the official runtime (the `ptts` crate, src/{flow_lm,tts_model,
transformer,mlp,conditioners,mimi,seanet,conv}.rs) operation by operation:

  voice   speaker_wavs [1,512,125] -> transpose -> speaker_proj (512x512, no bias) -> [1,125,512]
  prime   transformer(voice) then transformer(text_embed(tokens)), both written to the KV cache
  step    x = input_linear(prev latent or bos_emb) + cond      (cond = num_speakers("1") LUT
          + duration_delta(0.0) sinusoid, each through its output_proj)
          h = out_norm(transformer(x))[:, -1]; eos_logit = out_eos(h)
          latent = noise + flow_net(h, s=0, t=1, noise)          (LSD, 1 step; noise ~ N(0, temp))
  decode  (latent * emb_std + emb_mean) -> 1x1 conv 32->512 -> depthwise convtr x16 -> 2-layer
          transformer (window 250, layer scale) -> SEANet decoder (ratios 6,5,4) -> 1920 samples

Export programs (WebGPU):
  FlowLMProgram   forward(embeds[1,S,512] S<=128, input_pos[1] i64, noise[1,32])
                  -> latent[1,32], eos_logit[1,1], next_embeds[1,1,512]  (KV cache in the program)
  MimiStepProgram forward(latent[1,1,32], *state) -> pcm[1,1,1920], *new_state  (state explicit)

Rewrites for the public WebGPU runtime (same function, checked against the eager reference):
  * RoPE: q/k rows of in_proj permuted from interleaved pairs to halves, so the runtime's fused
    HF RoPE (et_vk.apply_rotary_emb_hf) applies; q.k is unchanged.
  * flow-LM LayerNorm (dynamic S): mean as a 1/512 linear + the rms_norm pattern, because
    native_layer_norm has no resize hook and mean.dim's resize drops a keepdim'ed dim.
  * ELU = relu(x) + exp(min(x, 0)) - 1 and SiLU = x * sigmoid(x) (no aten.elu / aten.silu).
  * ConvTranspose1d as a 4-D (H = 1) conv_transpose2d (3-D convolution has no transposed path).
  * Mimi streaming state as explicit tensors: mutable buffers other than the llama KV cache are
    not kept live by the delegate.
"""

import math

import torch
import torch.nn.functional as F
from torch import nn

D = 512
H = 8
HD = 64
LDIM = 32
FLOW_LAYERS = 6
MIMI_LAYERS = 2
MIMI_CONTEXT = 250
FRAME = 1920  # samples per latent frame at 24 kHz
UPS = 16  # mimi transformer tokens per latent frame
NEG = -1.0e30


# ---------------------------------------------------------------------------------------------
# Weights: a module whose state_dict keys are exactly the checkpoint's keys.
# ---------------------------------------------------------------------------------------------


def _p(*shape):
    return nn.Parameter(torch.zeros(*shape), requires_grad=False)


class _Lin(nn.Module):
    def __init__(self, i, o, bias):
        super().__init__()
        self.weight = _p(o, i)
        if bias:
            self.bias = _p(o)


class _LN(nn.Module):
    def __init__(self, n):
        super().__init__()
        self.weight = _p(n)
        self.bias = _p(n)


class _Attn(nn.Module):
    def __init__(self):
        super().__init__()
        self.in_proj = _Lin(D, 3 * D, False)
        self.out_proj = _Lin(D, D, False)


class _Scale(nn.Module):
    def __init__(self):
        super().__init__()
        self.scale = _p(D)


class _Layer(nn.Module):
    def __init__(self, layer_scale):
        super().__init__()
        self.self_attn = _Attn()
        self.norm1 = _LN(D)
        self.norm2 = _LN(D)
        self.linear1 = _Lin(D, 2048, False)
        self.linear2 = _Lin(2048, D, False)
        if layer_scale:
            self.layer_scale_1 = _Scale()
            self.layer_scale_2 = _Scale()


class _Transformer(nn.Module):
    def __init__(self, n, layer_scale):
        super().__init__()
        self.layers = nn.ModuleList(_Layer(layer_scale) for _ in range(n))


class _Seq(nn.Module):
    """nn.Sequential-style numbered children with only the indices that hold weights."""

    def __init__(self, items):
        super().__init__()
        for k, m in items.items():
            self.add_module(str(k), m)

    def __getitem__(self, i):
        return getattr(self, str(i))


class _TimeEmbed(nn.Module):
    def __init__(self):
        super().__init__()
        self.freqs = _p(128)
        self.mlp = _Seq({0: _Lin(256, D, True), 2: _Lin(D, D, True), 3: _Alpha()})


class _Alpha(nn.Module):
    def __init__(self):
        super().__init__()
        self.alpha = _p(D)


class _ResBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.in_ln = _LN(D)
        self.mlp = _Seq({0: _Lin(D, D, True), 2: _Lin(D, D, True)})
        self.adaLN_modulation = _Seq({1: _Lin(D, 3 * D, True)})


class _Final(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = _Lin(D, LDIM, True)
        self.adaLN_modulation = _Seq({1: _Lin(D, 2 * D, True)})


class _FlowNet(nn.Module):
    def __init__(self, depth=6):
        super().__init__()
        self.time_embed = nn.ModuleList([_TimeEmbed(), _TimeEmbed()])
        self.cond_embed = _Lin(D, D, True)
        self.input_proj = _Lin(LDIM, D, True)
        self.res_blocks = nn.ModuleList(_ResBlock() for _ in range(depth))
        self.final_layer = _Final()


class _Embed(nn.Module):
    def __init__(self, n, d):
        super().__init__()
        self.weight = _p(n, d)


class _Cond(nn.Module):
    def __init__(self, embed=None, proj_in=None):
        super().__init__()
        self.learnt_padding = _p(1, 1, D)
        if embed is not None:
            self.embed = _Embed(*embed)
        if proj_in is not None:
            self.output_proj = _Lin(proj_in, D, False)


class _Conditioners(nn.Module):
    def __init__(self):
        super().__init__()
        self.duration_delta = _Cond(proj_in=128)
        self.num_speakers = _Cond(embed=(32, 16), proj_in=16)
        self.speaker_wavs = _Cond(proj_in=512)
        self.transcript_in_segment = _Cond(embed=(4001, D))


class _CondProvider(nn.Module):
    def __init__(self):
        super().__init__()
        self.conditioners = _Conditioners()


class _WST(nn.Module):
    """flow.w_s_t: a training-time loss-weighting net; the runtime drops it (loader.rs remap_key)."""

    def __init__(self):
        super().__init__()
        self.w_s_t = _Seq(
            {
                0: _Lin(2, 32, True),
                3: _Lin(32, 32, True),
                6: _Lin(32, 32, True),
                9: _Lin(32, 1, True),
            }
        )


class FlowLMWeights(nn.Module):
    def __init__(self):
        super().__init__()
        self.bos_emb = _p(LDIM)
        self.emb_mean = _p(LDIM)
        self.emb_std = _p(LDIM)
        self.condition_provider = _CondProvider()
        self.flow = _WST()
        self.flow_net = _FlowNet()
        self.input_linear = _Lin(LDIM, D, False)
        self.out_eos = _Lin(D, 1, True)
        self.out_norm = _LN(D)
        self.transformer = _Transformer(FLOW_LAYERS, False)


class _Conv(nn.Module):
    def __init__(self, i, o, k):
        super().__init__()
        self.conv = _ConvW(i, o, k, bias=True)


class _ConvW(nn.Module):
    def __init__(self, i, o, k, bias):
        super().__init__()
        self.weight = _p(o, i, k)
        if bias:
            self.bias = _p(o)


class _ConvTr(nn.Module):
    def __init__(self, i, o, k):
        super().__init__()
        self.convtr = _ConvTrW(i, o, k, True)


class _ConvTrW(nn.Module):
    def __init__(self, i, o_per_group, k, bias):
        super().__init__()
        self.weight = _p(i, o_per_group, k)
        if bias:
            self.bias = _p(o_per_group)


class _Res(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.block = _Seq({1: _Conv(dim, dim // 2, 3), 3: _Conv(dim // 2, dim, 1)})


class _SEANetDecoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = _Seq(
            {
                0: _Conv(512, 512, 7),
                2: _ConvTr(512, 256, 12),
                3: _Res(256),
                5: _ConvTr(256, 128, 10),
                6: _Res(128),
                8: _ConvTr(128, 64, 8),
                9: _Res(64),
                11: _Conv(64, 1, 3),
            }
        )


class _ProjT(nn.Module):
    def __init__(self):
        super().__init__()
        self.transformer = _Transformer(MIMI_LAYERS, True)


class _Quantizer(nn.Module):
    def __init__(self):
        super().__init__()
        self.output_proj = _ConvW(LDIM, D, 1, bias=False)


class _Upsample(nn.Module):
    def __init__(self):
        super().__init__()
        self.convtr = _Seq({})
        self.convtr.convtr = _ConvTrW(D, 1, 2 * UPS, False)


class MimiWeights(nn.Module):
    def __init__(self):
        super().__init__()
        self.decoder = _SEANetDecoder()
        self.decoder_transformer = _ProjT()
        self.quantizer = _Quantizer()
        self.upsample = _Upsample()


class PhononWeights(nn.Module):
    def __init__(self):
        super().__init__()
        self.flow_lm = FlowLMWeights()
        self.mimi = MimiWeights()


def load_weights(safetensors_path):
    """Every checkpoint tensor, as float32 (bf16 -> f32 is exact), loaded with strict=True."""
    from safetensors.torch import load_file

    sd = {k: v.float() for k, v in load_file(str(safetensors_path)).items()}
    w = PhononWeights()
    w.load_state_dict(sd, strict=True)
    return w.eval()


def load_voice(path):
    """speaker_wavs [1, 512, T] (stored speaker-Mimi latents) from a voice .safetensors file."""
    from safetensors.torch import load_file

    t = load_file(str(path))
    assert list(t) == ["speaker_wavs"], list(t)
    return t["speaker_wavs"].float()


# ---------------------------------------------------------------------------------------------
# Shared math
# ---------------------------------------------------------------------------------------------


def lin(m, x):
    """nn.Linear on a 2-D view (the public WebGPU runtime's fp32 aten.linear takes 2-D input)."""
    shp = x.shape
    y = F.linear(x.reshape(-1, shp[-1]), m.weight, getattr(m, "bias", None))
    return y.reshape(*shp[:-1], y.shape[-1])


def lin3(m, x):
    """nn.Linear on the tensor as it is (any rank); the Vulkan lowering adds its own 2-D views."""
    return F.linear(x, m.weight, getattr(m, "bias", None))


def layer_norm_dyn(m, x, eps, mean_w):
    """LayerNorm for the dynamic-length flow-LM program, from ops that follow a resized sequence
    length on the public WebGPU runtime (its native_layer_norm has no resize hook, and its mean.dim
    resize drops a keepdim'ed dim): the mean is a linear with weight 1/512 (exact: a power of two),
    the rest matches the Vulkan rms_norm pattern (fused into et_vk.rms_norm)."""
    xc = x - F.linear(x, mean_w)
    return xc * torch.rsqrt((xc * xc).mean(-1, keepdim=True) + eps) * m.weight + m.bias


def layer_norm(m, x, eps):
    return F.layer_norm(x, (x.shape[-1],), m.weight, m.bias, eps)


def silu(x):
    return x * torch.sigmoid(x)


def elu(x):
    """ELU(alpha=1) from ops the public WebGPU runtime registers (it has no aten.elu)."""
    return torch.relu(x) + torch.exp(torch.clamp(x, max=0.0)) - 1.0


def gelu(x):
    return F.gelu(x)  # erf


def inv_freq():
    """RoPE frequencies as the runtime computes them (rope.rs): 1 / max_period^(i / half) in f32."""
    half = HD // 2
    i = torch.arange(half, dtype=torch.float32)
    return 1.0 / torch.pow(torch.tensor(10000.0, dtype=torch.float32), i / half)


def rope_interleaved(x, pos):
    """rope_i on [B, T, H, HD] with absolute positions pos [T] (the runtime's convention:
    pairs (2i, 2i+1) rotate by pos * inv_freq[i])."""
    ang = pos.to(torch.float32)[:, None] * inv_freq()[None, :]  # [T, half]
    c, s = torch.cos(ang)[None, :, None, :], torch.sin(ang)[None, :, None, :]
    xr, xi = x[..., 0::2], x[..., 1::2]
    return torch.stack([xr * c - xi * s, xr * s + xi * c], dim=-1).flatten(-2)


def time_embed_rms(te, t):
    """TimestepEmbedder with TimeRmsNorm::Rms (phonon config time_rms_norm: rms)."""
    args = t * te.freqs
    emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
    x = lin(te.mlp[2], silu(lin(te.mlp[0], emb)))
    ms = (x * x).mean(-1, keepdim=True)
    return x * torch.rsqrt(ms + 1e-5) * te.mlp[3].alpha


def flow_net(fn, c, x, t_comb):
    """SimpleMLPAdaLN: c [N, 512] backbone output, x [N, 32] noise, t_comb [N or 1, 512]."""
    x = lin(fn.input_proj, x)
    y = t_comb + lin(fn.cond_embed, c)
    sy = silu(y)
    for b in fn.res_blocks:
        ada = lin(b.adaLN_modulation[1], sy)
        shift, scale, gate = ada[..., :D], ada[..., D : 2 * D], ada[..., 2 * D :]
        h = layer_norm(b.in_ln, x, 1e-6) * (scale + 1.0) + shift
        h = lin(b.mlp[2], silu(lin(b.mlp[0], h)))
        x = x + gate * h
    fl = fn.final_layer
    ada = lin(fl.adaLN_modulation[1], sy)
    shift, scale = ada[..., :D], ada[..., D:]
    x = F.layer_norm(x, (D,), None, None, 1e-6) * (scale + 1.0) + shift
    return lin(fl.linear, x)


def lsd_time(fn):
    """(time_embed[0](s=0) + time_embed[1](t=1)) / 2, the flow head's constant time condition
    for lsd_decode_steps = 1."""
    s = torch.zeros(1, 1)
    t = torch.ones(1, 1)
    return (
        time_embed_rms(fn.time_embed[0], s) + time_embed_rms(fn.time_embed[1], t)
    ) / 2


def sin_embedding(position, dim, max_period):
    half = dim // 2
    i = torch.arange(half, dtype=torch.float32)
    ph = position / torch.pow(
        torch.tensor(max_period, dtype=torch.float32), i / (half - 1)
    )
    return torch.cat([torch.cos(ph), torch.sin(ph)])


def summed_conditions(flw, num_speakers="1", duration_delta=0.0):
    """conditioners.rs Conditioners::sum with the runtime's defaults, [1, 1, 512]."""
    c = flw.condition_provider.conditioners
    idx = ["0", "1", "2"].index(num_speakers)
    ns = F.linear(
        c.num_speakers.embed.weight[idx : idx + 1], c.num_speakers.output_proj.weight
    )
    dd = F.linear(
        sin_embedding(1000.0 * duration_delta, 128, 10000.0)[None],
        c.duration_delta.output_proj.weight,
    )
    return (ns + dd).view(1, 1, D)


def voice_embedding(flw, speaker_wavs):
    """[1, 512, T] speaker latents -> [1, T, 512] voice prompt (loader.rs project_latents)."""
    w = flw.condition_provider.conditioners.speaker_wavs.output_proj.weight
    return F.linear(speaker_wavs.transpose(1, 2), w)


def text_embedding(flw, tokens):
    emb = flw.condition_provider.conditioners.transcript_in_segment.embed.weight
    return emb[torch.as_tensor(tokens, dtype=torch.long)][None]


# ---------------------------------------------------------------------------------------------
# Eager reference (Python-side state, mirrors the runtime)
# ---------------------------------------------------------------------------------------------


class FlowLMRef:
    def __init__(self, w: PhononWeights, num_speakers="1", duration_delta=0.0):
        self.w = w.flow_lm
        self.cond = summed_conditions(self.w, num_speakers, duration_delta)
        self.t_comb = lsd_time(self.w.flow_net)
        self.reset()

    def reset(self):
        self.k = [None] * FLOW_LAYERS
        self.v = [None] * FLOW_LAYERS
        self.pos = 0

    @torch.no_grad()
    def backbone(self, x):
        B, T, _ = x.shape
        pos = torch.arange(self.pos, self.pos + T)
        for i, L in enumerate(self.w.transformer.layers):
            h = layer_norm(L.norm1, x, 1e-5)
            qkv = lin(L.self_attn.in_proj, h).view(B, T, 3, H, HD)
            q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
            q, k = rope_interleaved(q, pos), rope_interleaved(k, pos)
            self.k[i] = k if self.k[i] is None else torch.cat([self.k[i], k], 1)
            self.v[i] = v if self.v[i] is None else torch.cat([self.v[i], v], 1)
            kk, vv = self.k[i].transpose(1, 2), self.v[i].transpose(1, 2)
            qq = q.transpose(1, 2)
            Tk = kk.shape[2]
            qpos = torch.arange(Tk - T, Tk)[:, None]
            mask = torch.where(torch.arange(Tk)[None, :] <= qpos, 0.0, float("-inf"))
            att = torch.softmax(qq @ kk.transpose(-1, -2) / math.sqrt(HD) + mask, -1)
            o = (att @ vv).transpose(1, 2).reshape(B, T, D)
            x = x + lin(L.self_attn.out_proj, o)
            h = layer_norm(L.norm2, x, 1e-5)
            x = x + lin(L.linear2, gelu(lin(L.linear1, h)))
        self.pos += T
        return x

    def prompt(self, embeds):
        self.backbone(embeds)

    @torch.no_grad()
    def step(self, prev_latent, noise):
        """prev_latent None (BOS) or [1, 1, 32]; noise [1, 32]. -> latent [1,1,32], eos_logit."""
        seq = self.w.bos_emb.view(1, 1, LDIM) if prev_latent is None else prev_latent
        x = lin(self.w.input_linear, seq) + self.cond
        h = layer_norm(self.w.out_norm, self.backbone(x), 1e-5)[:, -1]
        eos = lin(self.w.out_eos, h)
        latent = noise + flow_net(self.w.flow_net, h, noise, self.t_comb)
        return latent.view(1, 1, LDIM), eos


def _conv1d_stream(m, x, prev, dilation=1):
    xx = torch.cat([prev, x], 2) if prev is not None else x
    y = F.conv1d(xx, m.conv.weight, m.conv.bias, dilation=dilation)
    tp = prev.shape[2] if prev is not None else 0
    return y, (xx[:, :, xx.shape[2] - tp :] if tp else None)


def _convtr_stream(w, bias, x, partial, stride, groups=1):
    y = F.conv_transpose1d(x, w, bias, stride=stride, groups=groups)
    pt = partial.shape[2]
    y = torch.cat([y[:, :, :pt] + partial, y[:, :, pt:]], 2)
    new_partial = y[:, :, y.shape[2] - pt :]
    if bias is not None:
        new_partial = new_partial - bias.view(1, -1, 1)
    return y[:, :, : y.shape[2] - pt], new_partial


class MimiRef:
    """Streaming Mimi decoder (MimiDecoder::decode_from_latent_step), one or more frames per call."""

    def __init__(self, w: PhononWeights):
        self.w = w.mimi
        self.emb_std = w.flow_lm.emb_std
        self.emb_mean = w.flow_lm.emb_mean
        self.reset()

    def reset(self):
        z = lambda *s: torch.zeros(1, *s)  # noqa: E731
        self.up = z(D, UPS)
        self.kv = [[None, None] for _ in range(MIMI_LAYERS)]
        self.offset = 0
        self.prev = {
            "init": z(512, 6),
            "r1": z(256, 2),
            "r2": z(128, 2),
            "r3": z(64, 2),
            "final": z(64, 2),
        }
        self.partial = {"u1": z(256, 6), "u2": z(128, 5), "u3": z(64, 4)}

    @torch.no_grad()
    def transformer(self, x):  # x [1, T, 512]
        T = x.shape[1]
        pos = torch.arange(self.offset, self.offset + T)
        for i, L in enumerate(self.w.decoder_transformer.transformer.layers):
            h = layer_norm(L.norm1, x, 1e-5)
            qkv = lin(L.self_attn.in_proj, h).view(1, T, 3, H, HD)
            q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
            q, k = rope_interleaved(q, pos), rope_interleaved(k, pos)
            q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
            pk, pv = self.kv[i]
            kk = k if pk is None else torch.cat([pk, k], 2)
            vv = v if pv is None else torch.cat([pv, v], 2)
            kvlen = kk.shape[2] - T
            qi = torch.arange(T)[:, None] + kvlen
            kj = torch.arange(kk.shape[2])[None, :]
            ok = (kj <= qi) & (kj >= qi - MIMI_CONTEXT)
            mask = torch.where(ok, 0.0, float("-inf"))
            att = torch.softmax(q @ kk.transpose(-1, -2) / math.sqrt(HD) + mask, -1)
            o = (att @ vv).transpose(1, 2).reshape(1, T, D)
            self.kv[i] = [kk[:, :, -MIMI_CONTEXT:], vv[:, :, -MIMI_CONTEXT:]]
            x = x + L.layer_scale_1.scale * lin(L.self_attn.out_proj, o)
            h = layer_norm(L.norm2, x, 1e-5)
            x = x + L.layer_scale_2.scale * lin(L.linear2, gelu(lin(L.linear1, h)))
        self.offset += T
        return x

    def _res(self, blk, x, key):
        v, self.prev[key] = _conv1d_stream(blk.block[1], elu(x), self.prev[key])
        v, _ = _conv1d_stream(blk.block[3], elu(v), None)
        return x + v

    @torch.no_grad()
    def decode(self, latent):
        """latent [1, T, 32] (flow-LM space) -> pcm [1, 1, T * 1920]."""
        x = (latent * self.emb_std + self.emb_mean).transpose(1, 2)
        x = F.conv1d(x, self.w.quantizer.output_proj.weight)
        x, self.up = _convtr_stream(
            self.w.upsample.convtr.convtr.weight, None, x, self.up, UPS, groups=D
        )
        x = self.transformer(x.transpose(1, 2)).transpose(1, 2)
        m = self.w.decoder.model
        x, self.prev["init"] = _conv1d_stream(m[0], x, self.prev["init"])
        for up, res, s, ku, kr in [
            (m[2], m[3], 6, "u1", "r1"),
            (m[5], m[6], 5, "u2", "r2"),
            (m[8], m[9], 4, "u3", "r3"),
        ]:
            x, self.partial[ku] = _convtr_stream(
                up.convtr.weight, up.convtr.bias, elu(x), self.partial[ku], s
            )
            x = self._res(res, x, kr)
        x, self.prev["final"] = _conv1d_stream(m[11], elu(x), self.prev["final"])
        return x


# ---------------------------------------------------------------------------------------------
# Generation policy (plan.rs, tts_model.rs)
# ---------------------------------------------------------------------------------------------


def prepare_text_prompt(text):
    text = " ".join(text.strip().replace("\n", " ").replace("\r", " ").split())
    if not text:
        return text, 3
    frames_after_eos = 3 if len(text.split()) <= 4 else 1
    text = text[0].upper() + text[1:]
    if text[-1].isalnum():
        text += "."
    return text, frames_after_eos


def frame_budget(num_tokens, frame_rate=12.5):
    return math.ceil((num_tokens / 3.0 + 2.0) * frame_rate)


class EosPolicy:
    def __init__(self, frames_after_eos):
        self.n, self.countdown = frames_after_eos, None

    def should_stop(self, is_eos):
        if is_eos and self.countdown is None:
            self.countdown = self.n
        if self.countdown is None:
            return False
        if self.countdown == 0:
            return True
        self.countdown -= 1
        return False


def generate(step_fn, tokens, frames_after_eos, noise_fn, eos_threshold=4.0):
    """The runtime's loop (synth.rs run_inline): BOS first, the EOS frame and its tail are kept.
    step_fn(prev_latent_or_None, noise) -> (latent [1,1,32], eos_logit). Returns latents, eos logits."""
    eos = EosPolicy(frames_after_eos)
    prev, lats, logits = None, [], []
    for i in range(frame_budget(len(tokens))):
        lat, e = step_fn(prev, noise_fn(i))
        lats.append(lat)
        logits.append(float(e.reshape(-1)[0]))
        if eos.should_stop(logits[-1] > eos_threshold):
            break
        prev = lat
    return lats, logits


def tokenizer(tokenizer_json):
    from tokenizers import Tokenizer

    return Tokenizer.from_file(str(tokenizer_json))


def tokenize(tok, text):
    prepared, frames_after_eos = prepare_text_prompt(text)
    return (
        tok.encode(prepared, add_special_tokens=False).ids,
        frames_after_eos,
        prepared,
    )


# ---------------------------------------------------------------------------------------------
# Export programs
# ---------------------------------------------------------------------------------------------

MAX_CTX = 512  # flow-LM KV cache positions: voice prompt + text tokens + generated frames
PREFILL_MAX = 128  # most positions one flow-LM call may write (the 125-frame voice prompt fits)


def _half_split_perm():
    """Within each head, reorder q/k channels from interleaved pairs (2i, 2i+1) to halves
    (i, i + 32): rope_i on the old layout equals rotate_half RoPE on the new one, and q.k is
    unchanged because q and k are permuted alike."""
    per = torch.cat([torch.arange(0, HD, 2), torch.arange(1, HD, 2)])
    return torch.cat([h * HD + per for h in range(H)])


def _permuted_in_proj(w):
    """in_proj weight [3*512, 512] with the q and k row blocks permuted, v untouched."""
    perm = _half_split_perm()
    return torch.cat([w[:D][perm], w[D : 2 * D][perm], w[2 * D :]], 0).contiguous()


def rope_tables(n):
    """cos/sin [n, 64] in the half-split layout, angle = f32(pos) * f32(inv_freq) as rope.rs."""
    ang = torch.arange(n, dtype=torch.float32)[:, None] * inv_freq()[None, :]
    c, s = torch.cos(ang), torch.sin(ang)
    return torch.cat([c, c], -1).contiguous(), torch.cat([s, s], -1).contiguous()


def _rope_hf(q, k, cos, sin):
    from executorch.backends.vulkan.patterns.rope_hf import HfRotaryEmbeddingPattern

    return HfRotaryEmbeddingPattern()(q, k, cos, sin)


class _LinP(nn.Module):
    def __init__(self, weight, bias=None):
        super().__init__()
        self.weight = nn.Parameter(weight.detach().clone().contiguous(), requires_grad=False)
        if bias is not None:
            self.bias = nn.Parameter(bias.detach().clone().contiguous(), requires_grad=False)


class FlowLMProgram(nn.Module):
    """forward(embeds[1,S,512], input_pos[1] i64, noise[1,32])
    -> latent[1,32], eos_logit[1,1], next_embeds[1,1,512].

    Writes the S positions [input_pos, input_pos + S) of the in-program KV cache (llama
    sdpa_with_kv_cache, causal), then runs the head on the last position: eos_logit = out_eos(h),
    latent = noise + flow_net(h, s=0, t=1, noise) and next_embeds = input_linear(latent) + cond,
    the embedding of the next step's input. Prompt calls (voice, text) ignore the outputs."""

    def __init__(self, w: PhononWeights, num_speakers="1", duration_delta=0.0, attn="sdpa_kv"):
        super().__init__()
        from executorch.extension.llm.custom_ops import custom_ops  # noqa: F401

        f = w.flow_lm
        self.attn = attn
        self.layers = nn.ModuleList()
        for L in f.transformer.layers:
            m = nn.Module()
            m.norm1, m.norm2 = L.norm1, L.norm2
            m.in_proj = _LinP(_permuted_in_proj(L.self_attn.in_proj.weight))
            m.out_proj, m.linear1, m.linear2 = L.self_attn.out_proj, L.linear1, L.linear2
            self.layers.append(m)
        self.out_norm, self.out_eos, self.input_linear = f.out_norm, f.out_eos, f.input_linear
        self.fn = f.flow_net
        with torch.no_grad():
            self.register_buffer("cond", summed_conditions(f, num_speakers, duration_delta).clone(), persistent=False)
            self.register_buffer("t_comb", lsd_time(f.flow_net).clone(), persistent=False)
        self.register_buffer("mean_w", torch.full((1, D), 1.0 / D), persistent=False)
        cos, sin = rope_tables(MAX_CTX)
        self.register_buffer("rope_cos", cos, persistent=False)
        self.register_buffer("rope_sin", sin, persistent=False)
        for i in range(FLOW_LAYERS):
            self.register_buffer(f"k_cache_{i}", torch.zeros(1, MAX_CTX, H, HD), persistent=False)
            self.register_buffer(f"v_cache_{i}", torch.zeros(1, MAX_CTX, H, HD), persistent=False)

    def reset(self):
        for i in range(FLOW_LAYERS):
            getattr(self, f"k_cache_{i}").zero_()
            getattr(self, f"v_cache_{i}").zero_()

    def forward(self, embeds, input_pos, noise):
        S = embeds.shape[1]
        p = input_pos[0].item()
        torch._check_is_size(p)
        torch._check(p + S <= MAX_CTX)
        cos = self.rope_cos.narrow(0, p, S)
        sin = self.rope_sin.narrow(0, p, S)
        x = embeds
        for i, L in enumerate(self.layers):
            h = layer_norm_dyn(L.norm1, x, 1e-5, self.mean_w)
            qkv = lin3(L.in_proj, h)
            q = qkv[..., :D].reshape(1, S, H, HD)
            k = qkv[..., D : 2 * D].reshape(1, S, H, HD)
            v = qkv[..., 2 * D :].reshape(1, S, H, HD)
            q, k = _rope_hf(q, k, cos, sin)
            kc, vc = getattr(self, f"k_cache_{i}"), getattr(self, f"v_cache_{i}")
            o = torch.ops.llama.sdpa_with_kv_cache(q, k, v, kc, vc, p, S, None, 0.0, True, None)
            x = x + lin3(L.out_proj, o.reshape(1, S, D))
            h = layer_norm_dyn(L.norm2, x, 1e-5, self.mean_w)
            x = x + lin3(L.linear2, gelu(lin3(L.linear1, h)))
        h = layer_norm(self.out_norm, x[:, -1, :], 1e-5)
        eos = lin3(self.out_eos, h)
        latent = noise + flow_net(self.fn, h, noise, self.t_comb)
        nxt = lin3(self.input_linear, latent).reshape(1, 1, D) + self.cond
        return latent, eos, nxt


MIMI_STATE = [
    ("upsample_partial", (1, D, UPS)),
    ("k0", (1, H, MIMI_CONTEXT, HD)),
    ("v0", (1, H, MIMI_CONTEXT, HD)),
    ("k1", (1, H, MIMI_CONTEXT, HD)),
    ("v1", (1, H, MIMI_CONTEXT, HD)),
    ("kv_seen", (1, 1, 1, MIMI_CONTEXT)),
    ("rope_offset", (1,)),
    ("init_conv_prev", (1, 512, 6)),
    ("up1_partial", (1, 256, 6)),
    ("res1_prev", (1, 256, 2)),
    ("up2_partial", (1, 128, 5)),
    ("res2_prev", (1, 128, 2)),
    ("up3_partial", (1, 64, 4)),
    ("res3_prev", (1, 64, 2)),
    ("final_conv_prev", (1, 64, 2)),
]


def mimi_zero_state():
    return [torch.zeros(s) for _, s in MIMI_STATE]


def _convtr(x, w, b, s, groups=1):
    """ConvTranspose1d as a 4-D (H = 1) conv_transpose2d: the public runtime's 3-D convolution
    has no transposed path, its 4-D one does (with groups)."""
    y = F.conv_transpose2d(x.unsqueeze(2), w.unsqueeze(2), b, stride=(1, s), groups=groups)
    return y.squeeze(2)


def _convtr_step(x, w, b, partial, s, groups=1):
    y = _convtr(x, w, b, s, groups)
    pt = partial.shape[2]
    n = y.shape[2]
    head = y[:, :, :pt] + partial
    out = torch.cat([head, y[:, :, pt : n - pt]], 2)
    tail = y[:, :, n - pt :]
    if b is not None:
        tail = tail - b.view(1, -1, 1)
    return out, tail


def _conv_step(m, x, prev):
    xx = torch.cat([prev, x], 2)
    y = F.conv1d(xx, m.conv.weight, m.conv.bias)
    return y, xx[:, :, xx.shape[2] - prev.shape[2] :]


class MimiStepProgram(nn.Module):
    """forward(latent[1,1,32], *state) -> pcm[1,1,1920], *new_state (MIMI_STATE order).

    One latent frame per call, exactly the runtime's streaming decoder: every piece of streaming
    state (conv tails, transposed-conv partial sums, the 2-layer transformer's last 250 keys and
    values, which of them are filled, and the RoPE offset) is an explicit input and output. The
    first call takes all-zero state."""

    def __init__(self, w: PhononWeights):
        super().__init__()
        m = w.mimi
        self.emb_std = nn.Parameter(w.flow_lm.emb_std.detach().clone().view(1, 1, LDIM), requires_grad=False)
        self.emb_mean = nn.Parameter(w.flow_lm.emb_mean.detach().clone().view(1, 1, LDIM), requires_grad=False)
        self.qproj = _LinP(m.quantizer.output_proj.weight.view(D, LDIM))
        self.up_w = nn.Parameter(m.upsample.convtr.convtr.weight.detach().clone(), requires_grad=False)
        self.layers = nn.ModuleList()
        for L in m.decoder_transformer.transformer.layers:
            q = nn.Module()
            q.norm1, q.norm2 = L.norm1, L.norm2
            q.in_proj = _LinP(_permuted_in_proj(L.self_attn.in_proj.weight))
            q.out_proj, q.linear1, q.linear2 = L.self_attn.out_proj, L.linear1, L.linear2
            q.ls1 = nn.Parameter(L.layer_scale_1.scale.detach().clone(), requires_grad=False)
            q.ls2 = nn.Parameter(L.layer_scale_2.scale.detach().clone(), requires_grad=False)
            self.layers.append(q)
        self.dec = m.decoder.model
        self.register_buffer("inv_freq", inv_freq().view(1, HD // 2), persistent=False)
        self.register_buffer("steps", torch.arange(UPS, dtype=torch.float32).view(UPS, 1), persistent=False)
        i = torch.arange(UPS)[:, None]
        j = torch.arange(MIMI_CONTEXT + UPS)[None, :]
        band = torch.where((j >= i) & (j <= i + MIMI_CONTEXT), 0.0, NEG)
        self.register_buffer("band", band.view(1, 1, UPS, MIMI_CONTEXT + UPS), persistent=False)
        self.register_buffer("ones_seen", torch.ones(1, 1, 1, UPS), persistent=False)

    def _attn(self, L, x, kc, vc, seen, cos, sin):
        T = UPS
        h = layer_norm(L.norm1, x, 1e-5)
        qkv = lin(L.in_proj, h)
        q = qkv[..., :D].reshape(1, T, H, HD)
        k = qkv[..., D : 2 * D].reshape(1, T, H, HD)
        v = qkv[..., 2 * D :].reshape(1, T, H, HD)
        rh = lambda t: torch.cat([-t[..., HD // 2 :], t[..., : HD // 2]], -1)  # noqa: E731
        q = q * cos + rh(q) * sin
        k = k * cos + rh(k) * sin
        q = q.permute(0, 2, 1, 3)  # [1, H, T, HD]
        k = k.permute(0, 2, 1, 3)
        v = v.permute(0, 2, 1, 3)
        kk = torch.cat([kc, k], 2)  # [1, H, 266, HD]
        vv = torch.cat([vc, v], 2)
        seen_all = torch.cat([seen, self.ones_seen], 3)
        mask = self.band + (seen_all - 1.0) * 1.0e30
        sc = torch.bmm(q.reshape(H, T, HD), kk.reshape(H, MIMI_CONTEXT + T, HD).transpose(1, 2)) * (HD**-0.5)
        pr = torch.softmax(sc.reshape(1, H, T, MIMI_CONTEXT + T) + mask, -1)
        o = torch.bmm(pr.reshape(H, T, MIMI_CONTEXT + T), vv.reshape(H, MIMI_CONTEXT + T, HD))
        o = o.reshape(1, H, T, HD).permute(0, 2, 1, 3).reshape(1, T, D)
        x = x + L.ls1 * lin(L.out_proj, o)
        h = layer_norm(L.norm2, x, 1e-5)
        x = x + L.ls2 * lin(L.linear2, gelu(lin(L.linear1, h)))
        return x, kk[:, :, UPS:], vv[:, :, UPS:], seen_all[..., UPS:]

    def _res(self, blk, x, prev):
        v, prev = _conv_step(blk.block[1], elu(x), prev)
        v = F.conv1d(elu(v), blk.block[3].conv.weight, blk.block[3].conv.bias)
        return x + v, prev

    def forward(self, latent, up_p, k0, v0, k1, v1, seen, off, c0, u1, r1, u2, r2, u3, r3, cf):
        x = latent * self.emb_std + self.emb_mean  # [1, 1, 32]
        x = lin(self.qproj, x).reshape(1, D, 1)
        x, up_p = _convtr_step(x, self.up_w, None, up_p, UPS, groups=D)  # [1, 512, 16]
        pos = self.steps + off  # [16, 1]
        ang = pos * self.inv_freq  # [16, 32]
        cos = torch.cat([torch.cos(ang), torch.cos(ang)], -1).reshape(1, UPS, 1, HD)
        sin = torch.cat([torch.sin(ang), torch.sin(ang)], -1).reshape(1, UPS, 1, HD)
        x = x.permute(0, 2, 1)  # [1, 16, 512]
        x, k0, v0, seen0 = self._attn(self.layers[0], x, k0, v0, seen, cos, sin)
        x, k1, v1, _ = self._attn(self.layers[1], x, k1, v1, seen, cos, sin)
        x = x.permute(0, 2, 1)  # [1, 512, 16]
        m = self.dec
        x, c0 = _conv_step(m[0], x, c0)
        x, u1 = _convtr_step(elu(x), m[2].convtr.weight, m[2].convtr.bias, u1, 6)
        x, r1 = self._res(m[3], x, r1)
        x, u2 = _convtr_step(elu(x), m[5].convtr.weight, m[5].convtr.bias, u2, 5)
        x, r2 = self._res(m[6], x, r2)
        x, u3 = _convtr_step(elu(x), m[8].convtr.weight, m[8].convtr.bias, u3, 4)
        x, r3 = self._res(m[9], x, r3)
        pcm, cf = _conv_step(m[11], elu(x), cf)
        return pcm, up_p, k0, v0, k1, v1, seen0, off + float(UPS), c0, u1, r1, u2, r2, u3, r3, cf