# test_depth_checkpoint.py: Frame-by-frame decoding with the depth module (per-level sampling, cached state,
# plan schedule) and the SilenceGuard.

import argparse
import json
import os
import sys
from pathlib import Path

import torch

_REPO_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(_REPO_ROOT))
sys.path.insert(0, str(_REPO_ROOT / "scripts"))

NUM_LEVELS = int(os.environ.get("MVC_NUM_CODEC_LEVELS", "8"))
CODEBOOK_SIZE = 2048

DEFAULT_PROMPTS = [
    "The quick brown fox jumps over the lazy dog.",
    "It was the best of times, it was the worst of times.",
    "Good morning. I hope you slept well last night.",
]


@torch.no_grad()
def trim_trailing_silence(wav, sr, thresh_db=-40.0, keep_ms=120):
    import numpy as np
    w = np.asarray(wav, dtype=np.float32)
    mono = w.mean(axis=1) if w.ndim > 1 else w
    peak = float(np.max(np.abs(mono))) + 1e-9
    thr = peak * (10.0 ** (thresh_db / 20.0))
    above = np.abs(mono) > thr
    if not above.any():
        return wav
    last = int(np.nonzero(above)[0][-1])
    keep = min(len(mono), last + int(sr * keep_ms / 1000))
    return w[:keep]


@torch.no_grad()
def think_generate(model, prefix, registry, tokenizer, max_think_tokens=200,
                   temp=0.7, speaker_id=None, device="cuda"):
    backbone = model.backbone.backbone
    emb = backbone.embedding
    inject = None
    speaker_emb = None
    if speaker_id is not None and getattr(model, "use_speaker_conditioning", False):
        sid = torch.tensor([speaker_id], dtype=torch.long, device=device)
        speaker_emb = model.speaker_encoder(sid)
        if getattr(model, "use_speaker_input", False):
            inject = model.speaker_input_proj(speaker_emb)

    ids = prefix.clone().to(device)
    think_ids = []
    ban_lo, ban_hi = registry.speech_token_id_min, registry.speech_token_id_max
    banned_ctl = [registry.audio_start_id, registry.audio_end_id,
                  registry.eos_id, registry.think_start_id]
    for _ in range(max_think_tokens):
        if inject is not None:
            emb.set_speaker(inject)
        emb.set_codes(torch.full((1, ids.shape[1], NUM_LEVELS), -1,
                                 dtype=torch.long, device=device))
        h = backbone(ids)
        if speaker_emb is not None:
            h = model.adaln(h, speaker_emb)
        logits = model.backbone.lm_head(h[:, -1, :])[0].float()
        logits[ban_lo:ban_hi + 1] = float("-inf")
        for t in banned_ctl:
            logits[t] = float("-inf")
        if temp > 0.0:
            nxt = int(torch.multinomial(
                torch.softmax(logits / temp, dim=-1), 1).item())
        else:
            nxt = int(logits.argmax().item())
        if nxt == registry.think_end_id:
            break
        think_ids.append(nxt)
        ids = torch.cat([ids, torch.tensor([[nxt]], device=device)], dim=1)
    tail = torch.tensor([[registry.think_end_id, registry.audio_start_id]],
                        device=device)
    ids = torch.cat([ids, tail], dim=1)
    reasoning = tokenizer.decode(think_ids) if think_ids else ""
    return ids, reasoning


def _half_up(v) -> int:
    return int(float(v) + 0.5)


