from __future__ import annotations

import argparse
import json
import math
import random
import sys
import time
from dataclasses import asdict, dataclass
from pathlib import Path

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

SCRIPT_ROOT = Path(__file__).resolve().parent
PROJECT_ROOT = SCRIPT_ROOT.parents[0]
TINY_ROOT = PROJECT_ROOT / "third_party" / "tiny-tts"
sys.path = [str(TINY_ROOT), str(SCRIPT_ROOT)] + [p for p in sys.path if p]

from train_hifigan_oracle_v1 import HifiGanConfig, HifiGanGenerator, MelFrontend


@dataclass
class MicroFastSpeechConfig:
    vocab_size: int = 256
    tone_size: int = 16
    lang_size: int = 4
    n_mels: int = 80
    hidden: int = 168
    encoder_layers: int = 5
    decoder_layers: int = 6
    decoder_ff_mult: int = 3
    kernel_size: int = 7
    speaker_count: int = 2
    speaker_dim: int = 64
    dropout: float = 0.08
    sample_rate: int = 24000
    max_frames: int = 1400
    postnet_scale: float = 0.10
    use_frame_pitch: bool = True
    abs_frame_bins: int = 512
    use_contextual_predictors: bool = False
    use_group_duration_planner: bool = False


def count_parameters(model: nn.Module) -> int:
    return sum(p.numel() for p in model.parameters())


def load_rows(path: Path, max_rows: int = 0) -> list[dict]:
    rows = []
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            if line.strip():
                row = json.loads(line)
                if Path(str(row.get("target_audio") or "")).is_file():
                    rows.append(row)
                    if max_rows and len(rows) >= max_rows:
                        break
    if not rows:
        raise RuntimeError(f"No usable rows in {path}")
    return rows


def load_audio(path: str, sample_rate: int, max_seconds: float) -> torch.Tensor:
    wav, sr = torchaudio.load(path)
    if wav.shape[0] > 1:
        wav = wav.mean(dim=0, keepdim=True)
    if sr != sample_rate:
        wav = torchaudio.functional.resample(wav, sr, sample_rate)
    return wav[:, : int(sample_rate * max_seconds)].squeeze(0).clamp(-1.0, 1.0)


def fit_durations(durations: list[int], target_frames: int) -> list[int]:
    if sum(durations) == target_frames:
        return list(durations)
    total = max(1, sum(durations))
    raw = [max(0.0, d * target_frames / total) for d in durations]
    out = [int(math.floor(x)) for x in raw]
    order = sorted(((raw[i] - out[i], i) for i in range(len(out))), reverse=True)
    for _, idx in order[: max(0, target_frames - sum(out))]:
        out[idx] += 1
    while sum(out) > target_frames:
        idx = max(range(len(out)), key=lambda i: out[i])
        out[idx] -= 1
    return out


def pad_1d(items: list[torch.Tensor], value: float = 0.0) -> torch.Tensor:
    max_len = max(x.numel() for x in items)
    out = torch.full((len(items), max_len), value, dtype=items[0].dtype)
    for i, item in enumerate(items):
        out[i, : item.numel()] = item
    return out


def pad_2d(items: list[torch.Tensor], value: float = 0.0) -> torch.Tensor:
    max_len = max(x.shape[0] for x in items)
    dim = items[0].shape[1]
    out = torch.full((len(items), max_len, dim), value, dtype=items[0].dtype)
    for i, item in enumerate(items):
        out[i, : item.shape[0]] = item
    return out


def pad_mels(items: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
    max_len = max(x.shape[-1] for x in items)
    n_mels = items[0].shape[0]
    out = torch.zeros(len(items), n_mels, max_len, dtype=items[0].dtype)
    mask = torch.zeros(len(items), max_len, dtype=torch.bool)
    for i, mel in enumerate(items):
        frames = mel.shape[-1]
        out[i, :, :frames] = mel
        mask[i, :frames] = True
    return out, mask


def pad_wavs(items: list[torch.Tensor], frames: list[int], hop_size: int) -> torch.Tensor:
    max_len = max(max(1, int(frame_count)) * hop_size for frame_count in frames)
    out = torch.zeros(len(items), max_len, dtype=items[0].dtype)
    for i, (wav, frame_count) in enumerate(zip(items, frames)):
        length = max(1, int(frame_count)) * hop_size
        cropped = wav[:length]
        out[i, : cropped.numel()] = cropped
    return out


def aggregate_token_features(mel: torch.Tensor, durations: list[int]) -> tuple[torch.Tensor, torch.Tensor]:
    # mel: [80, frames], log-mel from the exact V2+ frontend.
    frames = mel.shape[-1]
    amp = torch.exp(mel).clamp_min(1e-5)
    energy_frame = mel.mean(dim=0)
    bins = torch.linspace(0.0, 1.0, mel.shape[0], device=mel.device).view(-1, 1)
    bright_frame = (amp * bins).sum(dim=0) / amp.sum(dim=0).clamp_min(1e-5)
    energies = []
    brights = []
    pos = 0
    for dur in durations:
        end = min(frames, pos + max(0, int(dur)))
        if end > pos:
            energies.append(energy_frame[pos:end].mean())
            brights.append(bright_frame[pos:end].mean())
        else:
            energies.append(torch.zeros((), device=mel.device, dtype=mel.dtype))
            brights.append(torch.zeros((), device=mel.device, dtype=mel.dtype))
        pos = end
    return torch.stack(energies), torch.stack(brights)


def aggregate_token_pitch(pitch_frame: torch.Tensor, durations: list[int]) -> torch.Tensor:
    # pitch_frame: [2, frames] with normalized log-f0 and voiced flag.
    frames = pitch_frame.shape[-1]
    out = []
    pos = 0
    for dur in durations:
        end = min(frames, pos + max(0, int(dur)))
        if end > pos:
            span = pitch_frame[:, pos:end]
            voiced = span[1].mean()
            voiced_mask = span[1] > 0.5
            if bool(voiced_mask.any()):
                log_f0 = span[0, voiced_mask].mean()
            else:
                log_f0 = torch.zeros((), dtype=pitch_frame.dtype)
            out.append(torch.stack([log_f0, voiced]))
        else:
            out.append(torch.zeros(2, dtype=pitch_frame.dtype))
        pos = end
    return torch.stack(out, dim=0)


def extract_pitch_features(wav: torch.Tensor, sample_rate: int, frames: int) -> torch.Tensor:
    # Returns [2, frames]: normalized log-f0 and voiced flag. The detector can
    # produce octave spikes, so clip to speech range and median-smooth lightly.
    pitch = torchaudio.functional.detect_pitch_frequency(
        wav.unsqueeze(0).cpu(),
        sample_rate,
        frame_time=256 / sample_rate,
    ).squeeze(0)
    if pitch.numel() < frames:
        pitch = F.pad(pitch, (0, frames - pitch.numel()), value=0.0)
    pitch = pitch[:frames]
    voiced = ((pitch >= 55.0) & (pitch <= 420.0)).float()
    pitch = pitch.clamp(55.0, 420.0)
    # Median filter over 5 frames to reduce spurious jumps.
    if pitch.numel() >= 5:
        padded = F.pad(pitch.view(1, 1, -1), (2, 2), mode="replicate")
        windows = padded.unfold(-1, 5, 1).squeeze(0).squeeze(0)
        pitch = windows.median(dim=-1).values
    log_f0 = (torch.log(pitch) - math.log(140.0)) / 0.45
    log_f0 = log_f0.clamp(-3.0, 3.0) * voiced
    return torch.stack([log_f0, voiced], dim=0)


class ConvFFNBlock(nn.Module):
    def __init__(self, hidden: int, kernel_size: int, dropout: float, ff_mult: int = 4) -> None:
        super().__init__()
        pad = kernel_size // 2
        self.norm1 = nn.LayerNorm(hidden)
        self.depth = nn.Conv1d(hidden, hidden * 2, kernel_size, padding=pad, groups=hidden)
        self.point = nn.Conv1d(hidden, hidden, 1)
        self.drop = nn.Dropout(dropout)
        self.norm2 = nn.LayerNorm(hidden)
        self.ff = nn.Sequential(
            nn.Linear(hidden, hidden * ff_mult),
            nn.SiLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden * ff_mult, hidden),
        )

    def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
        y = self.norm1(x).transpose(1, 2)
        a, b = self.depth(y).chunk(2, dim=1)
        y = self.point(a * torch.sigmoid(b)).transpose(1, 2)
        x = x + self.drop(y)
        x = x + self.drop(self.ff(self.norm2(x)))
        if mask is not None:
            x = x * mask.unsqueeze(-1)
        return x


