import torch
import torch.nn as nn
from transformers import WavLMModel, DistilBertModel, DistilBertTokenizer


class MultimodalFusionModel(nn.Module):
    """Frozen WavLM (audio) + DistilBERT (text) encoders with a
    BiLSTM + attention fusion network and multi-label classification heads.

    The encoders are only used to compute turn-level embeddings
    (`compute_embeddings` / `get_embeddings`). The fusion network
    operates on precomputed `(batch, num_turns, emb_dim)` tensors.
    """

    def __init__(self,
                 audio_emb_dim=768,
                 text_emb_dim=768,
                 fusion_hidden=256,
                 lstm_layers=2,
                 num_heads=4,
                 num_classes=4,
                 freeze_encoders=True):
        super().__init__()

        self.freeze_encoders = freeze_encoders
        self.audio_emb_dim = audio_emb_dim
        self.text_emb_dim = text_emb_dim
        self.fusion_hidden = fusion_hidden
        self.num_classes = num_classes

        self.audio_encoder = WavLMModel.from_pretrained("microsoft/wavlm-base")
        self.text_encoder = DistilBertModel.from_pretrained("distilbert-base-uncased")
        self.tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")

        self.audio_encoder.eval()
        self.text_encoder.eval()
        if freeze_encoders:
            for encoder in (self.audio_encoder, self.text_encoder):
                for p in encoder.parameters():
                    p.requires_grad = False

        self.audio_proj = nn.Linear(audio_emb_dim, fusion_hidden)
        self.text_proj = nn.Linear(text_emb_dim, fusion_hidden)

        self.bilstm = nn.LSTM(fusion_hidden * 2,
                              fusion_hidden,
                              num_layers=lstm_layers,
                              bidirectional=True,
                              batch_first=True,
                              dropout=0.2 if lstm_layers > 1 else 0)

        self.attention = nn.MultiheadAttention(embed_dim=fusion_hidden * 2,
                                               num_heads=num_heads,
                                               batch_first=True)

        self.classifiers = nn.ModuleList([
            nn.Linear(fusion_hidden * 2, 1) for _ in range(num_classes)
        ])

    def forward(self, audio_embeds=None, text_embeds=None, lengths=None):
        """Fuse precomputed turn embeddings.

        audio_embeds: (batch, num_turns, audio_emb_dim)
        text_embeds:  (batch, num_turns, text_emb_dim)
        lengths:      (batch,) ints, number of real turns (optional, used to
                      mask padding when mean-pooling).
        Returns logits: (batch, num_classes)
        """
        if audio_embeds is None or text_embeds is None:
            raise ValueError("forward() expects precomputed embeddings "
                             "(audio_embeds, text_embeds).")

        seq_len = audio_embeds.shape[1]
        if lengths is None:
            lengths = torch.full((audio_embeds.shape[0],),
                                 seq_len,
                                 dtype=torch.long,
                                 device=audio_embeds.device)

        audio_proj = self.audio_proj(audio_embeds)
        text_proj = self.text_proj(text_embeds)
        combined = torch.cat([audio_proj, text_proj], dim=-1)

        lstm_out, _ = self.bilstm(combined)

        attn_out, _ = self.attention(lstm_out, lstm_out, lstm_out)

        mask = torch.arange(seq_len, device=attn_out.device).unsqueeze(0) < lengths.unsqueeze(1)
        mask = mask.unsqueeze(-1).float()
        pooled = (attn_out * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1.0)

        logits = torch.cat([cls(pooled) for cls in self.classifiers], dim=1)
        return logits

    def get_embeddings(self, audio_waveforms, input_ids, attention_mask=None):
        """Compute per-turn audio and text embeddings from raw inputs.

        audio_waveforms: (batch, num_samples) float waveform (or 1/2 channel).
        input_ids:       (batch, seq_len) token ids.
        attention_mask:  (batch, seq_len) optional.
        Returns (audio_emb, text_emb):
            audio_emb: (batch, audio_emb_dim) mean-pooled over time
            text_emb:  (batch, text_emb_dim) [CLS] embedding
        """
        with torch.no_grad():
            audio_dev = next(self.audio_encoder.parameters()).device
            text_dev = next(self.text_encoder.parameters()).device

            audio_out = self.audio_encoder(audio_waveforms.to(audio_dev),
                                           output_hidden_states=True)
            audio_feat = audio_out.last_hidden_state
            if audio_feat.dim() == 4:
                audio_feat = audio_feat.squeeze(1)
            audio_emb = audio_feat.mean(dim=1)

            out = self.text_encoder(input_ids=input_ids.to(text_dev),
                                    attention_mask=attention_mask.to(text_dev))
            text_emb = out.last_hidden_state[:, 0, :]
        return audio_emb, text_emb

    def compute_embeddings(self, waveforms, texts, max_text_len=128, device=None):
        """Batch helper: raw waveforms + raw strings -> per-turn embeddings.

        waveforms: list/tensor of (num_samples,) waveforms.
        texts:     list of strings.
        Returns (audio_embeds, text_embeds): (batch, emb_dim) tensors.
        """
        self.eval()
        enc = self.tokenizer(texts, truncation=True, padding=True,
                             max_length=max_text_len, return_tensors="pt")
        input_ids = enc["input_ids"]
        attention_mask = enc["attention_mask"]
        device = device or next(self.parameters()).device

        if isinstance(waveforms, (list, tuple)):
            max_len = max(w.shape[0] for w in waveforms) if waveforms else 0
            if waveforms and waveforms[0].dim() == 1:
                wav = torch.zeros(len(waveforms), max_len)
            else:
                wav = torch.cat([w.unsqueeze(0) for w in waveforms], dim=0)
            for i, w in enumerate(waveforms):
                wav[i, :w.shape[-1]] = w
        else:
            wav = waveforms

        return self.get_embeddings(wav.to(device),
                                   input_ids.to(device),
                                   attention_mask.to(device))