"""
================================================================================
Text Classification Dataset - Active Learning Module
================================================================================
Project: Text Classification Dataset
Category: Text Data / NLP

Author: Molla Samser
Designer & Tester: Rima Khatun
Website: https://rskworld.in
Email: help@rskworld.in | support@rskworld.in
Phone: +91 93305 39277

Copyright (c) 2026 RSK World - All Rights Reserved
Content used for educational purposes only.

Features:
- Uncertainty Sampling (Least Confidence, Margin, Entropy)
- Query-by-Committee
- Diversity Sampling
- Expected Model Change
- Information Density
- Batch Active Learning
- Human-in-the-loop Interface
- Learning Curve Tracking

Created: December 2026
================================================================================
"""

import numpy as np
import pandas as pd
from typing import List, Tuple, Dict, Optional, Callable, Union
from abc import ABC, abstractmethod
from collections import Counter
import warnings
from datetime import datetime

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.naive_bayes import MultinomialNB
from sklearn.svm import LinearSVC
from sklearn.metrics import accuracy_score, f1_score
from sklearn.base import clone
import scipy.sparse as sp

# Project information
__author__ = "Molla Samser"
__website__ = "https://rskworld.in"
__email__ = "help@rskworld.in"

# Category mapping
CATEGORIES = {
    0: 'Technology', 1: 'Sports', 2: 'Politics',
    3: 'Entertainment', 4: 'Business', 5: 'Science'
}


class QueryStrategy(ABC):
    """
    Abstract base class for query strategies.
    
    Author: Molla Samser | RSK World (https://rskworld.in)
    """
    
    @abstractmethod
    def select_samples(
        self,
        X_pool: Union[np.ndarray, sp.spmatrix],
        model,
        n_samples: int = 1
    ) -> np.ndarray:
        """
        Select samples from the pool for labeling.
        
        Args:
            X_pool: Features of unlabeled samples
            model: Trained model with predict_proba
            n_samples: Number of samples to select
            
        Returns:
            Indices of selected samples
        """
        pass


class UncertaintySampling(QueryStrategy):
    """
    Uncertainty sampling selects samples where the model is least confident.
    
    Strategies:
    - least_confidence: 1 - max(P(y|x))
    - margin: P(y1|x) - P(y2|x)  (difference between top 2 classes)
    - entropy: -sum(P(y|x) * log(P(y|x)))
    
    Author: Molla Samser | RSK World (https://rskworld.in)
    """
    
    def __init__(self, strategy: str = 'entropy'):
        """
        Initialize uncertainty sampling.
        
        Args:
            strategy: One of 'least_confidence', 'margin', 'entropy'
        """
        valid_strategies = ['least_confidence', 'margin', 'entropy']
        if strategy not in valid_strategies:
            raise ValueError(f"Strategy must be one of {valid_strategies}")
        self.strategy = strategy
    
    def select_samples(
        self,
        X_pool: Union[np.ndarray, sp.spmatrix],
        model,
        n_samples: int = 1
    ) -> np.ndarray:
        """Select most uncertain samples."""
        
        # Get probability predictions
        if hasattr(model, 'predict_proba'):
            probs = model.predict_proba(X_pool)
        else:
            # For models without predict_proba (like LinearSVC)
            decision = model.decision_function(X_pool)
            # Convert to pseudo-probabilities using softmax
            exp_decision = np.exp(decision - np.max(decision, axis=1, keepdims=True))
            probs = exp_decision / exp_decision.sum(axis=1, keepdims=True)
        
        # Calculate uncertainty scores
        if self.strategy == 'least_confidence':
            scores = 1 - np.max(probs, axis=1)
        
        elif self.strategy == 'margin':
            # Sort probabilities for each sample
            sorted_probs = np.sort(probs, axis=1)
            # Margin = difference between top 2 classes
            scores = 1 - (sorted_probs[:, -1] - sorted_probs[:, -2])
        
        elif self.strategy == 'entropy':
            # Entropy = -sum(p * log(p))
            with np.errstate(divide='ignore', invalid='ignore'):
                log_probs = np.log(probs + 1e-10)
                scores = -np.sum(probs * log_probs, axis=1)
        
        # Select top n_samples with highest uncertainty
        return np.argsort(scores)[-n_samples:][::-1]


