"""Dynamic Accent - Arm A, written from scratch in PyTorch.

Backbone: Parallel Tacotron 2 (Elias et al. 2021, arXiv 2103.14574): non-autoregressive, token durations,
  differentiable learned upsampling (MLP over start/end offsets -> token->frame weights W), iterative mel loss.
Multilingual recipe: Zhang et al. 2019 (arXiv 1907.04448): separate speaker + language embeddings,
  domain-adversarial speaker classifier (gradient reversal) on text-encoder outputs, residual VAE.

Documented deviation (alignment): PT2 trains durations end-to-end with Soft-DTW. Soft-DTW is O(T*N)-sequential and
too slow for a 2 h smoke budget, so we use the internal unsupervised alignment-learning objective of
Badlani et al. 2022 ("One TTS Alignment To Rule Them All", arXiv 2108.10447): a small text/mel attention trained with
a CTC forward-sum loss + beta-binomial prior, hardened with monotonic alignment search (MAS). The hard durations
supervise the duration predictor and drive the SAME differentiable PT2 upsampler at train time; predicted durations
drive it at inference. No external aligner is used.
"""
import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from .text import VOCAB, LANGS

N_MELS = 100


# ---------------------------------------------------------------- utils
def seq_mask(lens, maxlen=None):
    maxlen = maxlen or int(lens.max())
    return torch.arange(maxlen, device=lens.device)[None] < lens[:, None]   # True = valid


def sinusoid(T, d, device):
    pos = torch.arange(T, device=device).float()[:, None]
    i = torch.arange(0, d, 2, device=device).float()[None]
    ang = pos / (10000 ** (i / d))
    pe = torch.zeros(T, d, device=device)
    pe[:, 0::2] = torch.sin(ang)
    pe[:, 1::2] = torch.cos(ang)
    return pe


class GradReverse(torch.autograd.Function):
    """Gradient reversal. As in Zhang et al. 2019 the reversed gradient sent into the text encoder is clipped
    (here: its L2 norm is capped at `clip`), which keeps the min-max game from blowing up the encoder."""

    @staticmethod
    def forward(ctx, x, lam, clip):
        ctx.lam, ctx.clip = lam, clip
        return x.view_as(x)

    @staticmethod
    def backward(ctx, g):
        g = -ctx.lam * g
        if ctx.clip is not None:
            n = g.norm()
            if n > ctx.clip:
                g = g * (ctx.clip / (n + 1e-12))
        return g, None, None


def grad_reverse(x, lam=1.0, clip=0.5):
    return GradReverse.apply(x, lam, clip)


