//! Integration tests for rtx-nas

use rtx_nas::{
    algorithms::{DARTS, DARTSConfig, RandomSearch, RandomSearchConfig},
    search_space::{CellConfig, DARTSSearchSpace, SearchSpace},
};
use rtx_tensor::Device;

#[test]
fn test_random_search_workflow() {
    // Create search space
    let search_space = DARTSSearchSpace::default().expect("Failed to create search space");

    // Configure random search
    let config = RandomSearchConfig::new(5);
    let mut search = RandomSearch::new(config).expect("Failed to create random search");

    // Sample architectures
    search
        .sample(&search_space)
        .expect("Failed to sample architectures");

    // Verify samples
    assert_eq!(search.num_samples(), 5);
    assert!(!search.samples().is_empty());

    // Test get_best
    let scores = vec![0.5, 0.3, 0.9, 0.4, 0.6];
    let best = search.get_best(&scores).expect("Failed to get best");
    assert_eq!(best.id, search.samples()[2].id);

    // Test get_topk
    let topk = search.get_topk(&scores, 3).expect("Failed to get topk");
    assert_eq!(topk.len(), 3);
}

#[test]
fn test_darts_workflow() {
    let device = Device::cuda(0).unwrap_or(Device::default());

    // Configure DARTS
    let config = DARTSConfig::default();
    let cell_configs = vec![CellConfig::default_darts()];

    // Create DARTS instance
    let mut darts = DARTS::new(config, cell_configs, &device).expect("Failed to create DARTS");

    // Check initial state
    assert_eq!(darts.current_epoch(), 0);
    assert!(darts.is_warmup());

    // Perform warmup step
    darts
        .step(0.5, 0.6, None, None)
        .expect("Failed to perform step");
    assert_eq!(darts.current_epoch(), 1);

    // Get probabilities
    let probs = darts.get_probabilities();
    assert_eq!(probs.len(), 1);

    // Check all probabilities sum to 1.0
    for prob_map in probs {
        for prob_vec in prob_map.values() {
            let sum: f32 = prob_vec.iter().sum();
            assert!((sum - 1.0).abs() < 1e-5);
        }
    }

    // Derive architecture
    let arch = darts
        .derive_architecture()
        .expect("Failed to derive architecture");
    assert!(arch.validate().is_ok());
}

#[test]
fn test_search_space_encode_decode() {
    let search_space = DARTSSearchSpace::default().expect("Failed to create search space");

    // Sample an architecture
    let arch = search_space
        .sample()
        .expect("Failed to sample architecture");

    // Encode
    let encoding = search_space
        .encode(&arch)
        .expect("Failed to encode architecture");
    assert!(!encoding.is_empty());

    // Decode
    let decoded = search_space
        .decode(&encoding)
        .expect("Failed to decode architecture");

    // Both should have the same structure
    assert_eq!(decoded.num_cells(), arch.num_cells());
}

#[test]
fn test_cell_structure() {
    use rtx_nas::search_space::{Cell, Edge, OperationType};

    // Create a cell
    let config = CellConfig::default_darts();
    let mut cell = Cell::new(config).expect("Failed to create cell");

    // Set operations on edges
    let edge = Edge::new(0, 2);
    cell.set_operation(edge, OperationType::Conv3x3)
        .expect("Failed to set operation");

    // Query cell structure
    assert!(cell.total_nodes() > 0);
    assert!(!cell.edges().is_empty());
    assert_eq!(cell.get_operation(&edge), Some(OperationType::Conv3x3));
}

#[test]
fn test_darts_multi_step() {
    let device = Device::cuda(0).unwrap_or(Device::default());

    // Create DARTS with no warmup
    let mut config = DARTSConfig::default();
    config.warmup_epochs = 0;

    let cell_configs = vec![CellConfig::default_darts()];
    let mut darts = DARTS::new(config, cell_configs, &device).expect("Failed to create DARTS");

    // Perform multiple steps
    for _ in 0..5 {
        darts
            .step(0.5, 0.6, None, None)
            .expect("Failed to perform step");
    }

    assert_eq!(darts.current_epoch(), 5);
    assert!(!darts.is_warmup());

    // Architecture should still be derivable
    let arch = darts
        .derive_architecture()
        .expect("Failed to derive architecture");
    assert!(arch.validate().is_ok());
}

#[test]
fn test_architecture_validation() {
    use rtx_nas::search_space::{Architecture, Cell};

    let cell = Cell::default_darts().expect("Failed to create cell");

    // Valid architecture
    let arch = Architecture::new("test_arch".into(), vec![cell.clone()], 16, 8);
    assert!(arch.validate().is_ok());

    // Invalid: no cells
    let arch = Architecture::new("test_arch".into(), vec![], 16, 8);
    assert!(arch.validate().is_err());

    // Invalid: zero channels
    let arch = Architecture::new("test_arch".into(), vec![cell.clone()], 0, 8);
    assert!(arch.validate().is_err());

    // Invalid: zero layers
    let arch = Architecture::new("test_arch".into(), vec![cell], 16, 0);
    assert!(arch.validate().is_err());
}

