"""Independent multirate Mini prototype; never modifies historical model factories.

The deployment graph includes text/duration, predicted frame prosody, a
three-resolution conditional flow field, and the separately pinned vocoder.
Reference durations/prosody are training or explicitly labelled diagnostic inputs.
"""
from __future__ import annotations

import copy
import math

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

from scripts.v3mini_padding_safe import boolean_sequence_mask, masked_flow_loss

SCHEMA = "inflect-mini-multirate-v1"
VOCODER_PARAMETERS = 14_015_041
CONDITION_CHANNELS = 196


def clean(x, mask):
    return x.masked_fill(~mask.bool(), 0)


def expand_tokens(hidden, durations, max_length=None):
    """Expand integer durations without silently truncating or inventing frames."""
    if hidden.ndim != 3 or durations.shape != (hidden.shape[0], hidden.shape[2]):
        raise ValueError("Expected BxCxN hidden and BxN durations")
    if not torch.isfinite(durations).all() or (durations < 0).any() or (durations != durations.round()).any():
        raise ValueError("Durations must be nonnegative finite integers")
    lengths = durations.long().sum(-1)
    if (lengths < 1).any():
        raise ValueError("Every utterance needs at least one acoustic frame")
    needed = int(lengths.max())
    size = needed if max_length is None else int(max_length)
    if size < needed or size > 20000:
        raise ValueError("Expansion length is truncated or exceeds safety limit")
    frames = torch.arange(size, device=hidden.device)[None, :]
    ends = durations.long().cumsum(-1).contiguous()
    indices = torch.searchsorted(ends, frames.expand(hidden.shape[0], -1).contiguous(), right=True)
    indices = indices.clamp_max(hidden.shape[-1] - 1)
    expanded = hidden.gather(2, indices[:, None].expand(-1, hidden.shape[1], -1))
    starts = F.pad(ends[:, :-1], (1, 0)).gather(1, indices)
    widths = durations.gather(1, indices).clamp_min(1)
    # Frame-centred phase preserves both one-frame tokens and long vowels.
    position = (2 * (frames - starts + 0.5) / widths - 1)[:, None].to(hidden.dtype)
    mask = (frames < lengths[:, None])[:, None]
    return clean(expanded, mask), clean(position, mask), mask


