from medcat.cat import CAT
import csv
import os
from huggingface_hub import hf_hub_download
from transformers import pipeline, AutoTokenizer, AutoModelForTokenClassification

class MedFlairProcessor:
    """
    MedFlairProcessor is a wrapper around MedCAT and a Hugging Face Medical NER model for medical text annotation.
    """

    def __init__(self, model_path: str, model_repo_id: str = None):
        """
        Initialize MedFlairProcessor with a given MedCAT model path and a Hugging Face Medical NER model.
        
        :param model_path: Path to the MedCAT model.
        :param model_repo_id: Hugging Face model repository ID (optional, used if downloading).
        """
        self.tui_to_label = self.load_tui_mapping('/home/delta/Documents/workspace/medflair/medflair/data.csv')
        print("TUI mapping loaded successfully.")

        # if model_repo_id:
        #     # Download model from Hugging Face if a repository ID is provided
        #     model_path = self.download_model_from_hub(model_repo_id, model_path)
        
        self.cat = CAT.load_model_pack(model_path)
        print("MedFlairProcessor CAT successfully.")
        # Initialize the Hugging Face Medical NER pipeline
        self.ner_pipeline = pipeline("token-classification", model="Clinical-AI-Apollo/Medical-NER", aggregation_strategy='simple')
        
        print("MedFlairProcessor initialized successfully.")

    def load_tui_mapping(self, csv_path):
        """
        Load the TUI-to-label mapping from a CSV file.

        :param csv_path: Path to the CSV file.
        :return: Dictionary mapping TUI codes to labels.
        """
        tui_mapping = {}
        with open(csv_path, newline='', encoding='utf-8') as csvfile:
            reader = csv.reader(csvfile)
            next(reader)  # Skip the header
            for row in reader:
                if len(row) == 2:
                    tui_mapping[row[0]] = row[1]
        return tui_mapping

    def download_model_from_hub(self, model_repo_id: str, model_path: str):
        """
        Downloads the model from Hugging Face Hub to the specified model path.

        :param model_repo_id: Hugging Face model repository ID.
        :param model_path: Local path where the model should be downloaded.
        :return: Path to the downloaded model.
        """
        model_file_path = hf_hub_download(repo_id=model_repo_id, filename="model_pack", cache_dir=model_path)
        print(f"Model downloaded to {model_file_path}")
        return model_file_path

    def process_text(self, text: str):
        """
        Annotate a given text using both MedCAT and Hugging Face Medical NER model.

        :param text: The medical text to process.
        :return: Combined extracted entities from both models.
        """
        # MedCAT entity extraction
        medcat_entities = self.cat.get_entities(text)
        formatted_medcat_entities = []
        for entity in medcat_entities.get("entities", {}).values():
            formatted_medcat_entities.append({
                "text": entity["source_value"],
                "start_pos": entity["start"],
                "end_pos": entity["end"],
                "type": self.tui_to_label.get(entity["type_ids"][0], entity["type_ids"][0]) if entity["type_ids"] else "Unknown",
                "confidence": entity["acc"],
                "detected_name": entity["detected_name"],
                "pretty_name": entity["pretty_name"],
                "context_similarity": entity["context_similarity"]
            })
        
        # Hugging Face Medical NER extraction
        ner_results = self.ner_pipeline(text)
        formatted_ner_entities = [
            {
                "text": entity["word"],
                "start_pos": entity["start"],
                "end_pos": entity["end"],
                "type": entity["entity_group"],
                "confidence": entity["score"]
            }
            for entity in ner_results
        ]
        
        return {
            "medcat_entities": formatted_medcat_entities,
            "ner_entities": formatted_ner_entities
        }

    def process_combined(self, text: str):
        """
        Process text using both MedCAT and Hugging Face Medical NER, combining outputs into a unified format.

        :param text: The medical text to process.
        :return: List of combined extracted entities.
        """
        results = self.process_text(text)
        combined_entities = results["medcat_entities"] + results["ner_entities"]
        combined_entities.sort(key=lambda x: x["start_pos"])  # Sort by start position
        return combined_entities