class QueryByCommittee(QueryStrategy):
    """
    Query-by-Committee uses an ensemble of models to identify
    samples where the committee disagrees most.
    
    Author: Molla Samser | RSK World (https://rskworld.in)
    """
    
    def __init__(self, n_committee: int = 5, base_estimator=None):
        """
        Initialize QBC.
        
        Args:
            n_committee: Number of committee members
            base_estimator: Base model to use (will be trained with bootstrap)
        """
        self.n_committee = n_committee
        self.base_estimator = base_estimator or LogisticRegression(max_iter=1000)
        self.committee = []
    
    def train_committee(self, X: Union[np.ndarray, sp.spmatrix], y: np.ndarray):
        """Train the committee with bootstrap sampling."""
        self.committee = []
        n_samples = X.shape[0]
        
        for _ in range(self.n_committee):
            # Bootstrap sampling
            indices = np.random.choice(n_samples, size=n_samples, replace=True)
            X_boot = X[indices] if isinstance(X, np.ndarray) else X[indices, :]
            y_boot = y[indices]
            
            # Train committee member
            member = clone(self.base_estimator)
            member.fit(X_boot, y_boot)
            self.committee.append(member)
    
    def select_samples(
        self,
        X_pool: Union[np.ndarray, sp.spmatrix],
        model=None,  # Not used, uses committee instead
        n_samples: int = 1
    ) -> np.ndarray:
        """Select samples with highest disagreement."""
        if not self.committee:
            raise ValueError("Committee not trained. Call train_committee first.")
        
        # Get predictions from each committee member
        predictions = np.array([
            member.predict(X_pool) for member in self.committee
        ])
        
        # Calculate vote entropy (disagreement)
        n_classes = len(np.unique(predictions))
        scores = np.zeros(X_pool.shape[0])
        
        for i in range(X_pool.shape[0]):
            votes = predictions[:, i]
            vote_counts = Counter(votes)
            vote_fracs = np.array([vote_counts.get(c, 0) / self.n_committee 
                                   for c in range(n_classes)])
            # Entropy of votes
            with np.errstate(divide='ignore', invalid='ignore'):
                log_fracs = np.log(vote_fracs + 1e-10)
                scores[i] = -np.sum(vote_fracs * log_fracs)
        
        return np.argsort(scores)[-n_samples:][::-1]


class DiversitySampling(QueryStrategy):
    """
    Diversity sampling selects a diverse batch of samples
    using k-means++ style selection.
    
    Author: Molla Samser | RSK World (https://rskworld.in)
    """
    
    def __init__(self, combine_with_uncertainty: bool = True):
        """
        Initialize diversity sampling.
        
        Args:
            combine_with_uncertainty: Also consider model uncertainty
        """
        self.combine_with_uncertainty = combine_with_uncertainty
    
    def select_samples(
        self,
        X_pool: Union[np.ndarray, sp.spmatrix],
        model=None,
        n_samples: int = 1
    ) -> np.ndarray:
        """Select diverse samples using k-means++ initialization."""
        
        # Convert sparse to dense if needed
        if sp.issparse(X_pool):
            X_dense = X_pool.toarray()
        else:
            X_dense = X_pool
        
        n_pool = X_dense.shape[0]
        selected = []
        
        # If combining with uncertainty, weight by uncertainty
        weights = np.ones(n_pool)
        if self.combine_with_uncertainty and model is not None:
            if hasattr(model, 'predict_proba'):
                probs = model.predict_proba(X_pool)
                weights = 1 - np.max(probs, axis=1)  # Uncertainty
        
        # Select first sample (highest uncertainty or random)
        first_idx = np.argmax(weights) if self.combine_with_uncertainty else np.random.randint(n_pool)
        selected.append(first_idx)
        
        # k-means++ style selection
        for _ in range(n_samples - 1):
            # Calculate distance to nearest selected sample
            min_distances = np.full(n_pool, np.inf)
            for sel_idx in selected:
                distances = np.linalg.norm(X_dense - X_dense[sel_idx], axis=1)
                min_distances = np.minimum(min_distances, distances)
            
            # Weight by distance and uncertainty
            selection_probs = min_distances ** 2 * weights
            selection_probs[selected] = 0  # Don't reselect
            selection_probs = selection_probs / selection_probs.sum()
            
            # Select next sample
            next_idx = np.random.choice(n_pool, p=selection_probs)
            selected.append(next_idx)
        
        return np.array(selected)


