//! EmbodiedSim - Robotics World Model Trainer.
//!
//! This demo showcases world models for sim-to-real robotics transfer,
//! implementing RSSM and JEPA architectures for learning robot dynamics.

pub mod actor_critic;
pub mod jepa;
pub mod replay_buffer;
pub mod rssm;
pub mod sample_data;
pub mod simulator;

use thiserror::Error;

use embodied_shared::{
    Action, EnvironmentConfig, EvaluationRequest, EvaluationResult, ImaginationRequest,
    ImaginedTrajectory, Observation, PolicyConfig, TrainingConfig, TrainingProgress, Transition,
    WorldModelConfig,
};

/// Errors that can occur in EmbodiedSim.
#[derive(Debug, Error)]
pub enum EmbodiedError {
    /// Invalid configuration.
    #[error("Invalid configuration: {0}")]
    InvalidConfig(String),

    /// Training failed.
    #[error("Training failed: {0}")]
    TrainingFailed(String),

    /// World model not initialized.
    #[error("World model not initialized")]
    ModelNotInitialized,

    /// Policy not initialized.
    #[error("Policy not initialized")]
    PolicyNotInitialized,

    /// Buffer empty.
    #[error("Replay buffer is empty")]
    BufferEmpty,

    /// Imagination failed.
    #[error("Imagination failed: {0}")]
    ImaginationFailed(String),
}

/// Main EmbodiedSim system.
#[derive(Debug)]
pub struct EmbodiedSim {
    /// World model configuration.
    wm_config: WorldModelConfig,
    /// Policy configuration.
    policy_config: PolicyConfig,
    /// Environment configuration.
    env_config: EnvironmentConfig,
    /// RSSM world model.
    rssm: rssm::RSSM,
    /// Actor-Critic policy.
    policy: actor_critic::ActorCritic,
    /// Replay buffer.
    buffer: replay_buffer::ReplayBuffer,
    /// Simulator.
    simulator: simulator::Simulator,
    /// Is initialized.
    initialized: bool,
}

impl Default for EmbodiedSim {
    fn default() -> Self {
        Self::new(
            WorldModelConfig::default(),
            PolicyConfig::default(),
            EnvironmentConfig::default(),
        )
    }
}

impl EmbodiedSim {
    /// Create a new EmbodiedSim instance.
    #[must_use]
    pub fn new(
        wm_config: WorldModelConfig,
        policy_config: PolicyConfig,
        env_config: EnvironmentConfig,
    ) -> Self {
        Self {
            rssm: rssm::RSSM::new(&wm_config),
            policy: actor_critic::ActorCritic::new(&policy_config),
            buffer: replay_buffer::ReplayBuffer::new(100_000),
            simulator: simulator::Simulator::new(&env_config),
            wm_config,
            policy_config,
            env_config,
            initialized: true,
        }
    }

    /// Add transition to replay buffer.
    pub fn add_transition(&mut self, transition: Transition) {
        self.buffer.add(transition);
    }

    /// Collect experience from environment.
    pub fn collect_experience(
        &mut self,
        num_steps: usize,
    ) -> Result<Vec<Transition>, EmbodiedError> {
        if !self.initialized {
            return Err(EmbodiedError::ModelNotInitialized);
        }

        let mut transitions = Vec::with_capacity(num_steps);
        let mut obs = self.simulator.reset();

        for _ in 0..num_steps {
            // Get action from policy
            let state = self.rssm.encode_observation(&obs);
            let action = self.policy.sample_action(&state);
            let action_struct = self.action_from_vec(&action);

            // Step environment
            let (next_obs, reward, done, info) = self.simulator.step(&action_struct);

            let transition = Transition {
                observation: obs.clone(),
                action: action_struct,
                reward,
                next_observation: next_obs.clone(),
                done,
                truncated: false,
                info,
            };

            transitions.push(transition.clone());
            self.buffer.add(transition);

            if done {
                obs = self.simulator.reset();
            } else {
                obs = next_obs;
            }
        }

        Ok(transitions)
    }