def build_decode_plan_schedule(durations, prefix_len, max_frames, *,
                               mode="stretch", target_frames=None, device=None):
    from delay_dataset import NO_WORD, build_plan_schedule as _frame_map
    from model import build_plan_schedule as _position_schedule

    if prefix_len < 1:
        raise ValueError(f"prefix_len must be >= 1 (the <AUDIO> token sits at "
                         f"prefix_len - 1), got {prefix_len}")
    if max_frames < 1:
        raise ValueError(f"max_frames must be >= 1, got {max_frames}")
    durs = [_half_up(d) for d in durations]
    if not durs:
        raise ValueError("plan has no words — nothing to schedule")
    if any(d < 0 for d in durs):
        raise ValueError(f"negative word duration in plan: {durs}")
    if sum(durs) <= 0:
        raise ValueError("plan durations sum to 0 frames — the keystone token "
                         "(§3.1) is degenerate; refuse rather than schedule "
                         "every frame onto word 0")

    if mode == "stretch":
        if target_frames is None:
            raise ValueError(
                "mode='stretch' needs target_frames (the expected TOTAL frame "
                "count). ΣD counts speech frames only, so the caller must "
                "supply ΣD̂ / ρ with ρ the corpus speech/total frame ratio — or "
                "use mode='pack', which needs no estimate.")
        target_frames = int(target_frames)
        if target_frames < 1:
            raise ValueError(f"target_frames must be >= 1, got {target_frames}")
        frame_to_word = _frame_map(durs, target_frames, "stretch")
        word_frames = [0] * len(durs)
        for w in frame_to_word:
            if w != NO_WORD:
                word_frames[w] += 1
    elif mode == "pack":
        word_frames = durs
    else:
        raise ValueError(f"plan schedule mode {mode!r} not in ('stretch', 'pack')")

    return _position_schedule([word_frames], audio_start=prefix_len - 1,
                              n_frames=max_frames,
                              seq_len=prefix_len + max_frames, device=device)


def _fresh_inference_cache(max_seqlen, max_batch_size=1):
    from mamba_ssm.utils.generation import InferenceParams
    return InferenceParams(max_seqlen=int(max_seqlen), max_batch_size=int(max_batch_size))


def _check_plan_decode(model, plan_bins, plan_schedule, prefix_len, max_frames):
    if not getattr(model, "use_plan_injection", False):
        raise RuntimeError(
            "plan_bins/plan_schedule were passed but this model has no plan "
            "injection channel — call model.enable_plan_injection() BEFORE "
            "load_checkpoint so the checkpoint's plan_injector.* weights land "
            "in it. Without the call the tensors would be dropped and the run "
            "would silently measure the prefix-only condition.")
    if model.training:
        raise RuntimeError(
            "plan-injected decode requires model.eval(): PlanInjector's "
            "scaffold dropout keys off self.training and would drop the "
            "injection on a random ~scaffold_dropout share of utterances.")
    if plan_bins.dim() != 3 or plan_bins.shape[0] != 1 or plan_bins.shape[-1] != 3:
        raise ValueError(f"plan_bins must be [1, W, 3] (pitch, duration, "
                         f"energy), got {tuple(plan_bins.shape)}")
    if plan_schedule.dim() != 2 or plan_schedule.shape[0] != 1:
        raise ValueError(f"plan_schedule must be [1, S], got "
                         f"{tuple(plan_schedule.shape)}")
    need = prefix_len + max_frames - 1
    if plan_schedule.shape[1] < need:
        raise ValueError(
            f"plan_schedule covers {plan_schedule.shape[1]} positions but the "
            f"decode can reach position {need - 1} (prefix {prefix_len} + "
            f"{max_frames} frames) — build it with build_decode_plan_schedule("
            f"..., prefix_len={prefix_len}, max_frames={max_frames}).")
    if bool((plan_schedule[0, :prefix_len - 1] >= 0).any()):
        raise ValueError(
            "plan_schedule injects at a position before the <AUDIO> token "
            f"(prefix_len - 1 = {prefix_len - 1}). The program conditions the "
            "AUDIO block only; a schedule reaching the prompt/<THINK> region "
            "was built against a different prefix length.")
    if int(plan_schedule[0, prefix_len - 1]) < 0:
        raise ValueError(
            "plan_schedule has no word on the <AUDIO> token at position "
            f"{prefix_len - 1}, which is where frame 0's word belongs "
            "(delay_dataset plan_lookahead=1). The schedule is off by at least "
            "one position relative to training.")


def flipk_rate(cb0, k_lo=2, k_hi=8):
    n = len(cb0)
    hits = tot = 0
    for k in range(k_lo, k_hi + 1):
        if n <= k:
            continue
        for t in range(k, n):
            if cb0[t] == cb0[t - k] and cb0[t] != cb0[t - 1]:
                hits += 1
        tot += n - k
    return 100.0 * hits / tot if tot else 0.0


