//! Sample data and configurations for EmbodiedSim demo.

use embodied_shared::{
    ActionDistribution, DomainRandomization, EnvironmentConfig, EnvironmentType, EvaluationRequest,
    ImageEncoderType, ImaginationRequest, Observation, PolicyConfig, PolicyType, RewardType,
    SafetyLimits, SimToRealConfig, TaskType, TrainingConfig, WorldModelConfig, WorldModelType,
};

// ============================================================================
// World Model Configurations
// ============================================================================

/// Create a small RSSM configuration for fast training.
#[must_use]
pub fn small_rssm_config() -> WorldModelConfig {
    WorldModelConfig {
        model_type: WorldModelType::RSSM,
        obs_embed_dim: 128,
        action_embed_dim: 32,
        deter_dim: 256,
        stoch_dim: 16,
        stoch_classes: 16,
        hidden_dim: 256,
        num_layers: 1,
        image_encoder: ImageEncoderType::MLP,
        predict_reward: true,
        predict_done: true,
        imagination_horizon: 10,
    }
}

/// Create a medium RSSM configuration.
#[must_use]
pub fn medium_rssm_config() -> WorldModelConfig {
    WorldModelConfig {
        model_type: WorldModelType::RSSM,
        obs_embed_dim: 256,
        action_embed_dim: 64,
        deter_dim: 512,
        stoch_dim: 32,
        stoch_classes: 32,
        hidden_dim: 512,
        num_layers: 2,
        image_encoder: ImageEncoderType::CNN,
        predict_reward: true,
        predict_done: true,
        imagination_horizon: 15,
    }
}

/// Create a JEPA configuration.
#[must_use]
pub fn jepa_config() -> WorldModelConfig {
    WorldModelConfig {
        model_type: WorldModelType::JEPA,
        obs_embed_dim: 256,
        action_embed_dim: 64,
        deter_dim: 512,
        stoch_dim: 0,
        stoch_classes: 0,
        hidden_dim: 512,
        num_layers: 4,
        image_encoder: ImageEncoderType::ViT,
        predict_reward: false,
        predict_done: false,
        imagination_horizon: 20,
    }
}

/// Create a transformer world model configuration.
#[must_use]
pub fn transformer_wm_config() -> WorldModelConfig {
    WorldModelConfig {
        model_type: WorldModelType::Transformer,
        obs_embed_dim: 256,
        action_embed_dim: 64,
        deter_dim: 512,
        stoch_dim: 32,
        stoch_classes: 32,
        hidden_dim: 512,
        num_layers: 6,
        image_encoder: ImageEncoderType::ViT,
        predict_reward: true,
        predict_done: true,
        imagination_horizon: 25,
    }
}

// ============================================================================
// Policy Configurations
// ============================================================================

/// Create a small actor-critic policy.
#[must_use]
pub fn small_actor_critic_config() -> PolicyConfig {
    PolicyConfig {
        policy_type: PolicyType::ActorCritic,
        state_dim: 256 + 16 * 16, // deter + stoch * classes
        action_dim: 7,
        actor_hidden: vec![128, 128],
        critic_hidden: vec![128, 128],
        gamma: 0.99,
        gae_lambda: 0.95,
        entropy_coef: 0.001,
        continuous_actions: true,
        action_dist: ActionDistribution::Normal,
    }
}

/// Create a medium actor-critic policy.
#[must_use]
pub fn medium_actor_critic_config() -> PolicyConfig {
    PolicyConfig {
        policy_type: PolicyType::ActorCritic,
        state_dim: 512 + 32 * 32,
        action_dim: 7,
        actor_hidden: vec![256, 256],
        critic_hidden: vec![256, 256],
        gamma: 0.997,
        gae_lambda: 0.95,
        entropy_coef: 3e-4,
        continuous_actions: true,
        action_dist: ActionDistribution::Normal,
    }
}

/// Create an MPC policy configuration.
#[must_use]
pub fn mpc_policy_config() -> PolicyConfig {
    PolicyConfig {
        policy_type: PolicyType::MPC,
        state_dim: 512,
        action_dim: 7,
        actor_hidden: vec![],
        critic_hidden: vec![],
        gamma: 0.99,
        gae_lambda: 0.0,
        entropy_coef: 0.0,
        continuous_actions: true,
        action_dist: ActionDistribution::Normal,
    }
}

