"""
Entity Recognition Module
Extracts named entities using spaCy and NLTK

Developer: RSK World
Website: https://rskworld.in
Email: help@rskworld.in
Phone: +91 93305 39277
Year: 2026
"""

import spacy
from collections import defaultdict

class EntityRecognizer:
    """
    Named Entity Recognition class using spaCy
    Developer: RSK World - https://rskworld.in
    """
    
    def __init__(self):
        """Initialize spaCy model"""
        try:
            # Try to load English model
            self.nlp = spacy.load("en_core_web_sm")
        except OSError:
            print("Warning: spaCy English model not found.")
            print("Please install it using: python -m spacy download en_core_web_sm")
            self.nlp = None
    
    def extract_entities(self, text):
        """
        Extract named entities from text
        
        Args:
            text (str): Input text
            
        Returns:
            dict: Extracted entities grouped by type
        """
        if self.nlp is None:
            return {
                'entities': [],
                'entity_types': {},
                'total_entities': 0
            }
        
        doc = self.nlp(text)
        
        entities = []
        entity_types = defaultdict(list)
        
        for ent in doc.ents:
            entity_info = {
                'text': ent.text,
                'label': ent.label_,
                'description': spacy.explain(ent.label_),
                'start': ent.start_char,
                'end': ent.end_char
            }
            entities.append(entity_info)
            entity_types[ent.label_].append(ent.text)
        
        # Convert defaultdict to regular dict
        entity_types = dict(entity_types)
        
        return {
            'entities': entities,
            'entity_types': entity_types,
            'total_entities': len(entities),
            'unique_entity_types': list(entity_types.keys())
        }
    
    def extract_organizations(self, text):
        """
        Extract organization entities
        
        Args:
            text (str): Input text
            
        Returns:
            list: List of organization names
        """
        entities = self.extract_entities(text)
        return entities['entity_types'].get('ORG', [])
    
    def extract_persons(self, text):
        """
        Extract person entities
        
        Args:
            text (str): Input text
            
        Returns:
            list: List of person names
        """
        entities = self.extract_entities(text)
        return entities['entity_types'].get('PERSON', [])
    
    def extract_locations(self, text):
        """
        Extract location entities
        
        Args:
            text (str): Input text
            
        Returns:
            list: List of location names
        """
        entities = self.extract_entities(text)
        locations = list(entities['entity_types'].get('GPE', []))
        locations.extend(entities['entity_types'].get('LOC', []))
        return locations