#[test]
fn test_search_space_num_choices() {
    let space = DARTSSearchSpace::default().expect("Failed to create search space");
    let num_choices = space.num_choices();

    // With default config: 2 cells × 14 edges × 9 operations = 252
    assert_eq!(num_choices, 252);
}

// =============================================================================
// PC-DARTS Integration Tests
// =============================================================================

#[test]
fn test_pcdarts_workflow() {
    use rtx_nas::algorithms::{PCDARTS, PCDARTSConfig};

    let device = Device::cuda(0).unwrap_or(Device::default());

    // Configure PC-DARTS with partial channels
    let config = PCDARTSConfig::default();
    let cell_configs = vec![CellConfig::default_darts()];

    // Create PC-DARTS instance
    let mut pcdarts =
        PCDARTS::new(config, cell_configs, &device).expect("Failed to create PC-DARTS");

    // Verify initial state
    assert_eq!(pcdarts.current_epoch(), 0);
    assert!(pcdarts.is_warmup());

    // Perform warmup step
    pcdarts
        .step(0.5, 0.6, None, None)
        .expect("Failed to perform PC-DARTS step");
    assert_eq!(pcdarts.current_epoch(), 1);

    // Derive architecture
    let arch = pcdarts
        .derive_architecture()
        .expect("Failed to derive PC-DARTS architecture");
    assert!(arch.validate().is_ok());
}

#[test]
fn test_pcdarts_channel_fraction() {
    use rtx_nas::algorithms::{DARTSConfig, PCDARTSConfig};

    // Test default channel fraction
    let config = PCDARTSConfig::default();
    assert_eq!(config.channel_fraction, 0.125);

    // Test custom channel fraction
    let base_config = DARTSConfig::default();
    let config = PCDARTSConfig::new(base_config, 0.25);
    assert_eq!(config.channel_fraction, 0.25);

    // Validate config
    assert!(config.validate().is_ok());
}

#[test]
fn test_pcdarts_edge_normalization() {
    use rtx_nas::algorithms::{PCDARTS, PCDARTSConfig};

    let device = Device::cuda(0).unwrap_or(Device::default());

    // Enable edge normalization
    let mut config = PCDARTSConfig::default();
    config.edge_normalization = true;

    let cell_configs = vec![CellConfig::default_darts()];
    let mut pcdarts =
        PCDARTS::new(config, cell_configs, &device).expect("Failed to create PC-DARTS");

    // Perform multiple steps
    for _ in 0..3 {
        pcdarts
            .step(0.5, 0.6, None, None)
            .expect("Failed to perform step");
    }

    // Architecture should be derivable
    let arch = pcdarts
        .derive_architecture()
        .expect("Failed to derive architecture");
    assert!(arch.validate().is_ok());
}

// =============================================================================
// Hardware-Aware NAS Integration Tests
// =============================================================================

#[test]
fn test_hardware_device_profiles() {
    use rtx_nas::hardware::device::{CommonDevices, DeviceType};

    // Test common device profiles
    let rtx3090 = CommonDevices::rtx_3090();
    assert_eq!(rtx3090.name, "RTX 3090");
    assert!(matches!(rtx3090.device_type, DeviceType::CUDA));
    assert!(rtx3090.peak_tflops_fp32 > 30.0);

    let a100 = CommonDevices::a100_40gb();
    assert_eq!(a100.name, "A100 40GB");
    assert!(a100.memory_gb >= 40.0);

    let mobile = CommonDevices::mobile_arm();
    assert!(matches!(mobile.device_type, DeviceType::Mobile));
}

#[test]
fn test_latency_prediction() {
    use rtx_nas::hardware::{
        device::CommonDevices,
        latency::{LatencyPredictor, LookupTablePredictor},
    };
    use rtx_nas::search_space::{Architecture, Cell};

    // Create test architecture
    let cell = Cell::default_darts().expect("Failed to create cell");
    let arch = Architecture::new("test_arch".to_string(), vec![cell], 16, 8);

    // Create predictor
    let predictor = LookupTablePredictor::new();
    let device = CommonDevices::rtx_3090();

    // Predict latency
    let latency = predictor
        .predict(&arch, &device)
        .expect("Failed to predict latency");
    assert!(latency >= 0.0);

    // Confidence should be medium for uncalibrated predictor
    assert_eq!(predictor.confidence(), 0.5);
}