/// Create a diffusion policy configuration.
#[must_use]
pub fn diffusion_policy_config() -> PolicyConfig {
    PolicyConfig {
        policy_type: PolicyType::DiffusionPolicy,
        state_dim: 512 + 32 * 32,
        action_dim: 7,
        actor_hidden: vec![256, 256, 256],
        critic_hidden: vec![256, 256],
        gamma: 0.99,
        gae_lambda: 0.95,
        entropy_coef: 0.0,
        continuous_actions: true,
        action_dist: ActionDistribution::Normal,
    }
}

// ============================================================================
// Environment Configurations
// ============================================================================

/// Create a Panda reach task configuration.
#[must_use]
pub fn panda_reach_config() -> EnvironmentConfig {
    EnvironmentConfig {
        env_type: EnvironmentType::Panda,
        task_type: TaskType::Reach,
        num_joints: 7,
        action_dim: 7,
        obs_dim: 21,
        use_camera: false,
        camera_size: (64, 64),
        max_episode_length: 100,
        action_repeat: 2,
        control_freq: 20.0,
        reward_type: RewardType::Dense,
    }
}

/// Create a Panda pick and place configuration.
#[must_use]
pub fn panda_pick_place_config() -> EnvironmentConfig {
    EnvironmentConfig {
        env_type: EnvironmentType::Panda,
        task_type: TaskType::PickPlace,
        num_joints: 7,
        action_dim: 8, // 7 joints + gripper
        obs_dim: 25,   // joint state + object state
        use_camera: true,
        camera_size: (84, 84),
        max_episode_length: 200,
        action_repeat: 2,
        control_freq: 20.0,
        reward_type: RewardType::Shaped,
    }
}

/// Create a UR5 push configuration.
#[must_use]
pub fn ur5_push_config() -> EnvironmentConfig {
    EnvironmentConfig {
        env_type: EnvironmentType::UR5,
        task_type: TaskType::Push,
        num_joints: 6,
        action_dim: 6,
        obs_dim: 18,
        use_camera: true,
        camera_size: (64, 64),
        max_episode_length: 150,
        action_repeat: 2,
        control_freq: 20.0,
        reward_type: RewardType::Dense,
    }
}

/// Create a peg insertion configuration.
#[must_use]
pub fn peg_insertion_config() -> EnvironmentConfig {
    EnvironmentConfig {
        env_type: EnvironmentType::Panda,
        task_type: TaskType::PegInsertion,
        num_joints: 7,
        action_dim: 8,
        obs_dim: 28,
        use_camera: true,
        camera_size: (84, 84),
        max_episode_length: 250,
        action_repeat: 1,
        control_freq: 50.0,
        reward_type: RewardType::Shaped,
    }
}

// ============================================================================
// Training Configurations
// ============================================================================

/// Create a fast training configuration for demos.
#[must_use]
pub fn fast_training_config() -> TrainingConfig {
    TrainingConfig {
        batch_size: 8,
        sequence_length: 32,
        wm_learning_rate: 1e-4,
        actor_learning_rate: 3e-5,
        critic_learning_rate: 3e-5,
        epochs: 10,
        buffer_size: 10_000,
        prefill_steps: 500,
        train_every: 5,
        train_ratio: 64.0,
        mixed_precision: false,
        gradient_clip: 100.0,
        seed: Some(42),
    }
}

/// Create a standard training configuration.
#[must_use]
pub fn standard_training_config() -> TrainingConfig {
    TrainingConfig {
        batch_size: 64,
        sequence_length: 64,
        wm_learning_rate: 1e-4,
        actor_learning_rate: 3e-5,
        critic_learning_rate: 3e-5,
        epochs: 100,
        buffer_size: 1_000_000,
        prefill_steps: 5000,
        train_every: 5,
        train_ratio: 512.0,
        mixed_precision: true,
        gradient_clip: 100.0,
        seed: Some(42),
    }
}

/// Create a DreamerV3-style training configuration.
#[must_use]
pub fn dreamerv3_config() -> TrainingConfig {
    TrainingConfig {
        batch_size: 16,
        sequence_length: 64,
        wm_learning_rate: 1e-4,
        actor_learning_rate: 3e-5,
        critic_learning_rate: 3e-5,
        epochs: 1000,
        buffer_size: 5_000_000,
        prefill_steps: 2500,
        train_every: 1,
        train_ratio: 512.0,
        mixed_precision: true,
        gradient_clip: 1000.0,
        seed: Some(42),
    }
}

