use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
use rtx_rl::standalone_rl::{Experience, Matrix, PPO, ReplayBuffer};

fn bench_matrix_operations(c: &mut Criterion) {
    let mut group = c.benchmark_group("matrix_operations");

    for size in [10, 100, 1000].iter() {
        group.throughput(Throughput::Elements(*size as u64));

        group.bench_with_input(
            BenchmarkId::new("matrix_creation", size),
            size,
            |b, &size| {
                b.iter(|| Matrix::randn(vec![size, size]));
            },
        );

        let matrix_a = Matrix::randn(vec![*size, *size]);
        let matrix_b = Matrix::randn(vec![*size, *size]);

        group.bench_with_input(BenchmarkId::new("matrix_addition", size), size, |b, _| {
            b.iter(|| matrix_a.add(&matrix_b));
        });

        group.bench_with_input(BenchmarkId::new("matrix_scalar_mul", size), size, |b, _| {
            b.iter(|| matrix_a.mul(0.5));
        });

        group.bench_with_input(BenchmarkId::new("matrix_mean", size), size, |b, _| {
            b.iter(|| matrix_a.mean());
        });
    }

    group.finish();
}

fn bench_replay_buffer_operations(c: &mut Criterion) {
    let mut group = c.benchmark_group("replay_buffer");

    for capacity in [1000, 10000, 100000].iter() {
        group.throughput(Throughput::Elements(*capacity as u64));

        group.bench_with_input(
            BenchmarkId::new("buffer_fill", capacity),
            capacity,
            |b, &cap| {
                b.iter(|| {
                    let mut buffer = ReplayBuffer::new(cap);
                    for i in 0..cap {
                        let exp = Experience {
                            state: vec![i as f32, (i * 2) as f32],
                            action: vec![(i % 3) as f32],
                            reward: (i as f32) * 0.01,
                            next_state: vec![(i + 1) as f32, ((i + 1) * 2) as f32],
                            done: i == cap - 1,
                        };
                        buffer.push(exp);
                    }
                });
            },
        );

        let mut buffer = ReplayBuffer::new(*capacity);
        for i in 0..*capacity {
            let exp = Experience {
                state: vec![i as f32, (i * 2) as f32],
                action: vec![(i % 3) as f32],
                reward: (i as f32) * 0.01,
                next_state: vec![(i + 1) as f32, ((i + 1) * 2) as f32],
                done: false,
            };
            buffer.push(exp);
        }

        for batch_size in [32, 128, 512].iter() {
            if *batch_size <= *capacity {
                group.bench_with_input(
                    BenchmarkId::new("buffer_sample", format!("{}_{}", capacity, batch_size)),
                    &(*capacity, *batch_size),
                    |b, &(_cap, batch)| {
                        b.iter(|| buffer.sample(batch).unwrap());
                    },
                );
            }
        }
    }

    group.finish();
}

fn bench_policy_network(c: &mut Criterion) {
    let mut group = c.benchmark_group("policy_network");

    for (input_dim, output_dim) in [(4, 2), (16, 8), (64, 32), (256, 128)].iter() {
        group.throughput(Throughput::Elements(*input_dim as u64));

        let network = rtx_rl::standalone_rl::PolicyNetwork::new(*input_dim, *output_dim);
        let input: Vec<f32> = (0..*input_dim).map(|i| (i as f32) * 0.1).collect();

        group.bench_with_input(
            BenchmarkId::new("forward_pass", format!("{}x{}", input_dim, output_dim)),
            &(*input_dim, *output_dim),
            |b, _| {
                b.iter(|| network.forward(&input));
            },
        );
    }

    group.finish();
}

