"""Domain Adaptation System with LoRA Adapters.

This module provides domain-specific model adaptation using LoRA (Low-Rank Adaptation)
adapters for technical, medical, and academic domains to improve transcription accuracy.
"""

import logging
import os
from pathlib import Path
from typing import Dict, Optional, Any, List
from dataclasses import dataclass

import torch
import numpy as np
from transformers import WhisperForConditionalGeneration
from peft import LoraConfig, get_peft_model, PeftModel
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix

logger = logging.getLogger(__name__)


@dataclass
class DomainConfig:
    """Configuration for domain-specific adaptation."""
    name: str
    rank: int = 8
    alpha: int = 32
    dropout: float = 0.05
    target_modules: List[str] = None
    
    def __post_init__(self):
        if self.target_modules is None:
            self.target_modules = ["q_proj", "v_proj"]


class DomainAdapter:
    """LoRA adapter architecture for domain-specific model adaptation.
    
    Manages multiple LoRA adapters for different domains, allowing efficient
    switching between domain-specific models while sharing base model parameters.
    """
    
    def __init__(self, base_model_id: str = "openai/whisper-large-v2"):
        """Initialize the domain adapter with a base model.
        
        Args:
            base_model_id: HuggingFace model identifier for the base model
        """
        self.base_model_id = base_model_id
        self.base_model = None
        self.domain_adapters: Dict[str, PeftModel] = {}
        self._load_base_model()
        
    def _load_base_model(self):
        """Load the base Whisper model."""
        try:
            logger.info(f"Loading base model: {self.base_model_id}")
            self.base_model = WhisperForConditionalGeneration.from_pretrained(
                self.base_model_id,
                torch_dtype=torch.float16,
                device_map="auto"
            )
            logger.info("Base model loaded successfully")
        except Exception as e:
            logger.error(f"Failed to load base model: {e}")
            raise
    
    def create_adapter(self, domain_name: str, config: Optional[DomainConfig] = None) -> PeftModel:
        """Create a new LoRA adapter for a specific domain.
        
        Args:
            domain_name: Name of the domain (e.g., "technical", "medical")
            config: Domain-specific configuration, uses defaults if None
            
        Returns:
            Configured PeftModel with LoRA adapter
        """
        if config is None:
            config = DomainConfig(name=domain_name)
        
        logger.info(f"Creating LoRA adapter for domain: {domain_name}")
        
        # Create LoRA configuration
        lora_config = LoraConfig(
            r=config.rank,
            lora_alpha=config.alpha,
            target_modules=config.target_modules,
            lora_dropout=config.dropout,
            bias="none",
            task_type="SEQ_2_SEQ_LM"
        )
        
        # Create PEFT model
        adapter_model = get_peft_model(self.base_model, lora_config)
        
        # Store the adapter
        self.domain_adapters[domain_name] = adapter_model
        
        logger.info(f"LoRA adapter created for domain: {domain_name}")
        return adapter_model
    
    def load_adapter(self, domain_name: str, adapter_path: str) -> PeftModel:
        """Load a pre-trained adapter from disk.
        
        Args:
            domain_name: Name of the domain
            adapter_path: Path to the saved adapter weights
            
        Returns:
            Loaded PeftModel with adapter weights
        """
        if not os.path.exists(adapter_path):
            raise FileNotFoundError(f"Adapter path not found: {adapter_path}")
        
        logger.info(f"Loading adapter for domain {domain_name} from {adapter_path}")
        
        # Create adapter if it doesn't exist
        if domain_name not in self.domain_adapters:
            self.create_adapter(domain_name)
        
        # Load the adapter weights
        adapter_model = self.domain_adapters[domain_name]
        adapter_model.load_adapter(adapter_path)
        
        logger.info(f"Adapter loaded successfully for domain: {domain_name}")
        return adapter_model
    
    def save_adapter(self, domain_name: str, save_path: str) -> None:
        """Save an adapter to disk.
        
        Args:
            domain_name: Name of the domain
            save_path: Path where to save the adapter
        """
        if domain_name not in self.domain_adapters:
            raise ValueError(f"Domain adapter '{domain_name}' not found")
        
        logger.info(f"Saving adapter for domain {domain_name} to {save_path}")
        
        # Ensure directory exists
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        
        # Save the adapter
        adapter_model = self.domain_adapters[domain_name]
        adapter_model.save_adapter(save_path)
        
        logger.info(f"Adapter saved successfully for domain: {domain_name}")
    
    def switch_adapter(self, domain_name: str) -> PeftModel:
        """Switch to a specific domain adapter.
        
        Args:
            domain_name: Name of the domain to switch to
            
        Returns:
            Active PeftModel for the specified domain
        """
        if domain_name not in self.domain_adapters:
            raise ValueError(f"Domain adapter '{domain_name}' not found")
        
        logger.info(f"Switching to domain adapter: {domain_name}")
        return self.domain_adapters[domain_name]
    
    def list_adapters(self) -> List[str]:
        """List all available domain adapters.
        
        Returns:
            List of domain names with available adapters
        """
        return list(self.domain_adapters.keys())
    
    def remove_adapter(self, domain_name: str) -> None:
        """Remove a domain adapter from memory.
        
        Args:
            domain_name: Name of the domain to remove
        """
        if domain_name in self.domain_adapters:
            logger.info(f"Removing adapter for domain: {domain_name}")
            del self.domain_adapters[domain_name]
            # Force garbage collection
            torch.cuda.empty_cache() if torch.cuda.is_available() else None
            logger.info(f"Adapter removed for domain: {domain_name}")
        else:
            logger.warning(f"Adapter for domain '{domain_name}' not found")
    
    def get_adapter_info(self, domain_name: str) -> Dict[str, Any]:
        """Get information about a specific adapter.
        
        Args:
            domain_name: Name of the domain
            
        Returns:
            Dictionary with adapter information
        """
        if domain_name not in self.domain_adapters:
            raise ValueError(f"Domain adapter '{domain_name}' not found")
        
        adapter_model = self.domain_adapters[domain_name]
        
        # Count trainable parameters
        trainable_params = sum(p.numel() for p in adapter_model.parameters() if p.requires_grad)
        total_params = sum(p.numel() for p in adapter_model.parameters())
        
        return {
            "domain": domain_name,
            "trainable_parameters": trainable_params,
            "total_parameters": total_params,
            "trainable_ratio": trainable_params / total_params if total_params > 0 else 0,
            "model_type": type(adapter_model).__name__
        }