#[test]
fn test_cost_model() {
    use rtx_nas::hardware::cost_model::compute_cost;
    use rtx_nas::search_space::{Architecture, Cell, Edge, OperationType};

    // Create test architecture with operations set
    let mut cell = Cell::default_darts().expect("Failed to create cell");

    // Set operations on some edges
    cell.set_operation(Edge::new(0, 2), OperationType::Conv3x3)
        .expect("Failed to set operation");
    cell.set_operation(Edge::new(1, 2), OperationType::Conv3x3)
        .expect("Failed to set operation");
    cell.set_operation(Edge::new(0, 3), OperationType::MaxPool3x3)
        .expect("Failed to set operation");

    let arch = Architecture::new("test_arch".to_string(), vec![cell.clone(), cell], 32, 8);

    // Compute cost
    let cost = compute_cost(&arch).expect("Failed to compute cost");

    // Verify cost components - with operations set, we should have non-zero costs
    assert!(cost.flops > 0, "Expected non-zero FLOPs");
    assert!(cost.params > 0, "Expected non-zero params");
    assert!(cost.memory_mb >= 0.0);
    assert!(cost.model_size_mb >= 0.0);
}

#[test]
fn test_cost_model_fields() {
    use rtx_nas::hardware::cost_model::ArchitectureCost;

    let cost = ArchitectureCost {
        flops: 1000,
        params: 100,
        memory_mb: 10.0,
        model_size_mb: 1.0,
        estimated_latency_ms: Some(5.0),
        cell_costs: vec![],
    };

    // Test that cost fields are accessible
    assert_eq!(cost.flops, 1000);
    assert_eq!(cost.params, 100);
    assert_eq!(cost.memory_mb, 10.0);
    assert_eq!(cost.model_size_mb, 1.0);
    assert_eq!(cost.estimated_latency_ms, Some(5.0));
}

// =============================================================================
// Multi-Objective Search Integration Tests
// =============================================================================

#[test]
fn test_multi_objective_scoring() {
    use rtx_nas::hardware::cost_model::ArchitectureCost;
    use rtx_nas::search::{MultiObjective, ObjectiveScorer};

    // Create multi-objective config
    let objectives = MultiObjective::mobile_optimized();
    let scorer = ObjectiveScorer::new(objectives).expect("Failed to create scorer");

    // Create cost
    let cost = ArchitectureCost {
        flops: 500_000_000,
        params: 5_000_000,
        memory_mb: 500.0,
        model_size_mb: 20.0,
        estimated_latency_ms: Some(50.0),
        cell_costs: vec![],
    };

    // Score architecture (scorer uses internal reference values)
    let accuracy = 0.9;
    let score = scorer.score(accuracy, &cost);

    // Score should be positive
    assert!(score > 0.0);
}

#[test]
fn test_pareto_frontier() {
    use rtx_nas::hardware::cost_model::ArchitectureCost;
    use rtx_nas::search::{ParetoEntry, ParetoFrontier};
    use rtx_nas::search_space::{Architecture, Cell};

    let mut frontier = ParetoFrontier::new();

    // Add entries with different accuracy/cost tradeoffs
    let cell = Cell::default_darts().expect("Failed to create cell");

    for i in 0..5 {
        let arch = Architecture::new(format!("arch_{}", i), vec![cell.clone()], 16, 8);
        let cost = ArchitectureCost {
            flops: 1000 * (5 - i) as u64,
            params: 100 * (5 - i) as u64,
            memory_mb: 10.0 * (5 - i) as f32,
            model_size_mb: 1.0 * (5 - i) as f32,
            estimated_latency_ms: Some(5.0 * (5 - i) as f32),
            cell_costs: vec![],
        };

        let accuracy = 0.5 + 0.1 * i as f32;
        let entry = ParetoEntry::new(arch, cost, accuracy);
        frontier.add(entry);
    }

    // Frontier should contain non-dominated solutions
    assert!(!frontier.is_empty());
    assert!(frontier.len() <= 5);
}

#[test]
fn test_pareto_dominance() {
    use rtx_nas::hardware::cost_model::ArchitectureCost;
    use rtx_nas::search::ParetoEntry;
    use rtx_nas::search_space::{Architecture, Cell};

    let cell = Cell::default_darts().expect("Failed to create cell");

    // Create two entries
    let arch1 = Architecture::new("arch_1".to_string(), vec![cell.clone()], 16, 8);
    let cost1 = ArchitectureCost {
        flops: 1000,
        params: 100,
        memory_mb: 10.0,
        model_size_mb: 1.0,
        estimated_latency_ms: Some(5.0),
        cell_costs: vec![],
    };
    let entry1 = ParetoEntry::new(arch1, cost1, 0.9);

    let arch2 = Architecture::new("arch_2".to_string(), vec![cell], 16, 8);
    let cost2 = ArchitectureCost {
        flops: 2000,     // Worse
        params: 200,     // Worse
        memory_mb: 20.0, // Worse
        model_size_mb: 2.0,
        estimated_latency_ms: Some(10.0),
        cell_costs: vec![],
    };
    let entry2 = ParetoEntry::new(arch2, cost2, 0.8); // Lower accuracy too

    // entry1 dominates entry2
    assert!(entry1.dominates(&entry2));
    assert!(!entry2.dominates(&entry1));
}