class ActiveLearner:
    """
    Main active learning class that orchestrates the learning loop.
    
    Features:
    - Multiple query strategies
    - Automatic model retraining
    - Learning curve tracking
    - Human-in-the-loop interface
    
    Author: Molla Samser | RSK World (https://rskworld.in)
    """
    
    def __init__(
        self,
        model=None,
        query_strategy: QueryStrategy = None,
        vectorizer=None,
        random_state: int = 42
    ):
        """
        Initialize active learner.
        
        Args:
            model: Classifier model
            query_strategy: Strategy for selecting samples
            vectorizer: Text vectorizer (TfidfVectorizer)
            random_state: Random seed
        """
        self.model = model or LogisticRegression(max_iter=1000, random_state=random_state)
        self.query_strategy = query_strategy or UncertaintySampling(strategy='entropy')
        self.vectorizer = vectorizer or TfidfVectorizer(max_features=10000, ngram_range=(1, 2))
        self.random_state = random_state
        np.random.seed(random_state)
        
        # Data storage
        self.X_labeled = None
        self.y_labeled = None
        self.texts_labeled = []
        self.X_pool = None
        self.pool_indices = None
        self.texts_pool = []
        
        # Tracking
        self.learning_curve = []
        self.query_history = []
        self.is_fitted = False
    
    def initialize(
        self,
        texts_pool: List[str],
        texts_initial: List[str] = None,
        labels_initial: List[int] = None,
        n_initial: int = 10
    ):
        """
        Initialize the active learner with pool and initial labeled data.
        
        Args:
            texts_pool: List of unlabeled texts
            texts_initial: Initial labeled texts
            labels_initial: Labels for initial texts
            n_initial: If no initial data, randomly select this many
        """
        print(f"\n{'='*60}")
        print("Initializing Active Learner - RSK World")
        print(f"Author: {__author__} | Website: {__website__}")
        print(f"{'='*60}\n")
        
        self.texts_pool = list(texts_pool)
        
        if texts_initial is not None and labels_initial is not None:
            # Use provided initial data
            self.texts_labeled = list(texts_initial)
            self.y_labeled = np.array(labels_initial)
        else:
            # Random initial selection
            indices = np.random.choice(len(self.texts_pool), size=n_initial, replace=False)
            self.texts_labeled = [self.texts_pool[i] for i in indices]
            # Remove from pool
            for i in sorted(indices, reverse=True):
                self.texts_pool.pop(i)
            # Placeholder labels (need to be set)
            self.y_labeled = np.zeros(n_initial, dtype=int)
            print(f"⚠️  Selected {n_initial} random samples for initial labeling.")
            print("   Please provide labels using set_initial_labels() method.")
            return
        
        # Fit vectorizer on all texts
        all_texts = self.texts_labeled + self.texts_pool
        self.vectorizer.fit(all_texts)
        
        # Transform data
        self.X_labeled = self.vectorizer.transform(self.texts_labeled)
        self.X_pool = self.vectorizer.transform(self.texts_pool)
        self.pool_indices = np.arange(len(self.texts_pool))
        
        # Initial training
        self._train_model()
        
        print(f"✓ Initialized with {len(self.texts_labeled)} labeled samples")
        print(f"✓ Pool size: {len(self.texts_pool)} unlabeled samples")
        print(f"✓ Vocabulary size: {len(self.vectorizer.vocabulary_)}")
    
    def set_initial_labels(self, labels: List[int]):
        """Set labels for initially selected samples."""
        self.y_labeled = np.array(labels)
        
        # Fit vectorizer
        all_texts = self.texts_labeled + self.texts_pool
        self.vectorizer.fit(all_texts)
        
        # Transform
        self.X_labeled = self.vectorizer.transform(self.texts_labeled)
        self.X_pool = self.vectorizer.transform(self.texts_pool)
        self.pool_indices = np.arange(len(self.texts_pool))
        
        # Train
        self._train_model()
        
        print(f"✓ Labels set. Model trained on {len(labels)} samples.")
    
    def _train_model(self):
        """Train the model on labeled data."""
        self.model.fit(self.X_labeled, self.y_labeled)
        self.is_fitted = True
    
    def query(self, n_samples: int = 1) -> List[Tuple[int, str]]:
        """
        Query for samples to label.
        
        Args:
            n_samples: Number of samples to query
            
        Returns:
            List of (index, text) tuples
        """
        if not self.is_fitted:
            raise ValueError("Model not fitted. Call initialize() first.")
        
        if len(self.texts_pool) == 0:
            print("⚠️  No more samples in pool!")
            return []
        
        n_samples = min(n_samples, len(self.texts_pool))
        
        # Get indices using query strategy
        if isinstance(self.query_strategy, QueryByCommittee):
            self.query_strategy.train_committee(self.X_labeled, self.y_labeled)
            selected_indices = self.query_strategy.select_samples(
                self.X_pool, n_samples=n_samples
            )
        else:
            selected_indices = self.query_strategy.select_samples(
                self.X_pool, self.model, n_samples=n_samples
            )
        
        # Get texts
        queries = [(int(idx), self.texts_pool[idx]) for idx in selected_indices]
        
        # Store query history
        self.query_history.append({
            'iteration': len(self.learning_curve),
            'indices': selected_indices.tolist(),
            'timestamp': datetime.now().isoformat()
        })
        
        return queries
    
    def teach(
        self,
        indices: List[int],
        labels: List[int],
        X_test: np.ndarray = None,
        y_test: np.ndarray = None
    ):
        """
        Add newly labeled samples and retrain.
        
        Args:
            indices: Indices of labeled samples (from pool)
            labels: Labels for the samples
            X_test: Optional test features for tracking
            y_test: Optional test labels for tracking
        """
        # Add to labeled set
        for idx, label in zip(indices, labels):
            self.texts_labeled.append(self.texts_pool[idx])
            
        # Update labels
        self.y_labeled = np.concatenate([self.y_labeled, np.array(labels)])
        
        # Update labeled features
        new_texts = [self.texts_pool[idx] for idx in indices]
        new_features = self.vectorizer.transform(new_texts)
        self.X_labeled = sp.vstack([self.X_labeled, new_features])
        
        # Remove from pool (in reverse order to maintain indices)
        for idx in sorted(indices, reverse=True):
            self.texts_pool.pop(idx)
        
        # Rebuild pool
        if len(self.texts_pool) > 0:
            self.X_pool = self.vectorizer.transform(self.texts_pool)
            self.pool_indices = np.arange(len(self.texts_pool))
        else:
            self.X_pool = None
            self.pool_indices = np.array([])
        
        # Retrain
        self._train_model()
        
        # Track performance
        metrics = {
            'n_labeled': len(self.texts_labeled),
            'pool_size': len(self.texts_pool),
            'train_accuracy': accuracy_score(
                self.y_labeled, 
                self.model.predict(self.X_labeled)
            )
        }
        
        if X_test is not None and y_test is not None:
            y_pred = self.model.predict(X_test)
            metrics['test_accuracy'] = accuracy_score(y_test, y_pred)
            metrics['test_f1'] = f1_score(y_test, y_pred, average='macro')
        
        self.learning_curve.append(metrics)
        
        print(f"✓ Added {len(indices)} samples. Total labeled: {len(self.texts_labeled)}")
        if 'test_accuracy' in metrics:
            print(f"  Test accuracy: {metrics['test_accuracy']:.4f}")
    
    def interactive_labeling(
        self,
        n_iterations: int = 10,
        batch_size: int = 5,
        X_test: np.ndarray = None,
        y_test: np.ndarray = None
    ):
        """
        Run interactive labeling session.
        
        Args:
            n_iterations: Number of labeling iterations
            batch_size: Samples per iteration
            X_test: Test features for tracking
            y_test: Test labels for tracking
        """
        print(f"\n{'='*60}")
        print("Interactive Labeling Session")
        print(f"Categories: {CATEGORIES}")
        print(f"{'='*60}\n")
        
        for iteration in range(n_iterations):
            print(f"\n--- Iteration {iteration + 1}/{n_iterations} ---")
            
            queries = self.query(n_samples=batch_size)
            
            if not queries:
                print("No more samples to label!")
                break
            
            labels = []
            indices = []
            
            for idx, text in queries:
                print(f"\nText: {text[:200]}...")
                print(f"Categories: {CATEGORIES}")
                
                while True:
                    try:
                        label = int(input(f"Enter label (0-{len(CATEGORIES)-1}): "))
                        if 0 <= label < len(CATEGORIES):
                            break
                        print(f"Invalid label. Must be 0-{len(CATEGORIES)-1}")
                    except ValueError:
                        print("Please enter a number.")
                
                labels.append(label)
                indices.append(idx)
            
            self.teach(indices, labels, X_test, y_test)
        
        print(f"\n{'='*60}")
        print("Labeling session complete!")
        print(f"Total labeled: {len(self.texts_labeled)}")
        print(f"{'='*60}")
    
    def simulate_labeling(
        self,
        true_labels: List[int],
        n_iterations: int = 20,
        batch_size: int = 5,
        X_test: np.ndarray = None,
        y_test: np.ndarray = None
    ) -> List[Dict]:
        """
        Simulate labeling with known ground truth.
        
        Args:
            true_labels: True labels for pool samples
            n_iterations: Number of iterations
            batch_size: Samples per iteration
            X_test: Test features
            y_test: Test labels
            
        Returns:
            Learning curve data
        """
        print(f"\n{'='*60}")
        print("Simulated Active Learning - RSK World")
        print(f"Iterations: {n_iterations}, Batch size: {batch_size}")
        print(f"{'='*60}\n")
        
        labels_pool = list(true_labels)
        
        for iteration in range(n_iterations):
            if len(self.texts_pool) == 0:
                print("Pool exhausted!")
                break
            
            queries = self.query(n_samples=batch_size)
            
            if not queries:
                break
            
            indices = [q[0] for q in queries]
            labels = [labels_pool[idx] for idx in indices]
            
            self.teach(indices, labels, X_test, y_test)
            
            # Remove used labels from pool (in reverse order)
            for idx in sorted(indices, reverse=True):
                labels_pool.pop(idx)
            
            if (iteration + 1) % 5 == 0:
                acc = self.learning_curve[-1].get('test_accuracy', 'N/A')
                print(f"Iteration {iteration + 1}: {len(self.texts_labeled)} labeled, acc={acc}")
        
        return self.learning_curve
    
    def get_learning_curve(self) -> pd.DataFrame:
        """Get learning curve as DataFrame."""
        return pd.DataFrame(self.learning_curve)
    
    def plot_learning_curve(self, save_path: str = None):
        """Plot the learning curve."""
        try:
            import matplotlib.pyplot as plt
            
            df = self.get_learning_curve()
            
            fig, ax = plt.subplots(figsize=(10, 6))
            
            ax.plot(df['n_labeled'], df['train_accuracy'], 'b-o', label='Train Accuracy')
            if 'test_accuracy' in df.columns:
                ax.plot(df['n_labeled'], df['test_accuracy'], 'r-o', label='Test Accuracy')
            
            ax.set_xlabel('Number of Labeled Samples')
            ax.set_ylabel('Accuracy')
            ax.set_title('Active Learning Curve - RSK World')
            ax.legend()
            ax.grid(True, alpha=0.3)
            
            if save_path:
                plt.savefig(save_path, dpi=150, bbox_inches='tight')
            
            plt.tight_layout()
            plt.show()
            
        except ImportError:
            print("matplotlib not available for plotting")


