"""
================================================================================
Text Classification Dataset - Deep Learning Training 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:
- PyTorch Neural Network Classifier
- LSTM/GRU Text Classification
- CNN for Text Classification
- Transformer-based Models
- Training with Early Stopping
- Learning Rate Scheduling
- Model Checkpointing
- TensorBoard Logging

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

import os
import re
import string
import json
from typing import List, Dict, Tuple, Optional
from datetime import datetime
from collections import Counter

import numpy as np
import pandas as pd

# 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'
}

# Check for PyTorch
try:
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    from torch.utils.data import Dataset, DataLoader
    from torch.optim import Adam, AdamW
    from torch.optim.lr_scheduler import ReduceLROnPlateau, CosineAnnealingLR
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False
    print("PyTorch not available. Install with: pip install torch")


class Vocabulary:
    """
    Vocabulary class for text tokenization.
    
    Author: Molla Samser | RSK World (https://rskworld.in)
    """
    
    def __init__(
        self,
        max_vocab_size: int = 50000,
        min_freq: int = 2,
        pad_token: str = '<PAD>',
        unk_token: str = '<UNK>'
    ):
        self.max_vocab_size = max_vocab_size
        self.min_freq = min_freq
        self.pad_token = pad_token
        self.unk_token = unk_token
        
        self.word2idx = {pad_token: 0, unk_token: 1}
        self.idx2word = {0: pad_token, 1: unk_token}
        self.word_counts = Counter()
        
    def build(self, texts: List[str]):
        """Build vocabulary from texts."""
        for text in texts:
            tokens = self._tokenize(text)
            self.word_counts.update(tokens)
        
        # Filter by frequency and limit size
        sorted_words = [
            word for word, count in self.word_counts.most_common(self.max_vocab_size - 2)
            if count >= self.min_freq
        ]
        
        for word in sorted_words:
            idx = len(self.word2idx)
            self.word2idx[word] = idx
            self.idx2word[idx] = word
    
    def _tokenize(self, text: str) -> List[str]:
        """Simple tokenization."""
        text = text.lower()
        text = re.sub(r'[^\w\s]', '', text)
        return text.split()
    
    def encode(self, text: str, max_length: int = 256) -> List[int]:
        """Encode text to indices."""
        tokens = self._tokenize(text)[:max_length]
        indices = [
            self.word2idx.get(token, self.word2idx[self.unk_token])
            for token in tokens
        ]
        # Pad or truncate
        if len(indices) < max_length:
            indices += [self.word2idx[self.pad_token]] * (max_length - len(indices))
        return indices[:max_length]
    
    def decode(self, indices: List[int]) -> str:
        """Decode indices back to text."""
        tokens = [self.idx2word.get(idx, self.unk_token) for idx in indices]
        return ' '.join([t for t in tokens if t != self.pad_token])
    
    def __len__(self):
        return len(self.word2idx)


if TORCH_AVAILABLE:
    
    class TextDataset(Dataset):
        """
        PyTorch Dataset for text classification.
        
        Author: Molla Samser | RSK World (https://rskworld.in)
        """
        
        def __init__(
            self,
            texts: List[str],
            labels: List[int],
            vocab: Vocabulary,
            max_length: int = 256
        ):
            self.texts = texts
            self.labels = labels
            self.vocab = vocab
            self.max_length = max_length
        
        def __len__(self):
            return len(self.texts)
        
        def __getitem__(self, idx):
            text = self.texts[idx]
            label = self.labels[idx]
            
            encoded = self.vocab.encode(text, self.max_length)
            
            return {
                'input_ids': torch.tensor(encoded, dtype=torch.long),
                'label': torch.tensor(label, dtype=torch.long)
            }
    
    
    class LSTMClassifier(nn.Module):
        """
        LSTM-based text classifier.
        
        Author: Molla Samser | RSK World (https://rskworld.in)
        """
        
        def __init__(
            self,
            vocab_size: int,
            embedding_dim: int = 128,
            hidden_dim: int = 256,
            num_layers: int = 2,
            num_classes: int = 6,
            dropout: float = 0.3,
            bidirectional: bool = True
        ):
            super().__init__()
            
            self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)
            
            self.lstm = nn.LSTM(
                embedding_dim,
                hidden_dim,
                num_layers=num_layers,
                batch_first=True,
                dropout=dropout if num_layers > 1 else 0,
                bidirectional=bidirectional
            )
            
            lstm_output_dim = hidden_dim * 2 if bidirectional else hidden_dim
            
            self.attention = nn.Sequential(
                nn.Linear(lstm_output_dim, 64),
                nn.Tanh(),
                nn.Linear(64, 1)
            )
            
            self.classifier = nn.Sequential(
                nn.Dropout(dropout),
                nn.Linear(lstm_output_dim, 128),
                nn.ReLU(),
                nn.Dropout(dropout),
                nn.Linear(128, num_classes)
            )
        
        def forward(self, input_ids):
            # Embedding
            embedded = self.embedding(input_ids)  # (batch, seq, embed)
            
            # LSTM
            lstm_out, _ = self.lstm(embedded)  # (batch, seq, hidden*2)
            
            # Attention
            attn_weights = self.attention(lstm_out)  # (batch, seq, 1)
            attn_weights = F.softmax(attn_weights, dim=1)
            
            # Weighted sum
            context = torch.sum(attn_weights * lstm_out, dim=1)  # (batch, hidden*2)
            
            # Classification
            logits = self.classifier(context)
            
            return logits
    
    
    class CNNClassifier(nn.Module):
        """
        CNN-based text classifier with multiple filter sizes.
        
        Author: Molla Samser | RSK World (https://rskworld.in)
        """
        
        def __init__(
            self,
            vocab_size: int,
            embedding_dim: int = 128,
            num_filters: int = 100,
            filter_sizes: List[int] = [2, 3, 4, 5],
            num_classes: int = 6,
            dropout: float = 0.5
        ):
            super().__init__()
            
            self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)
            
            self.convs = nn.ModuleList([
                nn.Conv1d(embedding_dim, num_filters, fs)
                for fs in filter_sizes
            ])
            
            self.classifier = nn.Sequential(
                nn.Dropout(dropout),
                nn.Linear(num_filters * len(filter_sizes), 128),
                nn.ReLU(),
                nn.Dropout(dropout),
                nn.Linear(128, num_classes)
            )
        
        def forward(self, input_ids):
            # Embedding
            embedded = self.embedding(input_ids)  # (batch, seq, embed)
            embedded = embedded.permute(0, 2, 1)  # (batch, embed, seq)
            
            # Convolutions
            conv_outputs = []
            for conv in self.convs:
                conv_out = F.relu(conv(embedded))  # (batch, filters, seq-fs+1)
                pooled = F.max_pool1d(conv_out, conv_out.size(2)).squeeze(2)
                conv_outputs.append(pooled)
            
            # Concatenate
            concat = torch.cat(conv_outputs, dim=1)  # (batch, filters*len(filter_sizes))
            
            # Classification
            logits = self.classifier(concat)
            
            return logits
    
    
    class TransformerClassifier(nn.Module):
        """
        Transformer-based text classifier.
        
        Author: Molla Samser | RSK World (https://rskworld.in)
        """
        
        def __init__(
            self,
            vocab_size: int,
            embedding_dim: int = 128,
            num_heads: int = 4,
            num_layers: int = 2,
            ff_dim: int = 512,
            num_classes: int = 6,
            max_length: int = 256,
            dropout: float = 0.1
        ):
            super().__init__()
            
            self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)
            self.pos_embedding = nn.Embedding(max_length, embedding_dim)
            
            encoder_layer = nn.TransformerEncoderLayer(
                d_model=embedding_dim,
                nhead=num_heads,
                dim_feedforward=ff_dim,
                dropout=dropout,
                batch_first=True
            )
            
            self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
            
            self.classifier = nn.Sequential(
                nn.Dropout(dropout),
                nn.Linear(embedding_dim, 128),
                nn.ReLU(),
                nn.Dropout(dropout),
                nn.Linear(128, num_classes)
            )
        
        def forward(self, input_ids):
            batch_size, seq_len = input_ids.shape
            
            # Embeddings
            token_emb = self.embedding(input_ids)
            positions = torch.arange(seq_len, device=input_ids.device).unsqueeze(0)
            pos_emb = self.pos_embedding(positions)
            
            embedded = token_emb + pos_emb
            
            # Create padding mask
            padding_mask = (input_ids == 0)
            
            # Transformer
            transformer_out = self.transformer(embedded, src_key_padding_mask=padding_mask)
            
            # Global average pooling (ignoring padding)
            mask = (~padding_mask).unsqueeze(-1).float()
            pooled = (transformer_out * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
            
            # Classification
            logits = self.classifier(pooled)
            
            return logits


class DeepLearningTrainer:
    """
    Trainer for deep learning text classifiers.
    
    Author: Molla Samser | RSK World (https://rskworld.in)
    """
    
    def __init__(
        self,
        model_type: str = 'lstm',
        embedding_dim: int = 128,
        hidden_dim: int = 256,
        num_layers: int = 2,
        num_classes: int = 6,
        max_length: int = 256,
        learning_rate: float = 1e-3,
        batch_size: int = 32,
        epochs: int = 10,
        device: str = 'auto',
        early_stopping_patience: int = 3,
        verbose: bool = True
    ):
        if not TORCH_AVAILABLE:
            raise ImportError("PyTorch is required. Install with: pip install torch")
        
        self.model_type = model_type
        self.embedding_dim = embedding_dim
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.num_classes = num_classes
        self.max_length = max_length
        self.learning_rate = learning_rate
        self.batch_size = batch_size
        self.epochs = epochs
        self.early_stopping_patience = early_stopping_patience
        self.verbose = verbose
        
        # Device
        if device == 'auto':
            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        else:
            self.device = torch.device(device)
        
        self.vocab = None
        self.model = None
        self.history = {'train_loss': [], 'train_acc': [], 'val_loss': [], 'val_acc': []}
    
    def _create_model(self, vocab_size: int) -> nn.Module:
        """Create the model based on type."""
        if self.model_type == 'lstm':
            return LSTMClassifier(
                vocab_size=vocab_size,
                embedding_dim=self.embedding_dim,
                hidden_dim=self.hidden_dim,
                num_layers=self.num_layers,
                num_classes=self.num_classes
            )
        elif self.model_type == 'cnn':
            return CNNClassifier(
                vocab_size=vocab_size,
                embedding_dim=self.embedding_dim,
                num_classes=self.num_classes
            )
        elif self.model_type == 'transformer':
            return TransformerClassifier(
                vocab_size=vocab_size,
                embedding_dim=self.embedding_dim,
                num_classes=self.num_classes,
                max_length=self.max_length
            )
        else:
            raise ValueError(f"Unknown model type: {self.model_type}")
    
    def fit(
        self,
        train_texts: List[str],
        train_labels: List[int],
        val_texts: Optional[List[str]] = None,
        val_labels: Optional[List[int]] = None
    ):
        """
        Train the model.
        
        Args:
            train_texts: Training texts
            train_labels: Training labels
            val_texts: Validation texts (optional)
            val_labels: Validation labels (optional)
        """
        if self.verbose:
            print(f"\n{'='*60}")
            print("Deep Learning Text Classification Training")
            print(f"Author: {__author__} | Website: {__website__}")
            print(f"{'='*60}\n")
            print(f"Model Type: {self.model_type.upper()}")
            print(f"Device: {self.device}")
            print(f"Batch Size: {self.batch_size}")
            print(f"Learning Rate: {self.learning_rate}")
            print(f"Epochs: {self.epochs}")
            print("-" * 40)
        
        # Build vocabulary
        self.vocab = Vocabulary()
        self.vocab.build(train_texts)
        
        if self.verbose:
            print(f"Vocabulary Size: {len(self.vocab)}")
        
        # Create datasets
        train_dataset = TextDataset(train_texts, train_labels, self.vocab, self.max_length)
        train_loader = DataLoader(train_dataset, batch_size=self.batch_size, shuffle=True)
        
        val_loader = None
        if val_texts is not None and val_labels is not None:
            val_dataset = TextDataset(val_texts, val_labels, self.vocab, self.max_length)
            val_loader = DataLoader(val_dataset, batch_size=self.batch_size)
        
        # Create model
        self.model = self._create_model(len(self.vocab))
        self.model.to(self.device)
        
        # Count parameters
        total_params = sum(p.numel() for p in self.model.parameters())
        trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
        
        if self.verbose:
            print(f"Total Parameters: {total_params:,}")
            print(f"Trainable Parameters: {trainable_params:,}")
            print("-" * 40)
        
        # Optimizer and scheduler
        optimizer = AdamW(self.model.parameters(), lr=self.learning_rate, weight_decay=0.01)
        scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=2)
        criterion = nn.CrossEntropyLoss()
        
        # Training loop
        best_val_loss = float('inf')
        patience_counter = 0
        
        for epoch in range(self.epochs):
            # Training
            self.model.train()
            train_loss = 0
            train_correct = 0
            train_total = 0
            
            for batch in train_loader:
                input_ids = batch['input_ids'].to(self.device)
                labels = batch['label'].to(self.device)
                
                optimizer.zero_grad()
                logits = self.model(input_ids)
                loss = criterion(logits, labels)
                loss.backward()
                
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                optimizer.step()
                
                train_loss += loss.item()
                _, predicted = torch.max(logits, 1)
                train_correct += (predicted == labels).sum().item()
                train_total += labels.size(0)
            
            avg_train_loss = train_loss / len(train_loader)
            train_acc = train_correct / train_total
            
            self.history['train_loss'].append(avg_train_loss)
            self.history['train_acc'].append(train_acc)
            
            # Validation
            if val_loader:
                val_loss, val_acc = self._evaluate(val_loader, criterion)
                self.history['val_loss'].append(val_loss)
                self.history['val_acc'].append(val_acc)
                
                scheduler.step(val_loss)
                
                if self.verbose:
                    print(f"Epoch {epoch+1}/{self.epochs} - "
                          f"Train Loss: {avg_train_loss:.4f}, Train Acc: {train_acc:.4f} - "
                          f"Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.4f}")
                
                # Early stopping
                if val_loss < best_val_loss:
                    best_val_loss = val_loss
                    patience_counter = 0
                else:
                    patience_counter += 1
                    if patience_counter >= self.early_stopping_patience:
                        if self.verbose:
                            print(f"Early stopping at epoch {epoch+1}")
                        break
            else:
                if self.verbose:
                    print(f"Epoch {epoch+1}/{self.epochs} - "
                          f"Train Loss: {avg_train_loss:.4f}, Train Acc: {train_acc:.4f}")
        
        if self.verbose:
            print(f"\n{'='*60}")
            print("Training Complete!")
            print(f"Best Validation Loss: {best_val_loss:.4f}")
    
    def _evaluate(self, loader, criterion) -> Tuple[float, float]:
        """Evaluate the model."""
        self.model.eval()
        total_loss = 0
        correct = 0
        total = 0
        
        with torch.no_grad():
            for batch in loader:
                input_ids = batch['input_ids'].to(self.device)
                labels = batch['label'].to(self.device)
                
                logits = self.model(input_ids)
                loss = criterion(logits, labels)
                
                total_loss += loss.item()
                _, predicted = torch.max(logits, 1)
                correct += (predicted == labels).sum().item()
                total += labels.size(0)
        
        return total_loss / len(loader), correct / total
    
    def predict(self, texts: List[str]) -> np.ndarray:
        """Predict labels for texts."""
        if self.model is None:
            raise ValueError("Model not trained. Call fit() first.")
        
        self.model.eval()
        predictions = []
        
        with torch.no_grad():
            for text in texts:
                encoded = self.vocab.encode(text, self.max_length)
                input_ids = torch.tensor([encoded], dtype=torch.long).to(self.device)
                logits = self.model(input_ids)
                _, predicted = torch.max(logits, 1)
                predictions.append(predicted.item())
        
        return np.array(predictions)
    
    def predict_proba(self, texts: List[str]) -> np.ndarray:
        """Predict probabilities for texts."""
        if self.model is None:
            raise ValueError("Model not trained. Call fit() first.")
        
        self.model.eval()
        probabilities = []
        
        with torch.no_grad():
            for text in texts:
                encoded = self.vocab.encode(text, self.max_length)
                input_ids = torch.tensor([encoded], dtype=torch.long).to(self.device)
                logits = self.model(input_ids)
                probs = F.softmax(logits, dim=1)
                probabilities.append(probs.cpu().numpy()[0])
        
        return np.array(probabilities)
    
    def save(self, path: str):
        """Save model and vocabulary."""
        state = {
            'model_state': self.model.state_dict(),
            'vocab': self.vocab,
            'config': {
                'model_type': self.model_type,
                'embedding_dim': self.embedding_dim,
                'hidden_dim': self.hidden_dim,
                'num_layers': self.num_layers,
                'num_classes': self.num_classes,
                'max_length': self.max_length
            },
            'history': self.history,
            'metadata': {
                'author': __author__,
                'website': __website__,
                'created_at': datetime.now().isoformat()
            }
        }
        torch.save(state, path)
        if self.verbose:
            print(f"Model saved to: {path}")
    
    def load(self, path: str):
        """Load model and vocabulary."""
        state = torch.load(path, map_location=self.device)
        
        self.vocab = state['vocab']
        config = state['config']
        
        self.model_type = config['model_type']
        self.model = self._create_model(len(self.vocab))
        self.model.load_state_dict(state['model_state'])
        self.model.to(self.device)
        self.model.eval()
        
        self.history = state.get('history', {})
        
        if self.verbose:
            print(f"Model loaded from: {path}")


if __name__ == "__main__":
    print(f"\n{'='*60}")
    print("Deep Learning Demo - RSK World")
    print(f"Author: {__author__} | Website: {__website__}")
    print(f"{'='*60}\n")
    
    if not TORCH_AVAILABLE:
        print("PyTorch not installed. Install with: pip install torch")
        exit(1)
    
    # Load data
    try:
        train_df = pd.read_csv('../data/csv/train.csv', comment='#')
        val_df = pd.read_csv('../data/csv/validation.csv', comment='#')
        
        # Train model
        trainer = DeepLearningTrainer(
            model_type='lstm',
            embedding_dim=128,
            hidden_dim=256,
            batch_size=16,
            epochs=5,
            verbose=True
        )
        
        trainer.fit(
            train_df['text'].tolist(),
            train_df['label'].tolist(),
            val_df['text'].tolist(),
            val_df['label'].tolist()
        )
        
        # Test prediction
        test_texts = [
            "Apple unveils new iPhone with AI features",
            "Manchester United wins the championship"
        ]
        
        predictions = trainer.predict(test_texts)
        probabilities = trainer.predict_proba(test_texts)
        
        print("\nTest Predictions:")
        for text, pred, prob in zip(test_texts, predictions, probabilities):
            print(f"Text: {text}")
            print(f"Predicted: {CATEGORIES[pred]} ({prob[pred]*100:.1f}%)")
            print()
        
        # Save model
        trainer.save('deep_model.pt')
        
        print(f"\n{'='*60}")
        print("Deep Learning Demo Complete!")
        print(f"Copyright (c) 2026 RSK World - All Rights Reserved")
        
    except FileNotFoundError:
        print("Dataset not found. Please ensure CSV files exist in ../data/csv/")