"""NER annotation module using GLiNER models."""

from typing import List, Dict, Union, Optional
import torch
import random
from gliner import GLiNER
from ..utils.text_processing import tokenize_text

class AutoAnnotator:
    """A class for automatic NER annotation using GLiNER models."""
    
    def __init__(
        self,
        model: str = "BookingCare/gliner-multi-healthcare",
        device: Optional[torch.device] = None
    ) -> None:
        """Initialize the annotator with a GLiNER model.
        
        Args:
            model: Name or path of the GLiNER model to use
            device: Device to run the model on (CPU/GPU)
        """
        if device is None:
            device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

        # Set PyTorch memory management settings
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.set_per_process_memory_fraction(0.8)  # Use 80% of available GPU memory

        self.model = GLiNER.from_pretrained(model).to(device)
        self.annotated_data = []
        self.stat = {
            "total": None,
            "current": -1
        }

    def auto_annotate(
        self,
        data: List[str],
        labels: List[str],
        prompt: Optional[Union[str, List[str]]] = None,
        threshold: float = 0.5,
        nested_ner: bool = False
    ) -> List[Dict]:
        """Annotate a list of texts with NER labels.
        
        Args:
            data: List of texts to annotate
            labels: List of entity labels to detect
            prompt: Optional prompt or list of prompts to use
            threshold: Confidence threshold for entity detection
            nested_ner: Whether to allow nested entities
            
        Returns:
            List of annotated examples
        """
        self.stat["total"] = len(data)
        self.stat["current"] = -1
        
        # Process texts in batches
        processed_data = []
        batch_size = 8  # Reduced batch size to prevent OOM errors
        
        for i in range(0, len(data), batch_size):
            batch_texts = data[i:i + batch_size]
            batch_with_prompts = []
            
            # Add prompts to batch texts
            for text in batch_texts:
                if isinstance(prompt, list):
                    prompt_text = random.choice(prompt)
                else:
                    prompt_text = prompt
                text_with_prompt = f"{prompt_text}\n{text}" if prompt_text else text
                batch_with_prompts.append(text_with_prompt)
            
            # Process batch
            batch_results = self._batch_annotate_text(
                batch_with_prompts,
                labels,
                threshold,
                nested_ner
            )
            processed_data.extend(batch_results)
            
            # Clear CUDA cache after each batch
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            
            # Update progress
            self.stat["current"] = min(i + batch_size, len(data))
        
        self.annotated_data = processed_data
        return self.annotated_data

    def _batch_annotate_text(
        self,
        texts: List[str],
        labels: List[str],
        threshold: float,
        nested_ner: bool
    ) -> List[Dict]:
        """Annotate multiple texts in batch.
        
        Args:
            texts: List of texts to annotate
            labels: List of entity labels
            threshold: Confidence threshold
            nested_ner: Whether to allow nested entities
            
        Returns:
            List of annotated examples
        """
        batch_entities = self.model.batch_predict_entities(
            texts,
            labels,
            flat_ner=not nested_ner,
            threshold=threshold
        )
        
        results = []
        for text, entities in zip(texts, batch_entities):
            r = {
                "text": text,
                "entities": [
                    {
                        "entity": entity["label"],
                        "word": entity["text"],
                        "start": entity["start"],
                        "end": entity["end"],
                        "score": 0,
                    }
                    for entity in entities
                ],
            }
            r["entities"] = self._merge_entities(r["entities"])
            results.append(self._transform_data(r))
        return results

    def _merge_entities(self, entities: List[Dict]) -> List[Dict]:
        """Merge adjacent entities of the same type.
        
        Args:
            entities: List of entity dictionaries
            
        Returns:
            List of merged entities
        """
        if not entities:
            return []
        merged = []
        current = entities[0]
        for next_entity in entities[1:]:
            if (next_entity['entity'] == current['entity'] and
                (next_entity['start'] == current['end'] + 1 or
                 next_entity['start'] == current['end'])):
                current['word'] += ' ' + next_entity['word']
                current['end'] = next_entity['end']
            else:
                merged.append(current)
                current = next_entity
        merged.append(current)
        return merged

    def _transform_data(self, data: Dict) -> Dict:
        """Transform raw annotation data into tokenized format.
        
        Args:
            data: Raw annotation data
            
        Returns:
            Transformed data with tokenized text and NER spans
        """
        tokens = tokenize_text(data['text'])
        spans = []

        for entity in data['entities']:
            entity_tokens = tokenize_text(entity['word'])
            entity_length = len(entity_tokens)

            # Find the start and end indices of each entity in the tokenized text
            for i in range(len(tokens) - entity_length + 1):
                if tokens[i:i + entity_length] == entity_tokens:
                    spans.append([i, i + entity_length - 1, entity['entity']])
                    break

        return {
            "tokenized_text": tokens,
            "ner": spans,
            "validated": False
        }