"""
Text Classification Module
Classifies text into predefined categories

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

from transformers import pipeline
import re

class TextClassifier:
    """
    Text classification class
    Developer: RSK World - https://rskworld.in
    """
    
    def __init__(self):
        """Initialize text classifier"""
        try:
            # Zero-shot classification model
            self.zero_shot_classifier = pipeline(
                "zero-shot-classification",
                model="facebook/bart-large-mnli",
                device=-1  # Use CPU
            )
        except Exception as e:
            print(f"Warning: Could not load zero-shot classifier: {e}")
            self.zero_shot_classifier = None
    
    def classify(self, text, categories=None):
        """
        Classify text into categories
        
        Args:
            text (str): Input text
            categories (list): List of category labels
            
        Returns:
            dict: Classification results
        """
        if categories is None:
            categories = [
                "technology",
                "sports",
                "politics",
                "business",
                "entertainment",
                "health",
                "science",
                "education",
                "news",
                "opinion"
            ]
        
        if self.zero_shot_classifier is None:
            # Fallback classification based on keywords
            return self.keyword_based_classification(text, categories)
        
        try:
            # Truncate if too long
            max_length = 512
            if len(text) > max_length:
                text = text[:max_length]
            
            result = self.zero_shot_classifier(text, categories)
            
            return {
                'primary_category': result['labels'][0],
                'primary_score': result['scores'][0],
                'all_categories': [
                    {'label': label, 'score': score}
                    for label, score in zip(result['labels'], result['scores'])
                ],
                'top_3': [
                    {'label': label, 'score': score}
                    for label, score in zip(result['labels'][:3], result['scores'][:3])
                ]
            }
        except Exception as e:
            print(f"Error in classification: {e}")
            return self.keyword_based_classification(text, categories)
    
    def keyword_based_classification(self, text, categories):
        """
        Fallback keyword-based classification
        
        Args:
            text (str): Input text
            categories (list): List of categories
            
        Returns:
            dict: Classification results
        """
        text_lower = text.lower()
        
        # Keyword mapping
        keyword_map = {
            'technology': ['computer', 'software', 'tech', 'digital', 'internet', 'ai', 'app', 'code'],
            'sports': ['game', 'player', 'team', 'match', 'sport', 'championship', 'league'],
            'politics': ['government', 'president', 'election', 'political', 'policy', 'vote'],
            'business': ['company', 'market', 'business', 'economy', 'financial', 'profit', 'stock'],
            'entertainment': ['movie', 'music', 'film', 'celebrity', 'entertainment', 'show'],
            'health': ['health', 'medical', 'doctor', 'hospital', 'disease', 'treatment', 'medicine'],
            'science': ['research', 'study', 'scientific', 'experiment', 'discovery', 'data'],
            'education': ['school', 'student', 'education', 'university', 'learn', 'teacher'],
            'news': ['news', 'report', 'announcement', 'breaking', 'update'],
            'opinion': ['think', 'believe', 'opinion', 'view', 'feel', 'consider']
        }
        
        scores = {}
        for category in categories:
            keywords = keyword_map.get(category, [])
            score = sum(1 for keyword in keywords if keyword in text_lower)
            scores[category] = score / len(keywords) if keywords else 0
        
        # Normalize scores
        total_score = sum(scores.values())
        if total_score > 0:
            scores = {k: v / total_score for k, v in scores.items()}
        
        sorted_categories = sorted(scores.items(), key=lambda x: x[1], reverse=True)
        
        return {
            'primary_category': sorted_categories[0][0] if sorted_categories else categories[0],
            'primary_score': sorted_categories[0][1] if sorted_categories else 0.0,
            'all_categories': [
                {'label': cat, 'score': score}
                for cat, score in sorted_categories
            ],
            'top_3': [
                {'label': cat, 'score': score}
                for cat, score in sorted_categories[:3]
            ]
        }