processors/domain_classifier.py

# DEPENDENCIES
import os
import torch
from typing import Dict
from typing import List
from typing import Tuple
from loguru import logger
from typing import Optional
from config.enums import Domain
import torch.nn.functional as F
from config.schemas import DomainPrediction
from models.model_manager import get_model_manager
from config.constants import domain_classification_params
from config.threshold_config import interpolate_thresholds
from config.threshold_config import get_threshold_for_domain

# Device-agnostic safety
os.environ["TOKENIZERS_PARALLELISM"] = "false"


class DomainClassifier:
    """
    Classifies text into domains using zero-shot classification
    """
    # Use constants from config - map string keys to Domain enum
    DOMAIN_LABELS = {Domain.ACADEMIC      : domain_classification_params.DOMAIN_LABELS["academic"],
                     Domain.CREATIVE      : domain_classification_params.DOMAIN_LABELS["creative"],
                     Domain.AI_ML         : domain_classification_params.DOMAIN_LABELS["ai_ml"],
                     Domain.SOFTWARE_DEV  : domain_classification_params.DOMAIN_LABELS["software_dev"],
                     Domain.TECHNICAL_DOC : domain_classification_params.DOMAIN_LABELS["technical_doc"],
                     Domain.ENGINEERING   : domain_classification_params.DOMAIN_LABELS["engineering"],
                     Domain.SCIENCE       : domain_classification_params.DOMAIN_LABELS["science"],
                     Domain.BUSINESS      : domain_classification_params.DOMAIN_LABELS["business"],
                     Domain.JOURNALISM    : domain_classification_params.DOMAIN_LABELS["journalism"],
                     Domain.SOCIAL_MEDIA  : domain_classification_params.DOMAIN_LABELS["social_media"],
                     Domain.BLOG_PERSONAL : domain_classification_params.DOMAIN_LABELS["blog_personal"],
                     Domain.LEGAL         : domain_classification_params.DOMAIN_LABELS["legal"],
                     Domain.MEDICAL       : domain_classification_params.DOMAIN_LABELS["medical"],
                     Domain.MARKETING     : domain_classification_params.DOMAIN_LABELS["marketing"],
                     Domain.TUTORIAL      : domain_classification_params.DOMAIN_LABELS["tutorial"],
                     Domain.GENERAL       : domain_classification_params.DOMAIN_LABELS["general"],
                    }
    

    def __init__(self):
        self.model_manager       = get_model_manager()
        self.primary_classifier  = None
        self.fallback_classifier = None
        self.is_initialized      = False

    
    def initialize(self) -> bool:
        """
        Initialize the domain classifier with zero-shot models
        """
        try:
            logger.info("Initializing domain classifier...")
            
            # Load primary domain classifier (zero-shot)
            self.primary_classifier = self.model_manager.load_model(model_name = "content_domain_classifier")
            
            # Load fallback classifier
            try:
                self.fallback_classifier = self.model_manager.load_model(model_name = "domain_classifier_fallback")
                logger.info("Fallback classifier loaded successfully")

            except Exception as e:
                logger.warning(f"Could not load fallback classifier: {repr(e)}")
                self.fallback_classifier = None
            
            self.is_initialized = True
            logger.success("Domain classifier initialized successfully")
            return True
            
        except Exception as e:
            logger.error(f"Failed to initialize domain classifier: {repr(e)}")
            return False
    

    def classify(self, text: str, top_k: int = domain_classification_params.TOP_K_DOMAINS, min_confidence: float = domain_classification_params.MIN_CONFIDENCE_THRESHOLD) -> DomainPrediction:
        """
        Classify text into domain using zero-shot classification
        
        Arguments:
        ----------
            text           { str }   : Input text

            top_k          { int }   : Number of top domains to consider
            
            min_confidence { float } : Minimum confidence threshold
            
        Returns:
        --------
            { DomainPrediction }     : DomainPrediction object
        """
        if not self.is_initialized:
            logger.warning("Domain classifier not initialized, initializing now...")
            if not self.initialize():
                return self._get_default_prediction()
        
        try:
            # First try with primary classifier
            primary_result = self._classify_with_model(text       = text, 
                                                       classifier = self.primary_classifier, 
                                                       model_type = "primary",
                                                      )
            
            # Save it as best result
            best_result    = primary_result
            
            # If primary is low confidence but we have fallback, try fallback
            if (self.fallback_classifier and (primary_result.evidence_strength < domain_classification_params.HIGH_CONFIDENCE_THRESHOLD)):
                logger.info("Primary classifier shows low confidence, trying fallback model...")
                fallback_result = self._classify_with_model(text       = text, 
                                                            classifier = self.fallback_classifier, 
                                                            model_type = "fallback",
                                                           )
                
                # Use fallback if it has higher confidence
                if (fallback_result.evidence_strength > best_result.evidence_strength):
                    best_result = fallback_result

            # Hard Safety Gate: if not any solid evidence with great confidence to determine domain, fallback to General Domain
            if best_result.evidence_strength < domain_classification_params.ABS_DOMAIN_CONFIDENCE_THRESHOLD:
                logger.info(f"Domain confidence {best_result.evidence_strength:.3f} below hard threshold {domain_classification_params.ABS_DOMAIN_CONFIDENCE_THRESHOLD:.2f}; forcing GENERAL domain")
                return DomainPrediction(primary_domain    = Domain.GENERAL,
                                        secondary_domain  = None,
                                        evidence_strength = 0.5,
                                        domain_scores     = {Domain.GENERAL.value: 1.0},
                                       )
            
            # Return primary result even if low confidence
            return best_result
            
        except Exception as e:
            logger.error(f"Error in domain classification: {repr(e)}")
            
            # Try fallback classifier if primary failed
            if self.fallback_classifier:
                try:
                    logger.info("Trying fallback classifier after primary failure...")
                    return self._classify_with_model(text       = text, 
                                                     classifier = self.fallback_classifier, 
                                                     model_type = "fallback",
                                                    )
                
                except Exception as fallback_error:
                    logger.error(f"Fallback classifier also failed: {repr(fallback_error)}")
            
            # Both models failed, return default
            return self._get_default_prediction()
    

    # def _classify_with_model(self, text: str, classifier, model_type: str) -> DomainPrediction:
    #     """
    #     Classify using a zero-shot classification model
    #     """
    #     # Preprocess text
    #     processed_text  = self._preprocess_text(text)
        
    #     # Get all candidate labels
    #     all_labels      = list()
    #     label_to_domain = dict()

    #     for domain, labels in self.DOMAIN_LABELS.items():
    #         # Use the first label as the primary one for this domain
    #         primary_label                  = labels[0]
    #         all_labels.append(primary_label)

    #         label_to_domain[primary_label] = domain
        
    #     # Perform zero-shot classification
    #     result = classifier(processed_text,
    #                         candidate_labels    = all_labels,
    #                         multi_label         = False,
    #                         hypothesis_template = "This text is about {}.",
    #                        )
        
    #     # Convert to domain scores
    #     domain_scores = dict()

    #     for label, score in zip(result['labels'], result['scores']):
    #         domain     = label_to_domain[label]
    #         domain_key = domain.value

    #         if (domain_key not in domain_scores):
    #             domain_scores[domain_key] = list()

    #         domain_scores[domain_key].append(score)
        
    #     # Average scores for each domain
    #     avg_domain_scores                 = {domain: sum(scores) / len(scores) for domain, scores in domain_scores.items()}
        
    #     # Sort by score
    #     sorted_domains                    = sorted(avg_domain_scores.items(), key = lambda x: x[1], reverse = True)
        
    #     # Get primary and secondary domains
    #     primary_domain_str, primary_score = sorted_domains[0]
    #     primary_domain                    = Domain(primary_domain_str)
        
    #     secondary_domain                  = None
    #     secondary_score                   = 0.0
        
    #     # Use constant for secondary domain minimum score
    #     secondary_min_score               = domain_classification_params.SECONDARY_DOMAIN_MIN_SCORE

    #     if ((len(sorted_domains) > 1) and (sorted_domains[1][1] >= secondary_min_score)):
    #         secondary_domain = Domain(sorted_domains[1][0])
    #         secondary_score  = sorted_domains[1][1]
        
    #     # Calculate evidence_strength
    #     evidence_strength    = primary_score
        
    #     # Use constants for mixed domain detection
    #     high_conf_threshold  = domain_classification_params.HIGH_CONFIDENCE_THRESHOLD
    #     mixed_secondary_min  = domain_classification_params.MIXED_DOMAIN_SECONDARY_MIN
    #     mixed_ratio_thresh   = domain_classification_params.MIXED_DOMAIN_RATIO_THRESHOLD
    #     mixed_conf_penalty   = domain_classification_params.MIXED_DOMAIN_CONFIDENCE_PENALTY
        
    #     # If we have mixed domains with close scores, adjust confidence
    #     if (secondary_domain and (primary_score < high_conf_threshold) and (secondary_score > mixed_secondary_min)):
            
    #         score_ratio = secondary_score / primary_score
            
    #         # Secondary is at least 60% of primary
    #         if (score_ratio > mixed_ratio_thresh):
    #             # Lower confidence for mixed domains
    #             evidence_strength = ((primary_score + secondary_score) / 2 * mixed_conf_penalty)
    #             logger.info(f"Mixed domain detected: {primary_domain.value} + {secondary_domain.value}, will use interpolated thresholds")
        
    #     # Use constant for low confidence threshold
    #     low_conf_threshold = domain_classification_params.LOW_CONFIDENCE_THRESHOLD
        
    #     # If primary score is low and we have a secondary, it's uncertain
    #     if ((primary_score < low_conf_threshold) and secondary_domain):
    #         # Reduce confidence using penalty
    #         evidence_strength *= mixed_conf_penalty
        
    #     logger.info(f"{model_type.capitalize()} model classified domain: {primary_domain.value} (confidence: {evidence_strength:.3f})")
        
    #     return DomainPrediction(primary_domain    = primary_domain,
    #                             secondary_domain  = secondary_domain,
    #                             evidence_strength = evidence_strength,
    #                             domain_scores     = avg_domain_scores,
    #                            )

    def _classify_with_model(self, text: str, classifier, model_type: str) -> DomainPrediction:
        """
        Classify using a manual NLI-style zero-shot classifier (NO pipelines)
        """

        model, tokenizer = classifier

        processed_text   = self._preprocess_text(text)

        # Build labels
        all_labels      = list()
        label_to_domain = dict()

        for domain, labels in self.DOMAIN_LABELS.items():
            for label in labels[:3]: 
                all_labels.append(label)
                label_to_domain[label] = domain

        # NLI formulation
        premises   = [processed_text] * len(all_labels)
        hypotheses = [f"This text is a {label}." for label in all_labels]

        # Tokenize safely (device-agnostic)
        inputs     = tokenizer(premises,
                               hypotheses,
                               return_tensors = "pt",
                               padding        = True,
                               truncation     = True,
                               max_length     = 1024,  # HARD SAFETY CAP
                              )

        inputs     = {k: v.to(model.device) for k, v in inputs.items()}

        # Forward pass
        with torch.no_grad():
            logits = model(**inputs).logits

        # Resolve entailment index robustly (works for ALL devices/models)
        label2id       = {k.lower(): v for k, v in model.config.label2id.items()}

        # fallback: last column
        entailment_idx = (label2id.get("entailment") or label2id.get("entails") or (logits.shape[-1] - 1))

        probs          = F.softmax(logits, dim = -1)
        scores         = probs[:, entailment_idx].detach().cpu().tolist()

        # Aggregate per-domain scores
        domain_scores  = dict()

        for label, score in zip(all_labels, scores):
            domain     = label_to_domain[label]
            domain_key = domain.value
            domain_scores.setdefault(domain_key, []).append(score)

        # Average scores
        avg_domain_scores                 = {domain: sum(vals) / len(vals) for domain, vals in domain_scores.items()}

        # Sort
        sorted_domains                    = sorted(avg_domain_scores.items(),
                                                   key     = lambda x: x[1],
                                                   reverse = True,
                                                  )

        primary_domain_str, primary_score = sorted_domains[0]
        primary_domain                    = Domain(primary_domain_str)
        secondary_domain                  = None
        secondary_score                   = 0.0
        secondary_min_score               = domain_classification_params.SECONDARY_DOMAIN_MIN_SCORE

        if (len(sorted_domains) > 1) and (sorted_domains[1][1] >= secondary_min_score):
            secondary_domain = Domain(sorted_domains[1][0])
            secondary_score  = sorted_domains[1][1]

        evidence_strength   = primary_score

        # Mixed-domain confidence adjustment 
        high_conf_threshold = domain_classification_params.HIGH_CONFIDENCE_THRESHOLD
        mixed_secondary_min = domain_classification_params.MIXED_DOMAIN_SECONDARY_MIN
        mixed_ratio_thresh  = domain_classification_params.MIXED_DOMAIN_RATIO_THRESHOLD
        mixed_conf_penalty  = domain_classification_params.MIXED_DOMAIN_CONFIDENCE_PENALTY

        if secondary_domain and primary_score < high_conf_threshold and secondary_score > mixed_secondary_min:
            score_ratio = secondary_score / max(primary_score, 1e-6)
            
            if (score_ratio > mixed_ratio_thresh):
                evidence_strength = ((primary_score + secondary_score) / 2) * mixed_conf_penalty
                logger.info(f"Mixed domain detected: {primary_domain.value} + {secondary_domain.value}")

        low_conf_threshold = domain_classification_params.LOW_CONFIDENCE_THRESHOLD

        if ((primary_score < low_conf_threshold) and secondary_domain):
            evidence_strength *= mixed_conf_penalty

        logger.info(f"{model_type.capitalize()} model classified domain: {primary_domain.value} (confidence: {evidence_strength:.3f})")

        return DomainPrediction(primary_domain    = primary_domain,
                                secondary_domain  = secondary_domain,
                                evidence_strength = evidence_strength,
                                domain_scores     = avg_domain_scores,
                               )
    

    def _preprocess_text(self, text: str) -> str:
        """
        Preprocess text for classification
        """
        # Truncate to reasonable length using constant
        max_words = domain_classification_params.MAX_WORDS_FOR_CLASSIFICATION
        words     = text.split()
        
        if (len(words) > max_words):
            text = ' '.join(words[:max_words])
        
        # Clean up text
        text = text.strip()
        if not text:
            return "general content"
        
        return text
    

    def _get_default_prediction(self) -> DomainPrediction:
        """
        Get default prediction when classification fails
        """
        return DomainPrediction(primary_domain    = Domain.GENERAL,
                                secondary_domain  = None,
                                evidence_strength = 0.5,
                                domain_scores     = {Domain.GENERAL.value: 1.0},
                               )
    

    def get_adaptive_thresholds(self, domain_prediction: DomainPrediction):
        """
        Get adaptive thresholds based on domain prediction
        """
        # Use constants for threshold decisions
        high_conf_threshold = domain_classification_params.HIGH_CONFIDENCE_THRESHOLD
        med_conf_threshold  = domain_classification_params.MEDIUM_CONFIDENCE_THRESHOLD
        
        # High confidence, single domain - use domain-specific thresholds
        if ((domain_prediction.evidence_strength > high_conf_threshold) and (not domain_prediction.secondary_domain)):
            return get_threshold_for_domain(domain_prediction.primary_domain)
        
        # Mixed domains - interpolate between primary and secondary
        if domain_prediction.secondary_domain:
            primary_score   = domain_prediction.domain_scores.get(domain_prediction.primary_domain.value, 0)
            secondary_score = domain_prediction.domain_scores.get(domain_prediction.secondary_domain.value, 0)
            
            if (primary_score + secondary_score > 0):
                weight1 = primary_score / (primary_score + secondary_score)
            
            else:
                weight1 = domain_prediction.evidence_strength
                
            return interpolate_thresholds(domain1 = domain_prediction.primary_domain,
                                          domain2 = domain_prediction.secondary_domain,
                                          weight1 = weight1,
                                         )
        
        # Low/medium confidence - blend with general domain
        if (domain_prediction.evidence_strength < med_conf_threshold):
            return interpolate_thresholds(domain1 = domain_prediction.primary_domain,
                                          domain2 = Domain.GENERAL,
                                          weight1 = domain_prediction.evidence_strength,
                                         )
        
        # Default: use domain-specific thresholds
        return get_threshold_for_domain(domain_prediction.primary_domain)
    
    
    def cleanup(self):
        """
        Clean up resources
        """
        self.primary_classifier  = None
        self.fallback_classifier = None
        self.is_initialized      = False


