//! Sample data for WorldGen demo.

use worldgen_shared::{
    GenerationRequest, LightingState, ModelConfig, ModelType, PhysicsConstraint,
    PhysicsConstraintType, PhysicsParameters, SchedulerConfig, SchedulerType, TrainingConfig,
    WorldObject, WorldState,
};

/// Create a sample generation request for a bouncing ball.
#[must_use]
pub fn bouncing_ball_request() -> GenerationRequest {
    GenerationRequest {
        prompt: "A red ball bouncing on a wooden floor in a sunlit room".to_string(),
        negative_prompt: Some("blurry, distorted, unrealistic physics".to_string()),
        num_frames: 24,
        width: 256,
        height: 256,
        fps: 12.0,
        physics_constraints: vec![
            PhysicsConstraint {
                constraint_type: PhysicsConstraintType::Gravity,
                weight: 1.0,
                parameters: PhysicsParameters {
                    gravity: Some([0.0, -9.81, 0.0]),
                    ..Default::default()
                },
            },
            PhysicsConstraint {
                constraint_type: PhysicsConstraintType::Collision,
                weight: 0.9,
                parameters: PhysicsParameters {
                    restitution: Some(0.8),
                    ..Default::default()
                },
            },
            PhysicsConstraint {
                constraint_type: PhysicsConstraintType::TemporalCoherence,
                weight: 0.7,
                parameters: Default::default(),
            },
        ],
        seed: Some(42),
        guidance_scale: 7.5,
        num_inference_steps: 50,
    }
}

/// Create a sample generation request for flowing water.
#[must_use]
pub fn flowing_water_request() -> GenerationRequest {
    GenerationRequest {
        prompt: "Crystal clear water flowing over smooth river rocks, sunlight sparkling"
            .to_string(),
        negative_prompt: Some("frozen, static, unrealistic".to_string()),
        num_frames: 32,
        width: 512,
        height: 288,
        fps: 24.0,
        physics_constraints: vec![
            PhysicsConstraint {
                constraint_type: PhysicsConstraintType::Fluid,
                weight: 1.0,
                parameters: PhysicsParameters {
                    viscosity: Some(0.001), // Water viscosity
                    ..Default::default()
                },
            },
            PhysicsConstraint {
                constraint_type: PhysicsConstraintType::Lighting,
                weight: 0.8,
                parameters: PhysicsParameters {
                    light_direction: Some([0.3, -0.8, 0.5]),
                    ..Default::default()
                },
            },
            PhysicsConstraint {
                constraint_type: PhysicsConstraintType::TemporalCoherence,
                weight: 0.9,
                parameters: Default::default(),
            },
        ],
        seed: Some(123),
        guidance_scale: 8.0,
        num_inference_steps: 75,
    }
}

/// Create a sample generation request for a pendulum.
#[must_use]
pub fn pendulum_request() -> GenerationRequest {
    GenerationRequest {
        prompt: "A brass pendulum swinging in a grandfather clock, smooth motion".to_string(),
        negative_prompt: Some("jerky, unnatural motion".to_string()),
        num_frames: 48,
        width: 256,
        height: 384,
        fps: 30.0,
        physics_constraints: vec![
            PhysicsConstraint {
                constraint_type: PhysicsConstraintType::Gravity,
                weight: 1.0,
                parameters: PhysicsParameters {
                    gravity: Some([0.0, -9.81, 0.0]),
                    ..Default::default()
                },
            },
            PhysicsConstraint {
                constraint_type: PhysicsConstraintType::RigidBody,
                weight: 1.0,
                parameters: Default::default(),
            },
            PhysicsConstraint {
                constraint_type: PhysicsConstraintType::Momentum,
                weight: 0.9,
                parameters: Default::default(),
            },
        ],
        seed: Some(456),
        guidance_scale: 7.0,
        num_inference_steps: 60,
    }
}

