from transformers import pipeline

class ClinicalNER:
    """
    A class for Named Entity Recognition using bert-base-uncased_clinical-ner model.
    """
    
    def __init__(self):
        """
        Initialize the NER pipeline with bert-base-uncased_clinical-ner model.
        Note: Using aggregation_strategy="simple" to merge subword tokens.
        """
        self.ner_pipeline = pipeline(
            "ner",
            model="samrawal/bert-base-uncased_clinical-ner",
            aggregation_strategy="simple"
        )
    
    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)