// ============================================================================
// Evaluation Configurations
// ============================================================================

/// Create a quick evaluation request.
#[must_use]
pub fn quick_eval_request() -> EvaluationRequest {
    EvaluationRequest {
        num_episodes: 5,
        render: false,
        deterministic: true,
        record_video: false,
        env_config: panda_reach_config(),
    }
}

/// Create a full evaluation request.
#[must_use]
pub fn full_eval_request() -> EvaluationRequest {
    EvaluationRequest {
        num_episodes: 100,
        render: false,
        deterministic: true,
        record_video: false,
        env_config: panda_reach_config(),
    }
}

/// Create an evaluation with video recording.
#[must_use]
pub fn video_eval_request() -> EvaluationRequest {
    EvaluationRequest {
        num_episodes: 10,
        render: true,
        deterministic: true,
        record_video: true,
        env_config: panda_reach_config(),
    }
}

// ============================================================================
// Imagination Configurations
// ============================================================================

/// Create a short imagination request.
#[must_use]
pub fn short_imagination_request() -> ImaginationRequest {
    ImaginationRequest {
        initial_observation: Observation::default(),
        horizon: 10,
        num_trajectories: 5,
        use_policy: true,
        decode: false,
    }
}

/// Create a long imagination request.
#[must_use]
pub fn long_imagination_request() -> ImaginationRequest {
    ImaginationRequest {
        initial_observation: Observation::default(),
        horizon: 50,
        num_trajectories: 20,
        use_policy: true,
        decode: true,
    }
}

// ============================================================================
// Sim-to-Real Configurations
// ============================================================================

/// Create a conservative sim-to-real configuration.
#[must_use]
pub fn conservative_sim2real_config() -> SimToRealConfig {
    SimToRealConfig {
        domain_randomization: DomainRandomization {
            dynamics: true,
            visuals: true,
            lighting: true,
            mass_range: (0.9, 1.1),
            friction_range: (0.7, 1.3),
            damping_range: (0.9, 1.1),
            action_noise: 0.005,
            observation_noise: 0.005,
        },
        use_calibration: true,
        action_scale: 0.8,
        safety_limits: SafetyLimits {
            max_joint_velocity: 1.0,
            max_joint_acceleration: 3.0,
            max_ee_velocity: 0.3,
            max_force: 30.0,
            workspace: [[-0.4, 0.4], [-0.4, 0.4], [0.1, 0.7]],
        },
    }
}

/// Create an aggressive sim-to-real configuration.
#[must_use]
pub fn aggressive_sim2real_config() -> SimToRealConfig {
    SimToRealConfig {
        domain_randomization: DomainRandomization {
            dynamics: true,
            visuals: true,
            lighting: true,
            mass_range: (0.5, 2.0),
            friction_range: (0.2, 3.0),
            damping_range: (0.5, 2.0),
            action_noise: 0.02,
            observation_noise: 0.02,
        },
        use_calibration: false,
        action_scale: 1.0,
        safety_limits: SafetyLimits::default(),
    }
}

// ============================================================================
// Tests
// ============================================================================

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

    #[test]
    fn test_world_model_configs() {
        let small = small_rssm_config();
        let medium = medium_rssm_config();

        assert!(small.hidden_dim < medium.hidden_dim);
        assert!(small.deter_dim < medium.deter_dim);
    }

    #[test]
    fn test_policy_configs() {
        let small = small_actor_critic_config();
        let medium = medium_actor_critic_config();

        assert!(small.actor_hidden[0] < medium.actor_hidden[0]);
    }

    #[test]
    fn test_env_configs() {
        let reach = panda_reach_config();
        let pick = panda_pick_place_config();

        assert_eq!(reach.task_type, TaskType::Reach);
        assert_eq!(pick.task_type, TaskType::PickPlace);
        assert!(pick.use_camera);
    }

    #[test]
    fn test_training_configs() {
        let fast = fast_training_config();
        let standard = standard_training_config();

        assert!(fast.epochs < standard.epochs);
        assert!(fast.buffer_size < standard.buffer_size);
    }

    #[test]
    fn test_sim2real_configs() {
        let conservative = conservative_sim2real_config();
        let aggressive = aggressive_sim2real_config();

        assert!(conservative.action_scale < aggressive.action_scale);
        assert!(
            conservative.domain_randomization.action_noise
                < aggressive.domain_randomization.action_noise
        );
    }
}