//! Hardware-Aware Neural Architecture Search Example
//!
//! This example demonstrates how to perform hardware-aware NAS using rtx-nas,
//! including:
//!
//! - Device profiling and latency prediction
//! - Multi-objective optimization (accuracy vs latency/memory)
//! - Pareto frontier construction
//! - FairNAS constraints for balanced architecture evaluation
//!
//! # Running
//!
//! ```bash
//! cargo run --example hardware_aware_search -p rtx-nas
//! ```

use rtx_nas::{
    Result,
    algorithms::{
        PCDARTS, PCDARTSConfig, RandomSearch, RandomSearchConfig,
        fairness::{FairnessConfig, FairnessTracker},
    },
    hardware::{
        cost_model::compute_cost,
        device::{CommonDevices, DeviceProfile},
        latency::{LatencyPredictor, LookupTablePredictor},
    },
    search::{MultiObjective, ObjectiveScorer, ParetoEntry, ParetoFrontier},
    search_space::{CellConfig, DARTSSearchSpace, SearchSpace},
};
use rtx_tensor::Device;
use std::time::Instant;

fn main() -> Result<()> {
    println!("=== Hardware-Aware Neural Architecture Search ===\n");

    // 1. Select target devices
    println!("Step 1: Setting up target devices...");
    let devices = vec![
        CommonDevices::rtx_3090(),
        CommonDevices::t4(),
        CommonDevices::mobile_arm(),
    ];

    for device in &devices {
        println!(
            "  - {}: {:.1} TFLOPS, {:.0} GB memory",
            device.name, device.peak_tflops_fp32, device.memory_gb
        );
    }
    println!();

    // 2. Create search space
    println!("Step 2: Creating DARTS search space...");
    let search_space = DARTSSearchSpace::default()?;
    println!(
        "  Search space: {} total choices\n",
        search_space.num_choices()
    );

    // 3. Initialize latency predictor
    println!("Step 3: Setting up latency prediction...");
    let predictor = LookupTablePredictor::new();
    println!(
        "  Predictor confidence: {:.0}%\n",
        predictor.confidence() * 100.0
    );

    // 4. Configure multi-objective optimization
    println!("Step 4: Configuring multi-objective optimization...");
    let objectives = MultiObjective::mobile_optimized();
    println!("  Accuracy weight: {:.2}", objectives.accuracy_weight);
    println!("  Latency weight:  {:.2}", objectives.latency_weight);
    println!("  Memory weight:   {:.2}", objectives.memory_weight);
    println!("  Params weight:   {:.2}\n", objectives.params_weight);

    let scorer = ObjectiveScorer::new(objectives)?;

    // 5. Initialize fairness tracking
    println!("Step 5: Setting up fairness tracking...");
    let fairness_config = FairnessConfig::default()
        .with_tracking_window(500)
        .with_threshold(0.85);
    let mut fairness_tracker = FairnessTracker::new(fairness_config);
    println!("  Fairness threshold: 85%\n");

    // 6. Sample architectures using random search
    println!("Step 6: Sampling architectures with random search...");
    let start = Instant::now();

    let config = RandomSearchConfig::new(20);
    let mut search = RandomSearch::new(config)?;
    search.sample(&search_space)?;

    println!(
        "  Sampled {} architectures in {:?}\n",
        search.num_samples(),
        start.elapsed()
    );

    // 7. Evaluate architectures and build Pareto frontier
    println!("Step 7: Evaluating architectures...\n");
    let mut frontier = ParetoFrontier::with_max_size(10);

    let primary_device = &devices[0]; // Use RTX 3090 as primary evaluation device

    for (i, arch) in search.samples().iter().enumerate() {
        // Compute cost
        let mut cost = compute_cost(arch)?;

        // Predict latency
        let latency = predictor.predict(arch, primary_device)?;
        cost.estimated_latency_ms = Some(latency);

        // Simulate accuracy (in practice, this comes from actual training)
        let accuracy = simulate_accuracy(arch.channels, arch.num_layers);

        // Track fairness (simulate operation usage from the architecture)
        for cell in &arch.cells {
            for edge in cell.edges() {
                if let Some(op_type) = cell.get_operation(&edge) {
                    fairness_tracker.track_optimization(edge, op_type);
                }
            }
        }

        // Compute weighted score
        let score = scorer.score(accuracy, &cost);

        // Add to Pareto frontier
        let entry = ParetoEntry::new(arch.clone(), cost.clone(), accuracy);
        let added = frontier.add(entry);

        println!(
            "  Arch {:2}: acc={:.3}, latency={:.2}ms, params={:>6}, FLOPs={:>8}, score={:.3} {}",
            i + 1,
            accuracy,
            latency,
            format_number(cost.params),
            format_number(cost.flops),
            score,
            if added { "[Pareto]" } else { "" }
        );
    }

    // 8. Display Pareto frontier
    println!("\n=== Pareto Frontier ===");
    println!("Non-dominated architectures that offer best tradeoffs:\n");

    for (i, entry) in frontier.entries().iter().enumerate() {
        println!(
            "  {}. {} - Accuracy: {:.3}, Latency: {:.2}ms, Params: {}",
            i + 1,
            entry.arch_id(),
            entry.accuracy,
            entry.cost.estimated_latency_ms.unwrap_or(0.0),
            format_number(entry.cost.params)
        );
    }

    // 9. Display fairness report
    println!("\n=== Fairness Analysis ===");
    let report = fairness_tracker.generate_report();
    println!(
        "  Overall fairness score: {:.1}%",
        report.overall_score * 100.0
    );
    println!("  Total operations tracked: {}", report.total_iterations);

    if report.underrepresented.is_empty() {
        println!("  All operations fairly represented!");
    } else {
        println!("  Underrepresented operations:");
        for op in &report.underrepresented {
            println!("    - {:?}", op);
        }
    }

    // 10. Demonstrate PC-DARTS with hardware awareness
    println!("\n=== PC-DARTS with Hardware Constraints ===\n");
    demonstrate_pcdarts(primary_device, &predictor)?;

    println!("\n=== Hardware-Aware Search Complete ===");
    Ok(())
}