    /// Convert action vector to Action struct.
    fn action_from_vec(&self, action: &[f64]) -> Action {
        Action {
            joint_positions: Some(action.iter().map(|&x| x as f32).collect()),
            joint_velocities: None,
            joint_torques: None,
            ee_delta: None,
            gripper: if action.len() > self.env_config.num_joints {
                Some(action[self.env_config.num_joints] as f32)
            } else {
                None
            },
        }
    }

    /// Train the world model.
    pub fn train_world_model(
        &mut self,
        config: &TrainingConfig,
        progress_callback: Option<Box<dyn Fn(TrainingProgress) + Send>>,
    ) -> Result<(), EmbodiedError> {
        if !self.initialized {
            return Err(EmbodiedError::ModelNotInitialized);
        }

        if self.buffer.len() < config.batch_size * config.sequence_length {
            return Err(EmbodiedError::BufferEmpty);
        }

        let start_time = std::time::Instant::now();
        let total_steps = config.epochs * 100; // 100 batches per epoch

        for epoch in 0..config.epochs {
            for batch_idx in 0..100 {
                // Sample batch from buffer
                let batch = self
                    .buffer
                    .sample_sequence(config.batch_size, config.sequence_length);

                // Train world model on batch
                let loss = self.rssm.train_step(&batch, config.wm_learning_rate);

                let step = epoch * 100 + batch_idx + 1;

                if let Some(ref callback) = progress_callback {
                    callback(TrainingProgress {
                        epoch: epoch + 1,
                        total_epochs: config.epochs,
                        step,
                        total_steps,
                        wm_loss: loss,
                        image_loss: Some(loss * 0.6),
                        reward_loss: Some(loss * 0.2),
                        kl_loss: Some(loss * 0.2),
                        actor_loss: None,
                        critic_loss: None,
                        entropy: None,
                        learning_rate: config.wm_learning_rate,
                        env_return: None,
                        episode_length: None,
                        elapsed_seconds: start_time.elapsed().as_secs_f64(),
                    });
                }
            }
        }

        Ok(())
    }

    /// Train the policy using imagined rollouts.
    pub fn train_policy(
        &mut self,
        config: &TrainingConfig,
        progress_callback: Option<Box<dyn Fn(TrainingProgress) + Send>>,
    ) -> Result<(), EmbodiedError> {
        if !self.initialized {
            return Err(EmbodiedError::ModelNotInitialized);
        }

        let start_time = std::time::Instant::now();
        let total_steps = config.epochs * 50; // 50 updates per epoch

        for epoch in 0..config.epochs {
            for update_idx in 0..50 {
                // Sample starting states from buffer
                let batch = self.buffer.sample(config.batch_size);

                // Imagine trajectories
                let trajectories = self.imagine_batch(&batch, self.wm_config.imagination_horizon);

                // Compute returns and advantages
                let (returns, advantages) = self.compute_advantages(&trajectories);

                // Update policy
                let (actor_loss, critic_loss, entropy) =
                    self.policy
                        .update(&trajectories, &returns, &advantages, config);

                let step = epoch * 50 + update_idx + 1;

                if let Some(ref callback) = progress_callback {
                    callback(TrainingProgress {
                        epoch: epoch + 1,
                        total_epochs: config.epochs,
                        step,
                        total_steps,
                        wm_loss: 0.0,
                        image_loss: None,
                        reward_loss: None,
                        kl_loss: None,
                        actor_loss: Some(actor_loss),
                        critic_loss: Some(critic_loss),
                        entropy: Some(entropy),
                        learning_rate: config.actor_learning_rate,
                        env_return: None,
                        episode_length: None,
                        elapsed_seconds: start_time.elapsed().as_secs_f64(),
                    });
                }
            }
        }

        Ok(())
    }

    /// Imagine trajectories from a batch of observations.
    fn imagine_batch(&self, batch: &[Transition], horizon: usize) -> Vec<ImaginedTrajectory> {
        batch
            .iter()
            .map(|trans| self.imagine_from_obs(&trans.observation, horizon))
            .collect()
    }

