import os
from typing import List, Tuple, Union
import numpy as np
import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel, AutoConfig


class SapBertCrossEncoder(nn.Module):
    """
    Stage-2 Cross-Encoder for clinical entity linking.
    Accepts text pairs (mention, candidate_name) and outputs relevance probabilities
    computed via deep cross-attention layers.
    """

    def __init__(
        self,
        model_name_or_path: str = "checkpoints/sapbert_challenge_finetuned/best_model",
        num_labels: int = 1,
        dropout_prob: float = 0.1,
    ):
        super().__init__()
        self.config = AutoConfig.from_pretrained(model_name_or_path)
        self.encoder = AutoModel.from_pretrained(model_name_or_path, config=self.config)
        self.dropout = nn.Dropout(dropout_prob)
        self.classifier = nn.Linear(self.config.hidden_size, num_labels)
        
    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        token_type_ids: torch.Tensor = None,
    ) -> torch.Tensor:
        """
        Forward pass returning raw logits (shape: [batch_size, 1]).
        """
        kwargs = {"input_ids": input_ids, "attention_mask": attention_mask}
        if token_type_ids is not None:
            kwargs["token_type_ids"] = token_type_ids

        outputs = self.encoder(**kwargs)
        # [CLS] representation
        cls_rep = outputs.last_hidden_state[:, 0, :]
        cls_rep = self.dropout(cls_rep)
        logits = self.classifier(cls_rep)
        return logits.squeeze(-1)

    def predict_scores(
        self,
        pairs: List[Tuple[str, str]],
        tokenizer: AutoTokenizer,
        batch_size: int = 32,
        max_length: int = 64,
        device: Union[str, torch.device] = "cpu",
    ) -> np.ndarray:
        """
        Computes sigmoid-normalized match probabilities for a list of (query, candidate) pairs.
        """
        self.eval()
        self.to(device)
        all_scores = []

        for i in range(0, len(pairs), batch_size):
            batch_pairs = pairs[i : i + batch_size]
            queries = [p[0] for p in batch_pairs]
            candidates = [p[1] for p in batch_pairs]

            inputs = tokenizer(
                queries,
                candidates,
                padding=True,
                truncation=True,
                max_length=max_length,
                return_tensors="pt",
            ).to(device)

            with torch.no_grad():
                logits = self(**inputs)
                probs = torch.sigmoid(logits).cpu().numpy()
                all_scores.extend(probs.flatten().tolist())

        return np.array(all_scores, dtype=np.float32)

    def save_pretrained(self, save_directory: str):
        """Saves model weights and config to directory."""
        os.makedirs(save_directory, exist_ok=True)
        torch.save(self.state_dict(), os.path.join(save_directory, "pytorch_model.bin"))
        self.config.save_pretrained(save_directory)

    @classmethod
    def from_pretrained(cls, load_directory: str, **kwargs):
        """Loads a saved Cross-Encoder model."""
        model = cls(model_name_or_path=load_directory, **kwargs)
        weight_path = os.path.join(load_directory, "pytorch_model.bin")
        if os.path.exists(weight_path):
            state_dict = torch.load(weight_path, map_location="cpu")
            model.load_state_dict(state_dict, strict=False)
        return model