use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
use rtx_rl::{
    ActorLearner, ActorLearnerConfig, Environment, Experience, PrioritizedReplayBuffer,
    ReplayBuffer, Step,
    algorithms::{PPO, PPOConfig, SAC, SACConfig},
};
use rtx_tensor::{DType, Device, Tensor};
use std::collections::HashMap;

#[derive(Clone, Debug)]
struct BenchExperience {
    state: Vec<f32>,
    action: Vec<f32>,
    reward: f32,
    next_state: Vec<f32>,
    done: bool,
}

impl Experience for BenchExperience {
    type State = Vec<f32>;
    type Action = Vec<f32>;
    type Reward = f32;

    fn state(&self) -> &Self::State {
        &self.state
    }
    fn action(&self) -> &Self::Action {
        &self.action
    }
    fn reward(&self) -> Self::Reward {
        self.reward
    }
    fn next_state(&self) -> &Self::State {
        &self.next_state
    }
    fn done(&self) -> bool {
        self.done
    }
}

#[derive(Clone)]
struct BenchEnv {
    state_dim: usize,
    action_dim: usize,
    step_count: usize,
    max_steps: usize,
}

impl Environment for BenchEnv {
    type State = Vec<f32>;
    type Action = Vec<f32>;
    type Reward = f32;

    fn reset(&mut self) -> Self::State {
        self.step_count = 0;
        vec![0.0; self.state_dim]
    }

    fn step(&mut self, action: &Self::Action) -> Step<Self::State, Self::Reward> {
        self.step_count += 1;
        Step {
            state: vec![0.5; self.state_dim],
            reward: 1.0,
            done: self.step_count >= self.max_steps,
            info: HashMap::new(),
        }
    }

    fn action_space(&self) -> (Vec<f32>, Vec<f32>) {
        (vec![-1.0; self.action_dim], vec![1.0; self.action_dim])
    }

    fn observation_space(&self) -> (Vec<f32>, Vec<f32>) {
        (vec![-10.0; self.state_dim], vec![10.0; self.state_dim])
    }
}

fn bench_replay_buffer_operations(c: &mut Criterion) {
    let rt = tokio::runtime::Runtime::new().unwrap();
    let device = Device::cuda(0);

    let mut group = c.benchmark_group("replay_buffer");

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

        group.bench_with_input(
            BenchmarkId::new("add_experience", buffer_size),
            buffer_size,
            |b, &size| {
                b.iter(|| {
                    rt.block_on(async {
                        let mut buffer = ReplayBuffer::<BenchExperience>::new(size, device.clone());

                        for i in 0..size {
                            let exp = BenchExperience {
                                state: vec![i as f32; 8],
                                action: vec![0.5; 4],
                                reward: 1.0,
                                next_state: vec![i as f32 + 1.0; 8],
                                done: false,
                            };
                            buffer.add(exp);
                        }
                    })
                });
            },
        );

        group.bench_with_input(
            BenchmarkId::new("sample_batch", buffer_size),
            buffer_size,
            |b, &size| {
                let mut buffer = rt.block_on(async {
                    let mut buf = ReplayBuffer::<BenchExperience>::new(size, device.clone());
                    for i in 0..size {
                        let exp = BenchExperience {
                            state: vec![i as f32; 8],
                            action: vec![0.5; 4],
                            reward: 1.0,
                            next_state: vec![i as f32 + 1.0; 8],
                            done: false,
                        };
                        buf.add(exp);
                    }
                    buf
                });

                b.iter(|| {
                    rt.block_on(async {
                        let _batch = buffer.sample(256).await.unwrap();
                    })
                });
            },
        );
    }

    group.finish();
}

