"""
Multi-Task Learning Classifier for Hateful Content Detection.

This module implements a multi-task learning classifier with shared representation
and task-specific heads for each hate category. Compatible with Hugging Face Hub
for easy loading with `trust_remote_code=True`.

Usage:
    from transformers import AutoModel
    
    model = AutoModel.from_pretrained(
        "Amirhossein75/clip-vit-base-mmhs150k-mtl",
        trust_remote_code=True
    )
"""

from typing import List, Optional, Dict, Any
import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers import (
    CLIPTextModel,
    CLIPVisionModel,
    AutoModel,
    PreTrainedModel,
    PretrainedConfig,
    AutoImageProcessor,
    AutoTokenizer,
)


class MultiTaskClassifierConfig(PretrainedConfig):
    """Configuration class for MultiTaskClassifier."""
    
    model_type = "clip-mtl"
    
    def __init__(
        self,
        encoder_name: str = "openai/clip-vit-base-patch32",
        task_names: List[str] = None,
        fusion_dim: int = 512,
        backend: str = "clip",
        threshold: float = 0.5,
        thresholds: List[float] = None,
        head_hidden_dim: Optional[int] = None,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.encoder_name = encoder_name
        self.task_names = task_names or ["racist", "sexist", "homophobe", "religion", "otherhate"]
        self.fusion_dim = fusion_dim
        self.backend = backend
        self.threshold = threshold
        self.thresholds = thresholds or [0.5] * len(self.task_names)
        self.head_hidden_dim = head_hidden_dim


class MultiTaskClassifier(PreTrainedModel):
    """
    Multi-task classifier with shared projection + fusion, one binary head per task.
    
    This architecture is useful when predicting multiple related binary tasks
    (e.g., racist, sexist, homophobe, religion, otherhate).
    
    Backends:
      - "clip": Uses separate CLIP text & vision towers
      - "auto": Uses AutoModel for SigLIP/CLIP-like models
    
    Args:
        config: MultiTaskClassifierConfig with model parameters.
    """
    
    config_class = MultiTaskClassifierConfig
    
    def __init__(self, config: MultiTaskClassifierConfig):
        super().__init__(config)
        
        self.task_names = list(config.task_names)
        self.num_tasks = len(self.task_names)
        self.threshold = config.threshold
        self.thresholds = config.thresholds
        self.backend = config.backend.lower()
        self.fusion_dim = config.fusion_dim

        # Load backbone encoder
        if self.backend == "clip":
            self.tower_txt = CLIPTextModel.from_pretrained(config.encoder_name)
            self.tower_img = CLIPVisionModel.from_pretrained(config.encoder_name)
            tdim = self.tower_txt.config.hidden_size
            idim = self.tower_img.config.hidden_size
            self.backbone = None
        else:
            # Auto dual-encoder (e.g., SigLIP/SigLIP2)
            self.backbone = AutoModel.from_pretrained(config.encoder_name)
            tdim = getattr(getattr(self.backbone, "text_config", None), "hidden_size", None)
            idim = getattr(getattr(self.backbone, "vision_config", None), "hidden_size", None)
            if tdim is None or idim is None:
                pd = getattr(self.backbone.config, "projection_dim", None)
                tdim = tdim or pd
                idim = idim or pd
            assert tdim is not None and idim is not None, \
                "Could not infer hidden sizes for AutoModel backend."

        # Shared projection + gated fusion
        self.proj_t = nn.Linear(tdim, config.fusion_dim)
        self.proj_i = nn.Linear(idim, config.fusion_dim)
        self.g_t = nn.Linear(config.fusion_dim, config.fusion_dim)
        self.g_i = nn.Linear(config.fusion_dim, config.fusion_dim)
        self.gate = nn.Linear(config.fusion_dim * 2 + 2, config.fusion_dim)

        # Shared feature extraction head
        self.shared_head = nn.Sequential(
            nn.Dropout(0.2),
            nn.Linear(config.fusion_dim, config.fusion_dim),
            nn.GELU(),
            nn.Dropout(0.2),
        )

        # Task-specific heads
        head_hidden_dim = config.head_hidden_dim
        def make_head():
            if head_hidden_dim and head_hidden_dim > 0:
                return nn.Sequential(
                    nn.Linear(config.fusion_dim, head_hidden_dim),
                    nn.GELU(),
                    nn.Dropout(0.1),
                    nn.Linear(head_hidden_dim, 1),
                )
            else:
                return nn.Linear(config.fusion_dim, 1)
        
        self.heads = nn.ModuleList([make_head() for _ in range(self.num_tasks)])

        # Register buffers for pos_weight and log_vars (will be loaded from checkpoint)
        self.register_buffer("pos_weight", torch.ones(self.num_tasks))
        self.log_vars = nn.Parameter(torch.zeros(self.num_tasks))

    def _encode_text(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        """Encode text input to feature vector."""
        if self.backend == "clip":
            out = self.tower_txt(input_ids=input_ids, attention_mask=attention_mask)
            if getattr(out, "pooler_output", None) is not None:
                return out.pooler_output
            return out.last_hidden_state[:, 0]
        else:
            out = self.backbone.text_model(input_ids=input_ids, attention_mask=attention_mask)
            if getattr(out, "pooler_output", None) is not None:
                return out.pooler_output
            return out.last_hidden_state.mean(dim=1)

    def _encode_image(self, pixel_values: torch.Tensor) -> torch.Tensor:
        """Encode image input to feature vector."""
        if self.backend == "clip":
            out = self.tower_img(pixel_values=pixel_values)
            if getattr(out, "pooler_output", None) is not None:
                return out.pooler_output
            return out.last_hidden_state.mean(dim=1)
        else:
            out = self.backbone.vision_model(pixel_values=pixel_values)
            if getattr(out, "pooler_output", None) is not None:
                return out.pooler_output
            return out.last_hidden_state.mean(dim=1)

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        pixel_values: torch.Tensor,
        text_present: Optional[torch.Tensor] = None,
        image_present: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None
    ) -> Dict[str, Any]:
        """
        Forward pass.
        
        Args:
            input_ids: Tokenized text input [B, seq_len].
            attention_mask: Attention mask for text [B, seq_len].
            pixel_values: Preprocessed image tensor [B, C, H, W].
            text_present: Binary flag indicating text presence [B].
            image_present: Binary flag indicating image presence [B].
            labels: Ground truth labels [B, num_tasks].
            
        Returns:
            Dictionary with 'loss' (if labels provided) and 'logits'.
        """
        batch_size = input_ids.shape[0]
        device = input_ids.device
        
        # Handle missing presence flags
        if text_present is None:
            text_present = torch.ones(batch_size, device=device)
        if image_present is None:
            image_present = torch.ones(batch_size, device=device)
        
        # Encode modalities
        tfeat_raw = self._encode_text(input_ids, attention_mask)
        vfeat_raw = self._encode_image(pixel_values)

        # Project to fusion dimension
        tfeat = self.proj_t(tfeat_raw)
        vfeat = self.proj_i(vfeat_raw)

        # Gated fusion
        presence = torch.stack([text_present, image_present], dim=1)
        zt = torch.tanh(self.g_t(tfeat))
        zi = torch.tanh(self.g_i(vfeat))
        g = torch.sigmoid(self.gate(torch.cat([tfeat, vfeat, presence], dim=1)))
        
        # Conditional fusion based on modality presence
        fused = torch.where(
            (image_present < 0.5).unsqueeze(1), zt,
            torch.where((text_present < 0.5).unsqueeze(1), zi, g * zt + (1.0 - g) * zi)
        )

        # Shared head
        shared = self.shared_head(fused)

        # Task-specific predictions
        logits_per_task = []
        for head in self.heads:
            logit = head(shared).squeeze(-1)
            logits_per_task.append(logit)
        logits = torch.stack(logits_per_task, dim=1)

        # Compute loss if labels provided
        loss = None
        if labels is not None:
            per_task_losses = []
            for j in range(self.num_tasks):
                pw = self.pos_weight[j] if self.pos_weight is not None else None
                lj = F.binary_cross_entropy_with_logits(
                    logits[:, j], labels[:, j], pos_weight=pw, reduction="mean"
                )
                if self.log_vars is not None:
                    per_task_losses.append(
                        torch.exp(-self.log_vars[j]) * lj + 0.5 * self.log_vars[j]
                    )
                else:
                    per_task_losses.append(lj)
            loss = torch.stack(per_task_losses, dim=0).mean()

        return {"loss": loss, "logits": logits}

    def predict(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        pixel_values: torch.Tensor,
        text_present: Optional[torch.Tensor] = None,
        image_present: Optional[torch.Tensor] = None,
    ) -> Dict[str, Any]:
        """
        Make predictions with threshold application.
        
        Returns:
            Dictionary with predictions per task and probabilities.
        """
        self.eval()
        with torch.no_grad():
            outputs = self.forward(
                input_ids=input_ids,
                attention_mask=attention_mask,
                pixel_values=pixel_values,
                text_present=text_present,
                image_present=image_present,
            )
            logits = outputs["logits"]
            probs = torch.sigmoid(logits)
            
            # Apply per-class thresholds
            thresholds = torch.tensor(self.thresholds, device=probs.device)
            predictions = (probs > thresholds).int()
            
            return {
                "predictions": predictions,
                "probabilities": probs,
                "logits": logits,
                "task_names": self.task_names,
            }


# Register the model for auto classes
MultiTaskClassifierConfig.register_for_auto_class()
MultiTaskClassifier.register_for_auto_class("AutoModel")