class FrameNorm(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.norm = nn.LayerNorm(channels)

    def forward(self, x):
        return self.norm(x.transpose(1, 2)).transpose(1, 2)


class ProsodyHead(nn.Module):
    """Predict normalized log-F0, voicing logits, normalized log-RMS from text."""
    def __init__(self, channels=192, width=256):
        super().__init__()
        self.input = nn.Conv1d(channels, width, 1)
        self.layers = nn.ModuleList([nn.ModuleList([
            FrameNorm(width), nn.Conv1d(width, width * 2, 1),
            nn.Conv1d(width * 2, width * 2, 7, padding=3, groups=width * 2),
            nn.Conv1d(width * 2, width, 1),
        ]) for _ in range(2)])
        self.output = nn.Conv1d(width, 3, 1)

    def forward(self, hidden, mask):
        mask = boolean_sequence_mask(mask, hidden, prefix=True)
        x = clean(self.input(clean(hidden, mask)), mask)
        for norm, expand, conv, project in self.layers:
            h = clean(F.silu(expand(norm(x))), mask)
            h = clean(F.silu(conv(h)), mask)
            x = clean(x + project(h), mask)
        return clean(self.output(x), mask)


def prosody_loss(prediction, target, mask):
    """Never regress pitch on unvoiced frames; NaN pitch there is ignored."""
    valid = boolean_sequence_mask(mask, prediction, prefix=True)[:, 0]
    if prediction.shape != target.shape or prediction.shape[1] != 3:
        raise ValueError("Expected matching Bx3xT prosody")
    voice = target[:, 1].masked_select(valid)
    if not torch.all((voice == 0) | (voice == 1)):
        raise ValueError("Voicing targets must be binary")
    voiced = valid & (target[:, 1] == 1)
    pitch = F.smooth_l1_loss(prediction[:, 0][voiced], target[:, 0][voiced], reduction="sum") / voiced.sum().clamp_min(1)
    energy = F.smooth_l1_loss(prediction[:, 2][valid], target[:, 2][valid])
    uv = F.binary_cross_entropy_with_logits(prediction[:, 1][valid], voice)
    return {"pitch": pitch, "energy": energy, "voicing": uv}


def pool_pairs(x, mask):
    """Masked pair average: adding right padding cannot change the last pair."""
    pad = x.shape[-1] % 2
    values = F.pad(clean(x, mask), (0, pad)).reshape(x.shape[0], x.shape[1], -1, 2)
    valid = F.pad(mask, (0, pad)).reshape(mask.shape[0], 1, -1, 2)
    count = valid.sum(-1)
    return values.sum(-1) / count.clamp_min(1), count > 0


def rope(x):
    # B,H,T,D. Absolute frame indices are independent of batch/padded length.
    half = x.shape[-1] // 2
    frequency = torch.exp(-math.log(10000) * torch.arange(half, device=x.device, dtype=torch.float32) / half)
    angle = torch.arange(x.shape[-2], device=x.device, dtype=torch.float32)[:, None] * frequency[None]
    c, s = angle.cos().to(x.dtype), angle.sin().to(x.dtype)
    a, b = x.chunk(2, dim=-1)
    return torch.cat((a * c - b * s, a * s + b * c), dim=-1)


class FieldBlock(nn.Module):
    def __init__(self, width, attention=False):
        super().__init__()
        self.norm = FrameNorm(width)
        self.time = nn.Linear(256, 2 * width)
        self.condition = nn.Conv1d(CONDITION_CHANNELS, width, 1)
        self.expand = nn.Conv1d(width, 2 * width, 1)
        self.local = nn.Conv1d(2 * width, 2 * width, 7, padding=3, groups=2 * width)
        self.project = nn.Conv1d(2 * width, width, 1)
        self.attention = attention
        if attention:
            self.attn_norm = FrameNorm(width)
            self.qkv = nn.Linear(width, 3 * width)
            self.attn_output = nn.Linear(width, width)

    def forward(self, x, mask, condition, time):
        scale, shift = self.time(time).chunk(2, dim=1)
        h = self.norm(x) * (1 + scale[:, :, None]) + shift[:, :, None]
        h = clean(h + self.condition(clean(condition, mask)), mask)
        h = clean(F.silu(self.expand(h)), mask)
        h = clean(F.silu(self.local(h)), mask)
        x = clean(x + self.project(h), mask)
        if self.attention:
            b, c, length = x.shape
            qkv = self.qkv(self.attn_norm(x).transpose(1, 2)).reshape(b, length, 3, 4, c // 4)
            q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0)
            h = F.scaled_dot_product_attention(rope(q), rope(k), v, attn_mask=mask[:, :, None, :])
            h = self.attn_output(h.transpose(1, 2).reshape(b, length, c)).transpose(1, 2)
            x = clean(x + h, mask)
        return x


class MultirateEstimator(nn.Module):
    def __init__(self):
        super().__init__()
        self.time = nn.Sequential(nn.Linear(128, 256), nn.SiLU(), nn.Linear(256, 256))
        self.speaker = nn.Linear(64, 256)
        self.input = nn.Conv1d(200, 192, 1)
        self.down1 = nn.Conv1d(192, 256, 1)
        self.down2 = nn.Conv1d(256, 320, 1)
        self.up2 = nn.Conv1d(320 + 256, 256, 1)
        self.up1 = nn.Conv1d(256 + 192, 192, 1)
        self.stacks = nn.ModuleList([
            nn.ModuleList([FieldBlock(w, a) for _ in range(n)])
            for w, a, n in [(192, False, 2), (256, True, 2), (320, True, 4), (256, True, 2), (192, False, 2)]
        ])
        self.output_norm = FrameNorm(192)
        self.output = nn.Conv1d(192, 100, 1)

    def forward(self, x, mask, mu, t, spks=None, cond=None):
        mask = boolean_sequence_mask(mask, x, prefix=True)
        if mu.shape != x.shape or x.shape[1] != 100:
            raise ValueError("Expected matching 100-bin flow inputs")
        if cond is None or cond.shape != (x.shape[0], CONDITION_CHANNELS, x.shape[-1]):
            raise ValueError("Missing rich text/prosody conditioning")
        if spks is None or spks.shape != (x.shape[0], 64):
            raise ValueError("Expected Bx64 speaker embedding")
        t = t.reshape(-1)
        if t.numel() != x.shape[0]:
            raise ValueError("Flow time must have one value per example")
        frequencies = torch.exp(-math.log(10000) * torch.arange(64, device=x.device, dtype=torch.float32) / 63)
        angle = t.float()[:, None] * 1000 * frequencies
        time = self.time(torch.cat((angle.sin(), angle.cos()), -1).to(x.dtype)) + self.speaker(spks)
        masks, conditions, skips = [mask], [clean(cond, mask)], []
        for _ in range(2):
            c, m = pool_pairs(conditions[-1], masks[-1])
            conditions.append(c)
            masks.append(m)
        h = clean(self.input(torch.cat((clean(x, mask), clean(mu, mask)), 1)), mask)
        for stage, level in enumerate([0, 1, 2, 1, 0]):
            for block in self.stacks[stage]:
                h = block(h, masks[level], conditions[level], time)
            if stage < 2:
                skips.append(h)
                h, _ = pool_pairs(h, masks[level])
                h = clean((self.down1 if stage == 0 else self.down2)(h), masks[level + 1])
            elif stage < 4:
                skip = skips.pop()
                h = h.repeat_interleave(2, dim=-1)[..., :skip.shape[-1]]
                h = clean((self.up2 if stage == 2 else self.up1)(torch.cat((h, skip), 1)), masks[level - 1])
        return clean(self.output(self.output_norm(h)), mask)


class MiniAcoustic(nn.Module):
    """Own deployment/training wrapper, sharing only compatible donor modules."""
    def __init__(self, encoder, speaker, estimator, statistics, arm):
        super().__init__()
        self.encoder, self.spk_emb, self.estimator = encoder, speaker, estimator
        self.prosody = ProsodyHead() if arm == "multirate" else None
        self.arm = arm
        self.register_buffer("mel_mean", torch.tensor(float(statistics["mel_mean"])))
        self.register_buffer("mel_std", torch.tensor(float(statistics["mel_std"])))
        if not torch.isfinite(self.mel_mean) or not torch.isfinite(self.mel_std) or self.mel_std <= 0:
            raise ValueError("Mel statistics must be finite with positive standard deviation")

    def encode(self, ids, lengths, speakers):
        enc = self.encoder
        speaker = self.spk_emb(speakers)
        mask = (torch.arange(ids.shape[-1], device=ids.device)[None] < lengths[:, None])[:, None]
        x = enc.emb(ids).transpose(1, 2) * math.sqrt(enc.n_channels)
        x = enc.prenet(x, mask.to(x.dtype))
        x = torch.cat((x, speaker[:, :, None].expand(-1, -1, ids.shape[-1])), dim=1)
        hidden = enc.encoder(x, mask.to(x.dtype))
        mu = clean(enc.proj_m(hidden), mask)
        log_duration = enc.proj_w(hidden.detach(), mask.to(x.dtype))
        return hidden, mu, log_duration, mask, speaker

    def condition(self, ids, lengths, speakers, durations=None, max_length=None, length_scale=1.0):
        hidden, mu, log_duration, text_mask, speaker = self.encode(ids, lengths, speakers)
        if durations is None:
            if not 0.5 <= length_scale <= 2:
                raise ValueError("Invalid length scale")
            durations = torch.ceil(log_duration[:, 0].clamp(-7, 7).exp() * length_scale).long()
            durations = durations * text_mask[:, 0]
        elif durations.shape != ids.shape or (durations.masked_select(~text_mask[:, 0]) != 0).any():
            raise ValueError("Duration targets contain invalid tokens")
        expanded_mu, _, mask = expand_tokens(mu, durations, max_length)
        prediction, cond = None, None
        if self.prosody is not None:
            rich, position, _ = expand_tokens(hidden, durations, max_length)
            prediction = self.prosody(rich, mask)
            predicted = torch.cat((prediction[:, :1], prediction[:, 1:2].sigmoid(), prediction[:, 2:]), 1)
            cond = clean(torch.cat((rich, position, predicted), 1), mask)
        return {"mu": expanded_mu, "mask": mask, "speaker": speaker, "cond": cond,
                "prosody": prediction, "log_duration": log_duration,
                "text_mask": text_mask, "durations": durations}

    def flow_loss(self, target, conditioning, time, noise):
        t = time.reshape(-1, 1, 1)
        sigma = 1e-4
        x = (1 - (1 - sigma) * t) * noise + t * target
        u = target - (1 - sigma) * noise
        c = conditioning
        predicted = self.estimator(x, c["mask"], c["mu"], time, c["speaker"], c["cond"])
        return masked_flow_loss(predicted, u, c["mask"])

    @torch.inference_mode()
    def sample(self, ids, lengths, speakers, *, steps=16, temperature=0.667, seed=1234, durations=None):
        if self.training:
            raise ValueError("Sampling requires eval mode")
        if not 1 <= steps <= 128 or not 0 < temperature <= 2:
            raise ValueError("Invalid sampling settings")
        c = self.condition(ids, lengths, speakers, durations)
        rng = torch.Generator(device=ids.device).manual_seed(seed)
        x = torch.randn(c["mu"].shape, device=ids.device, generator=rng) * temperature
        x = clean(x, c["mask"])
        for step in range(steps):
            t = x.new_full((x.shape[0],), step / steps)
            x = clean(x + self.estimator(x, c["mask"], c["mu"], t, c["speaker"], c["cond"]) / steps, c["mask"])
        return x * self.mel_std + self.mel_mean, c


def create_from_donor(donor, config, arm):
    if arm not in ("multirate", "native240"):
        raise ValueError("Unknown architecture arm")
    for key in ("mel_mean", "mel_std"):
        actual = getattr(donor, key).detach().cpu()
        expected = torch.as_tensor(config["data_statistics"][key], dtype=actual.dtype)
        if actual.numel() != 1 or not torch.equal(actual.reshape(()), expected.reshape(())):
            raise ValueError("Donor/config mel statistics differ: " + key)
    if (donor.encoder.n_feats != 100 or donor.encoder.n_channels != 128
            or donor.encoder.n_vocab != 178 or tuple(donor.spk_emb.weight.shape) != (6, 64)):
        raise ValueError("Donor frontend/speaker dimensions differ from this architecture")
    if arm == "multirate":
        estimator = MultirateEstimator()
    else:
        from scripts.v3mini_padding_safe import native_decoder_types
        _, decoder_type = native_decoder_types()
        settings = copy.deepcopy(config["decoder"])
        settings["channels"] = [240, 240]
        estimator = decoder_type(in_channels=264, out_channels=100, **settings)
    model = MiniAcoustic(copy.deepcopy(donor.encoder), copy.deepcopy(donor.spk_emb), estimator,
                         config["data_statistics"], arm).float()
    model.requires_grad_(True)
    counts = parameter_counts(model)
    if not 24_000_000 <= counts["complete"] <= 30_000_000:
        raise ValueError(f"Deployment budget failed: {counts}")
    return model


def parameter_counts(model):
    values = {name: sum(p.numel() for p in module.parameters()) if module is not None else 0
              for name, module in [("text", model.encoder), ("speaker", model.spk_emb),
                                   ("prosody", model.prosody), ("estimator", model.estimator)]}
    values["acoustic"] = sum(p.numel() for p in model.parameters())
    values["vocoder"] = VOCODER_PARAMETERS
    values["complete"] = values["acoustic"] + VOCODER_PARAMETERS
    return values