trustformers-optim 0.1.1

#![allow(clippy::all)]
use std::collections::HashMap;
use std::time::Instant;
use trustformers_core::TrustformersError;
use trustformers_core::{traits::Optimizer, Tensor};
use trustformers_optim::*;

fn main() -> Result<(), TrustformersError> {
    println!("🚀 TrustformeRS Optimizer Visualization Tools");
    println!("===========================================");
    println!("📊 Generating performance analysis visualizations");

    generate_performance_comparison_chart()?;
    generate_convergence_analysis()?;
    generate_memory_usage_chart()?;
    generate_scaling_analysis()?;
    generate_optimizer_heatmap()?;

    println!("\n🎉 Visualization Tools Completed!");
    println!("   ✅ Performance comparison charts generated");
    println!("   📈 Convergence analysis available");
    println!("   💾 Memory usage visualizations ready");
    println!("   📊 Comprehensive optimizer analysis complete");

    Ok(())
}

fn generate_performance_comparison_chart() -> Result<(), TrustformersError> {
    println!("\n📊 Generating Performance Comparison Chart");
    println!("{}", "─".repeat(50));

    let param_sizes = vec![1000, 5000, 10000, 25000, 50000];
    let iterations = 50;

    // Data collection for visualization
    let mut performance_data = HashMap::new();

    for param_size in &param_sizes {
        println!("📈 Benchmarking {} parameters...", param_size);

        let mut params_adam = Tensor::randn(&[*param_size])?;
        let mut params_adamw = Tensor::randn(&[*param_size])?;
        let mut params_sgd = Tensor::randn(&[*param_size])?;
        let mut params_bge = Tensor::randn(&[*param_size])?;
        let gradients = Tensor::randn(&[*param_size])?;

        // Benchmark Adam
        let mut adam = Adam::new(0.001, (0.9, 0.999), 1e-8, 0.0);
        let start = Instant::now();
        for _ in 0..iterations {
            adam.update(&mut params_adam, &gradients)?;
            adam.step();
        }
        let adam_time = start.elapsed();

        // Benchmark AdamW
        let mut adamw = AdamW::new(0.001, (0.9, 0.999), 1e-8, 0.01);
        let start = Instant::now();
        for _ in 0..iterations {
            adamw.update(&mut params_adamw, &gradients)?;
            adamw.step();
        }
        let adamw_time = start.elapsed();

        // Benchmark SGD
        let mut sgd = SGD::new(0.01, 0.9, 0.0, false);
        let start = Instant::now();
        for _ in 0..iterations {
            sgd.update(&mut params_sgd, &gradients)?;
            sgd.step();
        }
        let sgd_time = start.elapsed();

        // Benchmark BGE-Adam
        let mut bge_adam = BGEAdam::new(0.001, (0.9, 0.999), 1e-8, 0.01, 0.1, 0.05, 0.05);
        let start = Instant::now();
        for _ in 0..iterations {
            bge_adam.update(&mut params_bge, &gradients)?;
            bge_adam.step();
        }
        let bge_time = start.elapsed();

        // Store data for visualization
        let size_data = vec![
            ("Adam", adam_time.as_nanos() as f64 / iterations as f64),
            ("AdamW", adamw_time.as_nanos() as f64 / iterations as f64),
            ("SGD", sgd_time.as_nanos() as f64 / iterations as f64),
            ("BGE-Adam", bge_time.as_nanos() as f64 / iterations as f64),
        ];
        performance_data.insert(*param_size, size_data);
    }

    // Generate ASCII chart
    println!("\n📊 Performance Comparison Chart (nanoseconds per iteration):");
    println!("{}", "─".repeat(80));

    // Header
    println!(
        "{:>8} | {:>12} | {:>12} | {:>12} | {:>12}",
        "Params", "Adam", "AdamW", "SGD", "BGE-Adam"
    );
    println!("{}", "─".repeat(80));

    for param_size in &param_sizes {
        if let Some(data) = performance_data.get(param_size) {
            let adam_ns = data
                .iter()
                .find(|(name, _)| *name == "Adam")
                .expect("Adam optimizer data should exist")
                .1;
            let adamw_ns = data
                .iter()
                .find(|(name, _)| *name == "AdamW")
                .expect("AdamW optimizer data should exist")
                .1;
            let sgd_ns = data
                .iter()
                .find(|(name, _)| *name == "SGD")
                .expect("SGD optimizer data should exist")
                .1;
            let bge_ns = data
                .iter()
                .find(|(name, _)| *name == "BGE-Adam")
                .expect("BGE-Adam optimizer data should exist")
                .1;

            println!(
                "{:>8} | {:>12.0} | {:>12.0} | {:>12.0} | {:>12.0}",
                param_size, adam_ns, adamw_ns, sgd_ns, bge_ns
            );
        }
    }

    println!("{}", "─".repeat(80));

    // Generate performance scaling visualization
    println!("\n📈 Performance Scaling Visualization:");
    for param_size in &param_sizes {
        if let Some(data) = performance_data.get(param_size) {
            let adam_ns = data
                .iter()
                .find(|(name, _)| *name == "Adam")
                .expect("Adam optimizer data should exist")
                .1;
            let scale = (adam_ns / 1000.0).min(50.0) as usize; // Scale for visualization
            let bar = "█".repeat(scale);
            println!(
                "{:>8} params: {} {:.1}µs",
                param_size,
                bar,
                adam_ns / 1000.0
            );
        }
    }

    println!("✅ Performance comparison chart generated");
    Ok(())
}