def compare_strategies(
    texts: List[str],
    labels: List[int],
    n_initial: int = 20,
    n_iterations: int = 15,
    batch_size: int = 5,
    test_size: float = 0.2
) -> Dict:
    """
    Compare different query strategies.
    
    Args:
        texts: All texts
        labels: All labels
        n_initial: Initial labeled samples
        n_iterations: Learning iterations
        batch_size: Samples per iteration
        test_size: Fraction for test set
        
    Returns:
        Comparison results
    """
    from sklearn.model_selection import train_test_split
    
    print(f"\n{'='*60}")
    print("Query Strategy Comparison - RSK World")
    print(f"Author: {__author__}")
    print(f"{'='*60}\n")
    
    # Split data
    texts_train, texts_test, y_train, y_test = train_test_split(
        texts, labels, test_size=test_size, random_state=42, stratify=labels
    )
    
    # Prepare test data
    vectorizer = TfidfVectorizer(max_features=10000, ngram_range=(1, 2))
    vectorizer.fit(texts_train + texts_test)
    X_test = vectorizer.transform(texts_test)
    
    strategies = {
        'Random': UncertaintySampling('least_confidence'),  # Will be random baseline
        'Least Confidence': UncertaintySampling('least_confidence'),
        'Margin': UncertaintySampling('margin'),
        'Entropy': UncertaintySampling('entropy'),
        'Diversity': DiversitySampling(combine_with_uncertainty=True)
    }
    
    results = {}
    
    for name, strategy in strategies.items():
        print(f"\nRunning {name}...")
        
        # Reset random state
        np.random.seed(42)
        
        # Initialize learner
        learner = ActiveLearner(
            query_strategy=strategy if name != 'Random' else None,
            random_state=42
        )
        
        # Random initial samples
        init_indices = np.random.choice(len(texts_train), size=n_initial, replace=False)
        texts_init = [texts_train[i] for i in init_indices]
        labels_init = [y_train[i] for i in init_indices]
        
        # Pool (remaining samples)
        pool_mask = np.ones(len(texts_train), dtype=bool)
        pool_mask[init_indices] = False
        texts_pool = [texts_train[i] for i, m in enumerate(pool_mask) if m]
        labels_pool = [y_train[i] for i, m in enumerate(pool_mask) if m]
        
        learner.initialize(texts_pool, texts_init, labels_init)
        
        # Simulate learning
        curve = learner.simulate_labeling(
            labels_pool,
            n_iterations=n_iterations,
            batch_size=batch_size,
            X_test=X_test,
            y_test=y_test
        )
        
        results[name] = learner.get_learning_curve()
    
    print(f"\n{'='*60}")
    print("Comparison Complete!")
    print(f"{'='*60}")
    
    return results


if __name__ == "__main__":
    print(f"\n{'='*60}")
    print("Active Learning Module Demo")
    print(f"Author: {__author__} | Website: {__website__}")
    print(f"{'='*60}\n")
    
    # Demo with sample data
    try:
        df = pd.read_csv('../data/csv/train.csv', comment='#')
        
        texts = df['text'].tolist()
        labels = df['label'].tolist()
        
        # Simple demo
        print("Running active learning demo...")
        
        # Initialize learner
        learner = ActiveLearner(
            query_strategy=UncertaintySampling('entropy')
        )
        
        # Split data
        n_init = 10
        init_texts = texts[:n_init]
        init_labels = labels[:n_init]
        pool_texts = texts[n_init:]
        pool_labels = labels[n_init:]
        
        learner.initialize(pool_texts, init_texts, init_labels)
        
        # Simulate a few iterations
        learner.simulate_labeling(
            pool_labels,
            n_iterations=5,
            batch_size=3
        )
        
        print("\nLearning curve:")
        print(learner.get_learning_curve())
        
    except FileNotFoundError:
        print("Dataset not found. Please ensure train.csv exists.")