class ConvNorm(nn.Module):
    def __init__(self, cin, cout, k=5, drop=0.1):
        super().__init__()
        self.conv = nn.Conv1d(cin, cout, k, padding=k // 2)
        self.ln = nn.LayerNorm(cout)
        self.drop = nn.Dropout(drop)

    def forward(self, x, mask):           # x [B,T,C], mask [B,T]
        x = x.masked_fill(~mask[..., None], 0)
        y = self.conv(x.transpose(1, 2)).transpose(1, 2)
        return self.drop(F.relu(self.ln(y)))


class FFTBlock(nn.Module):
    """Pre-LN self-attention + convolutional feed-forward (FastSpeech/PT2-style lightweight block)."""

    def __init__(self, d, heads=4, ff=4, k=3, drop=0.1):
        super().__init__()
        self.ln1 = nn.LayerNorm(d)
        self.att = nn.MultiheadAttention(d, heads, dropout=drop, batch_first=True)
        self.ln2 = nn.LayerNorm(d)
        self.c1 = nn.Conv1d(d, d * ff, k, padding=k // 2)
        self.c2 = nn.Conv1d(d * ff, d, 1)
        self.drop = nn.Dropout(drop)

    def forward(self, x, mask):
        h = self.ln1(x)
        a, _ = self.att(h, h, h, key_padding_mask=~mask, need_weights=False)
        x = x + self.drop(a)
        h = self.ln2(x).masked_fill(~mask[..., None], 0).transpose(1, 2)
        h = self.c2(F.gelu(self.c1(h))).transpose(1, 2)
        x = x + self.drop(h)
        return x.masked_fill(~mask[..., None], 0)


# ---------------------------------------------------------------- modules
class TextEncoder(nn.Module):
    def __init__(self, d, n_lang, layers=6):
        super().__init__()
        self.emb = nn.Embedding(VOCAB, d, padding_idx=0)
        self.lang_in = nn.Embedding(n_lang, d)          # per-token language id at the input
        self.convs = nn.ModuleList([ConvNorm(d, d, 5) for _ in range(3)])
        self.blocks = nn.ModuleList([FFTBlock(d) for _ in range(layers)])

    def forward(self, tok, lang, mask):
        x = self.emb(tok) + self.lang_in(lang)
        for c in self.convs:
            x = x + c(x, mask)
        x = x + sinusoid(x.shape[1], x.shape[2], x.device)[None]
        for b in self.blocks:
            x = b(x, mask)
        return x


class ResidualVAE(nn.Module):
    """Global residual (prosody) encoder: mel -> q(z|x); z broadcast to every token (Zhang 2019 / Hsu 2019)."""

    def __init__(self, zdim=16, h=256):
        super().__init__()
        self.convs = nn.Sequential(
            nn.Conv1d(N_MELS, h, 3, padding=1), nn.ReLU(),
            nn.Conv1d(h, h, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv1d(h, h, 3, stride=2, padding=1), nn.ReLU())
        self.gru = nn.GRU(h, h // 2, batch_first=True, bidirectional=True)
        self.mu = nn.Linear(h, zdim)
        self.logvar = nn.Linear(h, zdim)

    def forward(self, mel, mel_lens):
        h = self.convs(mel.transpose(1, 2)).transpose(1, 2)
        lens = torch.clamp((mel_lens + 3) // 4, min=1)
        h, _ = self.gru(h)
        m = seq_mask(lens, h.shape[1]).float()[..., None]
        h = (h * m).sum(1) / m.sum(1)
        mu, lv = self.mu(h), self.logvar(h).clamp(-8, 8)
        z = mu + torch.randn_like(mu) * torch.exp(0.5 * lv) if self.training else mu
        kl = 0.5 * (mu.pow(2) + lv.exp() - 1 - lv).sum(-1).mean()
        return z, kl


class Aligner(nn.Module):
    """Unsupervised text-mel alignment (Badlani et al. 2022)."""

    def __init__(self, d, h=256, temp=0.0005):
        super().__init__()
        self.temp = temp
        self.key = nn.Sequential(nn.Conv1d(d, h * 2, 3, padding=1), nn.ReLU(), nn.Conv1d(h * 2, 80, 1))
        self.query = nn.Sequential(nn.Conv1d(N_MELS, h * 2, 3, padding=1), nn.ReLU(),
                                   nn.Conv1d(h * 2, h, 1), nn.ReLU(), nn.Conv1d(h, 80, 1))

    def forward(self, txt, mel, tmask, mmask, prior_log):
        k = self.key(txt.transpose(1, 2))                         # B,80,N
        q = self.query(mel.transpose(1, 2))                       # B,80,T
        q, k = q.float(), k.float()
        dist = (q.pow(2).sum(1)[:, :, None] + k.pow(2).sum(1)[:, None, :]
                - 2 * torch.bmm(q.transpose(1, 2), k)).clamp_min(0)        # B,T,N  (memory-light ||q-k||^2)
        logits = -self.temp * dist
        logits = logits.masked_fill(~tmask[:, None, :], -1e4)
        logp = F.log_softmax(logits, -1) + prior_log                # prior in log domain
        logp = logp.masked_fill(~tmask[:, None, :], -1e4)
        soft = F.softmax(logp, -1) * mmask[..., None]
        return logits, logp, soft


def forward_sum_loss(logits, tlens, mlens, blank_logprob=-1.0):
    """CTC forward-sum: every token must be visited monotonically. logits [B,T,N] (pre-softmax, masked)."""
    B, T, N = logits.shape
    lp = F.pad(logits, (1, 0), value=blank_logprob)
    lp = F.log_softmax(lp, -1).transpose(0, 1)                      # T,B,N+1
    tgt = torch.arange(1, N + 1, device=logits.device)[None].expand(B, N)
    return F.ctc_loss(lp.float(), tgt, mlens, tlens, blank=0, reduction="mean", zero_infinity=True)


def _mas_np(logp, N, T):
    """Monotonic alignment search on one [T,N] log-prob matrix -> hard [T,N]."""
    Q = np.full((T, N), -np.inf, dtype=np.float32)
    Q[0, 0] = logp[0, 0]
    for t in range(1, T):
        for n in range(min(N, t + 1)):
            stay = Q[t - 1, n]
            move = Q[t - 1, n - 1] if n > 0 else -np.inf
            Q[t, n] = max(stay, move) + logp[t, n]
    path = np.zeros((T, N), dtype=np.float32)
    n = N - 1
    for t in range(T - 1, -1, -1):
        path[t, n] = 1
        if t > 0 and n > 0 and (n == t or Q[t - 1, n - 1] > Q[t - 1, n]):
            n -= 1
    return path


try:
    import numba

    _mas_np = numba.njit(cache=True)(_mas_np)  # noqa: F811
except Exception:  # numba optional (CPU unit tests)
    pass


@torch.no_grad()
def mas(logp, tlens, mlens):
    out = torch.zeros_like(logp)
    lp = logp.float().cpu().numpy()
    for b in range(lp.shape[0]):
        T, N = int(mlens[b]), int(tlens[b])
        out[b, :T, :N] = torch.from_numpy(_mas_np(np.ascontiguousarray(lp[b, :T, :N]), N, T))
    return out


_prior_cache = {}


def beta_binomial_prior(N, T, scale=1.0):
    from scipy.stats import betabinom
    key = (N, T)
    if key not in _prior_cache:
        x = np.arange(N)[None]
        t = np.arange(1, T + 1)[:, None]
        pmf = betabinom.pmf(x, N - 1, scale * t, scale * (T + 1 - t))
        _prior_cache[key] = torch.from_numpy(np.log(pmf + 1e-8)).float()
        if len(_prior_cache) > 20000:
            _prior_cache.clear()
    return _prior_cache[key]


def batch_prior(tlens, mlens, N, T, device):
    p = torch.full((len(tlens), T, N), -1e4)
    for b in range(len(tlens)):
        n, t = int(tlens[b]), int(mlens[b])
        p[b, :t, :n] = beta_binomial_prior(n, t)
    return p.to(device)


class DurationPredictor(nn.Module):
    def __init__(self, d, h=256):
        super().__init__()
        self.convs = nn.ModuleList([ConvNorm(d, h, 3), ConvNorm(h, h, 3), ConvNorm(h, h, 3)])
        self.out = nn.Linear(h, 1)

    def forward(self, x, mask):
        for c in self.convs:
            x = c(x, mask)
        return self.out(x).squeeze(-1).masked_fill(~mask, 0)       # log(1 + d)


class LearnedUpsampler(nn.Module):
    """PT2 differentiable upsampling: W[t,k] = softmax_k MLP(S_tk, E_tk, h_k);  O = W H + proj(sum_k W C)."""

    def __init__(self, d, hw=16, hc=8):
        super().__init__()
        self.se_w = nn.Linear(2, hw)
        self.h_w = nn.Linear(d, hw)
        self.w_out = nn.Linear(hw, 1)
        self.se_c = nn.Linear(2, hc)
        self.h_c = nn.Linear(d, hc)
        self.c_out = nn.Linear(hc, d)

    def forward(self, H, dur, tmask, T):
        dur = dur.masked_fill(~tmask, 0)
        e = torch.cumsum(dur, 1)
        s = e - dur
        t = torch.arange(T, device=H.device).float()[None, :, None] + 0.5
        S = (t - s[:, None, :]) / 10.0
        E = (e[:, None, :] - t) / 10.0
        se = torch.stack([S, E], -1)                                   # B,T,N,2
        lw = self.w_out(F.silu(self.se_w(se) + self.h_w(H)[:, None])).squeeze(-1)
        lw = lw.masked_fill(~tmask[:, None, :], -1e4)
        W = torch.softmax(lw, -1)                                      # B,T,N
        C = F.silu(self.se_c(se) + self.h_c(H)[:, None])               # B,T,N,hc
        out = torch.bmm(W, H) + self.c_out(torch.einsum("btn,btnc->btc", W, C))
        return out, W


class MelDecoder(nn.Module):
    def __init__(self, d, layers=6, n_iter=3):
        super().__init__()
        self.blocks = nn.ModuleList([FFTBlock(d) for _ in range(layers)])
        self.n_iter = n_iter
        self.outs = nn.ModuleList([nn.Linear(d, N_MELS) for _ in range(n_iter)])

    def forward(self, x, mask):
        x = x + sinusoid(x.shape[1], x.shape[2], x.device)[None]
        mels = []
        L = len(self.blocks)
        for i, b in enumerate(self.blocks):
            x = b(x, mask)
            j = i - (L - self.n_iter)
            if j >= 0:
                mels.append(self.outs[j](x))
        return mels


# ---------------------------------------------------------------- full model
class DynamicAccentTTS(nn.Module):
    def __init__(self, n_spk, n_lang=len(LANGS), d=384, enc_layers=6, dec_layers=6, spk_dim=128, lang_dim=64,
                 zdim=16, adv_hidden=256, mel_mean=None, mel_std=None, align_temp=0.0005):
        super().__init__()
        self.n_spk = n_spk
        self.encoder = TextEncoder(d, n_lang, enc_layers)
        self.spk_emb = nn.Embedding(n_spk, spk_dim)
        self.lang_out = nn.Embedding(n_lang, lang_dim)                 # per-token, concatenated to encoder outputs
        self.vae = ResidualVAE(zdim)
        self.cond = nn.Linear(d + spk_dim + lang_dim + zdim, d)
        self.adv = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, adv_hidden), nn.ReLU(), nn.Linear(adv_hidden, n_spk))
        self.aligner = Aligner(d, temp=align_temp)
        self.dur = DurationPredictor(d)
        self.up = LearnedUpsampler(d)
        self.dec_spk = nn.Linear(spk_dim, d)
        self.decoder = MelDecoder(d, dec_layers)
        self.register_buffer("mel_mean", torch.tensor(mel_mean if mel_mean is not None else -2.5))
        self.register_buffer("mel_std", torch.tensor(mel_std if mel_std is not None else 2.2))

    def _cond(self, Ht, spk, lang, z):
        B, N, _ = Ht.shape
        s = self.spk_emb(spk)[:, None].expand(B, N, -1)
        zz = z[:, None].expand(B, N, -1)
        return self.cond(torch.cat([Ht, s, self.lang_out(lang), zz], -1)), s

    def forward(self, tok, lang, tlens, mel, mlens, spk, step=0, adv_lambda=1.0, bin_start=4000, prior=None):
        tmask, mmask = seq_mask(tlens, tok.shape[1]), seq_mask(mlens, mel.shape[1])
        meln = (mel - self.mel_mean) / self.mel_std
        Ht = self.encoder(tok, lang, tmask)
        # adversarial speaker classifier on text encodings (Zhang 2019): GRL removes speaker info from text path
        adv_logits = self.adv(grad_reverse(Ht, adv_lambda))
        adv_ce = F.cross_entropy(adv_logits[tmask], spk[:, None].expand_as(tok)[tmask])
        adv_acc = (adv_logits[tmask].argmax(-1) == spk[:, None].expand_as(tok)[tmask]).float().mean()
        z, kl = self.vae(meln, mlens)
        H, _ = self._cond(Ht, spk, lang, z)
        # alignment (internal, unsupervised)
        if prior is None:
            prior = batch_prior(tlens, mlens, tok.shape[1], mel.shape[1], tok.device)
        logits, logp, soft = self.aligner(self.encoder.emb(tok) + Ht.detach(), meln, tmask, mmask, prior)
        l_fs = forward_sum_loss(logits.masked_fill(~mmask[..., None], -1e4) + prior, tlens, mlens)
        hard = mas(logp, tlens, mlens)
        l_bin = -(hard * torch.log(soft.clamp_min(1e-8))).sum() / hard.sum() if step >= bin_start else soft.new_zeros(())
        dur_t = hard.sum(1)                                            # B,N frames
        # durations
        logd = self.dur(H.detach() if step < 1000 else H, tmask)
        l_dur = F.mse_loss(logd[tmask], torch.log1p(dur_t[tmask]))
        # upsample with target durations (train) -> decoder
        up, W = self.up(H, dur_t, tmask, mel.shape[1])
        x = up + self.dec_spk(self.spk_emb(spk))[:, None]
        mels = self.decoder(x, mmask)
        m = mmask[..., None].float()
        l_mel = sum((F.l1_loss(o, meln, reduction="none") * m).sum() / (m.sum() * N_MELS) for o in mels) / len(mels)
        return dict(mel=l_mel, dur=l_dur, fs=l_fs, bin=l_bin, kl=kl, adv=adv_ce, adv_acc=adv_acc.detach()), mels[-1]

    @torch.no_grad()
    def infer(self, tok, lang, spk, pace=1.0, z=None):
        tlens = torch.full((tok.shape[0],), tok.shape[1], device=tok.device)
        tmask = seq_mask(tlens, tok.shape[1])
        Ht = self.encoder(tok, lang, tmask)
        if z is None:
            z = torch.zeros(tok.shape[0], self.vae.mu.out_features, device=tok.device)
        H, _ = self._cond(Ht, spk, lang, z)
        dur = (torch.expm1(self.dur(H, tmask)).clamp(min=0) * pace).masked_fill(~tmask, 0)
        T = max(int(dur.sum(1).max().round().item()), 1)
        up, W = self.up(H, dur, tmask, T)
        x = up + self.dec_spk(self.spk_emb(spk))[:, None]
        mel = self.decoder(x, torch.ones(tok.shape[0], T, dtype=torch.bool, device=tok.device))[-1]
        return mel * self.mel_std + self.mel_mean, dur


def count_params(m):
    return sum(p.numel() for p in m.parameters())