fn generate_convergence_analysis() -> Result<(), TrustformersError> {
    println!("\n📊 Generating Convergence Analysis");
    println!("{}", "─".repeat(50));

    // Simulate training loss convergence for different optimizers
    let total_steps = 200;
    let mut loss_history = HashMap::new();

    // Simulate different convergence patterns
    let optimizers = vec![
        ("Adam", generate_adam_convergence(total_steps)),
        ("AdamW", generate_adamw_convergence(total_steps)),
        ("SGD", generate_sgd_convergence(total_steps)),
        ("BGE-Adam", generate_bge_convergence(total_steps)),
    ];

    for (name, losses) in optimizers {
        loss_history.insert(name, losses);
    }

    // Generate convergence visualization
    println!("\n📈 Loss Convergence Analysis (simulated training):");
    println!("{}", "─".repeat(70));

    // Show key convergence milestones
    let milestones = vec![0, 25, 50, 100, 150, 199];

    println!(
        "{:>8} | {:>8} | {:>8} | {:>8} | {:>8}",
        "Step", "Adam", "AdamW", "SGD", "BGE-Adam"
    );
    println!("{}", "─".repeat(50));

    for &step in &milestones {
        let adam_loss = loss_history.get("Adam").expect("Adam loss history should exist")[step];
        let adamw_loss = loss_history.get("AdamW").expect("AdamW loss history should exist")[step];
        let sgd_loss = loss_history.get("SGD").expect("SGD loss history should exist")[step];
        let bge_loss =
            loss_history.get("BGE-Adam").expect("BGE-Adam loss history should exist")[step];

        println!(
            "{:>8} | {:>8.4} | {:>8.4} | {:>8.4} | {:>8.4}",
            step, adam_loss, adamw_loss, sgd_loss, bge_loss
        );
    }

    // Generate ASCII plot for Adam convergence
    println!("\n📉 Adam Loss Curve (ASCII plot):");
    let adam_losses = loss_history.get("Adam").expect("Adam loss history should exist");
    let max_loss = adam_losses.iter().fold(0.0f32, |a, &b| a.max(b));

    for (i, &loss) in adam_losses.iter().enumerate() {
        if i % 20 == 0 {
            // Show every 20th step
            let normalized = ((1.0 - loss / max_loss) * 40.0) as usize;
            let spaces = " ".repeat(normalized);
            let marker = "●";
            println!("Step {:>3}: {}{}  ({:.4})", i, spaces, marker, loss);
        }
    }

    // Convergence speed analysis
    println!("\n🎯 Convergence Speed Analysis:");
    for (optimizer, losses) in &loss_history {
        let initial_loss = losses[0];
        let final_loss = losses[losses.len() - 1];
        let improvement = ((initial_loss - final_loss) / initial_loss) * 100.0;

        // Find step where loss drops below 50% of initial
        let target_loss = initial_loss * 0.5;
        let convergence_step =
            losses.iter().position(|&loss| loss < target_loss).unwrap_or(total_steps);

        println!(
            "   {} | {:>6.1}% improvement | 50% reduction at step {}",
            optimizer, improvement, convergence_step
        );
    }

    println!("✅ Convergence analysis generated");
    Ok(())
}