// =============================================================================
// FairNAS Integration Tests
// =============================================================================

#[test]
fn test_fairness_tracking() {
    use rtx_nas::algorithms::fairness::{FairnessConfig, FairnessTracker};
    use rtx_nas::search_space::{Edge, OperationType};

    // Create fairness tracker
    let config = FairnessConfig::default();
    let mut tracker = FairnessTracker::new(config);

    // Simulate fair training
    let edge = Edge::new(0, 2);
    for _ in 0..100 {
        tracker.track_optimization(edge, OperationType::Conv3x3);
        tracker.track_optimization(edge, OperationType::Conv5x5);
        tracker.track_optimization(edge, OperationType::MaxPool3x3);
    }

    // Check fairness
    let score = tracker.compute_fairness_score();
    assert!(score > 0.9, "Expected high fairness for balanced training");
    assert!(tracker.is_fair());
}

#[test]
fn test_fairness_reweighting() {
    use rtx_nas::algorithms::fairness::{FairnessConfig, FairnessTracker};
    use rtx_nas::search_space::{Edge, OperationType};

    // Create tracker with auto-rebalancing
    let mut config = FairnessConfig::default();
    config.auto_rebalance = true;
    config.rebalance_interval = 10;

    let mut tracker = FairnessTracker::new(config);

    // Create imbalanced training
    let edge = Edge::new(0, 2);
    for _ in 0..100 {
        tracker.track_optimization(edge, OperationType::Conv3x3);
    }
    for _ in 0..10 {
        tracker.track_optimization(edge, OperationType::Conv5x5);
    }

    // Get weights - underrepresented ops should have higher weight
    let weight_conv3 = tracker.get_sampling_weight(OperationType::Conv3x3);
    let weight_conv5 = tracker.get_sampling_weight(OperationType::Conv5x5);

    assert!(
        weight_conv5 > weight_conv3,
        "Conv5x5 should have higher weight"
    );
}

#[test]
fn test_fairness_report() {
    use rtx_nas::algorithms::fairness::{FairnessConfig, FairnessTracker};
    use rtx_nas::search_space::{Edge, OperationType};

    let config = FairnessConfig::default();
    let mut tracker = FairnessTracker::new(config);

    // Add some tracking
    let edge = Edge::new(0, 2);
    for _ in 0..50 {
        tracker.track_optimization(edge, OperationType::Conv3x3);
        tracker.track_optimization(edge, OperationType::Conv5x5);
    }

    // Generate report
    let report = tracker.generate_report();

    assert_eq!(report.total_iterations, 100);
    assert!(report.overall_score > 0.9);
    assert!(report.is_fair(0.8));
}

// =============================================================================
// End-to-End Integration Tests
// =============================================================================

#[test]
fn test_hardware_aware_search() {
    use rtx_nas::algorithms::{RandomSearch, RandomSearchConfig};
    use rtx_nas::hardware::{
        cost_model::compute_cost,
        device::CommonDevices,
        latency::{LatencyPredictor, LookupTablePredictor},
    };
    use rtx_nas::search::{MultiObjective, ObjectiveScorer, ParetoEntry, ParetoFrontier};

    // Create search space
    let search_space = DARTSSearchSpace::default().expect("Failed to create search space");

    // Configure random search
    let config = RandomSearchConfig::new(10);
    let mut search = RandomSearch::new(config).expect("Failed to create random search");

    // Sample architectures
    search.sample(&search_space).expect("Failed to sample");

    // Set up hardware evaluation
    let device = CommonDevices::rtx_3090();
    let predictor = LookupTablePredictor::new();

    // Set up multi-objective scoring
    let objectives = MultiObjective::default();
    let _scorer = ObjectiveScorer::new(objectives).expect("Failed to create scorer");

    // Build Pareto frontier
    let mut frontier = ParetoFrontier::new();

    for arch in search.samples() {
        // Compute cost
        let mut cost = compute_cost(arch).expect("Failed to compute cost");

        // Predict latency
        let latency = predictor
            .predict(arch, &device)
            .expect("Failed to predict latency");
        cost.estimated_latency_ms = Some(latency);

        // Simulate accuracy (in real use, this would come from validation)
        let accuracy = 0.5 + rand::random::<f32>() * 0.4;

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

    // Frontier should have some entries
    assert!(!frontier.is_empty());

    // Get best by accuracy (objective index 0)
    let best = frontier.best_for_objective(0);
    assert!(best.is_some());
}