class LoopBreaker:
    def __init__(self, threshold=5.0, win=32, temp=0.6, rep=2.0,
                 stop_after=0):
        self.threshold = float(threshold)
        self.win = int(win)
        self.temp = float(temp)
        self.rep = float(rep)
        self.stop_after = int(stop_after)
        self.hist = []
        self.active = False
        self.metric = 0.0
        self.max_metric = 0.0
        self.events = 0
        self.frames_intervened = 0
        self.run = 0
        self.stopped = False

    def update(self, cb0_code):
        self.hist.append(int(cb0_code))
        if len(self.hist) > self.win:
            del self.hist[0]
        if len(self.hist) < self.win:
            return
        self.metric = flipk_rate(self.hist)
        self.max_metric = max(self.max_metric, self.metric)
        was = self.active
        self.active = self.metric >= self.threshold
        if self.active:
            if not was:
                self.events += 1
            self.run += 1
            self.frames_intervened += 1
        else:
            self.run = 0

    @property
    def force_stop(self):
        return bool(self.stop_after) and self.run >= self.stop_after

    def rewind(self, n):
        if n > 0:
            del self.hist[-n:]
        self.run = 0
        self.active = False

    def stats(self):
        return {"events": self.events,
                "frames": self.frames_intervened,
                "max_flipk": round(self.max_metric, 2),
                "stopped": self.stopped}


class SilenceGuard:
    def __init__(self, codes, budget=10, stop_after=10, stop_without_plan=False,
                 retries=0, retry_temp=0.6):
        self.codes = frozenset(int(c) for c in codes)
        if not self.codes:
            raise ValueError("SilenceGuard needs a non-empty pause-code set")
        self.budget = int(budget)
        self.stop_after = int(stop_after)
        self.stop_without_plan = bool(stop_without_plan)
        self.retries = int(retries)
        self.retry_temp = float(retry_temp)
        self.run = 0
        self.max_run = 0
        self.mask_events = 0
        self.frames_masked = 0
        self.stopped = False
        self.retries_used = 0
        self.rewinds = 0
        self.frames_rewound = 0
        self.mask_left = 0
        self.spoke_at_last = False
        self._masking = False
        self._idx = {}

    def update(self, cb0_code, at_last=None):
        if int(cb0_code) in self.codes:
            self.run += 1
            self.max_run = max(self.max_run, self.run)
        else:
            self.run = 0
            if at_last is True:
                self.spoke_at_last = True

    def plan_rewind(self, at_last):
        if (self.budget <= 0 or self.retries_used >= self.retries
                or at_last is True or self.run < self.budget):
            return 0
        n = self.run
        self.retries_used += 1
        self.rewinds += 1
        self.frames_rewound += n
        self.run = 0
        self.mask_left = self.budget
        return n

    @property
    def kicking(self):
        return self.mask_left > 0

    def mask_now(self, at_last):
        if self.mask_left > 0:
            self.mask_left -= 1
            if not self._masking:
                self.mask_events += 1
            self.frames_masked += 1
            self._masking = True
            return True
        want = (self.budget > 0 and self.run >= self.budget
                and at_last is not True)
        if want:
            if not self._masking:
                self.mask_events += 1
            self.frames_masked += 1
        self._masking = want
        return want

    def force_stop(self, at_last):
        if not self.stop_after or self.run < self.stop_after:
            return False
        if at_last is True:
            exhausted = self.retries > 0 and self.retries_used >= self.retries
            return self.spoke_at_last or exhausted
        return at_last is None and self.stop_without_plan

    def mask_index(self, device):
        key = str(device)
        if key not in self._idx:
            self._idx[key] = torch.tensor(sorted(self.codes), dtype=torch.long,
                                          device=device)
        return self._idx[key]

    def stats(self):
        return {"mask_events": self.mask_events,
                "frames_masked": self.frames_masked,
                "max_run": self.max_run,
                "stopped": self.stopped,
                "rewinds": self.rewinds,
                "frames_rewound": self.frames_rewound,
                "retries_used": self.retries_used,
                "spoke_at_last": self.spoke_at_last}