fn generate_memory_usage_chart() -> Result<(), TrustformersError> {
    println!("\n📊 Generating Memory Usage Chart");
    println!("{}", "─".repeat(50));

    let param_counts = vec![1000, 10000, 100000, 500000, 1000000];

    println!("\n💾 Memory Usage Comparison (MB):");
    println!("{}", "─".repeat(70));
    println!(
        "{:>10} | {:>10} | {:>10} | {:>10} | {:>10}",
        "Parameters", "Adam", "Adam-8bit", "AdamW", "ZeRO-3"
    );
    println!("{}", "─".repeat(70));

    for &param_count in &param_counts {
        // Memory calculations (in MB)
        let param_memory = (param_count * 4) as f64 / 1_048_576.0; // 4 bytes per f32

        // Regular Adam: params + momentum + variance
        let adam_memory = param_memory * 3.0;

        // Adam-8bit: params + quantized states (1 byte each) + overhead
        let adam_8bit_memory = param_memory + (param_count * 2) as f64 / 1_048_576.0 + 0.001; // 1MB overhead

        // AdamW: same as Adam
        let adamw_memory = adam_memory;

        // ZeRO-3: everything sharded across 8 GPUs
        let zero3_memory = adam_memory / 8.0;

        println!(
            "{:>10} | {:>10.2} | {:>10.2} | {:>10.2} | {:>10.2}",
            param_count, adam_memory, adam_8bit_memory, adamw_memory, zero3_memory
        );
    }

    println!("{}", "─".repeat(70));

    // Memory efficiency visualization
    println!("\n📊 Memory Efficiency Bars (1M parameters):");
    let param_count = 1_000_000;
    let base_memory = (param_count * 4 * 3) as f64 / 1_048_576.0; // Adam baseline

    let optimizers = vec![
        ("Adam", base_memory, "████████████████████"),
        ("Adam-8bit", base_memory * 0.25, "█████"),
        ("AdamW", base_memory, "████████████████████"),
        ("ZeRO-1", base_memory * 0.6, "████████████"),
        ("ZeRO-2", base_memory * 0.35, "███████"),
        ("ZeRO-3", base_memory * 0.125, "██"),
    ];

    for (name, memory, bar) in optimizers {
        println!("{:>8}: {} {:.1} MB", name, bar, memory);
    }

    println!("✅ Memory usage chart generated");
    Ok(())
}

fn generate_scaling_analysis() -> Result<(), TrustformersError> {
    println!("\n📊 Generating Scaling Analysis");
    println!("{}", "─".repeat(50));

    // Simulate distributed training scaling
    let node_counts = vec![1, 2, 4, 8, 16, 32];

    println!("\n🔗 Distributed Training Scaling Analysis:");
    println!("{}", "─".repeat(60));
    println!(
        "{:>6} | {:>12} | {:>12} | {:>12} | {:>8}",
        "Nodes", "Throughput", "Efficiency", "Comm Cost", "Speedup"
    );
    println!("{}", "─".repeat(60));

    let base_throughput = 1000.0; // samples/sec on single node

    for &nodes in &node_counts {
        // Simulate scaling efficiency with communication overhead
        let ideal_throughput = base_throughput * nodes as f64;
        let comm_overhead = if nodes == 1 {
            0.0
        } else {
            0.05 * (nodes as f64).log2() // Communication overhead grows with log(nodes)
        };
        let actual_throughput = ideal_throughput * (1.0 - comm_overhead);
        let efficiency = (actual_throughput / ideal_throughput) * 100.0;
        let comm_cost = comm_overhead * 100.0;
        let speedup = actual_throughput / base_throughput;

        println!(
            "{:>6} | {:>12.0} | {:>11.1}% | {:>11.1}% | {:>7.1}x",
            nodes, actual_throughput, efficiency, comm_cost, speedup
        );
    }

    println!("{}", "─".repeat(60));

    // Scaling visualization
    println!("\n📈 Scaling Efficiency Visualization:");
    for &nodes in &node_counts {
        let efficiency = if nodes == 1 {
            100.0
        } else {
            let comm_overhead = 0.05 * (nodes as f64).log2();
            (1.0 - comm_overhead) * 100.0
        };
        let bar_length = (efficiency / 5.0) as usize; // Scale to 20 chars max
        let bar = "█".repeat(bar_length);
        println!("{:>2} nodes: {} {:.1}%", nodes, bar, efficiency);
    }

    println!("✅ Scaling analysis generated");
    Ok(())
}

