"""
================================================================================
Text Classification Dataset - Advanced Visualization 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:
- Word Cloud Generation
- Category Distribution Charts
- Text Length Analysis
- Confusion Matrix Heatmaps
- Training History Plots
- Feature Importance Visualization
- t-SNE Embeddings Visualization

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

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

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.manifold import TSNE
from sklearn.decomposition import PCA

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

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

CATEGORY_COLORS = {
    'Technology': '#3b82f6',
    'Sports': '#22c55e',
    'Politics': '#8b5cf6',
    'Entertainment': '#ec4899',
    'Business': '#f59e0b',
    'Science': '#06b6d4'
}


def set_style():
    """Set consistent plotting style."""
    plt.style.use('seaborn-v0_8-darkgrid')
    plt.rcParams['figure.facecolor'] = '#0f0a1f'
    plt.rcParams['axes.facecolor'] = '#1a1333'
    plt.rcParams['axes.edgecolor'] = '#352d54'
    plt.rcParams['axes.labelcolor'] = '#f8fafc'
    plt.rcParams['text.color'] = '#f8fafc'
    plt.rcParams['xtick.color'] = '#a5a3b8'
    plt.rcParams['ytick.color'] = '#a5a3b8'
    plt.rcParams['grid.color'] = '#352d54'
    plt.rcParams['legend.facecolor'] = '#231d3a'
    plt.rcParams['legend.edgecolor'] = '#352d54'
    plt.rcParams['font.family'] = 'sans-serif'


def generate_wordcloud(
    texts: List[str],
    output_path: str = 'wordcloud.png',
    title: str = 'Word Cloud',
    width: int = 1200,
    height: int = 600,
    background_color: str = '#0f0a1f',
    colormap: str = 'Reds'
):
    """
    Generate word cloud from texts.
    
    Args:
        texts: List of text documents
        output_path: Path to save image
        title: Chart title
        width: Image width
        height: Image height
        background_color: Background color
        colormap: Matplotlib colormap
        
    Author: Molla Samser | RSK World (https://rskworld.in)
    """
    try:
        from wordcloud import WordCloud, STOPWORDS
    except ImportError:
        print("Please install wordcloud: pip install wordcloud")
        return
    
    # Combine all texts
    combined_text = ' '.join(texts)
    
    # Clean text
    combined_text = combined_text.lower()
    combined_text = re.sub(r'[^\w\s]', '', combined_text)
    
    # Generate word cloud
    wordcloud = WordCloud(
        width=width,
        height=height,
        background_color=background_color,
        colormap=colormap,
        stopwords=STOPWORDS,
        max_words=200,
        max_font_size=150,
        random_state=42
    ).generate(combined_text)
    
    # Plot
    set_style()
    fig, ax = plt.subplots(figsize=(15, 8))
    ax.imshow(wordcloud, interpolation='bilinear')
    ax.axis('off')
    ax.set_title(title, fontsize=20, fontweight='bold', color='#f8fafc', pad=20)
    
    # Add watermark
    fig.text(0.99, 0.01, 'RSK World | rskworld.in', fontsize=10, color='#6b6882',
             ha='right', va='bottom', style='italic')
    
    plt.tight_layout()
    plt.savefig(output_path, dpi=150, bbox_inches='tight', facecolor='#0f0a1f')
    plt.close()
    
    print(f"Word cloud saved to: {output_path}")


def generate_wordclouds_by_category(
    df: pd.DataFrame,
    text_column: str = 'text',
    label_column: str = 'label',
    output_dir: str = 'wordclouds'
):
    """
    Generate separate word clouds for each category.
    
    Args:
        df: DataFrame with texts and labels
        text_column: Column name for text
        label_column: Column name for labels
        output_dir: Output directory
        
    Author: Molla Samser | RSK World (https://rskworld.in)
    """
    os.makedirs(output_dir, exist_ok=True)
    
    for label, category in CATEGORIES.items():
        texts = df[df[label_column] == label][text_column].tolist()
        if texts:
            output_path = os.path.join(output_dir, f'wordcloud_{category.lower()}.png')
            generate_wordcloud(
                texts,
                output_path=output_path,
                title=f'{category} - Word Cloud',
                colormap='Blues' if category == 'Technology' else 
                         'Greens' if category == 'Sports' else
                         'Purples' if category == 'Politics' else
                         'RdPu' if category == 'Entertainment' else
                         'YlOrBr' if category == 'Business' else 'BuGn'
            )


def plot_category_distribution(
    df: pd.DataFrame,
    label_column: str = 'label',
    output_path: str = 'category_distribution.png'
):
    """
    Plot category distribution as pie and bar charts.
    
    Args:
        df: DataFrame with labels
        label_column: Column name for labels
        output_path: Path to save image
        
    Author: Molla Samser | RSK World (https://rskworld.in)
    """
    set_style()
    
    # Count categories
    counts = df[label_column].map(CATEGORIES).value_counts()
    colors = [CATEGORY_COLORS[cat] for cat in counts.index]
    
    fig, axes = plt.subplots(1, 2, figsize=(16, 7))
    
    # Bar chart
    bars = axes[0].bar(counts.index, counts.values, color=colors, edgecolor='white', linewidth=1.5)
    axes[0].set_xlabel('Category', fontsize=12)
    axes[0].set_ylabel('Number of Documents', fontsize=12)
    axes[0].set_title('Category Distribution', fontsize=16, fontweight='bold')
    axes[0].tick_params(axis='x', rotation=45)
    
    # Add value labels
    for bar, val in zip(bars, counts.values):
        axes[0].text(bar.get_x() + bar.get_width()/2, bar.get_height() + 2,
                    str(val), ha='center', va='bottom', fontsize=11, fontweight='bold')
    
    # Pie chart
    wedges, texts, autotexts = axes[1].pie(
        counts.values,
        labels=counts.index,
        colors=colors,
        autopct='%1.1f%%',
        startangle=90,
        explode=[0.02] * len(counts),
        shadow=True
    )
    axes[1].set_title('Category Proportions', fontsize=16, fontweight='bold')
    
    for autotext in autotexts:
        autotext.set_fontsize(10)
        autotext.set_fontweight('bold')
    
    plt.suptitle('Text Classification Dataset - Category Analysis\nRSK World | rskworld.in',
                 fontsize=18, fontweight='bold', y=1.02)
    
    plt.tight_layout()
    plt.savefig(output_path, dpi=150, bbox_inches='tight', facecolor='#0f0a1f')
    plt.close()
    
    print(f"Category distribution saved to: {output_path}")


def plot_text_length_distribution(
    df: pd.DataFrame,
    text_column: str = 'text',
    label_column: str = 'label',
    output_path: str = 'text_length_distribution.png'
):
    """
    Plot text length distribution by category.
    
    Args:
        df: DataFrame with texts
        text_column: Column name for text
        label_column: Column name for labels
        output_path: Path to save image
        
    Author: Molla Samser | RSK World (https://rskworld.in)
    """
    set_style()
    
    # Calculate lengths
    df = df.copy()
    df['word_count'] = df[text_column].str.split().str.len()
    df['char_count'] = df[text_column].str.len()
    df['category'] = df[label_column].map(CATEGORIES)
    
    fig, axes = plt.subplots(2, 2, figsize=(16, 12))
    
    # Word count histogram
    for cat in CATEGORIES.values():
        data = df[df['category'] == cat]['word_count']
        axes[0, 0].hist(data, bins=30, alpha=0.6, label=cat, color=CATEGORY_COLORS[cat])
    axes[0, 0].set_xlabel('Word Count', fontsize=12)
    axes[0, 0].set_ylabel('Frequency', fontsize=12)
    axes[0, 0].set_title('Word Count Distribution by Category', fontsize=14)
    axes[0, 0].legend(loc='upper right')
    
    # Character count histogram
    for cat in CATEGORIES.values():
        data = df[df['category'] == cat]['char_count']
        axes[0, 1].hist(data, bins=30, alpha=0.6, label=cat, color=CATEGORY_COLORS[cat])
    axes[0, 1].set_xlabel('Character Count', fontsize=12)
    axes[0, 1].set_ylabel('Frequency', fontsize=12)
    axes[0, 1].set_title('Character Count Distribution by Category', fontsize=14)
    axes[0, 1].legend(loc='upper right')
    
    # Box plot - word count
    colors = [CATEGORY_COLORS[CATEGORIES[i]] for i in range(6)]
    bp1 = df.boxplot(column='word_count', by='category', ax=axes[1, 0],
                     patch_artist=True, return_type='dict')
    for patch, color in zip(bp1['word_count']['boxes'], colors):
        patch.set_facecolor(color)
        patch.set_alpha(0.7)
    axes[1, 0].set_xlabel('Category', fontsize=12)
    axes[1, 0].set_ylabel('Word Count', fontsize=12)
    axes[1, 0].set_title('Word Count Box Plot', fontsize=14)
    plt.suptitle('')
    
    # Violin plot - character count
    violin_data = [df[df['category'] == cat]['char_count'].values for cat in CATEGORIES.values()]
    parts = axes[1, 1].violinplot(violin_data, positions=range(len(CATEGORIES)))
    for i, pc in enumerate(parts['bodies']):
        pc.set_facecolor(colors[i])
        pc.set_alpha(0.7)
    axes[1, 1].set_xticks(range(len(CATEGORIES)))
    axes[1, 1].set_xticklabels(CATEGORIES.values(), rotation=45)
    axes[1, 1].set_xlabel('Category', fontsize=12)
    axes[1, 1].set_ylabel('Character Count', fontsize=12)
    axes[1, 1].set_title('Character Count Violin Plot', fontsize=14)
    
    plt.suptitle('Text Length Analysis - RSK World | rskworld.in',
                 fontsize=18, fontweight='bold', y=1.02)
    
    plt.tight_layout()
    plt.savefig(output_path, dpi=150, bbox_inches='tight', facecolor='#0f0a1f')
    plt.close()
    
    print(f"Text length distribution saved to: {output_path}")


def plot_confusion_matrix(
    y_true: np.ndarray,
    y_pred: np.ndarray,
    output_path: str = 'confusion_matrix.png',
    title: str = 'Confusion Matrix',
    normalize: bool = True
):
    """
    Plot confusion matrix heatmap.
    
    Args:
        y_true: True labels
        y_pred: Predicted labels
        output_path: Path to save image
        title: Chart title
        normalize: Whether to normalize
        
    Author: Molla Samser | RSK World (https://rskworld.in)
    """
    from sklearn.metrics import confusion_matrix as cm_func
    
    set_style()
    
    cm = cm_func(y_true, y_pred)
    if normalize:
        cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
    
    fig, ax = plt.subplots(figsize=(12, 10))
    
    sns.heatmap(
        cm,
        annot=True,
        fmt='.2%' if normalize else 'd',
        cmap='Reds',
        xticklabels=CATEGORIES.values(),
        yticklabels=CATEGORIES.values(),
        ax=ax,
        linewidths=0.5,
        linecolor='#352d54',
        cbar_kws={'label': 'Proportion' if normalize else 'Count'}
    )
    
    ax.set_xlabel('Predicted Label', fontsize=14)
    ax.set_ylabel('True Label', fontsize=14)
    ax.set_title(f'{title}\nRSK World | rskworld.in', fontsize=16, fontweight='bold')
    
    plt.tight_layout()
    plt.savefig(output_path, dpi=150, bbox_inches='tight', facecolor='#0f0a1f')
    plt.close()
    
    print(f"Confusion matrix saved to: {output_path}")


def plot_training_history(
    history: Dict[str, List[float]],
    output_path: str = 'training_history.png'
):
    """
    Plot training history (loss and accuracy).
    
    Args:
        history: Dictionary with 'loss', 'val_loss', 'accuracy', 'val_accuracy'
        output_path: Path to save image
        
    Author: Molla Samser | RSK World (https://rskworld.in)
    """
    set_style()
    
    fig, axes = plt.subplots(1, 2, figsize=(16, 6))
    
    epochs = range(1, len(history.get('loss', [])) + 1)
    
    # Loss plot
    if 'loss' in history:
        axes[0].plot(epochs, history['loss'], 'o-', color='#dc2626', 
                    label='Training Loss', linewidth=2, markersize=6)
    if 'val_loss' in history:
        axes[0].plot(epochs, history['val_loss'], 's--', color='#f59e0b',
                    label='Validation Loss', linewidth=2, markersize=6)
    axes[0].set_xlabel('Epoch', fontsize=12)
    axes[0].set_ylabel('Loss', fontsize=12)
    axes[0].set_title('Training and Validation Loss', fontsize=14, fontweight='bold')
    axes[0].legend()
    axes[0].grid(True, alpha=0.3)
    
    # Accuracy plot
    if 'accuracy' in history:
        axes[1].plot(epochs, history['accuracy'], 'o-', color='#22c55e',
                    label='Training Accuracy', linewidth=2, markersize=6)
    if 'val_accuracy' in history:
        axes[1].plot(epochs, history['val_accuracy'], 's--', color='#3b82f6',
                    label='Validation Accuracy', linewidth=2, markersize=6)
    axes[1].set_xlabel('Epoch', fontsize=12)
    axes[1].set_ylabel('Accuracy', fontsize=12)
    axes[1].set_title('Training and Validation Accuracy', fontsize=14, fontweight='bold')
    axes[1].legend()
    axes[1].grid(True, alpha=0.3)
    axes[1].set_ylim(0, 1)
    
    plt.suptitle('Model Training History - RSK World | rskworld.in',
                 fontsize=18, fontweight='bold', y=1.02)
    
    plt.tight_layout()
    plt.savefig(output_path, dpi=150, bbox_inches='tight', facecolor='#0f0a1f')
    plt.close()
    
    print(f"Training history saved to: {output_path}")


def plot_feature_importance(
    feature_names: List[str],
    importances: np.ndarray,
    top_n: int = 20,
    output_path: str = 'feature_importance.png',
    title: str = 'Top Features'
):
    """
    Plot feature importance bar chart.
    
    Args:
        feature_names: List of feature names
        importances: Feature importance scores
        top_n: Number of top features to show
        output_path: Path to save image
        title: Chart title
        
    Author: Molla Samser | RSK World (https://rskworld.in)
    """
    set_style()
    
    # Get top features
    indices = np.argsort(importances)[-top_n:][::-1]
    top_features = [feature_names[i] for i in indices]
    top_importances = importances[indices]
    
    fig, ax = plt.subplots(figsize=(12, 8))
    
    colors = plt.cm.Reds(np.linspace(0.4, 0.9, len(top_features)))
    bars = ax.barh(range(len(top_features)), top_importances[::-1], color=colors[::-1])
    
    ax.set_yticks(range(len(top_features)))
    ax.set_yticklabels(top_features[::-1])
    ax.set_xlabel('Importance Score', fontsize=12)
    ax.set_title(f'{title}\nRSK World | rskworld.in', fontsize=16, fontweight='bold')
    
    # Add value labels
    for bar, val in zip(bars, top_importances[::-1]):
        ax.text(bar.get_width() + 0.001, bar.get_y() + bar.get_height()/2,
               f'{val:.4f}', va='center', fontsize=9)
    
    plt.tight_layout()
    plt.savefig(output_path, dpi=150, bbox_inches='tight', facecolor='#0f0a1f')
    plt.close()
    
    print(f"Feature importance saved to: {output_path}")


def plot_tsne_embeddings(
    embeddings: np.ndarray,
    labels: np.ndarray,
    output_path: str = 'tsne_embeddings.png',
    perplexity: int = 30
):
    """
    Plot t-SNE visualization of text embeddings.
    
    Args:
        embeddings: Document embeddings (n_samples, n_features)
        labels: Category labels
        output_path: Path to save image
        perplexity: t-SNE perplexity parameter
        
    Author: Molla Samser | RSK World (https://rskworld.in)
    """
    set_style()
    
    # Reduce dimensions with t-SNE
    print("Computing t-SNE embeddings...")
    tsne = TSNE(n_components=2, perplexity=perplexity, random_state=42, n_iter=1000)
    embeddings_2d = tsne.fit_transform(embeddings)
    
    fig, ax = plt.subplots(figsize=(14, 10))
    
    for label, category in CATEGORIES.items():
        mask = labels == label
        ax.scatter(
            embeddings_2d[mask, 0],
            embeddings_2d[mask, 1],
            c=CATEGORY_COLORS[category],
            label=category,
            alpha=0.7,
            s=50,
            edgecolors='white',
            linewidth=0.5
        )
    
    ax.set_xlabel('t-SNE Dimension 1', fontsize=12)
    ax.set_ylabel('t-SNE Dimension 2', fontsize=12)
    ax.set_title('t-SNE Visualization of Document Embeddings\nRSK World | rskworld.in',
                 fontsize=16, fontweight='bold')
    ax.legend(loc='best', framealpha=0.9)
    
    plt.tight_layout()
    plt.savefig(output_path, dpi=150, bbox_inches='tight', facecolor='#0f0a1f')
    plt.close()
    
    print(f"t-SNE visualization saved to: {output_path}")


def generate_all_visualizations(data_dir: str, output_dir: str = 'visualizations'):
    """
    Generate all visualizations from the dataset.
    
    Args:
        data_dir: Path to data directory
        output_dir: Output directory for visualizations
        
    Author: Molla Samser | RSK World (https://rskworld.in)
    """
    os.makedirs(output_dir, exist_ok=True)
    
    print(f"\n{'='*60}")
    print("Generating Visualizations - RSK World")
    print(f"Author: {__author__} | Website: {__website__}")
    print(f"{'='*60}\n")
    
    # Load data
    train_df = pd.read_csv(os.path.join(data_dir, 'csv', 'train.csv'), comment='#')
    
    # Generate visualizations
    print("1. Generating category distribution...")
    plot_category_distribution(
        train_df,
        output_path=os.path.join(output_dir, 'category_distribution.png')
    )
    
    print("2. Generating text length analysis...")
    plot_text_length_distribution(
        train_df,
        output_path=os.path.join(output_dir, 'text_length_distribution.png')
    )
    
    print("3. Generating word clouds...")
    generate_wordcloud(
        train_df['text'].tolist(),
        output_path=os.path.join(output_dir, 'wordcloud_all.png'),
        title='Text Classification Dataset - All Categories'
    )
    
    generate_wordclouds_by_category(
        train_df,
        output_dir=os.path.join(output_dir, 'wordclouds_by_category')
    )
    
    print(f"\n{'='*60}")
    print("All visualizations generated successfully!")
    print(f"Output directory: {output_dir}")


if __name__ == "__main__":
    import sys
    
    if len(sys.argv) > 1:
        data_dir = sys.argv[1]
    else:
        data_dir = '../data'
    
    generate_all_visualizations(data_dir)