/// Create a sample generation request for a car driving.
#[must_use]
pub fn driving_car_request() -> GenerationRequest {
    GenerationRequest {
        prompt: "A silver sports car driving along a coastal highway at sunset".to_string(),
        negative_prompt: Some("floating, unrealistic shadows, glitchy".to_string()),
        num_frames: 64,
        width: 512,
        height: 288,
        fps: 24.0,
        physics_constraints: vec![
            PhysicsConstraint {
                constraint_type: PhysicsConstraintType::RigidBody,
                weight: 1.0,
                parameters: Default::default(),
            },
            PhysicsConstraint {
                constraint_type: PhysicsConstraintType::Shadows,
                weight: 0.8,
                parameters: Default::default(),
            },
            PhysicsConstraint {
                constraint_type: PhysicsConstraintType::Lighting,
                weight: 0.9,
                parameters: PhysicsParameters {
                    light_direction: Some([-0.5, -0.3, 0.8]),
                    ..Default::default()
                },
            },
            PhysicsConstraint {
                constraint_type: PhysicsConstraintType::TemporalCoherence,
                weight: 0.95,
                parameters: Default::default(),
            },
        ],
        seed: Some(789),
        guidance_scale: 8.5,
        num_inference_steps: 80,
    }
}

/// Create a sample generation request for a jellyfish.
#[must_use]
pub fn jellyfish_request() -> GenerationRequest {
    GenerationRequest {
        prompt: "A bioluminescent jellyfish pulsing through dark ocean water".to_string(),
        negative_prompt: Some("static, rigid, unrealistic".to_string()),
        num_frames: 40,
        width: 384,
        height: 384,
        fps: 20.0,
        physics_constraints: vec![
            PhysicsConstraint {
                constraint_type: PhysicsConstraintType::SoftBody,
                weight: 1.0,
                parameters: Default::default(),
            },
            PhysicsConstraint {
                constraint_type: PhysicsConstraintType::Fluid,
                weight: 0.7,
                parameters: PhysicsParameters {
                    viscosity: Some(0.01),
                    ..Default::default()
                },
            },
            PhysicsConstraint {
                constraint_type: PhysicsConstraintType::Lighting,
                weight: 0.6,
                parameters: Default::default(),
            },
        ],
        seed: Some(999),
        guidance_scale: 9.0,
        num_inference_steps: 70,
    }
}

/// Create a sample generation request for falling leaves.
#[must_use]
pub fn falling_leaves_request() -> GenerationRequest {
    GenerationRequest {
        prompt: "Autumn leaves gently falling from a maple tree in golden afternoon light"
            .to_string(),
        negative_prompt: Some("unnatural motion, frozen".to_string()),
        num_frames: 60,
        width: 384,
        height: 512,
        fps: 24.0,
        physics_constraints: vec![
            PhysicsConstraint {
                constraint_type: PhysicsConstraintType::Gravity,
                weight: 0.6,
                parameters: PhysicsParameters {
                    gravity: Some([0.0, -2.0, 0.0]), // Slower falling
                    ..Default::default()
                },
            },
            PhysicsConstraint {
                constraint_type: PhysicsConstraintType::Fluid,
                weight: 0.4, // Air resistance
                parameters: PhysicsParameters {
                    viscosity: Some(0.1),
                    ..Default::default()
                },
            },
            PhysicsConstraint {
                constraint_type: PhysicsConstraintType::Lighting,
                weight: 0.8,
                parameters: PhysicsParameters {
                    light_direction: Some([0.4, -0.6, 0.7]),
                    ..Default::default()
                },
            },
        ],
        seed: Some(1111),
        guidance_scale: 7.5,
        num_inference_steps: 55,
    }
}

/// Create a minimal generation request for testing.
#[must_use]
pub fn minimal_test_request() -> GenerationRequest {
    GenerationRequest {
        prompt: "A simple scene".to_string(),
        negative_prompt: None,
        num_frames: 4,
        width: 64,
        height: 64,
        fps: 8.0,
        physics_constraints: vec![],
        seed: Some(1),
        guidance_scale: 5.0,
        num_inference_steps: 5,
    }
}

/// Create a small DiT model configuration for demos.
#[must_use]
pub fn small_dit_config() -> ModelConfig {
    ModelConfig {
        model_type: ModelType::DiT,
        hidden_dim: 256,
        num_layers: 4,
        num_heads: 4,
        patch_size: 4,
        latent_channels: 4,
        temporal_attention: true,
        cross_attention: true,
    }
}