fn generate_optimizer_heatmap() -> Result<(), TrustformersError> {
    println!("\n📊 Generating Optimizer Performance Heatmap");
    println!("{}", "─".repeat(50));

    // Create a performance heatmap across different scenarios
    let scenarios = [
        ("Small Model", "1M params"),
        ("Medium Model", "100M params"),
        ("Large Model", "1B+ params"),
        ("Vision Task", "CNN training"),
        ("NLP Task", "Transformer"),
        ("Memory Limited", "8GB GPU"),
    ];

    let optimizers = vec!["Adam", "AdamW", "SGD", "LAMB", "BGE-Adam", "8bit-Adam"];

    println!("\n🔥 Optimizer Performance Heatmap:");
    println!("   Legend: ██ Excellent  ▓▓ Good  ░░ Fair  ·· Poor");
    println!("{}", "─".repeat(70));

    // Header
    print!("{:>15} |", "Scenario");
    for opt in &optimizers {
        print!(" {:^8} |", opt);
    }
    println!();
    println!("{}", "─".repeat(70));

    // Performance ratings (simulated based on typical use cases)
    let ratings: &[&[&str]] = &[
        &["██", "██", "▓▓", "▓▓", "░░", "▓▓"], // Small Model
        &["██", "██", "▓▓", "██", "▓▓", "██"], // Medium Model
        &["▓▓", "██", "░░", "██", "▓▓", "██"], // Large Model
        &["██", "██", "██", "▓▓", "▓▓", "▓▓"], // Vision Task
        &["██", "██", "▓▓", "██", "██", "▓▓"], // NLP Task
        &["░░", "░░", "▓▓", "▓▓", "░░", "██"], // Memory Limited
    ];

    for (i, (scenario, description)) in scenarios.iter().enumerate() {
        print!("{:>15} |", scenario);
        for &rating in ratings[i].iter() {
            print!(" {:^8} |", rating);
        }
        println!(" {}", description);
    }

    println!("{}", "─".repeat(70));

    // Recommendations
    println!("\n💡 Optimizer Recommendations:");
    println!("   🎯 General Purpose: Adam/AdamW (reliable, well-tested)");
    println!("   🚀 Large Models: LAMB (better scaling), 8bit-Adam (memory efficient)");
    println!("   💾 Memory Constrained: 8bit-Adam, ZeRO optimizers");
    println!("   ⚡ Fast Convergence: BGE-Adam (entropy-weighted), AdamW (decoupled weight decay)");
    println!("   📱 Mobile/Edge: SGD (lightweight), quantized optimizers");

    println!("✅ Optimizer heatmap generated");
    Ok(())
}

// Helper functions for convergence simulation
fn generate_adam_convergence(steps: usize) -> Vec<f32> {
    let mut losses = Vec::new();
    let mut loss = 2.0;
    for i in 0..steps {
        // Adam: fast initial convergence, then slower
        let rate = 0.02 * (1.0 - (i as f32 / steps as f32).powf(0.5));
        loss *= 1.0 - rate;
        losses.push(loss);
    }
    losses
}

fn generate_adamw_convergence(steps: usize) -> Vec<f32> {
    let mut losses = Vec::new();
    let mut loss = 2.0;
    for i in 0..steps {
        // AdamW: similar to Adam but slightly better final convergence
        let rate = 0.022 * (1.0 - (i as f32 / steps as f32).powf(0.5));
        loss *= 1.0 - rate;
        losses.push(loss);
    }
    losses
}

fn generate_sgd_convergence(steps: usize) -> Vec<f32> {
    let mut losses = Vec::new();
    let mut loss = 2.0;
    for i in 0..steps {
        // SGD: slower initial convergence, steady improvement
        let rate = 0.015 * (1.0 - (i as f32 / steps as f32).powf(0.3));
        loss *= 1.0 - rate;
        losses.push(loss);
    }
    losses
}

fn generate_bge_convergence(steps: usize) -> Vec<f32> {
    let mut losses = Vec::new();
    let mut loss = 2.0;
    for i in 0..steps {
        // BGE-Adam: adaptive convergence with entropy weighting
        let rate = 0.025 * (1.0 - (i as f32 / steps as f32).powf(0.6));
        loss *= 1.0 - rate;
        losses.push(loss);
    }
    losses
}