"""
FIGMA inference model (inference-only).

Loads the packaged FIGMA checkpoint (frozen MuQ audio encoder + frozen E5 text
encoder + two Transformer projection heads) and produces the GLOBAL audio/text
embeddings used for retrieval. Retrieval score = cosine similarity between the
global audio and text embeddings.

Note: retrieval uses only the global embeddings (mean-pooled audio frames + E5
[CLS]). The frame/token projection path is used during training only.

License: this model bundles MuQ weights (CC BY-NC 4.0), so the composite model is
released for NON-COMMERCIAL research use (CC BY-NC 4.0). MuQ and E5 are attributed.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import TransformerEncoder, TransformerEncoderLayer
from transformers import AutoModel, AutoTokenizer
from muq import MuQ

TEXT_ENCODER = "intfloat/multilingual-e5-large-instruct"
AUDIO_ENCODER = "OpenMuQ/MuQ-large-msd-iter"
SAMPLE_RATE = 24000
CLIP_SECONDS = 10.0


class ProjectionHead(nn.Module):
    def __init__(self, input_dim, hidden_dim=512, output_dim=512):
        super().__init__()
        layer = TransformerEncoderLayer(
            d_model=input_dim, nhead=8, dim_feedforward=hidden_dim,
            dropout=0.1, batch_first=True,
        )
        self.transformer = TransformerEncoder(layer, num_layers=2)
        self.output_proj = nn.Linear(input_dim, output_dim)

    def forward(self, x):
        if x.dim() == 2:
            x = x.unsqueeze(1)
            x = self.transformer(x)
            x = x.squeeze(1)
        else:
            x = self.transformer(x)
        return self.output_proj(x)


class Figma(nn.Module):
    """Submodule names (muq / text_encoder / audio_proj / text_proj) match the
    training checkpoint, so its state_dict loads directly."""

    def __init__(self, audio_feat_dim=1024):
        super().__init__()
        self.muq = MuQ.from_pretrained(AUDIO_ENCODER).eval().requires_grad_(False)
        self.text_encoder = AutoModel.from_pretrained(TEXT_ENCODER)
        for p in self.text_encoder.parameters():
            p.requires_grad = False
        self.audio_proj = ProjectionHead(input_dim=audio_feat_dim)
        self.text_proj = ProjectionHead(input_dim=self.text_encoder.config.hidden_size)

    @torch.no_grad()
    def encode_audio(self, wavs):
        """wavs: [B, 1, samples] float tensor at 24 kHz -> [B, 512] L2-normalized."""
        wavs = torch.nan_to_num(wavs, nan=0.0, posinf=0.0, neginf=0.0)
        with torch.amp.autocast("cuda", enabled=False):
            a_seq = self.muq(wavs).last_hidden_state
        a_clip = a_seq.mean(dim=1)
        return F.normalize(self.audio_proj(a_clip), dim=-1, eps=1e-8)

    @torch.no_grad()
    def encode_text(self, text_inputs):
        """text_inputs: dict of tokenized tensors -> [B, 512] L2-normalized."""
        t_seq = self.text_encoder(**text_inputs).last_hidden_state
        t_cls = t_seq[:, 0, :]
        return F.normalize(self.text_proj(t_cls), dim=-1, eps=1e-8)

    @classmethod
    def from_checkpoint(cls, ckpt_path, device="cpu"):
        model = cls().to(device).eval()
        ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
        sd = ckpt.get("state_dict", ckpt)
        missing, _ = model.load_state_dict(sd, strict=False)
        bad = [k for k in missing if k.startswith(("audio_proj", "text_proj"))]
        if bad:
            raise RuntimeError(f"Projection-head weights missing from checkpoint: {bad}")
        return model


def get_tokenizer():
    return AutoTokenizer.from_pretrained(TEXT_ENCODER)