/// Create a medium DiT model configuration.
#[must_use]
pub fn medium_dit_config() -> ModelConfig {
    ModelConfig {
        model_type: ModelType::DiT,
        hidden_dim: 512,
        num_layers: 8,
        num_heads: 8,
        patch_size: 2,
        latent_channels: 4,
        temporal_attention: true,
        cross_attention: true,
    }
}

/// Create a large DiT model configuration.
#[must_use]
pub fn large_dit_config() -> ModelConfig {
    ModelConfig {
        model_type: ModelType::DiT,
        hidden_dim: 1024,
        num_layers: 16,
        num_heads: 16,
        patch_size: 2,
        latent_channels: 8,
        temporal_attention: true,
        cross_attention: true,
    }
}

/// Create a DDIM scheduler configuration.
#[must_use]
pub fn ddim_scheduler_config() -> SchedulerConfig {
    SchedulerConfig {
        scheduler_type: SchedulerType::DDIM,
        num_timesteps: 1000,
        beta_start: 0.0001,
        beta_end: 0.02,
        beta_schedule: "linear".to_string(),
    }
}

/// Create a cosine scheduler configuration.
#[must_use]
pub fn cosine_scheduler_config() -> SchedulerConfig {
    SchedulerConfig {
        scheduler_type: SchedulerType::DDPM,
        num_timesteps: 1000,
        beta_start: 0.0001,
        beta_end: 0.02,
        beta_schedule: "cosine".to_string(),
    }
}

/// Create a fast scheduler configuration for testing.
#[must_use]
pub fn fast_scheduler_config() -> SchedulerConfig {
    SchedulerConfig {
        scheduler_type: SchedulerType::Euler,
        num_timesteps: 100,
        beta_start: 0.001,
        beta_end: 0.02,
        beta_schedule: "linear".to_string(),
    }
}

/// Create a sample training configuration.
#[must_use]
pub fn sample_training_config() -> TrainingConfig {
    TrainingConfig {
        epochs: 100,
        batch_size: 4,
        learning_rate: 1e-4,
        weight_decay: 0.01,
        gradient_accumulation: 4,
        mixed_precision: true,
        physics_weight: 0.1,
        seed: Some(42),
    }
}

/// Create a fast training configuration for demos.
#[must_use]
pub fn fast_training_config() -> TrainingConfig {
    TrainingConfig {
        epochs: 10,
        batch_size: 2,
        learning_rate: 1e-4,
        weight_decay: 0.01,
        gradient_accumulation: 1,
        mixed_precision: false,
        physics_weight: 0.1,
        seed: Some(42),
    }
}

/// Create a sample world state for physics simulation.
#[must_use]
pub fn sample_world_state() -> WorldState {
    WorldState {
        objects: vec![
            WorldObject {
                id: "ball".to_string(),
                object_type: "sphere".to_string(),
                position: [0.0, 2.0, 0.0],
                velocity: [1.0, 0.0, 0.0],
                rotation: [0.0, 0.0, 0.0],
                scale: [0.5, 0.5, 0.5],
                mass: 1.0,
            },
            WorldObject {
                id: "floor".to_string(),
                object_type: "plane".to_string(),
                position: [0.0, 0.0, 0.0],
                velocity: [0.0, 0.0, 0.0],
                rotation: [0.0, 0.0, 0.0],
                scale: [10.0, 1.0, 10.0],
                mass: f32::INFINITY,
            },
            WorldObject {
                id: "wall".to_string(),
                object_type: "box".to_string(),
                position: [5.0, 1.0, 0.0],
                velocity: [0.0, 0.0, 0.0],
                rotation: [0.0, 0.0, 0.0],
                scale: [0.5, 2.0, 5.0],
                mass: f32::INFINITY,
            },
        ],
        camera_position: [5.0, 3.0, 5.0],
        camera_rotation: [-0.5, 0.8, 0.0],
        lighting: LightingState {
            ambient: [0.1, 0.1, 0.12],
            sun_direction: [0.5, -0.8, 0.3],
            sun_color: [1.0, 0.95, 0.9],
            time_of_day: 14.0,
        },
        time: 0.0,
    }
}

