from transformers import pipeline
import spacy

class ClinicalNER:
    """
    A class for Named Entity Recognition and POS tagging using bert-base-uncased_clinical-ner model.
    """
    
    def __init__(self, use_pos=True):
        """
        Initialize the NER pipeline with bert-base-uncased_clinical-ner model.
        
        Args:
            use_pos (bool): Whether to load the POS tagger. Default is True.
        """
        self.ner_pipeline = pipeline(
            "ner",
            model="samrawal/bert-base-uncased_clinical-ner",
            aggregation_strategy="simple"
        )
        
        # Load spaCy model for POS tagging
        self.nlp = None
        if use_pos:
            try:
                self.nlp = spacy.load("en_core_web_sm")
            except OSError:
                print("Warning: spaCy model 'en_core_web_sm' not found.")
                print("Install it with: python -m spacy download en_core_web_sm")
    
    def _merge_subwords(self, entities):
        """
        Merge subword tokens (those starting with ##) into complete words.
        
        Args:
            entities (list): List of entity dictionaries from the pipeline
            
        Returns:
            list: Merged entities with complete words
        """
        if not entities:
            return []
        
        merged = []
        i = 0
        
        while i < len(entities):
            current = entities[i].copy()
            word = current['word']
            end = current['end']
            
            # Look ahead for subword tokens (starting with ##)
            j = i + 1
            while j < len(entities):
                next_entity = entities[j]
                
                # Check if it's a subword of the same entity type
                if (next_entity['word'].startswith('##') and 
                    next_entity['entity_group'] == current['entity_group']):
                    # Remove ## prefix and append
                    word += next_entity['word'][2:]
                    end = next_entity['end']
                    j += 1
                else:
                    break
            
            # Update the merged entity
            current['word'] = word
            current['end'] = end
            merged.append(current)
            
            # Skip the merged tokens
            i = j
        
        return merged
    
    def basic_ner(self, text):
        """
        Performs NER on the input text and returns annotations with merged subwords.
        
        Args:
            text (str): Input text to analyze
            
        Returns:
            list: List of dictionaries containing entity annotations
                  Each dict has: entity_group, score, word, start, end
        """
        entities = self.ner_pipeline(text)
        return self._merge_subwords(entities)
    
    def prolog_ner(self, text):
        """
        Performs NER and returns results as Prolog facts compatible with Tau Prolog.
        Subword tokens are automatically merged.
        
        Args:
            text (str): Input text to analyze
            
        Returns:
            str: Prolog facts as a string, one per line
        """
        entities = self.ner_pipeline(text)
        merged_entities = self._merge_subwords(entities)
        
        prolog_facts = []
        for i, entity in enumerate(merged_entities):
            # Escape single quotes in words for Prolog
            word = entity['word'].replace("'", "\\'")
            
            # Format: entity(Id, Type, Word, Start, End, Score)
            fact = (
                f"entity({i}, '{entity['entity_group']}', "
                f"'{word}', {entity['start']}, "
                f"{entity['end']}, {entity['score']:.4f})."
            )
            prolog_facts.append(fact)
        
        return "\n".join(prolog_facts)
    
    def pos_tagging(self, text):
        """
        Performs Part-of-Speech tagging on the input text.
        
        Args:
            text (str): Input text to analyze
            
        Returns:
            list: List of dictionaries with token, pos, tag, and description
                  Each dict has: token, pos (universal), tag (fine-grained), dep, lemma
        """
        if self.nlp is None:
            raise RuntimeError("POS tagger not initialized. Install spaCy model: python -m spacy download en_core_web_sm")
        
        doc = self.nlp(text)
        
        pos_results = []
        for token in doc:
            pos_results.append({
                'token': token.text,
                'lemma': token.lemma_,
                'pos': token.pos_,  # Universal POS tag
                'tag': token.tag_,  # Fine-grained POS tag
                'dep': token.dep_,  # Dependency relation
                'start': token.idx,
                'end': token.idx + len(token.text)
            })
        
        return pos_results
    
    def prolog_pos(self, text):
        """
        Performs POS tagging and returns results as Prolog facts.
        
        Args:
            text (str): Input text to analyze
            
        Returns:
            str: Prolog facts as a string, one per line
        """
        if self.nlp is None:
            raise RuntimeError("POS tagger not initialized. Install spaCy model: python -m spacy download en_core_web_sm")
        
        pos_results = self.pos_tagging(text)
        
        prolog_facts = []
        for i, token_info in enumerate(pos_results):
            # Escape single quotes in tokens for Prolog
            token = token_info['token'].replace("'", "\\'")
            lemma = token_info['lemma'].replace("'", "\\'")
            
            # Format: pos(Id, Token, Lemma, POS, Tag, Dep, Start, End)
            fact = (
                f"pos({i}, '{token}', '{lemma}', '{token_info['pos']}', "
                f"'{token_info['tag']}', '{token_info['dep']}', "
                f"{token_info['start']}, {token_info['end']})."
            )
            prolog_facts.append(fact)
        
        return "\n".join(prolog_facts)
    
    def combined_analysis(self, text):
        """
        Performs both NER and POS tagging on the same text.
        
        Args:
            text (str): Input text to analyze
            
        Returns:
            dict: Dictionary with 'entities' and 'pos_tags' keys
        """
        return {
            'entities': self.basic_ner(text),
            'pos_tags': self.pos_tagging(text) if self.nlp else []
        }
    
    def prolog_combined(self, text):
        """
        Performs both NER and POS tagging and returns combined Prolog facts.
        
        Args:
            text (str): Input text to analyze
            
        Returns:
            str: Combined Prolog facts as a string
        """
        ner_facts = self.prolog_ner(text)
        pos_facts = self.prolog_pos(text) if self.nlp else ""
        
        if ner_facts and pos_facts:
            return f"{ner_facts}\n\n% POS Tags\n{pos_facts}"
        elif ner_facts:
            return ner_facts
        else:
            return pos_facts