//! Simplified robot simulator for the demo.

use embodied_shared::{Action, EnvironmentConfig, Observation, TaskType, TransitionInfo};

/// Robot simulator.
#[derive(Debug)]
pub struct Simulator {
    /// Environment configuration.
    config: EnvironmentConfig,
    /// Current joint positions.
    joint_positions: Vec<f32>,
    /// Current joint velocities.
    joint_velocities: Vec<f32>,
    /// Target position.
    target_position: [f32; 3],
    /// Current timestep.
    timestep: usize,
    /// RNG state.
    rng_state: u64,
}

impl Simulator {
    /// Create a new simulator.
    pub fn new(config: &EnvironmentConfig) -> Self {
        let mut sim = Self {
            config: config.clone(),
            joint_positions: vec![0.0; config.num_joints],
            joint_velocities: vec![0.0; config.num_joints],
            target_position: [0.3, 0.0, 0.5],
            timestep: 0,
            rng_state: 42,
        };
        sim.reset_state();
        sim
    }

    /// Reset the environment.
    pub fn reset(&mut self) -> Observation {
        self.reset_state();
        self.get_observation()
    }

    /// Reset internal state.
    fn reset_state(&mut self) {
        // Reset to home position
        self.joint_positions = match self.config.num_joints {
            7 => vec![0.0, -0.785, 0.0, -2.356, 0.0, 1.571, 0.785],
            6 => vec![0.0, -1.57, 1.57, 0.0, 1.57, 0.0],
            _ => vec![0.0; self.config.num_joints],
        };

        self.joint_velocities = vec![0.0; self.config.num_joints];
        self.timestep = 0;

        // Randomize target
        self.target_position = [
            0.3 + self.random() * 0.2 - 0.1,
            0.0 + self.random() * 0.4 - 0.2,
            0.5 + self.random() * 0.2 - 0.1,
        ];
    }

    /// Random number.
    fn random(&mut self) -> f32 {
        self.rng_state = self
            .rng_state
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        (self.rng_state >> 11) as f32 / (1u64 << 53) as f32
    }

    /// Step the environment.
    pub fn step(&mut self, action: &Action) -> (Observation, f32, bool, TransitionInfo) {
        self.timestep += 1;

        // Apply action
        if let Some(ref target_pos) = action.joint_positions {
            for (i, &target) in target_pos.iter().enumerate() {
                if i < self.joint_positions.len() {
                    // Simple P controller
                    let error = target - self.joint_positions[i];
                    self.joint_velocities[i] = error * 5.0; // Gain
                    self.joint_positions[i] += self.joint_velocities[i] * 0.02; // dt = 0.02s

                    // Clamp velocities
                    self.joint_velocities[i] = self.joint_velocities[i].clamp(-2.0, 2.0);
                }
            }
        } else if let Some(ref target_vel) = action.joint_velocities {
            for (i, &vel) in target_vel.iter().enumerate() {
                if i < self.joint_velocities.len() {
                    self.joint_velocities[i] = vel.clamp(-2.0, 2.0);
                    self.joint_positions[i] += self.joint_velocities[i] * 0.02;
                }
            }
        }

        // Compute end-effector position (simplified forward kinematics)
        let ee_pos = self.compute_end_effector_position();

        // Compute reward based on task
        let (reward, success) = self.compute_reward(&ee_pos);

        // Check if done
        let done = success
            || self.timestep >= self.config.max_episode_length
            || self.is_out_of_bounds(&ee_pos);

        let info = TransitionInfo {
            success: Some(success),
            distance_to_goal: Some(self.distance_to_target(&ee_pos)),
            contact_forces: None,
        };

        (self.get_observation(), reward, done, info)
    }

    /// Get current observation.
    fn get_observation(&self) -> Observation {
        let ee_pos = self.compute_end_effector_position();

        Observation {
            joint_positions: self.joint_positions.clone(),
            joint_velocities: self.joint_velocities.clone(),
            end_effector_pose: [ee_pos[0], ee_pos[1], ee_pos[2], 1.0, 0.0, 0.0, 0.0],
            camera_image: None,
            image_dims: None,
            timestamp: self.timestep as f64 * 0.02,
        }
    }