def depth_generate(model, prefix, registry, max_frames=200, level_temps=None,
                   top_k=0, top_p=0.0, allow_stop=True, min_frames=20,
                   stop_threshold=0.0, sticky_bias=0.0,
                   speaker_id=None, feedback="all", cross_frame_decode=True,
                   device="cuda", antifan_bias=0.0, antifan_levels=None,
                   onset_frames=0, onset_temp_scale=1.0,
                   rep_penalty=1.0, rep_win=16, ras_tau=0.0, ras_win=10,
                   stop_prob_out=None, plan_bins=None, plan_schedule=None,
                   loop_breaker=None, silence_guard=None, cached=False):
    placeholder = registry.audio_start_id
    audio_end = registry.audio_end_id
    eos = registry.eos_id
    backbone = model.backbone.backbone
    emb = backbone.embedding
    speaker_emb = None
    inject = None
    if speaker_id is not None and getattr(model, "use_speaker_conditioning", False):
        sid = torch.tensor([speaker_id], dtype=torch.long, device=device)
        speaker_emb = model.speaker_encoder(sid)
        if getattr(model, "use_speaker_input", False):
            inject = model.speaker_input_proj(speaker_emb)

    if (plan_bins is None) != (plan_schedule is None):
        raise ValueError(
            "plan_bins and plan_schedule must be supplied together (got "
            f"plan_bins={'set' if plan_bins is not None else 'None'}, "
            f"plan_schedule={'set' if plan_schedule is not None else 'None'}) "
            "— half a plan is a wiring bug whose failure mode is a silent null.")
    if plan_bins is not None:
        _check_plan_decode(model, plan_bins, plan_schedule, prefix.shape[1],
                           max_frames)

    ids = prefix.clone().to(device)
    codes = torch.full((1, ids.shape[1], NUM_LEVELS), -1,
                       dtype=torch.long, device=device)
    frames = []
    iters = max_frames
    if silence_guard is not None and silence_guard.retries > 0:
        iters += silence_guard.retries * (silence_guard.budget + 1)
    ip = None
    n_fed = 0
    if cached:
        ip = _fresh_inference_cache(ids.shape[1] + iters + 8)
    for _ in range(iters):
        if len(frames) >= max_frames:
            break
        if inject is not None:
            emb.set_speaker(inject)
        lo = n_fed if cached else 0
        emb.set_codes(codes[:, lo:])
        at_last = None
        if plan_bins is not None:
            plan_vec = model.plan_injector(plan_bins, plan_schedule[:, :ids.shape[1]])
            emb.set_plan(plan_vec[:, lo:])
            cur_word = int(plan_schedule[0, ids.shape[1] - 1].item())
            at_last = (cur_word >= plan_bins.shape[1] - 1) if cur_word >= 0 else False
        if silence_guard is not None:
            nb = silence_guard.plan_rewind(at_last)
            if nb:
                nb = min(nb, len(frames))
                if nb:
                    frames = frames[:-nb]
                    ids = ids[:, :-nb]
                    codes = codes[:, :-nb]
                    if loop_breaker is not None:
                        loop_breaker.rewind(nb)
                    if cached:
                        ip = _fresh_inference_cache(ids.shape[1] + iters + 8)
                        n_fed = 0
                continue
        if cached:
            ip.seqlen_offset = n_fed
            h = backbone(ids[:, lo:], inference_params=ip)
            n_fed = ids.shape[1]
        else:
            h = backbone(ids)
        if speaker_emb is not None:
            h = model.adaln(h, speaker_emb)
        last = h[:, -1, :]
        may_stop = allow_stop and len(frames) >= min_frames
        if may_stop or stop_prob_out is not None:
            logits = model.backbone.lm_head(last)[0]
            probs = torch.softmax(logits.float(), dim=-1)
            p_stop = float(probs[audio_end].item() + probs[eos].item())
            need_argmax = (stop_prob_out is not None) or stop_threshold <= 0.0
            argmax_stop = (int(logits.argmax().item()) in (audio_end, eos)
                           if need_argmax else False)
            if stop_prob_out is not None:
                stop_prob_out.append((p_stop, argmax_stop))
            if may_stop:
                if stop_threshold > 0.0:
                    if p_stop > stop_threshold:
                        break
                elif argmax_stop:
                    break
        lt = level_temps
        if onset_frames and lt is not None and len(frames) < onset_frames:
            lt = [lt[0]] + [t * onset_temp_scale for t in lt[1:]]
        rp_eff = rep_penalty
        if loop_breaker is not None and loop_breaker.active:
            if lt is not None:
                lt = [max(loop_breaker.temp, lt[0])] + list(lt[1:])
            else:
                lt = [loop_breaker.temp] + [0.0] * (NUM_LEVELS - 1)
            if rp_eff is None or isinstance(rp_eff, (int, float)):
                base = float(rp_eff) if rp_eff else 1.0
                rp_eff = ([max(loop_breaker.rep, base)]
                          + [base] * (NUM_LEVELS - 1))
            else:
                rp_eff = ([max(loop_breaker.rep, float(rp_eff[0]))]
                          + list(rp_eff[1:]))
        if silence_guard is not None and silence_guard.kicking:
            if lt is not None:
                lt = [max(silence_guard.retry_temp, lt[0])] + list(lt[1:])
            else:
                lt = [silence_guard.retry_temp] + [0.0] * (NUM_LEVELS - 1)
        rw = None
        _rp_on = (any(float(x) != 1.0 for x in rep_penalty)
                  if rep_penalty is not None
                  and not isinstance(rep_penalty, (int, float))
                  else bool(rep_penalty) and rep_penalty != 1.0)
        _lb_on = loop_breaker is not None and loop_breaker.active
        if (_rp_on or ras_tau or _lb_on) and frames:
            recent = frames[-rep_win:]
            rw = [[f[k] for f in recent] for k in range(NUM_LEVELS)]
        frame = model.depth_module.generate_frame(
            last, level_temps=lt, top_k=top_k, top_p=top_p,
            return_argmax=(feedback == "argmax"),
            rep_window=rw, rep_penalty=rp_eff,
            ras_tau=ras_tau, ras_win=ras_win,
            sticky_prev=(frames[-1] if ((sticky_bias or antifan_bias) and frames)
                         else None),
            sticky_bias=sticky_bias,
            prev_codes=(frames[-1] if (frames and cross_frame_decode)
                        else None),
            antifan_prev2=(frames[-2] if (antifan_bias and len(frames) >= 2)
                           else None),
            antifan_bias=antifan_bias, antifan_levels=antifan_levels,
            cb0_mask=(silence_guard.mask_index(device)
                      if (silence_guard is not None
                          and silence_guard.mask_now(at_last)) else None))
        frame_amax = None
        if feedback == "argmax":
            frame, frame_amax = frame
        frames.append(frame)
        if loop_breaker is not None:
            loop_breaker.update(frame[0])
            if loop_breaker.force_stop and len(frames) >= min_frames:
                loop_breaker.stopped = True
                break
        if silence_guard is not None:
            silence_guard.update(frame[0], at_last=at_last)
            if silence_guard.force_stop(at_last) and len(frames) >= min_frames:
                silence_guard.stopped = True
                break
        fb = list(frame)
        if feedback == "semantic":
            fb = [frame[0]] + [-1] * (NUM_LEVELS - 1)
        elif feedback == "argmax":
            fb = list(frame_amax)
        col = torch.tensor(fb, dtype=torch.long, device=device).view(1, 1, NUM_LEVELS)
        ids = torch.cat([ids, torch.tensor([[placeholder]], device=device)], dim=1)
        codes = torch.cat([codes, col], dim=1)
    if not frames:
        return None
    return [[f[k] for f in frames] for k in range(NUM_LEVELS)]


