"""
================================================================================
Text Classification Dataset - Batch Processing & Evaluation Tools
================================================================================
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:
- Batch text classification
- Model evaluation metrics
- Cross-validation support
- Export results to multiple formats
- Performance benchmarking

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

import os
import json
import time
import argparse
from datetime import datetime
from typing import List, Dict, Optional, Tuple, Any

import numpy as np
import pandas as pd
from sklearn.metrics import (
    accuracy_score, precision_score, recall_score, f1_score,
    classification_report, confusion_matrix, roc_auc_score
)
from sklearn.model_selection import cross_val_score, StratifiedKFold

# 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 BatchProcessor:
    """
    Batch text classification processor.
    
    Process large volumes of text data efficiently with
    progress tracking and error handling.
    
    Author: Molla Samser | RSK World (https://rskworld.in)
    """
    
    def __init__(
        self,
        model_path: str,
        batch_size: int = 100,
        verbose: bool = True
    ):
        """
        Initialize batch processor.
        
        Args:
            model_path: Path to trained model file
            batch_size: Number of texts per batch
            verbose: Print progress
        """
        import joblib
        
        self.model_path = model_path
        self.batch_size = batch_size
        self.verbose = verbose
        
        # Load model
        self.model_data = joblib.load(model_path)
        self.pipeline = self.model_data['pipeline']
        
        if self.verbose:
            print(f"Model loaded: {self.model_data.get('algorithm', 'unknown')}")
    
    def preprocess(self, text: str) -> str:
        """Basic text preprocessing."""
        import re
        import string
        
        text = text.lower()
        text = re.sub(r'https?://\S+|www\.\S+', '', text)
        text = text.translate(str.maketrans('', '', string.punctuation))
        text = ' '.join(text.split())
        return text
    
    def process_batch(self, texts: List[str]) -> List[Dict]:
        """
        Process a batch of texts.
        
        Args:
            texts: List of texts to classify
            
        Returns:
            List of prediction results
        """
        # Preprocess
        processed = [self.preprocess(t) for t in texts]
        
        # Predict
        predictions = self.pipeline.predict(processed)
        
        # Get probabilities if available
        probabilities = None
        if hasattr(self.pipeline, 'predict_proba'):
            try:
                probabilities = self.pipeline.predict_proba(processed)
            except:
                pass
        
        results = []
        for i, (text, pred) in enumerate(zip(texts, predictions)):
            result = {
                'text': text[:200] + '...' if len(text) > 200 else text,
                'predicted_label': int(pred),
                'predicted_category': CATEGORIES[pred]
            }
            
            if probabilities is not None:
                result['confidence'] = float(probabilities[i][pred])
                result['probabilities'] = {
                    CATEGORIES[j]: float(p) for j, p in enumerate(probabilities[i])
                }
            
            results.append(result)
        
        return results
    
    def process_file(
        self,
        input_path: str,
        output_path: str,
        text_column: str = 'text'
    ) -> Dict:
        """
        Process an entire file of texts.
        
        Args:
            input_path: Input CSV/JSON file
            output_path: Output file path
            text_column: Column name for text
            
        Returns:
            Processing summary
        """
        start_time = time.time()
        
        # Load data
        if input_path.endswith('.json'):
            df = pd.read_json(input_path)
        else:
            df = pd.read_csv(input_path, comment='#')
        
        texts = df[text_column].tolist()
        total = len(texts)
        
        if self.verbose:
            print(f"Processing {total} documents...")
        
        # Process in batches
        all_results = []
        for i in range(0, total, self.batch_size):
            batch = texts[i:i + self.batch_size]
            results = self.process_batch(batch)
            all_results.extend(results)
            
            if self.verbose:
                progress = min(i + self.batch_size, total)
                print(f"  Processed {progress}/{total} ({100*progress/total:.1f}%)")
        
        # Create output dataframe
        result_df = pd.DataFrame(all_results)
        
        # Save
        if output_path.endswith('.json'):
            result_df.to_json(output_path, orient='records', indent=2)
        else:
            result_df.to_csv(output_path, index=False)
        
        elapsed = time.time() - start_time
        
        summary = {
            'input_file': input_path,
            'output_file': output_path,
            'total_documents': total,
            'processing_time_seconds': round(elapsed, 2),
            'documents_per_second': round(total / elapsed, 2),
            'category_distribution': result_df['predicted_category'].value_counts().to_dict()
        }
        
        if self.verbose:
            print(f"\nProcessing complete!")
            print(f"  Time: {elapsed:.2f}s ({total/elapsed:.1f} docs/sec)")
            print(f"  Output: {output_path}")
        
        return summary


class ModelEvaluator:
    """
    Comprehensive model evaluation toolkit.
    
    Provides detailed metrics, confusion matrices, and
    cross-validation analysis.
    
    Author: Molla Samser | RSK World (https://rskworld.in)
    """
    
    def __init__(self, class_names: List[str] = None):
        """
        Initialize evaluator.
        
        Args:
            class_names: List of class names
        """
        self.class_names = class_names or list(CATEGORIES.values())
    
    def evaluate(
        self,
        y_true: np.ndarray,
        y_pred: np.ndarray,
        y_prob: Optional[np.ndarray] = None
    ) -> Dict:
        """
        Compute comprehensive evaluation metrics.
        
        Args:
            y_true: True labels
            y_pred: Predicted labels
            y_prob: Prediction probabilities (optional)
            
        Returns:
            Dictionary of metrics
        """
        metrics = {
            'accuracy': float(accuracy_score(y_true, y_pred)),
            'precision_macro': float(precision_score(y_true, y_pred, average='macro', zero_division=0)),
            'precision_weighted': float(precision_score(y_true, y_pred, average='weighted', zero_division=0)),
            'recall_macro': float(recall_score(y_true, y_pred, average='macro', zero_division=0)),
            'recall_weighted': float(recall_score(y_true, y_pred, average='weighted', zero_division=0)),
            'f1_macro': float(f1_score(y_true, y_pred, average='macro', zero_division=0)),
            'f1_weighted': float(f1_score(y_true, y_pred, average='weighted', zero_division=0)),
        }
        
        # Per-class metrics
        precision_per_class = precision_score(y_true, y_pred, average=None, zero_division=0)
        recall_per_class = recall_score(y_true, y_pred, average=None, zero_division=0)
        f1_per_class = f1_score(y_true, y_pred, average=None, zero_division=0)
        
        metrics['per_class'] = {
            self.class_names[i]: {
                'precision': float(precision_per_class[i]),
                'recall': float(recall_per_class[i]),
                'f1': float(f1_per_class[i])
            }
            for i in range(len(self.class_names)) if i < len(precision_per_class)
        }
        
        # Confusion matrix
        cm = confusion_matrix(y_true, y_pred)
        metrics['confusion_matrix'] = cm.tolist()
        
        # ROC-AUC if probabilities available
        if y_prob is not None:
            try:
                metrics['roc_auc_macro'] = float(
                    roc_auc_score(y_true, y_prob, multi_class='ovr', average='macro')
                )
                metrics['roc_auc_weighted'] = float(
                    roc_auc_score(y_true, y_prob, multi_class='ovr', average='weighted')
                )
            except:
                pass
        
        return metrics
    
    def cross_validate(
        self,
        model,
        X: np.ndarray,
        y: np.ndarray,
        cv: int = 5,
        scoring: str = 'accuracy'
    ) -> Dict:
        """
        Perform cross-validation evaluation.
        
        Args:
            model: Sklearn model/pipeline
            X: Features (can be text if pipeline includes vectorizer)
            y: Labels
            cv: Number of folds
            scoring: Scoring metric
            
        Returns:
            Cross-validation results
        """
        skf = StratifiedKFold(n_splits=cv, shuffle=True, random_state=42)
        
        scores = cross_val_score(model, X, y, cv=skf, scoring=scoring)
        
        return {
            'cv_folds': cv,
            'scoring': scoring,
            'mean_score': float(np.mean(scores)),
            'std_score': float(np.std(scores)),
            'fold_scores': [float(s) for s in scores],
            'min_score': float(np.min(scores)),
            'max_score': float(np.max(scores))
        }
    
    def generate_report(
        self,
        y_true: np.ndarray,
        y_pred: np.ndarray,
        output_path: str = 'evaluation_report.json'
    ) -> str:
        """
        Generate comprehensive evaluation report.
        
        Args:
            y_true: True labels
            y_pred: Predicted labels
            output_path: Path to save report
            
        Returns:
            Path to saved report
        """
        metrics = self.evaluate(y_true, y_pred)
        
        report = {
            'metadata': {
                'author': __author__,
                'website': __website__,
                'generated_at': datetime.now().isoformat(),
                'total_samples': len(y_true),
                'num_classes': len(self.class_names)
            },
            'overall_metrics': {
                'accuracy': metrics['accuracy'],
                'f1_macro': metrics['f1_macro'],
                'f1_weighted': metrics['f1_weighted']
            },
            'detailed_metrics': metrics,
            'classification_report': classification_report(
                y_true, y_pred,
                target_names=self.class_names[:len(np.unique(y_true))],
                output_dict=True
            )
        }
        
        with open(output_path, 'w') as f:
            json.dump(report, f, indent=2)
        
        return output_path


class PerformanceBenchmark:
    """
    Benchmark model performance across different configurations.
    
    Author: Molla Samser | RSK World (https://rskworld.in)
    """
    
    def __init__(self):
        """Initialize benchmark."""
        self.results = []
    
    def benchmark_inference(
        self,
        model,
        texts: List[str],
        iterations: int = 3
    ) -> Dict:
        """
        Benchmark inference speed.
        
        Args:
            model: Model/pipeline for predictions
            texts: Sample texts
            iterations: Number of iterations
            
        Returns:
            Benchmark results
        """
        times = []
        
        for _ in range(iterations):
            start = time.time()
            model.predict(texts)
            elapsed = time.time() - start
            times.append(elapsed)
        
        n = len(texts)
        
        return {
            'num_samples': n,
            'iterations': iterations,
            'total_time_mean': float(np.mean(times)),
            'total_time_std': float(np.std(times)),
            'time_per_sample_ms': float(np.mean(times) / n * 1000),
            'throughput_per_second': float(n / np.mean(times))
        }
    
    def compare_models(
        self,
        models: Dict[str, Any],
        X_test: np.ndarray,
        y_test: np.ndarray
    ) -> pd.DataFrame:
        """
        Compare multiple models.
        
        Args:
            models: Dictionary of name -> model
            X_test: Test features
            y_test: Test labels
            
        Returns:
            Comparison dataframe
        """
        evaluator = ModelEvaluator()
        results = []
        
        for name, model in models.items():
            print(f"Evaluating {name}...")
            
            # Predict
            start = time.time()
            y_pred = model.predict(X_test)
            inference_time = time.time() - start
            
            # Metrics
            metrics = evaluator.evaluate(y_test, y_pred)
            
            results.append({
                'model': name,
                'accuracy': metrics['accuracy'],
                'f1_macro': metrics['f1_macro'],
                'precision_macro': metrics['precision_macro'],
                'recall_macro': metrics['recall_macro'],
                'inference_time_s': inference_time,
                'samples_per_second': len(y_test) / inference_time
            })
        
        return pd.DataFrame(results).sort_values('f1_macro', ascending=False)


def main():
    """Main entry point for batch processing CLI."""
    parser = argparse.ArgumentParser(
        description='Batch Text Classification Processor - RSK World'
    )
    
    subparsers = parser.add_subparsers(dest='command', help='Commands')
    
    # Batch process command
    batch_parser = subparsers.add_parser('process', help='Process a file')
    batch_parser.add_argument('--input', '-i', required=True, help='Input file path')
    batch_parser.add_argument('--output', '-o', required=True, help='Output file path')
    batch_parser.add_argument('--model', '-m', required=True, help='Model file path')
    batch_parser.add_argument('--text-column', '-t', default='text', help='Text column name')
    batch_parser.add_argument('--batch-size', '-b', type=int, default=100)
    
    # Evaluate command
    eval_parser = subparsers.add_parser('evaluate', help='Evaluate predictions')
    eval_parser.add_argument('--predictions', '-p', required=True, help='Predictions file')
    eval_parser.add_argument('--ground-truth', '-g', required=True, help='Ground truth file')
    eval_parser.add_argument('--output', '-o', default='evaluation_report.json')
    
    args = parser.parse_args()
    
    print(f"\n{'='*60}")
    print("Batch Processor - RSK World")
    print(f"Author: {__author__} | Website: {__website__}")
    print(f"{'='*60}\n")
    
    if args.command == 'process':
        processor = BatchProcessor(args.model, args.batch_size)
        summary = processor.process_file(args.input, args.output, args.text_column)
        print(f"\nSummary: {json.dumps(summary, indent=2)}")
    
    elif args.command == 'evaluate':
        pred_df = pd.read_csv(args.predictions)
        truth_df = pd.read_csv(args.ground_truth, comment='#')
        
        evaluator = ModelEvaluator()
        report_path = evaluator.generate_report(
            truth_df['label'].values,
            pred_df['predicted_label'].values,
            args.output
        )
        print(f"Report saved to: {report_path}")
    
    else:
        parser.print_help()


if __name__ == '__main__':
    main()