class DomainDetector:
    """Domain detection system using TF-IDF and Random Forest classification.
    
    Automatically detects the domain of transcribed text to determine which
    domain-specific adapter should be used for optimal transcription quality.
    """
    
    def __init__(self, max_features: int = 5000, n_estimators: int = 100):
        """Initialize the domain detector.
        
        Args:
            max_features: Maximum number of features for TF-IDF vectorization
            n_estimators: Number of trees in the Random Forest classifier
        """
        self.vectorizer = TfidfVectorizer(
            max_features=max_features,
            stop_words='english',
            ngram_range=(1, 2),
            min_df=1,  # Allow single occurrence terms for small datasets
            max_df=1.0  # Allow all terms to appear in all documents
        )
        self.classifier = RandomForestClassifier(
            n_estimators=n_estimators,
            random_state=42,
            n_jobs=-1
        )
        self.domains = ["general", "technical", "medical", "academic"]
        self.is_trained = False
        
        logger.info("Domain detector initialized")
    
    def train(self, texts: List[str], domain_labels: List[str]) -> None:
        """Train the domain detector on labeled examples.
        
        Args:
            texts: List of text samples
            domain_labels: Corresponding domain labels
        """
        if len(texts) != len(domain_labels):
            raise ValueError("Number of texts and labels must match")
        
        logger.info(f"Training domain detector on {len(texts)} samples")
        
        # Validate domain labels
        valid_domains = set(self.domains)
        invalid_labels = set(domain_labels) - valid_domains
        if invalid_labels:
            raise ValueError(f"Invalid domain labels: {invalid_labels}")
        
        # Fit TF-IDF vectorizer and transform texts
        X = self.vectorizer.fit_transform(texts)
        
        # Train classifier
        self.classifier.fit(X, domain_labels)
        self.is_trained = True
        
        logger.info("Domain detector training completed")
    
    def detect_domain(self, text: str, threshold: float = 0.6) -> str:
        """Detect the domain of a given text.
        
        Args:
            text: Text to classify
            threshold: Confidence threshold for domain detection
            
        Returns:
            Detected domain name
        """
        if not self.is_trained:
            # Fallback to rule-based detection if ML model not trained
            logger.debug("ML model not trained, using rule-based detection")
            return self._rule_based_detection(text)
        
        if not text or not text.strip():
            return "general"
        
        # Transform text using fitted vectorizer
        X = self.vectorizer.transform([text])
        
        # Get prediction probabilities
        probabilities = self.classifier.predict_proba(X)[0]
        
        # Get highest probability domain
        max_prob_idx = np.argmax(probabilities)
        max_probability = probabilities[max_prob_idx]
        
        if max_probability >= threshold:
            detected_domain = self.domains[max_prob_idx]
            logger.debug(f"Detected domain: {detected_domain} (confidence: {max_probability:.3f})")
            return detected_domain
        else:
            logger.debug(f"Low confidence ({max_probability:.3f}), defaulting to general")
            return "general"
    
    def detect_domain_from_text(self, text: str, threshold: float = 0.6) -> Optional[str]:
        """Detect domain from transcript text.
        
        Args:
            text: Text to classify
            threshold: Confidence threshold for domain detection
            
        Returns:
            Detected domain name or None if detection fails
        """
        try:
            return self.detect_domain(text, threshold)
        except Exception as e:
            logger.error(f"Text-based domain detection failed: {e}")
            # Fallback to rule-based detection
            return self._rule_based_detection(text)
    
    def detect_domain_from_path(self, audio_path: Path) -> Optional[str]:
        """Detect domain from audio file path.
        
        Args:
            audio_path: Path to the audio file
            
        Returns:
            Detected domain name or None if detection fails
        """
        try:
            # Extract filename and path components for domain hints
            filename = audio_path.name.lower()
            path_parts = [part.lower() for part in audio_path.parts]
            
            # Check for domain indicators in filename and path
            domain_indicators = {
                "medical": ["medical", "health", "patient", "doctor", "hospital", "clinic", "diagnosis"],
                "technical": ["tech", "programming", "software", "hardware", "algorithm", "code", "development"],
                "academic": ["academic", "research", "lecture", "study", "university", "college", "thesis"],
                "legal": ["legal", "law", "court", "deposition", "testimony", "contract", "agreement"]
            }
            
            for domain, indicators in domain_indicators.items():
                for indicator in indicators:
                    if indicator in filename or any(indicator in part for part in path_parts):
                        logger.debug(f"Path-based domain detection: {domain} (indicator: {indicator})")
                        return domain
            
            # No clear indicators found
            logger.debug("No domain indicators found in audio path")
            return None
            
        except Exception as e:
            logger.error(f"Path-based domain detection failed: {e}")
            return None
    
    def _rule_based_detection(self, text: str) -> str:
        """Simple rule-based domain detection as fallback.
        
        Args:
            text: Text to classify
            
        Returns:
            Detected domain name
        """
        if not text or not text.strip():
            return "general"
        
        text = text.lower()
        
        # Medical domain keywords
        medical_terms = [
            'patient', 'diagnosis', 'treatment', 'symptom', 'clinical', 'medical', 'doctor',
            'nurse', 'hospital', 'clinic', 'disease', 'illness', 'medication', 'prescription',
            'surgery', 'therapy', 'recovery', 'health', 'care', 'emergency', 'ambulance'
        ]
        medical_score = sum(1 for term in medical_terms if term in text)
        
        # Technical domain keywords
        technical_terms = [
            'algorithm', 'system', 'software', 'hardware', 'implementation', 'code', 'programming',
            'development', 'database', 'network', 'server', 'client', 'api', 'interface',
            'function', 'class', 'method', 'variable', 'loop', 'condition', 'debug', 'test'
        ]
        technical_score = sum(1 for term in technical_terms if term in text)
        
        # Academic domain keywords
        academic_terms = [
            'research', 'study', 'analysis', 'theory', 'hypothesis', 'methodology', 'experiment',
            'data', 'results', 'conclusion', 'literature', 'citation', 'publication', 'journal',
            'conference', 'presentation', 'lecture', 'seminar', 'workshop', 'academic'
        ]
        academic_score = sum(1 for term in academic_terms if term in text)
        
        # Legal domain keywords
        legal_terms = [
            'contract', 'agreement', 'law', 'regulation', 'compliance', 'legal', 'court',
            'judge', 'attorney', 'lawyer', 'testimony', 'evidence', 'witness', 'deposition',
            'case', 'ruling', 'verdict', 'appeal', 'jurisdiction', 'statute', 'amendment'
        ]
        legal_score = sum(1 for term in legal_terms if term in text)
        
        # Calculate scores with weights
        scores = [
            ("general", 0),
            ("medical", medical_score * 2),  # Higher weight for medical terms
            ("technical", technical_score),
            ("academic", academic_score),
            ("legal", legal_score * 1.5)  # Higher weight for legal terms
        ]
        
        # Find domain with highest score
        best_domain, best_score = max(scores, key=lambda x: x[1])
        
        # Return general if no clear domain is detected
        if best_score == 0:
            logger.debug("Rule-based detection: No clear domain indicators, using general")
            return "general"
        
        logger.debug(f"Rule-based detection: {best_domain} (score: {best_score})")
        return best_domain
    
    def get_domain_probabilities(self, text: str) -> Dict[str, float]:
        """Get probability scores for all domains.
        
        Args:
            text: Text to classify
            
        Returns:
            Dictionary mapping domain names to probability scores
        """
        if not self.is_trained:
            # Fallback to rule-based detection for probability estimation
            logger.debug("ML model not trained, using rule-based probability estimation")
            detected_domain = self._rule_based_detection(text)
            # Create probability distribution with high confidence for detected domain
            probabilities = {domain: 0.1 for domain in self.domains}
            probabilities[detected_domain] = 0.7
            return probabilities
        
        if not text or not text.strip():
            return {domain: 1.0 if domain == "general" else 0.0 for domain in self.domains}
        
        X = self.vectorizer.transform([text])
        probabilities = self.classifier.predict_proba(X)[0]
        
        return dict(zip(self.domains, probabilities))
    
    def evaluate(self, test_texts: List[str], test_labels: List[str]) -> Dict[str, Any]:
        """Evaluate the domain detector performance.
        
        Args:
            test_texts: Test text samples
            test_labels: True domain labels
            
        Returns:
            Dictionary with evaluation metrics
        """
        if not self.is_trained:
            raise RuntimeError("Domain detector must be trained before evaluation")
        
        if len(test_texts) != len(test_labels):
            raise ValueError("Number of test texts and labels must match")
        
        # Transform test texts
        X_test = self.vectorizer.transform(test_texts)
        
        # Get predictions
        predictions = self.classifier.predict(X_test)
        probabilities = self.classifier.predict_proba(X_test)
        
        # Calculate metrics
        report = classification_report(test_labels, predictions, output_dict=True)
        conf_matrix = confusion_matrix(test_labels, predictions, labels=self.domains)
        
        # Calculate confidence statistics
        max_probabilities = np.max(probabilities, axis=1)
        avg_confidence = np.mean(max_probabilities)
        
        return {
            "classification_report": report,
            "confusion_matrix": conf_matrix.tolist(),
            "average_confidence": avg_confidence,
            "accuracy": report["accuracy"],
            "macro_avg_f1": report["macro avg"]["f1-score"]
        }
    
    def save_model(self, save_path: str) -> None:
        """Save the trained domain detector model.
        
        Args:
            save_path: Path where to save the model
        """
        if not self.is_trained:
            raise RuntimeError("Cannot save untrained model")
        
        import joblib
        
        logger.info(f"Saving domain detector to {save_path}")
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        
        model_data = {
            "vectorizer": self.vectorizer,
            "classifier": self.classifier,
            "domains": self.domains,
            "is_trained": self.is_trained
        }
        
        joblib.dump(model_data, save_path)
        logger.info("Domain detector saved successfully")
    
    def load_model(self, load_path: str) -> None:
        """Load a trained domain detector model.
        
        Args:
            load_path: Path to the saved model
        """
        import joblib
        
        logger.info(f"Loading domain detector from {load_path}")
        
        if not os.path.exists(load_path):
            raise FileNotFoundError(f"Model file not found: {load_path}")
        
        model_data = joblib.load(load_path)
        
        self.vectorizer = model_data["vectorizer"]
        self.classifier = model_data["classifier"]
        self.domains = model_data["domains"]
        self.is_trained = model_data["is_trained"]
        
        logger.info("Domain detector loaded successfully")