class MicroFastSpeech(nn.Module):
    def __init__(self, cfg: MicroFastSpeechConfig) -> None:
        super().__init__()
        self.cfg = cfg
        # Phone id 0 is a real inserted blank/silence token from TinyTTS, not
        # padding. Padding is tracked by duration masks instead.
        self.phone = nn.Embedding(cfg.vocab_size, cfg.hidden)
        self.tone = nn.Embedding(cfg.tone_size, cfg.hidden)
        self.lang = nn.Embedding(cfg.lang_size, cfg.hidden)
        self.speaker = nn.Embedding(cfg.speaker_count, cfg.speaker_dim)
        self.speaker_proj = nn.Linear(cfg.speaker_dim, cfg.hidden)
        self.encoder = nn.ModuleList([ConvFFNBlock(cfg.hidden, cfg.kernel_size, cfg.dropout) for _ in range(cfg.encoder_layers)])
        self.duration_head = nn.Sequential(nn.LayerNorm(cfg.hidden), nn.Linear(cfg.hidden, cfg.hidden), nn.SiLU(), nn.Linear(cfg.hidden, 1))
        self.energy_head = nn.Sequential(nn.LayerNorm(cfg.hidden), nn.Linear(cfg.hidden, cfg.hidden // 2), nn.SiLU(), nn.Linear(cfg.hidden // 2, 1))
        self.bright_head = nn.Sequential(nn.LayerNorm(cfg.hidden), nn.Linear(cfg.hidden, cfg.hidden // 2), nn.SiLU(), nn.Linear(cfg.hidden // 2, 1))
        self.pitch_head = nn.Sequential(nn.LayerNorm(cfg.hidden), nn.Linear(cfg.hidden, cfg.hidden), nn.SiLU(), nn.Linear(cfg.hidden, 2))
        self.group_duration_delta = nn.Linear(cfg.hidden, 1) if cfg.use_group_duration_planner else None
        if self.group_duration_delta is not None:
            nn.init.zeros_(self.group_duration_delta.weight)
            nn.init.zeros_(self.group_duration_delta.bias)
        self.predictor_context = (
            ConvFFNBlock(cfg.hidden, 5, cfg.dropout, 2) if cfg.use_contextual_predictors else nn.Identity()
        )
        self.duration_delta = nn.Linear(cfg.hidden, 1) if cfg.use_contextual_predictors else None
        self.energy_delta = nn.Linear(cfg.hidden, 1) if cfg.use_contextual_predictors else None
        self.bright_delta = nn.Linear(cfg.hidden, 1) if cfg.use_contextual_predictors else None
        self.pitch_delta = nn.Linear(cfg.hidden, 2) if cfg.use_contextual_predictors else None
        if cfg.use_contextual_predictors:
            for layer in (self.duration_delta, self.energy_delta, self.bright_delta, self.pitch_delta):
                nn.init.zeros_(layer.weight)
                nn.init.zeros_(layer.bias)
        self.energy_proj = nn.Linear(1, cfg.hidden)
        self.bright_proj = nn.Linear(1, cfg.hidden)
        self.pitch_proj = nn.Sequential(nn.Linear(2, cfg.hidden), nn.SiLU(), nn.Linear(cfg.hidden, cfg.hidden))
        self.abs_frame = nn.Embedding(cfg.abs_frame_bins, cfg.hidden)
        self.frame_proj = nn.Sequential(nn.Linear(8, cfg.hidden), nn.SiLU(), nn.Linear(cfg.hidden, cfg.hidden))
        self.local_ctx = nn.Sequential(
            nn.Linear(cfg.hidden * 3, cfg.hidden * 2),
            nn.SiLU(),
            nn.Linear(cfg.hidden * 2, cfg.hidden),
        )
        self.decoder = nn.ModuleList([ConvFFNBlock(cfg.hidden, cfg.kernel_size, cfg.dropout, cfg.decoder_ff_mult) for _ in range(cfg.decoder_layers)])
        self.frame_gru = nn.GRU(cfg.hidden, cfg.hidden // 2, num_layers=1, batch_first=True, bidirectional=True)
        self.mel_head = nn.Sequential(nn.LayerNorm(cfg.hidden), nn.Linear(cfg.hidden, cfg.hidden), nn.SiLU(), nn.Linear(cfg.hidden, cfg.n_mels))
        self.postnet = nn.Sequential(
            nn.Conv1d(cfg.n_mels, cfg.hidden, 5, padding=2),
            nn.Tanh(),
            nn.Conv1d(cfg.hidden, cfg.hidden, 5, padding=2),
            nn.Tanh(),
            nn.Conv1d(cfg.hidden, cfg.n_mels, 5, padding=2),
        )

    def encode(self, phone: torch.Tensor, tone: torch.Tensor, lang: torch.Tensor, speaker: torch.Tensor, token_mask: torch.Tensor) -> torch.Tensor:
        x = self.phone(phone) + self.tone(tone.clamp_max(self.cfg.tone_size - 1)) + self.lang(lang.clamp_max(self.cfg.lang_size - 1))
        x = x + self.speaker_proj(self.speaker(speaker)).unsqueeze(1)
        x = x * token_mask.unsqueeze(-1)
        for block in self.encoder:
            x = block(x, token_mask)
        return x

    def regulate(self, encoded: torch.Tensor, durations: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        batch_frames = []
        batch_meta = []
        lengths = []
        device = encoded.device
        for b in range(encoded.shape[0]):
            reps = []
            meta = []
            durs = durations[b].long().clamp_min(0)
            token_count = max(1, int((durs > 0).sum().item()))
            for i, dur_t in enumerate(durs.tolist()):
                dur = int(dur_t)
                if dur <= 0:
                    continue
                reps.append(encoded[b, i].view(1, -1).expand(dur, -1))
                rel = torch.linspace(0.0, 1.0, dur, device=device)
                token_pos = torch.full((dur,), i / max(1, token_count - 1), device=device)
                log_dur = torch.full((dur,), math.log1p(dur) / 6.0, device=device)
                inv_rel = 1.0 - rel
                center = 1.0 - torch.abs(rel * 2.0 - 1.0)
                meta.append(
                    torch.stack(
                        [
                            rel,
                            inv_rel,
                            center,
                            torch.sin(rel * math.pi),
                            torch.cos(rel * math.pi),
                            token_pos,
                            log_dur,
                            torch.full_like(rel, dur / 40.0),
                        ],
                        dim=-1,
                    )
                )
            if reps:
                frames = torch.cat(reps, dim=0)
                frame_meta = torch.cat(meta, dim=0)
            else:
                frames = encoded[b, :1]
                frame_meta = torch.zeros(1, 8, device=device)
            batch_frames.append(frames[: self.cfg.max_frames])
            batch_meta.append(frame_meta[: self.cfg.max_frames])
            lengths.append(min(frames.shape[0], self.cfg.max_frames))
        max_len = max(lengths)
        out = torch.zeros(encoded.shape[0], max_len, encoded.shape[-1], device=device)
        meta_out = torch.zeros(encoded.shape[0], max_len, 8, device=device)
        mask = torch.zeros(encoded.shape[0], max_len, dtype=torch.bool, device=device)
        for b, frames in enumerate(batch_frames):
            n = min(frames.shape[0], max_len)
            out[b, :n] = frames[:n]
            meta_out[b, :n] = batch_meta[b][:n]
            mask[b, :n] = True
        return out, meta_out, mask

    def add_local_context(self, encoded: torch.Tensor, durations: torch.Tensor) -> torch.Tensor:
        device = encoded.device
        batch_frames = []
        for b in range(encoded.shape[0]):
            reps = []
            durs = durations[b].long().clamp_min(0)
            for i, dur_t in enumerate(durs.tolist()):
                dur = int(dur_t)
                if dur <= 0:
                    continue
                prev_i = max(0, i - 1)
                next_i = min(encoded.shape[1] - 1, i + 1)
                ctx = torch.cat([encoded[b, prev_i], encoded[b, i], encoded[b, next_i]], dim=-1)
                reps.append(ctx.view(1, -1).expand(dur, -1))
            if reps:
                frames = torch.cat(reps, dim=0)
            else:
                frames = torch.zeros(1, encoded.shape[-1] * 3, device=device)
            batch_frames.append(frames[: self.cfg.max_frames])
        max_len = max(x.shape[0] for x in batch_frames)
        ctx_out = torch.zeros(encoded.shape[0], max_len, encoded.shape[-1] * 3, device=device)
        for b, frames in enumerate(batch_frames):
            ctx_out[b, : frames.shape[0]] = frames
        return self.local_ctx(ctx_out)

    def expand_token_feature(self, feature: torch.Tensor, durations: torch.Tensor) -> torch.Tensor:
        device = feature.device
        batch_frames = []
        for b in range(feature.shape[0]):
            reps = []
            durs = durations[b].long().clamp_min(0)
            for i, dur_t in enumerate(durs.tolist()):
                dur = int(dur_t)
                if dur <= 0:
                    continue
                reps.append(feature[b, i].view(1, -1).expand(dur, -1))
            if reps:
                frames = torch.cat(reps, dim=0)
            else:
                frames = torch.zeros(1, feature.shape[-1], device=device)
            batch_frames.append(frames[: self.cfg.max_frames])
        max_len = max(x.shape[0] for x in batch_frames)
        out = torch.zeros(feature.shape[0], max_len, feature.shape[-1], device=device)
        for b, frames in enumerate(batch_frames):
            out[b, : frames.shape[0]] = frames
        return out

    def forward(
        self,
        phone: torch.Tensor,
        tone: torch.Tensor,
        lang: torch.Tensor,
        speaker: torch.Tensor,
        durations: torch.Tensor,
        energy_target: torch.Tensor | None = None,
        bright_target: torch.Tensor | None = None,
        pitch_frame: torch.Tensor | None = None,
        predicted_prosody_mix: float = 0.0,
        detach_mixed_predictions: bool = True,
    ) -> dict[str, torch.Tensor]:
        token_mask = durations.gt(0)
        encoded = self.encode(phone, tone, lang, speaker, token_mask)
        log_dur, energy_pred, bright_pred, pitch_pred = self.predict_prosody(encoded, token_mask)
        mixed_energy_pred = energy_pred.detach() if detach_mixed_predictions else energy_pred
        mixed_bright_pred = bright_pred.detach() if detach_mixed_predictions else bright_pred
        if energy_target is not None:
            energy = torch.lerp(energy_target, mixed_energy_pred, predicted_prosody_mix)
        else:
            energy = energy_pred
        if bright_target is not None:
            bright = torch.lerp(bright_target, mixed_bright_pred, predicted_prosody_mix)
        else:
            bright = bright_pred
        conditioned = encoded + self.energy_proj(energy.unsqueeze(-1)) + self.bright_proj(bright.unsqueeze(-1))
        frames, frame_meta, frame_mask = self.regulate(conditioned, durations)
        x = frames + self.frame_proj(frame_meta) + self.add_local_context(conditioned, durations)
        pos = torch.arange(x.shape[1], device=x.device)
        pos = torch.div(pos * self.cfg.abs_frame_bins, max(1, self.cfg.max_frames), rounding_mode="floor").clamp_max(
            self.cfg.abs_frame_bins - 1
        )
        x = x + self.abs_frame(pos).unsqueeze(0)
        if self.cfg.use_frame_pitch:
            if pitch_frame is not None:
                pitch_frame = pitch_frame[:, :, : x.shape[1]].transpose(1, 2)
                if pitch_frame.shape[1] < x.shape[1]:
                    pitch_frame = F.pad(pitch_frame, (0, 0, 0, x.shape[1] - pitch_frame.shape[1]))
                if predicted_prosody_mix > 0.0:
                    mixed_pitch_pred = pitch_pred.detach() if detach_mixed_predictions else pitch_pred
                    predicted_pitch_frame = self.expand_token_feature(mixed_pitch_pred, durations)[:, : x.shape[1]]
                    pitch_frame = torch.lerp(pitch_frame, predicted_pitch_frame, predicted_prosody_mix)
            else:
                pitch_frame = self.expand_token_feature(pitch_pred, durations)[:, : x.shape[1]]
            x = x + self.pitch_proj(pitch_frame)
        for block in self.decoder:
            x = block(x, frame_mask)
        x = x + self.frame_gru(x)[0]
        mel = self.mel_head(x).transpose(1, 2)
        mel = mel + self.cfg.postnet_scale * self.postnet(mel)
        group_log_dur, group_mask = self.group_log_durations(phone, log_dur, encoded)
        return {
            "mel": mel,
            "frame_mask": frame_mask,
            "log_dur": log_dur,
            "group_log_dur": group_log_dur,
            "group_mask": group_mask,
            "energy": energy_pred,
            "bright": bright_pred,
            "pitch": pitch_pred,
            "token_mask": token_mask,
        }

    def predict_prosody(
        self, encoded: torch.Tensor, token_mask: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        log_dur = self.duration_head(encoded).squeeze(-1)
        energy = self.energy_head(encoded).squeeze(-1)
        bright = self.bright_head(encoded).squeeze(-1)
        pitch = self.pitch_head(encoded)
        if self.cfg.use_contextual_predictors:
            context = self.predictor_context(encoded, token_mask)
            log_dur = log_dur + self.duration_delta(context).squeeze(-1)
            energy = energy + self.energy_delta(context).squeeze(-1)
            bright = bright + self.bright_delta(context).squeeze(-1)
            pitch = pitch + self.pitch_delta(context)
        return log_dur, energy, bright, pitch

    def group_log_durations(
        self, phone: torch.Tensor, log_dur: torch.Tensor, encoded: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Predict stable blank-plus-phone region durations at visible phones."""
        base_dur = torch.expm1(log_dur).clamp_min(0.05)
        grouped = torch.zeros_like(log_dur)
        group_mask = torch.zeros_like(phone, dtype=torch.bool)
        delta = self.group_duration_delta(encoded).squeeze(-1) if self.group_duration_delta is not None else None
        for batch_index in range(phone.shape[0]):
            pending: list[torch.Tensor] = []
            last_visible: int | None = None
            for token_index in range(phone.shape[1]):
                pending.append(base_dur[batch_index, token_index])
                if int(phone[batch_index, token_index].item()) != 0:
                    value = torch.stack(pending).sum()
                    if delta is not None:
                        value = value * torch.exp(delta[batch_index, token_index].clamp(-1.5, 1.5))
                    grouped[batch_index, token_index] = torch.log1p(value)
                    group_mask[batch_index, token_index] = True
                    pending = []
                    last_visible = token_index
            if pending and last_visible is not None:
                value = torch.expm1(grouped[batch_index, last_visible]) + torch.stack(pending).sum()
                grouped[batch_index, last_visible] = torch.log1p(value)
        return grouped, group_mask

    def apply_group_duration_plan(
        self, phone: torch.Tensor, log_dur: torch.Tensor, encoded: torch.Tensor, length_scale: float, max_duration: int
    ) -> torch.Tensor:
        base = torch.expm1(log_dur).clamp_min(0.05)
        group_log, _ = self.group_log_durations(phone, log_dur, encoded)
        planned = torch.zeros_like(base, dtype=torch.long)
        for batch_index in range(phone.shape[0]):
            pending: list[int] = []
            last_visible: int | None = None
            for token_index in range(phone.shape[1]):
                pending.append(token_index)
                if int(phone[batch_index, token_index].item()) != 0:
                    target = max(len(pending), int(round(float(torch.expm1(group_log[batch_index, token_index]) * length_scale))))
                    weights = base[batch_index, pending]
                    remaining = target - len(pending)
                    raw = weights / weights.sum().clamp_min(1e-6) * remaining
                    allocated = torch.ones_like(raw, dtype=torch.long) + torch.floor(raw).long()
                    remainder = target - int(allocated.sum().item())
                    if remainder > 0:
                        order = torch.argsort(raw - torch.floor(raw), descending=True)
                        allocated[order[:remainder]] += 1
                    planned[batch_index, pending] = allocated
                    pending = []
                    last_visible = token_index
            if pending and last_visible is not None:
                planned[batch_index, last_visible] += max(1, int(round(float(base[batch_index, pending].sum() * length_scale))))
        return planned.clamp(0, max_duration)

    @torch.no_grad()
    def infer(
        self,
        phone: torch.Tensor,
        tone: torch.Tensor,
        lang: torch.Tensor,
        speaker: torch.Tensor,
        length_scale: float = 1.0,
        min_duration: int = 1,
        max_duration: int = 80,
        pitch_scale: float = 1.0,
        energy_scale: float = 1.0,
        smooth_predictors: bool = False,
    ) -> torch.Tensor:
        # In single-sample inference there is no padded tail; id 0 remains the
        # explicit blank/pause token and must keep duration.
        token_mask = torch.ones_like(phone, dtype=torch.bool)
        encoded = self.encode(phone, tone, lang, speaker, token_mask)
        log_dur, energy, bright, pitch = self.predict_prosody(encoded, token_mask)
        if self.group_duration_delta is not None:
            durations = self.apply_group_duration_plan(phone, log_dur, encoded, length_scale, max_duration)
            durations = durations.masked_fill(~token_mask, 0)
        else:
            pred_dur = torch.expm1(log_dur).clamp(0, max_duration) * length_scale
            durations = torch.round(pred_dur).long().clamp_min(min_duration).masked_fill(~token_mask, 0)
        energy = energy * energy_scale
        pitch = torch.stack([pitch[..., 0] * pitch_scale, pitch[..., 1].clamp(0.0, 1.0)], dim=-1)
        if smooth_predictors and phone.shape[1] >= 3:
            energy = F.avg_pool1d(energy.unsqueeze(1), 3, stride=1, padding=1).squeeze(1)
            bright = F.avg_pool1d(bright.unsqueeze(1), 3, stride=1, padding=1).squeeze(1)
            pitch_t = pitch.transpose(1, 2)
            pitch = F.avg_pool1d(pitch_t, 3, stride=1, padding=1).transpose(1, 2)
        conditioned = encoded + self.energy_proj(energy.unsqueeze(-1)) + self.bright_proj(bright.unsqueeze(-1))
        frames, frame_meta, frame_mask = self.regulate(conditioned, durations)
        x = frames + self.frame_proj(frame_meta) + self.add_local_context(conditioned, durations)
        pos = torch.arange(x.shape[1], device=x.device)
        pos = torch.div(pos * self.cfg.abs_frame_bins, max(1, self.cfg.max_frames), rounding_mode="floor").clamp_max(
            self.cfg.abs_frame_bins - 1
        )
        x = x + self.abs_frame(pos).unsqueeze(0)
        if self.cfg.use_frame_pitch:
            pitch_frame = self.expand_token_feature(pitch, durations)[:, : x.shape[1]]
            x = x + self.pitch_proj(pitch_frame)
        for block in self.decoder:
            x = block(x, frame_mask)
        x = x + self.frame_gru(x)[0]
        mel = self.mel_head(x).transpose(1, 2)
        mel = mel + self.cfg.postnet_scale * self.postnet(mel)
        return mel


def collate(batch: list[dict], cfg: MicroFastSpeechConfig, mel_frontend: MelFrontend, device: torch.device, max_seconds: float, hop_size: int):
    phones = [torch.LongTensor(x["phone_ids"]) for x in batch]
    tones = [torch.LongTensor(x["tone_ids"]) for x in batch]
    langs = [torch.LongTensor(x["lang_ids"]) for x in batch]
    durations_raw = [list(map(int, x["hifigan_durations"])) for x in batch]
    speakers = torch.LongTensor([int(x["speaker_id"]) for x in batch])
    phone = pad_1d(phones, 0).long()
    tone = pad_1d(tones, 0).long()
    lang = pad_1d(langs, 0).long()
    mels = []
    durations = []
    energies = []
    brights = []
    pitches = []
    token_pitches = []
    wavs = []
    frame_counts = []
    with torch.no_grad():
        for row, dur in zip(batch, durations_raw):
            wav_1d = load_audio(str(row["target_audio"]), cfg.sample_rate, max_seconds)
            wav = wav_1d.unsqueeze(0).to(device)
            mel = mel_frontend(wav).squeeze(0).detach().cpu()
            dur = fit_durations(dur[: len(row["phone_ids"])], min(mel.shape[-1], cfg.max_frames))
            mel = mel[:, : sum(dur)]
            energy, bright = aggregate_token_features(mel, dur)
            pitch = extract_pitch_features(wav_1d, cfg.sample_rate, mel.shape[-1])
            token_pitch = aggregate_token_pitch(pitch, dur)
            mels.append(mel)
            durations.append(torch.LongTensor(dur))
            energies.append(energy)
            brights.append(bright)
            pitches.append(pitch)
            token_pitches.append(token_pitch)
            wavs.append(wav_1d)
            frame_counts.append(mel.shape[-1])
    duration = pad_1d(durations, 0).long()
    energy = pad_1d(energies, 0.0).float()
    bright = pad_1d(brights, 0.0).float()
    token_pitch = pad_2d(token_pitches, 0.0).float()
    target_mel, frame_mask = pad_mels(mels)
    pitch_frame, _ = pad_mels(pitches)
    target_wav = pad_wavs(wavs, frame_counts, hop_size)
    return (
        phone.to(device),
        tone.to(device),
        lang.to(device),
        speakers.to(device),
        duration.to(device),
        energy.to(device),
        bright.to(device),
        token_pitch.to(device),
        target_mel.to(device),
        frame_mask.to(device),
        pitch_frame.to(device),
        target_wav.to(device),
    )


def prepare_row_features(
    row: dict,
    cfg: MicroFastSpeechConfig,
    mel_frontend: MelFrontend,
    device: torch.device,
    max_seconds: float,
) -> dict:
    dur = list(map(int, row["hifigan_durations"]))
    wav_1d = load_audio(str(row["target_audio"]), cfg.sample_rate, max_seconds)
    with torch.no_grad():
        wav = wav_1d.unsqueeze(0).to(device)
        mel = mel_frontend(wav).squeeze(0).detach().cpu()
    dur = fit_durations(dur[: len(row["phone_ids"])], min(mel.shape[-1], cfg.max_frames))
    mel = mel[:, : sum(dur)]
    energy, bright = aggregate_token_features(mel, dur)
    pitch = extract_pitch_features(wav_1d, cfg.sample_rate, mel.shape[-1])
    token_pitch = aggregate_token_pitch(pitch, dur)
    return {
        "phone": torch.LongTensor(row["phone_ids"]),
        "tone": torch.LongTensor(row["tone_ids"]),
        "lang": torch.LongTensor(row["lang_ids"]),
        "speaker": int(row["speaker_id"]),
        "duration": torch.LongTensor(dur),
        "energy": energy.float(),
        "bright": bright.float(),
        "token_pitch": token_pitch.float(),
        "target_mel": mel.float(),
        "pitch_frame": pitch.float(),
        "target_wav": wav_1d.float(),
        "frame_count": int(mel.shape[-1]),
    }


def collate_prepared(batch: list[dict], device: torch.device, hop_size: int):
    phone = pad_1d([x["phone"] for x in batch], 0).long()
    tone = pad_1d([x["tone"] for x in batch], 0).long()
    lang = pad_1d([x["lang"] for x in batch], 0).long()
    speakers = torch.LongTensor([int(x["speaker"]) for x in batch])
    duration = pad_1d([x["duration"] for x in batch], 0).long()
    energy = pad_1d([x["energy"] for x in batch], 0.0).float()
    bright = pad_1d([x["bright"] for x in batch], 0.0).float()
    token_pitch = pad_2d([x["token_pitch"] for x in batch], 0.0).float()
    target_mel, frame_mask = pad_mels([x["target_mel"] for x in batch])
    pitch_frame, _ = pad_mels([x["pitch_frame"] for x in batch])
    target_wav = pad_wavs([x["target_wav"] for x in batch], [int(x["frame_count"]) for x in batch], hop_size)
    return (
        phone.to(device),
        tone.to(device),
        lang.to(device),
        speakers.to(device),
        duration.to(device),
        energy.to(device),
        bright.to(device),
        token_pitch.to(device),
        target_mel.to(device),
        frame_mask.to(device),
        pitch_frame.to(device),
        target_wav.to(device),
    )


def masked_l1(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    common = min(pred.shape[-1], target.shape[-1], mask.shape[-1])
    pred = pred[..., :common]
    target = target[..., :common]
    mask = mask[:, :common].unsqueeze(1)
    return (torch.abs(pred - target) * mask).sum() / (mask.sum() * pred.shape[1]).clamp_min(1.0)


def masked_mse(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    common = min(pred.shape[-1], target.shape[-1], mask.shape[-1])
    pred = pred[..., :common]
    target = target[..., :common]
    mask = mask[:, :common].unsqueeze(1)
    return (((pred - target) ** 2) * mask).sum() / (mask.sum() * pred.shape[1]).clamp_min(1.0)


def masked_delta_loss(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    common = min(pred.shape[-1], target.shape[-1], mask.shape[-1])
    if common < 2:
        return torch.zeros((), device=pred.device)
    dp = pred[..., 1:common] - pred[..., : common - 1]
    dt = target[..., 1:common] - target[..., : common - 1]
    dm = (mask[:, 1:common] & mask[:, : common - 1]).unsqueeze(1)
    return (torch.abs(dp - dt) * dm).sum() / (dm.sum() * pred.shape[1]).clamp_min(1.0)


def masked_accel_loss(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    common = min(pred.shape[-1], target.shape[-1], mask.shape[-1])
    if common < 3:
        return torch.zeros((), device=pred.device)
    dp = pred[..., 2:common] - 2.0 * pred[..., 1 : common - 1] + pred[..., : common - 2]
    dt = target[..., 2:common] - 2.0 * target[..., 1 : common - 1] + target[..., : common - 2]
    dm = (mask[:, 2:common] & mask[:, 1 : common - 1] & mask[:, : common - 2]).unsqueeze(1)
    return (torch.abs(dp - dt) * dm).sum() / (dm.sum() * pred.shape[1]).clamp_min(1.0)


def token_mse(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    common = min(pred.shape[1], target.shape[1], mask.shape[1])
    pred = pred[:, :common]
    target = target[:, :common]
    mask = mask[:, :common]
    return (((pred - target) ** 2) * mask).sum() / mask.sum().clamp_min(1.0)


def token_mse_nd(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    common = min(pred.shape[1], target.shape[1], mask.shape[1])
    pred = pred[:, :common]
    target = target[:, :common]
    mask = mask[:, :common].unsqueeze(-1)
    return (((pred - target) ** 2) * mask).sum() / (mask.sum() * pred.shape[-1]).clamp_min(1.0)


def group_duration_targets(phone: torch.Tensor, durations: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    grouped = torch.zeros_like(durations, dtype=torch.float32)
    mask = torch.zeros_like(phone, dtype=torch.bool)
    for batch_index in range(phone.shape[0]):
        pending: list[torch.Tensor] = []
        last_visible: int | None = None
        for token_index in range(phone.shape[1]):
            pending.append(durations[batch_index, token_index].float())
            if int(phone[batch_index, token_index].item()) != 0:
                grouped[batch_index, token_index] = torch.log1p(torch.stack(pending).sum())
                mask[batch_index, token_index] = True
                pending = []
                last_visible = token_index
        if pending and last_visible is not None:
            value = torch.expm1(grouped[batch_index, last_visible]) + torch.stack(pending).sum()
            grouped[batch_index, last_visible] = torch.log1p(value)
    return grouped, mask


def masked_wav_l1(pred: torch.Tensor, target: torch.Tensor, frame_mask: torch.Tensor, hop_size: int) -> torch.Tensor:
    if pred.dim() == 3:
        pred = pred.squeeze(1)
    common = min(pred.shape[-1], target.shape[-1], frame_mask.shape[-1] * hop_size)
    pred = pred[:, :common]
    target = target[:, :common]
    sample_mask = frame_mask.repeat_interleave(hop_size, dim=1)[:, :common].to(pred.dtype)
    return (torch.abs(pred - target) * sample_mask).sum() / sample_mask.sum().clamp_min(1.0)


def load_frozen_vocoder(path: Path, device: torch.device) -> tuple[HifiGanGenerator, HifiGanConfig]:
    ckpt = torch.load(path, map_location=device, weights_only=False)
    cfg_payload = ckpt.get("config") or {"variant": "v2plus"}
    cfg = HifiGanConfig(**cfg_payload) if isinstance(cfg_payload, dict) else cfg_payload
    vocoder = HifiGanGenerator(cfg).to(device)
    vocoder.load_state_dict(ckpt["generator"])
    vocoder.eval()
    for param in vocoder.parameters():
        param.requires_grad_(False)
    return vocoder, cfg


def load_model_state_flexible(model: nn.Module, state: dict[str, torch.Tensor]) -> tuple[int, int]:
    current = model.state_dict()
    compatible = {key: value for key, value in state.items() if key in current and current[key].shape == value.shape}
    model.load_state_dict(compatible, strict=False)
    return len(compatible), len(state) - len(compatible)


def set_trainable_by_mode(model: MicroFastSpeech, mode: str) -> None:
    if mode == "all":
        for param in model.parameters():
            param.requires_grad_(True)
        return
    for param in model.parameters():
        param.requires_grad_(False)
    prefixes: tuple[str, ...]
    if mode == "duration":
        prefixes = ("phone.", "tone.", "lang.", "speaker.", "speaker_proj.", "encoder.", "duration_head.")
    elif mode == "predictors":
        prefixes = (
            "phone.",
            "tone.",
            "lang.",
            "speaker.",
            "speaker_proj.",
            "encoder.",
            "duration_head.",
            "energy_head.",
            "bright_head.",
            "pitch_head.",
        )
    elif mode == "heads":
        prefixes = (
            "duration_head.",
            "energy_head.",
            "bright_head.",
            "pitch_head.",
            "predictor_context.",
            "duration_delta.",
            "energy_delta.",
            "bright_delta.",
            "pitch_delta.",
        )
    elif mode == "contextual":
        prefixes = (
            "predictor_context.",
            "duration_delta.",
            "energy_delta.",
            "bright_delta.",
            "pitch_delta.",
        )
    elif mode == "group_duration":
        prefixes = ("group_duration_delta.",)
    elif mode == "decoder_adapt":
        prefixes = (
            "energy_proj.",
            "bright_proj.",
            "pitch_proj.",
            "abs_frame.",
            "frame_proj.",
            "local_ctx.",
            "decoder.",
            "frame_gru.",
            "mel_head.",
            "postnet.",
        )
    else:
        raise ValueError(f"Unknown trainable mode: {mode}")
    for name, param in model.named_parameters():
        if name.startswith(prefixes):
            param.requires_grad_(True)


def latest_checkpoint(out_dir: Path) -> Path | None:
    found = []
    for path in out_dir.glob("inflect-micro-fastspeech-*.pt"):
        tail = path.stem.rsplit("-", 1)[-1]
        if tail.isdigit():
            found.append((int(tail), path))
    return max(found)[1] if found else None


def save_checkpoint(path: Path, model: nn.Module, optim, cfg: MicroFastSpeechConfig, step: int, args, speakers: dict[str, int]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    torch.save(
        {
            "model": model.state_dict(),
            "optim": optim.state_dict(),
            "config": asdict(cfg),
            "step": step,
            "speakers": speakers,
            "args": {k: str(v) if isinstance(v, Path) else v for k, v in vars(args).items()},
            "params": count_parameters(model),
        },
        tmp,
    )
    tmp.replace(path)


def train(args: argparse.Namespace) -> None:
    device = torch.device(args.device)
    rows = load_rows(args.durations_jsonl, args.max_rows)
    speakers = {voice: idx for idx, voice in enumerate(sorted({str(r.get("voice_id") or "mark") for r in rows}))}
    max_phone_id = max(max(map(int, r["phone_ids"])) for r in rows)
    max_tone_id = max(max(map(int, r["tone_ids"])) for r in rows)
    max_lang_id = max(max(map(int, r["lang_ids"])) for r in rows)
    cfg = MicroFastSpeechConfig(
        vocab_size=max(256, max_phone_id + 1),
        tone_size=max(16, max_tone_id + 1),
        lang_size=max(4, max_lang_id + 1),
        speaker_count=max(2, len(speakers)),
        hidden=args.hidden,
        encoder_layers=args.encoder_layers,
        decoder_layers=args.decoder_layers,
        decoder_ff_mult=args.decoder_ff_mult,
        max_frames=args.max_frames,
        postnet_scale=args.postnet_scale,
        abs_frame_bins=args.abs_frame_bins,
        use_contextual_predictors=args.contextual_predictors,
        use_group_duration_planner=args.group_duration_planner,
    )
    for row in rows:
        row["speaker_id"] = speakers[str(row.get("voice_id") or "mark")]
    random.Random(args.seed).shuffle(rows)

    model = MicroFastSpeech(cfg).to(device)
    start_step = 0
    if args.init_checkpoint and not args.resume:
        ckpt = torch.load(args.init_checkpoint, map_location=device, weights_only=False)
        copied, skipped = load_model_state_flexible(model, ckpt["model"])
        print(f"Initialized model from {args.init_checkpoint} ({copied} tensors copied, {skipped} skipped)")
    set_trainable_by_mode(model, args.trainable)
    trainable_params = [param for param in model.parameters() if param.requires_grad]
    optim = torch.optim.AdamW(trainable_params, lr=args.lr, betas=(0.9, 0.98), weight_decay=args.weight_decay)
    if args.resume:
        ckpt_path = latest_checkpoint(args.out_dir)
        if ckpt_path:
            ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
            model.load_state_dict(ckpt["model"])
            optim.load_state_dict(ckpt["optim"])
            start_step = int(ckpt.get("step") or 0)
            print(f"Resumed {ckpt_path} at step {start_step}")

    hifi_cfg = HifiGanConfig(variant="v2plus")
    mel_frontend = MelFrontend(hifi_cfg).to(device)
    prepared_rows = None
    if args.preload_features:
        print("Preloading audio/mel/pitch features...", flush=True)
        prepared_rows = [prepare_row_features(row, cfg, mel_frontend, device, args.max_seconds) for row in rows]
        total_frames = sum(int(row["frame_count"]) for row in prepared_rows)
        print(f"Preloaded {len(prepared_rows)} rows ({total_frames:,} frames)", flush=True)
    consistency_vocoder = None
    if args.vocoder_checkpoint:
        consistency_vocoder, consistency_cfg = load_frozen_vocoder(args.vocoder_checkpoint, device)
        if consistency_cfg.hop_size != hifi_cfg.hop_size:
            raise RuntimeError(f"Vocoder hop mismatch: {consistency_cfg.hop_size} != {hifi_cfg.hop_size}")
        print(f"Loaded frozen vocoder consistency checkpoint: {args.vocoder_checkpoint}")
    if (args.vocoder_wav_weight > 0.0 or args.vocoder_mel_weight > 0.0) and consistency_vocoder is None:
        raise RuntimeError("--vocoder-checkpoint is required when vocoder consistency losses are enabled")
    args.out_dir.mkdir(parents=True, exist_ok=True)
    (args.out_dir / "config.json").write_text(
        json.dumps({"config": asdict(cfg), "speakers": speakers, "rows": len(rows), "params": count_parameters(model)}, indent=2),
        encoding="utf-8",
    )

    print(f"Rows: {len(rows)}")
    print(f"Speakers: {speakers}")
    print(f"Acoustic params: {count_parameters(model):,} ({count_parameters(model)/1_000_000:.3f}M)")
    print(f"Trainable params: {sum(p.numel() for p in model.parameters() if p.requires_grad):,} mode={args.trainable}")
    print(f"Total with V2+ vocoder: {(count_parameters(model)+1_426_842):,} ({(count_parameters(model)+1_426_842)/1_000_000:.3f}M)")

    rng = random.Random(args.seed + start_step)
    step = start_step
    started = time.time()
    while step < args.steps:
        source_rows = prepared_rows if prepared_rows is not None else rows
        batch = [source_rows[rng.randrange(len(source_rows))] for _ in range(args.batch_size)]
        if prepared_rows is not None:
            phone, tone, lang, speaker, durations, energy_t, bright_t, pitch_token_t, target_mel, frame_mask, pitch_frame, target_wav = collate_prepared(
                batch, device, hifi_cfg.hop_size
            )
        else:
            phone, tone, lang, speaker, durations, energy_t, bright_t, pitch_token_t, target_mel, frame_mask, pitch_frame, target_wav = collate(
                batch, cfg, mel_frontend, device, args.max_seconds, hifi_cfg.hop_size
            )
        out = model(phone, tone, lang, speaker, durations, energy_t, bright_t, pitch_frame)
        token_mask = out["token_mask"]
        log_dur_t = torch.log1p(durations.float())
        group_log_dur_t, group_mask = group_duration_targets(phone, durations)
        mel_l1 = masked_l1(out["mel"], target_mel, frame_mask)
        mel_mse = masked_mse(out["mel"], target_mel, frame_mask)
        delta = masked_delta_loss(out["mel"], target_mel, frame_mask)
        accel = masked_accel_loss(out["mel"], target_mel, frame_mask)
        dur_loss = token_mse(out["log_dur"], log_dur_t, token_mask)
        group_dur_loss = token_mse(out["group_log_dur"], group_log_dur_t, group_mask)
        energy_loss = token_mse(out["energy"], energy_t, token_mask)
        bright_loss = token_mse(out["bright"], bright_t, token_mask)
        pitch_loss = token_mse_nd(out["pitch"], pitch_token_t, token_mask)
        predicted_prosody_mel_loss = torch.zeros((), device=device)
        predicted_prosody_delta_loss = torch.zeros((), device=device)
        if args.predicted_prosody_mel_weight > 0.0 or args.predicted_prosody_delta_weight > 0.0:
            # Train the predictor heads against the acoustic result they produce at
            # inference, while retaining reference durations so this path remains
            # differentiable and isolates prosody exposure bias.
            predicted_conditioning = model(phone, tone, lang, speaker, durations)
            if args.predicted_prosody_mel_weight > 0.0:
                predicted_prosody_mel_loss = masked_l1(predicted_conditioning["mel"], target_mel, frame_mask)
            if args.predicted_prosody_delta_weight > 0.0:
                predicted_prosody_delta_loss = masked_delta_loss(predicted_conditioning["mel"], target_mel, frame_mask)
        robust_prosody_mel_loss = torch.zeros((), device=device)
        robust_prosody_delta_loss = torch.zeros((), device=device)
        if args.robust_prosody_mel_weight > 0.0 or args.robust_prosody_delta_weight > 0.0:
            robust_conditioning = model(
                phone,
                tone,
                lang,
                speaker,
                durations,
                energy_t,
                bright_t,
                pitch_frame,
                predicted_prosody_mix=args.robust_prosody_mix,
                detach_mixed_predictions=True,
            )
            if args.robust_prosody_mel_weight > 0.0:
                robust_prosody_mel_loss = masked_l1(robust_conditioning["mel"], target_mel, frame_mask)
            if args.robust_prosody_delta_weight > 0.0:
                robust_prosody_delta_loss = masked_delta_loss(robust_conditioning["mel"], target_mel, frame_mask)
        voc_wav_loss = torch.zeros((), device=device)
        voc_mel_loss = torch.zeros((), device=device)
        if consistency_vocoder is not None and (args.vocoder_wav_weight > 0.0 or args.vocoder_mel_weight > 0.0):
            pred_wav = consistency_vocoder(out["mel"].clamp(-12.0, 2.0))
            if args.vocoder_wav_weight > 0.0:
                voc_wav_loss = masked_wav_l1(pred_wav, target_wav, frame_mask, hifi_cfg.hop_size)
            if args.vocoder_mel_weight > 0.0:
                pred_recon_mel = mel_frontend(pred_wav.squeeze(1))
                voc_mel_loss = masked_l1(pred_recon_mel, target_mel, frame_mask)
        loss = (
            mel_l1
            + args.mse_weight * mel_mse
            + args.delta_weight * delta
            + args.accel_weight * accel
            + args.duration_weight * dur_loss
            + args.group_duration_weight * group_dur_loss
            + args.energy_weight * energy_loss
            + args.bright_weight * bright_loss
            + args.pitch_weight * pitch_loss
            + args.predicted_prosody_mel_weight * predicted_prosody_mel_loss
            + args.predicted_prosody_delta_weight * predicted_prosody_delta_loss
            + args.robust_prosody_mel_weight * robust_prosody_mel_loss
            + args.robust_prosody_delta_weight * robust_prosody_delta_loss
            + args.vocoder_wav_weight * voc_wav_loss
            + args.vocoder_mel_weight * voc_mel_loss
        )
        optim.zero_grad(set_to_none=True)
        loss.backward()
        grad = torch.nn.utils.clip_grad_norm_(trainable_params, args.grad_clip)
        optim.step()
        step += 1

        if step == 1 or step % args.log_interval == 0:
            elapsed = max(1e-6, time.time() - started)
            speed = (step - start_step) / elapsed
            eta = (args.steps - step) / max(1e-6, speed)
            print(
                f"step={step}/{args.steps} loss={loss.item():.4f} mel={mel_l1.item():.4f} "
                f"mse={mel_mse.item():.4f} delta={delta.item():.4f} accel={accel.item():.4f} dur={dur_loss.item():.4f} "
                f"gdur={group_dur_loss.item():.4f} "
                f"energy={energy_loss.item():.4f} bright={bright_loss.item():.4f} "
                f"pitch={pitch_loss.item():.4f} pmel={predicted_prosody_mel_loss.item():.4f} "
                f"pdelta={predicted_prosody_delta_loss.item():.4f} rmel={robust_prosody_mel_loss.item():.4f} "
                f"rdelta={robust_prosody_delta_loss.item():.4f} vwav={voc_wav_loss.item():.4f} "
                f"vmel={voc_mel_loss.item():.4f} grad={float(grad):.2f} "
                f"speed={speed:.3f} step/s eta={eta/60:.1f}m",
                flush=True,
            )
        if step % args.save_interval == 0 or step >= args.steps:
            save_checkpoint(args.out_dir / f"inflect-micro-fastspeech-{step}.pt", model, optim, cfg, step, args, speakers)
            save_checkpoint(args.out_dir / "inflect-micro-fastspeech-latest.pt", model, optim, cfg, step, args, speakers)

    print(f"Done. {args.out_dir}")


def main() -> None:
    ap = argparse.ArgumentParser(description="Train Inflect Micro duration-conditioned acoustic model.")
    ap.add_argument("--durations-jsonl", type=Path, required=True)
    ap.add_argument("--out-dir", type=Path, required=True)
    ap.add_argument("--max-rows", type=int, default=0)
    ap.add_argument("--steps", type=int, default=20000)
    ap.add_argument("--batch-size", type=int, default=6)
    ap.add_argument("--lr", type=float, default=2.0e-4)
    ap.add_argument("--weight-decay", type=float, default=1.0e-4)
    ap.add_argument("--hidden", type=int, default=168)
    ap.add_argument("--encoder-layers", type=int, default=5)
    ap.add_argument("--decoder-layers", type=int, default=6)
    ap.add_argument("--decoder-ff-mult", type=int, default=3)
    ap.add_argument("--max-seconds", type=float, default=12.0)
    ap.add_argument("--max-frames", type=int, default=1400)
    ap.add_argument("--mse-weight", type=float, default=0.25)
    ap.add_argument("--delta-weight", type=float, default=0.18)
    ap.add_argument("--accel-weight", type=float, default=0.0)
    ap.add_argument("--duration-weight", type=float, default=0.08)
    ap.add_argument("--group-duration-weight", type=float, default=0.0)
    ap.add_argument("--energy-weight", type=float, default=0.04)
    ap.add_argument("--bright-weight", type=float, default=0.04)
    ap.add_argument("--pitch-weight", type=float, default=0.04)
    ap.add_argument("--predicted-prosody-mel-weight", type=float, default=0.0)
    ap.add_argument("--predicted-prosody-delta-weight", type=float, default=0.0)
    ap.add_argument("--robust-prosody-mix", type=float, default=0.0)
    ap.add_argument("--robust-prosody-mel-weight", type=float, default=0.0)
    ap.add_argument("--robust-prosody-delta-weight", type=float, default=0.0)
    ap.add_argument("--grad-clip", type=float, default=5.0)
    ap.add_argument("--postnet-scale", type=float, default=0.10)
    ap.add_argument("--abs-frame-bins", type=int, default=512)
    ap.add_argument("--init-checkpoint", type=Path)
    ap.add_argument("--vocoder-checkpoint", type=Path)
    ap.add_argument("--vocoder-wav-weight", type=float, default=0.0)
    ap.add_argument("--vocoder-mel-weight", type=float, default=0.0)
    ap.add_argument("--save-interval", type=int, default=2000)
    ap.add_argument("--log-interval", type=int, default=50)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--resume", action="store_true")
    ap.add_argument("--preload-features", action="store_true", help="Cache decoded audio, mels, pitch, and token features in RAM before training.")
    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    ap.add_argument(
        "--trainable",
        choices=["all", "duration", "predictors", "heads", "contextual", "group_duration", "decoder_adapt"],
        default="all",
    )
    ap.add_argument("--contextual-predictors", action="store_true")
    ap.add_argument("--group-duration-planner", action="store_true")
    args = ap.parse_args()
    train(args)


if __name__ == "__main__":
    main()