    /// Compute end-effector position (simplified FK).
    fn compute_end_effector_position(&self) -> [f32; 3] {
        // Very simplified forward kinematics
        let base_height = 0.3;

        if self.joint_positions.len() >= 7 {
            // Panda-like robot
            let q = &self.joint_positions;

            let x = 0.3 + 0.1 * q[0].sin() + 0.1 * q[1].sin() + 0.1 * q[3].sin();
            let y = 0.1 * q[0].cos() * q[1].sin() + 0.1 * q[2].sin();
            let z = base_height + 0.3 * q[1].cos() + 0.2 * q[3].cos() + 0.1 * q[5].cos();

            [x, y, z]
        } else {
            // Generic
            let x = 0.3 + 0.1 * self.joint_positions.iter().sum::<f32>();
            let y = 0.0;
            let z = base_height + 0.3;

            [x, y, z]
        }
    }

    /// Compute reward.
    fn compute_reward(&self, ee_pos: &[f32; 3]) -> (f32, bool) {
        let distance = self.distance_to_target(ee_pos);
        let success = distance < 0.05; // 5cm threshold

        match self.config.task_type {
            TaskType::Reach => {
                if success {
                    (10.0, true)
                } else {
                    (-distance, false) // Dense reward
                }
            }
            TaskType::Push | TaskType::PickPlace => {
                // Simplified
                if success {
                    (20.0, true)
                } else {
                    (-distance * 2.0, false)
                }
            }
            _ => {
                // Default
                if success {
                    (10.0, true)
                } else {
                    (-distance, false)
                }
            }
        }
    }

    /// Distance to target.
    fn distance_to_target(&self, ee_pos: &[f32; 3]) -> f32 {
        let dx = ee_pos[0] - self.target_position[0];
        let dy = ee_pos[1] - self.target_position[1];
        let dz = ee_pos[2] - self.target_position[2];

        (dx * dx + dy * dy + dz * dz).sqrt()
    }

    /// Check if out of bounds.
    fn is_out_of_bounds(&self, ee_pos: &[f32; 3]) -> bool {
        ee_pos[0].abs() > 1.0 || ee_pos[1].abs() > 1.0 || ee_pos[2] < 0.0 || ee_pos[2] > 1.5
    }

    /// Get target position.
    #[must_use]
    pub fn target_position(&self) -> [f32; 3] {
        self.target_position
    }

    /// Get current timestep.
    #[must_use]
    pub fn timestep(&self) -> usize {
        self.timestep
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use embodied_shared::sample_env_config;

    #[test]
    fn test_simulator_creation() {
        let config = sample_env_config();
        let sim = Simulator::new(&config);

        assert_eq!(sim.joint_positions.len(), config.num_joints);
    }

    #[test]
    fn test_reset() {
        let config = sample_env_config();
        let mut sim = Simulator::new(&config);

        let obs = sim.reset();
        assert_eq!(obs.joint_positions.len(), config.num_joints);
        assert_eq!(sim.timestep(), 0);
    }

    #[test]
    fn test_step() {
        let config = sample_env_config();
        let mut sim = Simulator::new(&config);
        sim.reset();

        let action = Action {
            joint_positions: Some(vec![0.1; config.num_joints]),
            ..Default::default()
        };

        let (obs, reward, done, info) = sim.step(&action);

        assert_eq!(obs.joint_positions.len(), config.num_joints);
        assert!(reward.is_finite());
        assert!(info.distance_to_goal.is_some());
        assert_eq!(sim.timestep(), 1);
    }

    #[test]
    fn test_multiple_steps() {
        let config = sample_env_config();
        let mut sim = Simulator::new(&config);
        sim.reset();

        for _ in 0..10 {
            let action = Action {
                joint_positions: Some(vec![0.1; config.num_joints]),
                ..Default::default()
            };
            let _ = sim.step(&action);
        }

        assert_eq!(sim.timestep(), 10);
    }

    #[test]
    fn test_reach_target() {
        let config = sample_env_config();
        let mut sim = Simulator::new(&config);
        sim.reset();

        // Force end-effector to target
        sim.target_position = [0.3, 0.0, 0.5];

        // Try to move toward target
        for _ in 0..100 {
            let action = Action {
                joint_positions: Some(vec![0.0, -0.785, 0.0, -2.356, 0.0, 1.571, 0.785]),
                ..Default::default()
            };
            let (_, _, done, info) = sim.step(&action);

            if done && info.success == Some(true) {
                break;
            }
        }
    }

    #[test]
    fn test_velocity_control() {
        let config = sample_env_config();
        let mut sim = Simulator::new(&config);
        sim.reset();

        let action = Action {
            joint_velocities: Some(vec![0.1; config.num_joints]),
            ..Default::default()
        };

        let _ = sim.step(&action);

        // Velocities should be applied
        for &vel in &sim.joint_velocities {
            assert!(vel.abs() <= 2.0);
        }
    }
}