# Student MeanFlow pour le DiT CosyVoice3 : v(x,t) -> u(x,t,r) (vitesse moyenne sur [t,r])
# Init: poids du teacher + gate r zero-init => student(x,t,r) == teacher v(x,t) a l'init.
import torch
import torch.nn as nn

from cosyvoice.flow.DiT.dit import DiT, TimestepEmbedding
from cosyvoice.utils.mask import add_optional_chunk_mask


class MeanFlowDiT(DiT):
    def __init__(self, **kw):
        super().__init__(**kw)
        self.r_embed = TimestepEmbedding(self.dim)
        self.r_gate = nn.Linear(self.dim, self.dim)
        nn.init.zeros_(self.r_gate.weight)
        nn.init.zeros_(self.r_gate.bias)

    @classmethod
    def from_teacher(cls, teacher: DiT, **kw):
        m = cls(**kw)
        missing, unexpected = m.load_state_dict(teacher.state_dict(), strict=False)
        assert not unexpected, f"poids teacher inattendus: {unexpected}"
        assert all(k.startswith(("r_embed.", "r_gate.")) for k in missing), f"manquants hors student: {missing}"
        return m

    def forward(self, x, mask, mu, t, r=None, spks=None, cond=None, streaming=False):
        # copie de DiT.forward avec injection de l'embedding d'intervalle r
        x = x.transpose(1, 2)
        mu = mu.transpose(1, 2)
        cond = cond.transpose(1, 2)
        spks = spks.unsqueeze(dim=1)
        batch, seq_len = x.shape[0], x.shape[1]
        if t.ndim == 0:
            t = t.repeat(batch)
        if r is None:
            r = t
        if r.ndim == 0:
            r = r.repeat(batch)

        t = self.time_embed(t) + self.r_gate(self.r_embed(r))
        x = self.input_embed(x, cond, mu, spks.squeeze(1))

        rope = self.rotary_embed.forward_from_seq_len(seq_len)

        if self.long_skip_connection is not None:
            residual = x

        if streaming is True:
            attn_mask = add_optional_chunk_mask(x, mask.bool(), False, False, 0, self.static_chunk_size, -1).unsqueeze(dim=1)
        else:
            attn_mask = add_optional_chunk_mask(x, mask.bool(), False, False, 0, 0, -1).repeat(1, x.size(1), 1).unsqueeze(dim=1)

        for block in self.transformer_blocks:
            x = block(x, t, mask=attn_mask.bool(), rope=rope)

        if self.long_skip_connection is not None:
            x = self.long_skip_connection(torch.cat((x, residual), dim=-1))

        x = self.norm_out(x, t)
        output = self.proj_out(x).transpose(1, 2)
        return output


# kwargs exacts du DiT CosyVoice3-0.5B-2512 (cosyvoice3.yaml)
CV3_DIT_KWARGS = dict(
    dim=1024, depth=22, heads=16, dim_head=64, ff_mult=2,
    mel_dim=80, mu_dim=80, spk_dim=80, out_channels=80,
    static_chunk_size=25 * 2, num_decoding_left_chunks=-1,
)


@torch.inference_mode()
def student_infer(student, z, mu, mask, spks, cond, t_span):
    """Inference few-step : x_{i+1} = x_i + (r-t) * u(x_i, t, r) sur la grille t_span (deja cosine-warpee)."""
    x = z
    for i in range(len(t_span) - 1):
        t, r = t_span[i], t_span[i + 1]
        bt = t.expand(x.size(0)).to(x.dtype)
        br = r.expand(x.size(0)).to(x.dtype)
        u = student(x, mask, mu, bt, r=br, spks=spks, cond=cond)
        x = x + (r - t) * u
    return x