fn bench_ppo_algorithm(c: &mut Criterion) {
    let mut group = c.benchmark_group("ppo");

    for batch_size in [32, 128, 512].iter() {
        group.throughput(Throughput::Elements(*batch_size as u64));

        let mut ppo = PPO::new(8, 4);
        let mut experiences = Vec::new();

        for i in 0..*batch_size {
            let exp = Experience {
                state: vec![
                    (i as f32) * 0.01,
                    (i as f32) * 0.02,
                    (i as f32) * 0.03,
                    (i as f32) * 0.04,
                    (i as f32) * 0.05,
                    (i as f32) * 0.06,
                    (i as f32) * 0.07,
                    (i as f32) * 0.08,
                ],
                action: vec![
                    ((i % 4) as f32) * 0.25,
                    ((i % 3) as f32) * 0.33,
                    ((i % 2) as f32) * 0.5,
                    (i as f32) * 0.01,
                ],
                reward: if i % 2 == 0 { 1.0 } else { -0.5 },
                next_state: vec![
                    (i as f32 + 1.0) * 0.01,
                    (i as f32 + 1.0) * 0.02,
                    (i as f32 + 1.0) * 0.03,
                    (i as f32 + 1.0) * 0.04,
                    (i as f32 + 1.0) * 0.05,
                    (i as f32 + 1.0) * 0.06,
                    (i as f32 + 1.0) * 0.07,
                    (i as f32 + 1.0) * 0.08,
                ],
                done: i == *batch_size - 1,
            };
            experiences.push(exp);
        }

        group.bench_with_input(
            BenchmarkId::new("get_action", batch_size),
            batch_size,
            |b, _| {
                b.iter(|| {
                    for exp in &experiences {
                        ppo.get_action(&exp.state);
                    }
                });
            },
        );

        group.bench_with_input(
            BenchmarkId::new("get_value", batch_size),
            batch_size,
            |b, _| {
                b.iter(|| {
                    for exp in &experiences {
                        ppo.get_value(&exp.state);
                    }
                });
            },
        );

        group.bench_with_input(
            BenchmarkId::new("compute_advantages", batch_size),
            batch_size,
            |b, _| {
                b.iter(|| ppo.compute_advantages(&experiences));
            },
        );

        group.bench_with_input(
            BenchmarkId::new("update", batch_size),
            batch_size,
            |b, _| {
                b.iter(|| ppo.update(&experiences));
            },
        );
    }

    group.finish();
}

fn bench_complete_rl_episode(c: &mut Criterion) {
    let mut group = c.benchmark_group("complete_rl_episode");

    for episode_length in [10, 100, 1000].iter() {
        group.throughput(Throughput::Elements(*episode_length as u64));

        group.bench_with_input(
            BenchmarkId::new("full_episode", episode_length),
            episode_length,
            |b, &length| {
                b.iter(|| {
                    let mut ppo = PPO::new(4, 2);
                    let mut buffer = ReplayBuffer::new(length * 2);

                    // Simulate episode
                    let mut state = vec![0.0, 0.0, 0.0, 0.0];
                    let mut total_reward = 0.0;

                    for step in 0..length {
                        let action = ppo.get_action(&state);

                        // Simple environment dynamics
                        let next_state = vec![
                            state[0] + action[0] * 0.1,
                            state[1] + action[1] * 0.1,
                            state[2] - state[0] * 0.05,
                            state[3] - state[1] * 0.05,
                        ];

                        let distance = next_state.iter().map(|x| x.powi(2)).sum::<f32>().sqrt();
                        let reward = 1.0 - distance.min(1.0);
                        total_reward += reward;

                        let experience = Experience {
                            state: state.clone(),
                            action: action.clone(),
                            reward,
                            next_state: next_state.clone(),
                            done: step >= length - 1,
                        };

                        buffer.push(experience);
                        state = next_state;

                        // Train every 32 steps
                        if step % 32 == 31 {
                            if let Some(batch) = buffer.sample(32.min(buffer.len())) {
                                ppo.update(&batch);
                            }
                        }
                    }

                    total_reward
                });
            },
        );
    }

    group.finish();
}

fn bench_memory_usage(c: &mut Criterion) {
    let mut group = c.benchmark_group("memory_usage");

    for num_experiences in [1000, 10000, 100000].iter() {
        group.throughput(Throughput::Elements(*num_experiences as u64));

        group.bench_with_input(
            BenchmarkId::new("experience_creation", num_experiences),
            num_experiences,
            |b, &count| {
                b.iter(|| {
                    let mut experiences = Vec::with_capacity(count);
                    for i in 0..count {
                        let exp = Experience {
                            state: vec![i as f32; 8],
                            action: vec![(i % 4) as f32; 4],
                            reward: (i as f32) * 0.001,
                            next_state: vec![(i + 1) as f32; 8],
                            done: i == count - 1,
                        };
                        experiences.push(exp);
                    }
                    experiences
                });
            },
        );
    }

    group.finish();
}

criterion_group!(
    benches,
    bench_matrix_operations,
    bench_replay_buffer_operations,
    bench_policy_network,
    bench_ppo_algorithm,
    bench_complete_rl_episode,
    bench_memory_usage
);

criterion_main!(benches);