def quick_classify(text: str, **kwargs) -> DomainPrediction:
    """
    Quick domain classification with default settings
    
    Arguments:
    ----------
        text     { str } : Input text

        **kwargs         : Override settings
        
    Returns:
    --------
        { DomainPrediction } : DomainPrediction object
    """
    classifier = DomainClassifier()
    classifier.initialize()
    return classifier.classify(text, **kwargs)


def get_domain_name(domain: Domain) -> str:
    """
    Get human-readable domain name
    
    Arguments:
    ----------
        domain { Domain } : Domain enum value
        
    Returns:
    --------
        { str } : Human-readable domain name
    """
    domain_names = {Domain.ACADEMIC      : "Academic",
                    Domain.CREATIVE      : "Creative Writing",
                    Domain.AI_ML         : "AI/ML",
                    Domain.SOFTWARE_DEV  : "Software Development",
                    Domain.TECHNICAL_DOC : "Technical Documentation",
                    Domain.ENGINEERING   : "Engineering",
                    Domain.SCIENCE       : "Science",
                    Domain.BUSINESS      : "Business",
                    Domain.JOURNALISM    : "Journalism",
                    Domain.SOCIAL_MEDIA  : "Social Media",
                    Domain.BLOG_PERSONAL : "Personal Blog",
                    Domain.LEGAL         : "Legal",
                    Domain.MEDICAL       : "Medical",
                    Domain.MARKETING     : "Marketing",
                    Domain.TUTORIAL      : "Tutorial",
                    Domain.GENERAL       : "General",
                   }

    return domain_names.get(domain, "Unknown")