@torch.no_grad()
def flow_generate(model, prefix, registry, max_frames=200, allow_stop=True,
                  min_frames=20, stop_threshold=0.0, speaker_id=None,
                  feedback="all", flow_steps=16, flow_noise=1.0, cb0_temp=0.0,
                  device="cuda"):
    placeholder = registry.audio_start_id
    audio_end = registry.audio_end_id
    eos = registry.eos_id
    backbone = model.backbone.backbone
    emb = backbone.embedding
    speaker_emb = None
    inject = None
    if speaker_id is not None and getattr(model, "use_speaker_conditioning", False):
        sid = torch.tensor([speaker_id], dtype=torch.long, device=device)
        speaker_emb = model.speaker_encoder(sid)
        if getattr(model, "use_speaker_input", False):
            inject = model.speaker_input_proj(speaker_emb)

    ids = prefix.clone().to(device)
    codes = torch.full((1, ids.shape[1], NUM_LEVELS), -1,
                       dtype=torch.long, device=device)
    latents = []
    for _ in range(max_frames):
        if inject is not None:
            emb.set_speaker(inject)
        emb.set_codes(codes)
        h = backbone(ids)
        if speaker_emb is not None:
            h = model.adaln(h, speaker_emb)
        last = h[:, -1, :]
        if allow_stop and len(latents) >= min_frames:
            logits = model.backbone.lm_head(last)[0]
            if stop_threshold > 0.0:
                probs = torch.softmax(logits.float(), dim=-1)
                if float(probs[audio_end].item() + probs[eos].item()) > stop_threshold:
                    break
            elif int(logits.argmax().item()) in (audio_end, eos):
                break
        frame = model.depth_module.generate_frame(
            last, level_temps=[cb0_temp] + [0.0] * (NUM_LEVELS - 1))
        cb0 = torch.tensor([frame[0]], device=device)
        lat = model.flow_head.sample(last, cb0, steps=flow_steps,
                                     noise_scale=flow_noise)
        latents.append(lat[0])
        fb = list(frame)
        if feedback == "semantic":
            fb = [frame[0]] + [-1] * (NUM_LEVELS - 1)
        col = torch.tensor(fb, dtype=torch.long, device=device).view(1, 1, NUM_LEVELS)
        ids = torch.cat([ids, torch.tensor([[placeholder]], device=device)], dim=1)
        codes = torch.cat([codes, col], dim=1)
    if not latents:
        return None
    return torch.stack(latents, dim=0).transpose(0, 1).unsqueeze(0)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--checkpoint", required=True)
    ap.add_argument("--model_name", default="state-spaces/mamba2-1.3b")
    ap.add_argument("--lora_rank", type=int, default=0)
    ap.add_argument("--num_speakers", type=int, default=20000)
    ap.add_argument("--speaker_dim", type=int, default=256)
    ap.add_argument("--speaker_input_injection", action="store_true",
                    help="MUST match training (all depth recipes pass it)")
    ap.add_argument("--speaker_id", type=int, default=None)
    ap.add_argument("--depth_dim", type=int, default=1024)
    ap.add_argument("--depth_layers", type=int, default=2)
    ap.add_argument("--no_cross_frame_decode", action="store_true",
                    help="S32 residual-loop diag: disable the cf3 cross-frame "
                         "channel at DECODE time (every frame uses the learned "
                         "no-predecessor cross_start vector) — fully opens the "
                         "acoustic loop on a cross_frame-trained checkpoint")
    ap.add_argument("--depth_feedback", choices=["all", "semantic"], default="all",
                    help="MUST match the checkpoint's training mode")
    ap.add_argument("--depth_cross_frame", type=int, default=0,
                    help="MUST match the checkpoint's --depth_cross_frame "
                         "(S30 cross-frame acoustic channel; 3 = fan fix)")
    ap.add_argument("--temp_semantic", type=float, default=0.0)
    ap.add_argument("--temp_acoustic", type=float, default=None)
    ap.add_argument("--acoustic_temps", type=float, nargs=7, default=None)
    ap.add_argument("--top_k", type=int, default=0)
    ap.add_argument("--top_p", type=float, default=0.0)
    ap.add_argument("--best_of", type=int, default=1)
    ap.add_argument("--bestof_semantic_temps", type=float, nargs="+",
                    default=[0.1, 0.3, 0.5, 0.7])
    ap.add_argument("--max_frames", type=int, default=200)
    ap.add_argument("--min_frames", type=int, default=20)
    ap.add_argument("--min_frames_per_word", type=float, default=0.0,
                    help="per-prompt stop floor = max(min_frames, "
                         "words * this). ~4.5 (0.36 s/word) prevents the "
                         "first-sentence EOS truncation on multi-sentence "
                         "prompts (S2 arms learned utterance-final stop)")
    ap.add_argument("--allow_stop", action="store_true")
    ap.add_argument("--stop_threshold", type=float, default=0.0,
                    help=">0: honor AUDIO_END/EOS only when P(stop) exceeds this "
                         "(softmax over lm_head); 0 = bare-argmax stop (legacy)")
    ap.add_argument("--sticky_bias", type=float, default=0.0,
                    help="S28 fan mitigation: logit bonus for the previous "
                         "frame's code at each acoustic level (cb1-7) — "
                         "decode-side temporal texture persistence. Try 1-3")
    ap.add_argument("--trim_silence", action="store_true",
                    help="trim trailing near-silence from decoded audio")
    ap.add_argument("--whisper", default="openai/whisper-base.en")
    ap.add_argument("--prompts", nargs="*", default=None)
    ap.add_argument("--prompts_from", default=None,
                    help="read prompts from a JSONL (6-20 word rows) instead of "
                         "the 3 default sentences — for a larger naturalness n")
    ap.add_argument("--num_prompts", type=int, default=15)
    ap.add_argument("--outdir", default="out_depth_eval")
    ap.add_argument("--device", default="cuda")
    ap.add_argument("--flow", action="store_true",
                    help="PB-14 option A: render via the continuous flow head "
                         "(RVQ-recon latent -> decode_latents) instead of the "
                         "discrete codebooks. Requires a --flow_head checkpoint.")
    ap.add_argument("--flow_hidden", type=int, default=1024)
    ap.add_argument("--flow_layers", type=int, default=4)
    ap.add_argument("--flow_steps", type=int, default=16,
                    help="Euler ODE steps for flow sampling")
    ap.add_argument("--hybrid_attention_top_k", type=int, default=0,
                    help="Phase 8: load a hybrid checkpoint (top-K mixers are "
                         "attention). Must match the trained value (4).")
    ap.add_argument("--cot", action="store_true",
                    help="Stage-2 CoT decode: free-generate the <THINK> block "
                         "(speech/audio tokens banned), then constrained depth "
                         "audio. Requires a stage2_cot_* checkpoint.")
    ap.add_argument("--max_think_tokens", type=int, default=200)
    ap.add_argument("--think_temp", type=float, default=0.7)
    ap.add_argument("--flow_noise", type=float, default=1.0,
                    help="flow start-noise scale (1=full N(0,I); 0=deterministic "
                         "from origin -> endpoint nearer the RVQ lattice)")
    args = ap.parse_args()
    prompts = args.prompts or DEFAULT_PROMPTS
    if args.prompts_from:
        import json as _json
        pf = []
        with open(args.prompts_from, encoding="utf-8") as _fh:
            for _line in _fh:
                _t = _json.loads(_line).get("prompt", "").strip()
                if 6 <= len(_t.split()) <= 20:
                    pf.append(_t)
                if len(pf) >= args.num_prompts:
                    break
        prompts = pf or prompts

    import config as _cfg
    import soundfile as sf
    from model import MambaCoTModel
    from train import load_checkpoint, apply_lora
    from inference import build_inference_prefix
    from codec import load_codec
    from eval.intelligibility import IntelligibilityScorer
    assert _cfg.NUM_SPEECH_TOKENS == CODEBOOK_SIZE, "run with MVC_NUM_SPEECH_TOKENS=2048"

    outdir = Path(args.outdir)
    outdir.mkdir(parents=True, exist_ok=True)

    model = MambaCoTModel(model_name=args.model_name, device=args.device,
                          dtype=torch.bfloat16, mtp_num_heads=0)
    model.enable_speaker_conditioning(speaker_dim=args.speaker_dim,
                                      num_speakers=args.num_speakers,
                                      input_injection=args.speaker_input_injection)
    if args.lora_rank > 0:
        apply_lora(model, rank=args.lora_rank, alpha=args.lora_rank,
                   target_modules=["in_proj", "out_proj"], dropout=0.0)
    model.enable_depth_module(num_levels=NUM_LEVELS, codebook_size=CODEBOOK_SIZE,
                              d_depth=args.depth_dim,
                              depth_layers=args.depth_layers,
                              depth_feedback=args.depth_feedback,
                              cross_frame=args.depth_cross_frame)
    if args.hybrid_attention_top_k > 0:
        model.enable_hybrid_attention(top_k=args.hybrid_attention_top_k)
    if args.flow:
        model.enable_flow_head(latent_dim=512, cb0_size=CODEBOOK_SIZE,
                               d_hidden=args.flow_hidden, n_layers=args.flow_layers)
    load_checkpoint(args.checkpoint, model, device=args.device)
    model.eval()
    registry = model.token_registry
    codec = load_codec("mimi", device=args.device)
    scorer = IntelligibilityScorer(model_id=args.whisper, device=args.device)

    results = []
    for i, prompt in enumerate(prompts):
        print(f"\n[{i}] prompt: {prompt!r}", flush=True)
        if args.cot:
            p2 = build_inference_prefix(model.tokenizer, registry, prompt, 2)
            prefix, reasoning = think_generate(
                model, p2, registry, model.tokenizer,
                max_think_tokens=args.max_think_tokens, temp=args.think_temp,
                speaker_id=args.speaker_id, device=args.device)
            print(f"    think ({prefix.shape[1] - p2.shape[1] - 2} tok): "
                  f"{reasoning!r}", flush=True)
        else:
            prefix = build_inference_prefix(model.tokenizer, registry, prompt, 1)
        min_fr = max(args.min_frames,
                     int(len(prompt.split()) * args.min_frames_per_word))
        best = None
        for c in range(max(1, args.best_of)):
            sem_t = (args.temp_semantic if args.best_of <= 1
                     else args.bestof_semantic_temps[c % len(args.bestof_semantic_temps)])
            if args.acoustic_temps is not None:
                lts = [sem_t] + list(args.acoustic_temps)
            else:
                ac = args.temp_acoustic if args.temp_acoustic is not None else 0.0
                lts = [sem_t] + [ac] * (NUM_LEVELS - 1)
            if args.flow:
                latents = flow_generate(
                    model, prefix, registry, max_frames=args.max_frames,
                    allow_stop=args.allow_stop, min_frames=min_fr,
                    stop_threshold=args.stop_threshold, speaker_id=args.speaker_id,
                    feedback=args.depth_feedback, flow_steps=args.flow_steps,
                    flow_noise=args.flow_noise,
                    cb0_temp=sem_t, device=args.device)
                if latents is None:
                    continue
                wav, sr = codec.decode_latents(latents)
                nframes = latents.shape[-1]
            else:
                codes = depth_generate(
                    model, prefix, registry, max_frames=args.max_frames,
                    level_temps=lts, top_k=args.top_k, top_p=args.top_p,
                    allow_stop=args.allow_stop, min_frames=min_fr,
                    stop_threshold=args.stop_threshold,
                    sticky_bias=args.sticky_bias,
                    speaker_id=args.speaker_id, feedback=args.depth_feedback,
                    cross_frame_decode=not args.no_cross_frame_decode,
                    device=args.device)
                if codes is None:
                    continue
                wav, sr = codec.decode_multi(codes)
                nframes = len(codes[0])
            if args.trim_silence:
                wav = trim_trailing_silence(wav, sr)
            tmp = outdir / f"_cand_{i:02d}_{c}.wav"
            sf.write(str(tmp), wav, sr)
            score = scorer.score(str(tmp), prompt)
            tmp.unlink(missing_ok=True)
            wer = score["wer"]
            cand = {"prompt": prompt, "frames": nframes, "sem_temp": sem_t,
                    "asr": score["hypothesis"],
                    "wer": round(wer, 4) if wer == wer else None,
                    "_wav": wav, "_sr": sr}
            print(f"    cand {c} semT={sem_t} frames={cand['frames']} "
                  f"wer={cand['wer']} asr={cand['asr']!r}", flush=True)
            if best is None or (cand["wer"] is not None and
                                (best["wer"] is None or cand["wer"] < best["wer"])):
                best = cand
        if best is None:
            print("    (no frames)"); continue
        if args.cot:
            best["reasoning"] = reasoning
        sf.write(str(outdir / f"sample_{i:02d}.wav"), best.pop("_wav"), best.pop("_sr"))
        results.append(best)
        print(f"    BEST wer={best['wer']}")

    (outdir / "intelligibility_depth.json").write_text(json.dumps(results, indent=2))
    wers = [r["wer"] for r in results if r["wer"] is not None]
    if wers:
        print(f"\ndepth-module mean WER: {sum(wers)/len(wers):.3f}")
    print("TEST_DEPTH_DONE")

if __name__ == "__main__":
    main()