aip-sci 0.1.0

use aip::edm::core::{EmotionDataModel, EmotionDataModelTrainer};
use aip::edm::roguelite::core::RogueliteEdm;
use aip::director::core::{InteractionStrategy, InteractionState, EmotionStats, InteractionStrategyTrainer};
use aip::director::roguelite::core::RogueliteDirector;
use candle_core::Device;
use std::collections::HashMap;
use std::time::{Duration, Instant};

fn get_model_dir() -> std::path::PathBuf {
    let mut path = std::env::current_dir().unwrap();
    path.push("models");
    if !path.exists() {
        std::fs::create_dir_all(&path).unwrap();
    }
    path
}

fn create_features() -> HashMap<u32, f32> {
    let mut features = HashMap::new();
    for i in 0..15 {
        features.insert(i, 0.5);
    }
    features
}

fn create_state() -> InteractionState {
    let mut user_traits = HashMap::new();
    for i in 0..8 {
        user_traits.insert(i, 0.5);
    }
    
    let mut env_state = HashMap::new();
    for i in 0..6 {
        env_state.insert(i, 0.5);
    }
    
    InteractionState {
        user_traits,
        env_state,
        emotion: aip::edm::core::EmotionState::new(0.5, 0.5, 0.5),
        emotion_stats: EmotionStats::default(),
    }
}

#[test]
fn benchmark_edm_inference() {
    let device = Device::Cpu;
    let edm = RogueliteEdm::new(device).unwrap();
    let features = create_features();
    
    let warmup = 10;
    for _ in 0..warmup {
        let _ = edm.infer(&features);
    }
    
    let iterations = 1000;
    let start = Instant::now();
    
    for _ in 0..iterations {
        let _ = edm.infer(&features).unwrap();
    }
    
    let total_duration = start.elapsed();
    let avg_duration = total_duration / iterations;
    
    println!("\nEDM Inference Benchmark:");
    println!("  Total time: {:?}", total_duration);
    println!("  Average latency: {:?}", avg_duration);
    println!("  Throughput: {:.2} inferences/sec", iterations as f64 / total_duration.as_secs_f64());
    
    assert!(avg_duration < Duration::from_millis(10), 
        "EDM inference latency {:?} exceeds 10ms target", avg_duration);
}

#[test]
fn benchmark_director_decision() {
    let device = Device::Cpu;
    let director = RogueliteDirector::new(device).unwrap();
    let state = create_state();
    
    let warmup = 10;
    for _ in 0..warmup {
        let _ = director.decide(&state);
    }
    
    let iterations = 1000;
    let start = Instant::now();
    
    for _ in 0..iterations {
        let _ = director.decide(&state).unwrap();
    }
    
    let total_duration = start.elapsed();
    let avg_duration = total_duration / iterations;
    
    println!("\nDirector Decision Benchmark:");
    println!("  Total time: {:?}", total_duration);
    println!("  Average latency: {:?}", avg_duration);
    println!("  Throughput: {:.2} decisions/sec", iterations as f64 / total_duration.as_secs_f64());
    
    assert!(avg_duration < Duration::from_millis(10), 
        "Director decision latency {:?} exceeds 10ms target", avg_duration);
}

#[test]
fn benchmark_full_pipeline() {
    let device = Device::Cpu;
    let edm = RogueliteEdm::new(device.clone()).unwrap();
    let director = RogueliteDirector::new(device).unwrap();
    let features = create_features();
    
    let warmup = 10;
    for _ in 0..warmup {
        let emotion = edm.infer(&features).unwrap();
        let mut state = create_state();
        state.emotion = emotion;
        let _ = director.decide(&state);
    }
    
    let iterations = 1000;
    let start = Instant::now();
    
    for _ in 0..iterations {
        let emotion = edm.infer(&features).unwrap();
        let mut state = create_state();
        state.emotion = emotion;
        let _ = director.decide(&state).unwrap();
    }
    
    let total_duration = start.elapsed();
    let avg_duration = total_duration / iterations;
    
    println!("\nFull Pipeline Benchmark (EDM + Director):");
    println!("  Total time: {:?}", total_duration);
    println!("  Average latency: {:?}", avg_duration);
    println!("  Throughput: {:.2} pipelines/sec", iterations as f64 / total_duration.as_secs_f64());
    
    assert!(avg_duration < Duration::from_millis(20), 
        "Full pipeline latency {:?} exceeds 20ms target", avg_duration);
}

#[test]
fn benchmark_model_save_load() {
    let device = Device::Cpu;
    let edm = RogueliteEdm::new(device.clone()).unwrap();
    let director = RogueliteDirector::new(device.clone()).unwrap();
    
    let model_dir = get_model_dir();
    let edm_path = model_dir.join("bench_edm.safetensors");
    let director_path = model_dir.join("bench_director.safetensors");
    
    let edm_save_start = Instant::now();
    edm.save(&edm_path).unwrap();
    let edm_save_time = edm_save_start.elapsed();
    
    let director_save_start = Instant::now();
    director.save(&director_path).unwrap();
    let director_save_time = director_save_start.elapsed();
    
    let mut edm2 = RogueliteEdm::new(device.clone()).unwrap();
    let mut director2 = RogueliteDirector::new(device).unwrap();
    
    let edm_load_start = Instant::now();
    EmotionDataModel::load(&mut edm2, &edm_path).unwrap();
    let edm_load_time = edm_load_start.elapsed();
    
    let director_load_start = Instant::now();
    InteractionStrategy::load(&mut director2, &director_path).unwrap();
    let director_load_time = director_load_start.elapsed();
    
    println!("\nModel Save/Load Benchmark:");
    println!("  EDM save: {:?}", edm_save_time);
    println!("  EDM load: {:?}", edm_load_time);
    println!("  Director save: {:?}", director_save_time);
    println!("  Director load: {:?}", director_load_time);
}