def is_technical_domain(domain: Domain) -> bool:
    """
    Check if domain is technical in nature
    
    Arguments:
    ----------
        domain { Domain } : Domain enum value
        
    Returns:
    --------
        { bool } : True if technical domain
    """
    technical_domains = {Domain.AI_ML,
                         Domain.SOFTWARE_DEV,
                         Domain.TECHNICAL_DOC,
                         Domain.ENGINEERING,
                         Domain.SCIENCE,
                        }

    return domain in technical_domains


def is_creative_domain(domain: Domain) -> bool:
    """
    Check if domain is creative in nature
    
    Arguments:
    ----------
        domain { Domain } : Domain enum value
        
    Returns:
    --------
        { bool } : True if creative domain
    """
    creative_domains = {Domain.CREATIVE,
                        Domain.JOURNALISM,
                        Domain.SOCIAL_MEDIA,
                        Domain.BLOG_PERSONAL,
                        Domain.MARKETING,
                       }

    return domain in creative_domains


def is_formal_domain(domain: Domain) -> bool:
    """
    Check if domain is formal in nature
    
    Arguments:
    ----------
        domain { Domain } : Domain enum value
        
    Returns:
    --------
        { bool } : True if formal domain
    """
    formal_domains = {Domain.ACADEMIC,
                      Domain.LEGAL,
                      Domain.MEDICAL,
                      Domain.BUSINESS,
                     }

    return domain in formal_domains


# Export
__all__ = ["Domain",
           "DomainClassifier",
           "DomainPrediction",
           "quick_classify",
           "get_domain_name",
           "is_technical_domain",
           "is_creative_domain",
           "is_formal_domain",
          ]