import torch
import torch.nn as nn
import json
import os
from typing import Dict, Optional, Tuple

class UnifiedCellClassifier(nn.Module):
    def __init__(self, models_base_path: str, sub_classifier_names: list):
        """
        Args:
            models_base_path: Path base dove sono salvati i modelli
            sub_classifier_names: Lista nomi cartelle sub-classificatori
                Struttura attesa:
                - main_classifier/model.pth + id2label.json
                - B_cells_classifier/model.pth + id2label.json
                - T_cells_classifier/model.pth + id2label.json
                - ...
        """
        super().__init__()
        
        # Carica classificatore principale
        main_path = os.path.join(models_base_path, "main_classifier")
        self.main_classifier = torch.load(os.path.join(main_path, "complete_model.pth"), map_location='cpu', weights_only=False)
        with open(os.path.join(main_path, "id2label.json")) as f:
            self.main_labels = json.load(f)
        
        
        # Carica sub-classificatori
        self.sub_classifiers = nn.ModuleDict()
        self.sub_labels = {}
        
        for sub_name in sub_classifier_names:
            sub_path = os.path.join(models_base_path, sub_name)
            
            if os.path.exists(sub_path):
                self.sub_classifiers[sub_name] = torch.load(
                    os.path.join(sub_path, "complete_model.pth"), map_location='cpu', weights_only=False
                )
                with open(os.path.join(sub_path, "id2label.json")) as f:
                    self.sub_labels[sub_name] = json.load(f)
        
        # Mapping macrocategoria -> sub-classificatore (da configurare)
        self.macro_to_sub = self._build_macro_to_sub_mapping()
        
        # Imposta modalità eval
        self.eval()
    
    def _build_macro_to_sub_mapping(self):

        return {
            "0": "B_cells_classifier",      
            "1": "CD4plus_T_cells_classifier",     
            "4": "Myeloid_cells_classifier",     
            "5": "NK_cells_classifier",
            "7": "TRAV1_2_CD8plus_T_cells",
            "8": "gd_T_cells_classfier"
        }
    
    def forward(self, x: torch.Tensor, return_probabilities: bool = False):
        """
        Forward pass gerarchico
        
        Args:
            x: Input embeddings [batch_size, embedding_dim]
            return_probabilities: Se True, restituisce anche le probabilità
            
        Returns:
            Dict con macro_prediction, sub_prediction, final_prediction
        """
        # Classificazione principale
        with torch.no_grad():
            main_logits = self.main_classifier(x)
            main_probs = torch.softmax(main_logits, dim=-1)
            main_pred = torch.argmax(main_logits, dim=-1)
        
        # Classificazione secondaria
        batch_size = x.shape[0]
        sub_predictions = []
        sub_probabilities = [] if return_probabilities else None
        
        for i in range(batch_size):
            macro_idx = str(main_pred[i].item())
            macro_label = self.main_labels.get(macro_idx, f"unknown_{macro_idx}")
            
            # Controlla se esiste sub-classificatore per questa macro
            if macro_idx in self.macro_to_sub:
                sub_classifier_name = self.macro_to_sub[macro_idx]
                if sub_classifier_name in self.sub_classifiers:
                    # Usa sub-classificatore
                    with torch.no_grad():
                        sub_logits = self.sub_classifiers[sub_classifier_name](x[i:i+1])
                        sub_probs = torch.softmax(sub_logits, dim=-1)
                        sub_pred = torch.argmax(sub_logits, dim=-1)
                    
                    sub_idx = str(sub_pred.item())
                    sub_label = self.sub_labels[sub_classifier_name].get(sub_idx, f"unknown_{sub_idx}")
                    final_prediction = f"{macro_label}_{sub_label}"
                    
                    if return_probabilities:
                        sub_probabilities.append(sub_probs[0])
                else:
                    # Sub-classificatore non trovato, usa solo macro
                    final_prediction = macro_label
                    if return_probabilities:
                        sub_probabilities.append(None)
            else:
                # Nessun sub-classificatore per questa macro, usa solo macro
                final_prediction = macro_label
                if return_probabilities:
                    sub_probabilities.append(None)
            
            sub_predictions.append(final_prediction)
        
        result = {
            'macro_predictions': [self.main_labels.get(str(idx.item()), f"unknown_{idx.item()}") 
                                for idx in main_pred],
            'final_predictions': sub_predictions
        }
        
        if return_probabilities:
            result['macro_probabilities'] = main_probs
            result['sub_probabilities'] = sub_probabilities
        
        return result
    
    def predict(self, x: torch.Tensor):
        """Metodo semplificato per predizione"""
        return self.forward(x, return_probabilities=False)['final_predictions']
    
    @classmethod
    def from_pretrained(cls, repo_path: str, **kwargs):
        """
        carica la struttura dal repo di HF: aspetta
          - config.json
          - id2label_main.json
          - macro_to_sub.json
          - sub_classifiers/<name>.bin + id2label
        """
        # 1. leggi config
        with open(os.path.join(repo_path, "config.json")) as f:
            config = json.load(f)

        # 2. istanzia l'oggetto
        model = cls(**config)  

        # 3. carica main
        main_sd = torch.load(os.path.join(repo_path, "main_classifier.bin"), map_location="cpu")
        model.main_classifier.load_state_dict(main_sd)
        model.main_labels = json.load(open(os.path.join(repo_path, "id2label_main.json")))

        # 4. carica sub
        model.sub_classifiers = nn.ModuleDict()
        model.sub_labels = {}
        for name in model.sub_classifier_names:
            bin_path = os.path.join(repo_path, "sub_classifiers", f"{name}.bin")
            model.sub_classifiers[name] = model._build_submodule(name)  # metodo helper che crea l’istanza
            model.sub_classifiers[name].load_state_dict(torch.load(bin_path, map_location="cpu"))
            model.sub_labels[name] = json.load(open(
                os.path.join(repo_path, "sub_classifiers", f"{name}_id2label.json")))
        
        model.macro_to_sub = json.load(open(os.path.join(repo_path, "macro_to_sub.json")))
        model.eval()
        return model