fn bench_prioritized_replay_buffer(c: &mut Criterion) {
    let rt = tokio::runtime::Runtime::new().unwrap();
    let device = Device::cuda(0);

    let mut group = c.benchmark_group("prioritized_replay_buffer");

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

        group.bench_with_input(
            BenchmarkId::new("prioritized_sample", buffer_size),
            buffer_size,
            |b, &size| {
                let mut buffer = rt.block_on(async {
                    let mut buf = PrioritizedReplayBuffer::<BenchExperience>::new(
                        size,
                        0.6,
                        0.4,
                        device.clone(),
                    );
                    for i in 0..size {
                        let exp = BenchExperience {
                            state: vec![i as f32; 8],
                            action: vec![0.5; 4],
                            reward: 1.0,
                            next_state: vec![i as f32 + 1.0; 8],
                            done: false,
                        };
                        buf.add_with_priority(exp, 0.5);
                    }
                    buf
                });

                b.iter(|| {
                    rt.block_on(async {
                        let (_batch, _weights, indices) =
                            buffer.sample_with_weights(256).await.unwrap();
                        let new_priorities = vec![0.8; indices.len()];
                        buffer.update_priorities(&indices, &new_priorities);
                    })
                });
            },
        );
    }

    group.finish();
}

fn bench_ppo_algorithm(c: &mut Criterion) {
    let rt = tokio::runtime::Runtime::new().unwrap();
    let device = Device::cuda(0);

    let mut group = c.benchmark_group("ppo");

    for batch_size in [64, 256, 1024].iter() {
        group.throughput(Throughput::Elements(*batch_size as u64));

        group.bench_with_input(
            BenchmarkId::new("forward_pass", batch_size),
            batch_size,
            |b, &size| {
                let config = PPOConfig::default();
                let ppo = PPO::new(config, 8, 4, 256, device.clone());

                b.iter(|| {
                    rt.block_on(async {
                        let states = Tensor::randn(&[size, 8], DType::F32, device.clone());
                        let _output = ppo.forward(&states).await.unwrap();
                    })
                });
            },
        );

        group.bench_with_input(
            BenchmarkId::new("update_step", batch_size),
            batch_size,
            |b, &size| {
                let config = PPOConfig {
                    n_epochs: 1,
                    batch_size: size,
                    minibatch_size: size.min(64),
                    ..PPOConfig::default()
                };
                let mut ppo = PPO::new(config, 8, 4, 256, device.clone());

                b.iter(|| {
                    rt.block_on(async {
                        let states = Tensor::randn(&[size, 8], DType::F32, device.clone());
                        let actions = Tensor::randn(&[size, 4], DType::F32, device.clone());
                        let old_log_probs = Tensor::randn(&[size], DType::F32, device.clone());
                        let rewards = Tensor::ones(&[size], DType::F32, device.clone());
                        let values = Tensor::zeros(&[size], DType::F32, device.clone());
                        let dones = Tensor::zeros(&[size], DType::Bool, device.clone());

                        let _metrics = ppo
                            .update(&states, &actions, &old_log_probs, &rewards, &values, &dones)
                            .await
                            .unwrap();
                    })
                });
            },
        );
    }

    group.finish();
}