    /// Imagine a trajectory from an observation.
    fn imagine_from_obs(&self, obs: &Observation, horizon: usize) -> ImaginedTrajectory {
        let mut states = Vec::with_capacity(horizon);
        let mut actions = Vec::with_capacity(horizon);
        let mut rewards = Vec::with_capacity(horizon);
        let mut values = Vec::with_capacity(horizon);
        let mut dones = Vec::with_capacity(horizon);

        let mut state = self.rssm.encode_observation(obs);

        for _ in 0..horizon {
            // Sample action from policy
            let action = self.policy.sample_action(&state);

            // Predict value
            let value = self.policy.predict_value(&state);

            // Imagine next state
            let (next_state, reward, done) = self.rssm.imagine_step(&state, &action);

            states.push(state.clone());
            actions.push(action);
            rewards.push(reward);
            values.push(value);
            dones.push(done);

            state = next_state;
        }

        ImaginedTrajectory {
            states,
            actions,
            rewards,
            values,
            dones,
            decoded_observations: None,
        }
    }

    /// Compute returns and advantages for imagined trajectories.
    fn compute_advantages(
        &self,
        trajectories: &[ImaginedTrajectory],
    ) -> (Vec<Vec<f64>>, Vec<Vec<f64>>) {
        let gamma = self.policy_config.gamma as f64;
        let lambda = self.policy_config.gae_lambda as f64;

        let mut all_returns = Vec::with_capacity(trajectories.len());
        let mut all_advantages = Vec::with_capacity(trajectories.len());

        for traj in trajectories {
            let horizon = traj.rewards.len();
            let mut returns = vec![0.0; horizon];
            let mut advantages = vec![0.0; horizon];

            // Bootstrap from last value
            let mut gae = 0.0;

            for t in (0..horizon).rev() {
                let next_value = if t + 1 < horizon {
                    traj.values[t + 1]
                } else {
                    0.0
                };

                let delta =
                    traj.rewards[t] + gamma * next_value * (1.0 - traj.dones[t]) - traj.values[t];
                gae = delta + gamma * lambda * (1.0 - traj.dones[t]) * gae;
                advantages[t] = gae;
                returns[t] = advantages[t] + traj.values[t];
            }

            all_returns.push(returns);
            all_advantages.push(advantages);
        }

        (all_returns, all_advantages)
    }

    /// Imagine future trajectories.
    pub fn imagine(
        &self,
        request: &ImaginationRequest,
    ) -> Result<Vec<ImaginedTrajectory>, EmbodiedError> {
        if !self.initialized {
            return Err(EmbodiedError::ModelNotInitialized);
        }

        let mut trajectories = Vec::with_capacity(request.num_trajectories);

        for _ in 0..request.num_trajectories {
            let traj = self.imagine_from_obs(&request.initial_observation, request.horizon);
            trajectories.push(traj);
        }

        Ok(trajectories)
    }

    /// Evaluate the policy.
    pub fn evaluate(
        &mut self,
        request: &EvaluationRequest,
    ) -> Result<EvaluationResult, EmbodiedError> {
        if !self.initialized {
            return Err(EmbodiedError::ModelNotInitialized);
        }

        let start_time = std::time::Instant::now();
        let mut episode_returns = Vec::with_capacity(request.num_episodes);
        let mut episode_lengths = Vec::with_capacity(request.num_episodes);
        let mut successes = 0;

        for _ in 0..request.num_episodes {
            let mut obs = self.simulator.reset();
            let mut episode_return = 0.0;
            let mut episode_length = 0;

            loop {
                let state = self.rssm.encode_observation(&obs);
                let action = if request.deterministic {
                    self.policy.get_action_mean(&state)
                } else {
                    self.policy.sample_action(&state)
                };

                let action_struct = self.action_from_vec(&action);
                let (next_obs, reward, done, info) = self.simulator.step(&action_struct);

                episode_return += reward as f64;
                episode_length += 1;

                if info.success == Some(true) {
                    successes += 1;
                }

                if done || episode_length >= self.env_config.max_episode_length {
                    break;
                }

                obs = next_obs;
            }

            episode_returns.push(episode_return);
            episode_lengths.push(episode_length);
        }

        let mean_return = episode_returns.iter().sum::<f64>() / request.num_episodes as f64;
        let mean_length =
            episode_lengths.iter().sum::<usize>() as f64 / request.num_episodes as f64;

        let std_return = if request.num_episodes > 1 {
            let variance = episode_returns
                .iter()
                .map(|r| (r - mean_return).powi(2))
                .sum::<f64>()
                / (request.num_episodes - 1) as f64;
            variance.sqrt()
        } else {
            0.0
        };

        Ok(EvaluationResult {
            mean_return,
            std_return,
            mean_length,
            success_rate: successes as f64 / request.num_episodes as f64,
            episode_returns,
            episode_lengths,
            eval_time: start_time.elapsed().as_secs_f64(),
        })
    }

