import os
import json
import torch
import numpy as np
import torch.nn.functional as F
try:
    from .model import SimpleMCQModel
    from .tokenizer import MCQTokenizer
except ImportError:
    from model import SimpleMCQModel
    from tokenizer import MCQTokenizer

class MCQInferencePipeline:
    def __init__(self, model, tokenizer, label_map, max_len=128, device='cpu'):
        self.model = model.to(device)
        self.model.eval()
        self.tokenizer = tokenizer
        self.label_map = label_map
        self.max_len = max_len
        self.device = device

    @classmethod
    def load_from_dir(cls, model_dir, device='cpu'):
        with open(os.path.join(model_dir, 'config.json'), 'r', encoding='utf-8') as f:
            config = json.load(f)

        with open(os.path.join(model_dir, 'label_mapping.json'), 'r', encoding='utf-8') as f:
            raw_labels = json.load(f)
            label_map = {int(k): v for k, v in raw_labels.items()}

        tokenizer = MCQTokenizer.load_vocab(os.path.join(model_dir, 'vocab.json'), max_len=config.get('max_length', 128))
        model = SimpleMCQModel(vocab_size=config['vocab_size'], embed_dim=config['embedding_dim'], hidden_dim=config['hidden_dim'])
        model.load_state_dict(torch.load(os.path.join(model_dir, 'model.pt'), map_location=device))

        return cls(model=model, tokenizer=tokenizer, label_map=label_map, max_len=config.get('max_length', 128), device=device)

    def predict(self, prompt, options):
        """
        options: list of 5 text options [optA, optB, optC, optD, optE]
        """
        option_tensors = []
        for opt_text in options:
            combined_text = str(prompt) + " " + str(opt_text)
            tokens = self.tokenizer.tokenize(combined_text)
            option_tensors.append(tokens)

        # Batch size 1: (1, 5, max_len)
        x = torch.tensor([option_tensors], dtype=torch.long, device=self.device)
        with torch.no_grad():
            logits = self.model(x) # (1, 5)
            probs = F.softmax(logits, dim=1).squeeze(0).cpu().numpy()

        sorted_indices = np.argsort(probs)[::-1]
        top1_idx = sorted_indices[0]
        top1_label = self.label_map[top1_idx]
        top3_labels = [self.label_map[i] for i in sorted_indices[:3]]

        confidence_scores = {self.label_map[i]: float(probs[i]) for i in range(len(options))}
        
        return {
            "top1_label": top1_label,
            "top1_option_text": options[top1_idx],
            "top1_confidence": float(probs[top1_idx]),
            "top3_labels": top3_labels,
            "top3_str": " ".join(top3_labels),
            "confidence_scores": confidence_scores
        }