/// Create an underwater world state.
#[must_use]
pub fn underwater_world_state() -> WorldState {
    WorldState {
        objects: vec![
            WorldObject {
                id: "jellyfish".to_string(),
                object_type: "soft_body".to_string(),
                position: [0.0, 5.0, 0.0],
                velocity: [0.0, -0.2, 0.0],
                rotation: [0.0, 0.0, 0.0],
                scale: [1.0, 1.5, 1.0],
                mass: 0.5,
            },
            WorldObject {
                id: "seabed".to_string(),
                object_type: "terrain".to_string(),
                position: [0.0, 0.0, 0.0],
                velocity: [0.0, 0.0, 0.0],
                rotation: [0.0, 0.0, 0.0],
                scale: [20.0, 1.0, 20.0],
                mass: f32::INFINITY,
            },
        ],
        camera_position: [3.0, 5.0, 3.0],
        camera_rotation: [-0.3, 0.5, 0.0],
        lighting: LightingState {
            ambient: [0.05, 0.1, 0.15],
            sun_direction: [0.0, -1.0, 0.0],
            sun_color: [0.3, 0.5, 0.7],
            time_of_day: 12.0,
        },
        time: 0.0,
    }
}

/// Get all sample generation requests.
#[must_use]
pub fn all_sample_requests() -> Vec<GenerationRequest> {
    vec![
        bouncing_ball_request(),
        flowing_water_request(),
        pendulum_request(),
        driving_car_request(),
        jellyfish_request(),
        falling_leaves_request(),
    ]
}

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

    #[test]
    fn test_bouncing_ball_request() {
        let req = bouncing_ball_request();
        assert!(!req.prompt.is_empty());
        assert!(req.num_frames > 0);
        assert!(!req.physics_constraints.is_empty());
    }

    #[test]
    fn test_flowing_water_request() {
        let req = flowing_water_request();
        assert!(req.prompt.contains("water"));
        assert!(
            req.physics_constraints
                .iter()
                .any(|c| c.constraint_type == PhysicsConstraintType::Fluid)
        );
    }

    #[test]
    fn test_pendulum_request() {
        let req = pendulum_request();
        assert!(req.prompt.contains("pendulum"));
        assert!(
            req.physics_constraints
                .iter()
                .any(|c| c.constraint_type == PhysicsConstraintType::RigidBody)
        );
    }

    #[test]
    fn test_minimal_test_request() {
        let req = minimal_test_request();
        assert_eq!(req.num_frames, 4);
        assert_eq!(req.width, 64);
        assert_eq!(req.height, 64);
        assert!(req.physics_constraints.is_empty());
    }

    #[test]
    fn test_model_configs() {
        let small = small_dit_config();
        let medium = medium_dit_config();
        let large = large_dit_config();

        assert!(small.hidden_dim < medium.hidden_dim);
        assert!(medium.hidden_dim < large.hidden_dim);
        assert!(small.num_layers < medium.num_layers);
        assert!(medium.num_layers < large.num_layers);
    }

    #[test]
    fn test_scheduler_configs() {
        let ddim = ddim_scheduler_config();
        let cosine = cosine_scheduler_config();
        let fast = fast_scheduler_config();

        assert_eq!(ddim.scheduler_type, SchedulerType::DDIM);
        assert_eq!(cosine.beta_schedule, "cosine");
        assert!(fast.num_timesteps < ddim.num_timesteps);
    }

    #[test]
    fn test_training_configs() {
        let sample = sample_training_config();
        let fast = fast_training_config();

        assert!(fast.epochs < sample.epochs);
        assert!(fast.batch_size <= sample.batch_size);
    }

    #[test]
    fn test_world_states() {
        let world = sample_world_state();
        assert!(!world.objects.is_empty());
        assert!(world.objects.iter().any(|o| o.id == "ball"));
        assert!(world.objects.iter().any(|o| o.id == "floor"));

        let underwater = underwater_world_state();
        assert!(underwater.objects.iter().any(|o| o.id == "jellyfish"));
    }

    #[test]
    fn test_all_sample_requests() {
        let requests = all_sample_requests();
        assert!(!requests.is_empty());

        for req in requests {
            assert!(!req.prompt.is_empty());
            assert!(req.num_frames > 0);
            assert!(req.width > 0);
            assert!(req.height > 0);
        }
    }
}