    /// Get action for a given observation.
    pub fn get_action(
        &self,
        obs: &Observation,
        deterministic: bool,
    ) -> Result<Action, EmbodiedError> {
        if !self.initialized {
            return Err(EmbodiedError::PolicyNotInitialized);
        }

        let state = self.rssm.encode_observation(obs);
        let action = if deterministic {
            self.policy.get_action_mean(&state)
        } else {
            self.policy.sample_action(&state)
        };

        Ok(self.action_from_vec(&action))
    }

    /// Get world model config.
    #[must_use]
    pub fn world_model_config(&self) -> &WorldModelConfig {
        &self.wm_config
    }

    /// Get policy config.
    #[must_use]
    pub fn policy_config(&self) -> &PolicyConfig {
        &self.policy_config
    }

    /// Get environment config.
    #[must_use]
    pub fn env_config(&self) -> &EnvironmentConfig {
        &self.env_config
    }

    /// Get buffer size.
    #[must_use]
    pub fn buffer_size(&self) -> usize {
        self.buffer.len()
    }
}

/// Run the demo.
pub fn run_demo() -> Result<EvaluationResult, EmbodiedError> {
    let mut embodied = EmbodiedSim::default();

    // Collect some experience
    let _ = embodied.collect_experience(100)?;

    // Evaluate
    let request = EvaluationRequest {
        num_episodes: 5,
        ..Default::default()
    };
    embodied.evaluate(&request)
}

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

    #[test]
    fn test_embodied_creation() {
        let embodied = EmbodiedSim::default();
        assert!(embodied.initialized);
    }

    #[test]
    fn test_collect_experience() {
        let mut embodied = EmbodiedSim::default();
        let result = embodied.collect_experience(10);
        assert!(result.is_ok());

        let transitions = result.unwrap();
        assert_eq!(transitions.len(), 10);
        assert_eq!(embodied.buffer_size(), 10);
    }

    #[test]
    fn test_get_action() {
        let embodied = EmbodiedSim::default();
        let obs = embodied_shared::sample_observation();

        let action = embodied.get_action(&obs, true);
        assert!(action.is_ok());
    }

    #[test]
    fn test_imagine() {
        let mut embodied = EmbodiedSim::default();
        let _ = embodied.collect_experience(10);

        let request = ImaginationRequest {
            initial_observation: embodied_shared::sample_observation(),
            horizon: 5,
            num_trajectories: 3,
            use_policy: true,
            decode: false,
        };

        let result = embodied.imagine(&request);
        assert!(result.is_ok());

        let trajectories = result.unwrap();
        assert_eq!(trajectories.len(), 3);
        for traj in &trajectories {
            assert_eq!(traj.states.len(), 5);
            assert_eq!(traj.rewards.len(), 5);
        }
    }

    #[test]
    fn test_evaluate() {
        let mut embodied = EmbodiedSim::default();
        let _ = embodied.collect_experience(10);

        let request = EvaluationRequest {
            num_episodes: 2,
            ..Default::default()
        };

        let result = embodied.evaluate(&request);
        assert!(result.is_ok());

        let eval = result.unwrap();
        assert_eq!(eval.episode_returns.len(), 2);
    }

    #[test]
    fn test_run_demo() {
        let result = run_demo();
        assert!(result.is_ok());
    }
}