fn bench_sac_algorithm(c: &mut Criterion) {
    let rt = tokio::runtime::Runtime::new().unwrap();
    let device = Device::cuda(0);

    let mut group = c.benchmark_group("sac");

    for batch_size in [64, 256, 1024].iter() {
        group.throughput(Throughput::Elements(*batch_size as u64));

        group.bench_with_input(
            BenchmarkId::new("actor_forward", batch_size),
            batch_size,
            |b, &size| {
                let config = SACConfig::default();
                let sac = SAC::new(config, 8, 4, 256, device.clone());

                b.iter(|| {
                    rt.block_on(async {
                        let states = Tensor::randn(&[size, 8], DType::F32, device.clone());
                        let _output = sac.actor_forward(&states, false).await.unwrap();
                    })
                });
            },
        );

        group.bench_with_input(
            BenchmarkId::new("critic_forward", batch_size),
            batch_size,
            |b, &size| {
                let config = SACConfig::default();
                let sac = SAC::new(config, 8, 4, 256, device.clone());

                b.iter(|| {
                    rt.block_on(async {
                        let states = Tensor::randn(&[size, 8], DType::F32, device.clone());
                        let actions = Tensor::randn(&[size, 4], DType::F32, device.clone());
                        let _q_vals = sac.critic_forward(&states, &actions, 0).await.unwrap();
                    })
                });
            },
        );

        group.bench_with_input(
            BenchmarkId::new("full_update", batch_size),
            batch_size,
            |b, &size| {
                let config = SACConfig {
                    batch_size: size,
                    ..SACConfig::default()
                };
                let mut sac = SAC::new(config, 8, 4, 256, device.clone());

                b.iter(|| {
                    rt.block_on(async {
                        let states = Tensor::randn(&[size, 8], DType::F32, device.clone());
                        let actions = Tensor::randn(&[size, 4], DType::F32, device.clone());
                        let rewards = Tensor::randn(&[size], DType::F32, device.clone());
                        let next_states = Tensor::randn(&[size, 8], DType::F32, device.clone());
                        let dones = Tensor::zeros(&[size], DType::Bool, device.clone());

                        let _metrics = sac
                            .update(&states, &actions, &rewards, &next_states, &dones)
                            .await
                            .unwrap();
                    })
                });
            },
        );
    }

    group.finish();
}

fn bench_actor_learner(c: &mut Criterion) {
    let rt = tokio::runtime::Runtime::new().unwrap();
    let device = Device::cuda(0);

    let mut group = c.benchmark_group("actor_learner");

    for num_actors in [1, 4, 8].iter() {
        group.throughput(Throughput::Elements(*num_actors as u64));

        group.bench_with_input(
            BenchmarkId::new("rollout_collection", num_actors),
            num_actors,
            |b, &actors| {
                let config = ActorLearnerConfig {
                    num_actors: actors,
                    num_learners: 1,
                    rollout_length: 128,
                    batch_size: 256,
                    device: device.clone(),
                    ..Default::default()
                };
                let mut actor_learner = ActorLearner::new(config, 8, 4, 256);

                b.iter(|| {
                    rt.block_on(async {
                        let mut envs: Vec<BenchEnv> = (0..actors)
                            .map(|_| BenchEnv {
                                state_dim: 8,
                                action_dim: 4,
                                step_count: 0,
                                max_steps: 200,
                            })
                            .collect();

                        let _metrics = actor_learner.training_step(&mut envs).await.unwrap();
                    })
                });
            },
        );
    }

    group.finish();
}

fn bench_throughput_scaling(c: &mut Criterion) {
    let rt = tokio::runtime::Runtime::new().unwrap();
    let device = Device::cuda(0);

    let mut group = c.benchmark_group("throughput_scaling");

    // Test how throughput scales with different parameters
    for (state_dim, action_dim) in [(4, 2), (16, 8), (64, 32), (256, 128)].iter() {
        group.throughput(Throughput::Elements(1000));

        group.bench_with_input(
            BenchmarkId::new(
                "inference_throughput",
                format!("{}x{}", state_dim, action_dim),
            ),
            &(*state_dim, *action_dim),
            |b, &(s_dim, a_dim)| {
                let config = PPOConfig::default();
                let ppo = PPO::new(config, s_dim, a_dim, 256, device.clone());

                b.iter(|| {
                    rt.block_on(async {
                        let states = Tensor::randn(&[1000, s_dim], DType::F32, device.clone());
                        let _output = ppo.forward(&states).await.unwrap();
                    })
                });
            },
        );
    }

    group.finish();
}

criterion_group!(
    benches,
    bench_replay_buffer_operations,
    bench_prioritized_replay_buffer,
    bench_ppo_algorithm,
    bench_sac_algorithm,
    bench_actor_learner,
    bench_throughput_scaling
);

criterion_main!(benches);