/// Simulate accuracy based on architecture complexity
fn simulate_accuracy(channels: usize, num_layers: usize) -> f32 {
    // Simple heuristic: more channels and layers = higher accuracy (with noise)
    let base = 0.7 + (channels as f32 / 100.0).min(0.15) + (num_layers as f32 / 50.0).min(0.1);
    let noise = (rand::random::<f32>() - 0.5) * 0.05;
    (base + noise).clamp(0.5, 0.98)
}

/// Format large numbers for display
fn format_number(n: u64) -> String {
    if n >= 1_000_000_000 {
        format!("{:.1}B", n as f64 / 1_000_000_000.0)
    } else if n >= 1_000_000 {
        format!("{:.1}M", n as f64 / 1_000_000.0)
    } else if n >= 1_000 {
        format!("{:.1}K", n as f64 / 1_000.0)
    } else {
        n.to_string()
    }
}

/// Demonstrate PC-DARTS with hardware constraints
fn demonstrate_pcdarts(device: &DeviceProfile, predictor: &LookupTablePredictor) -> Result<()> {
    let compute_device = Device::cuda(0).unwrap_or(Device::default());

    println!("Creating PC-DARTS with partial channel connections...");
    println!("  Target device: {}", device.name);
    println!("  Channel fraction: 12.5% (memory efficient)");
    println!("  Edge normalization: enabled\n");

    // Create PC-DARTS configuration
    let mut config = PCDARTSConfig::default();
    config.edge_normalization = true;

    let cell_configs = vec![CellConfig::default_darts()];
    let mut pcdarts = PCDARTS::new(config, cell_configs, &compute_device)?;

    // Simulate a few search steps
    println!("Running architecture search steps...");
    for epoch in 0..3 {
        // In practice, train_loss and valid_loss come from actual training
        let train_loss = 0.5 - (epoch as f32 * 0.1);
        let valid_loss = 0.6 - (epoch as f32 * 0.08);

        pcdarts.step(train_loss, valid_loss, None, None)?;

        // Derive architecture and estimate cost
        let arch = pcdarts.derive_architecture()?;
        let cost = compute_cost(&arch)?;
        let latency = predictor.predict(&arch, device)?;

        println!(
            "  Epoch {}: train_loss={:.3}, valid_loss={:.3}, est_latency={:.2}ms",
            epoch + 1,
            train_loss,
            valid_loss,
            latency
        );
    }

    // Derive final architecture
    let final_arch = pcdarts.derive_architecture()?;
    let final_cost = compute_cost(&final_arch)?;
    let final_latency = predictor.predict(&final_arch, device)?;

    println!("\nDerived architecture:");
    println!("  ID: {}", final_arch.id);
    println!("  Cells: {}", final_arch.num_cells());
    println!("  Channels: {}", final_arch.channels);
    println!("  Estimated FLOPs: {}", format_number(final_cost.flops));
    println!("  Estimated params: {}", format_number(final_cost.params));
    println!(
        "  Estimated latency on {}: {:.2}ms",
        device.name, final_latency
    );

    Ok(())
}