//! Sample data generation for NeuralOp Studio demos.

use neuralop_studio_shared::{
    ActivationType, BCType, BoundaryCondition, Domain, InitialCondition, OperatorConfig,
    OperatorType, PDEDefinition, PDEParameters, PDEType, SchedulerType, TrainingConfig,
};

/// Create a complete demo configuration.
#[must_use]
pub fn create_demo_config() -> (PDEDefinition, OperatorConfig, TrainingConfig) {
    (
        create_poisson_2d(),
        create_fno_config(),
        create_training_config(),
    )
}

/// Create a 2D Poisson problem.
#[must_use]
pub fn create_poisson_2d() -> PDEDefinition {
    PDEDefinition {
        pde_type: PDEType::Poisson,
        equation: r"-\nabla^2 u = \sin(\pi x) \sin(\pi y)".to_string(),
        domain: Domain {
            dimensions: 2,
            bounds: vec![(0.0, 1.0), (0.0, 1.0)],
            resolution: vec![64, 64],
            time_bounds: None,
            time_resolution: None,
        },
        boundary_conditions: vec![BoundaryCondition {
            boundary: "all".to_string(),
            bc_type: BCType::Dirichlet,
            value: "0".to_string(),
            robin_coefficients: None,
        }],
        initial_condition: None,
        parameters: PDEParameters {
            source_term: Some("sin(pi*x)*sin(pi*y)".to_string()),
            ..Default::default()
        },
    }
}

/// Create a 1D heat equation problem.
#[must_use]
pub fn create_heat_1d() -> PDEDefinition {
    PDEDefinition {
        pde_type: PDEType::Heat,
        equation: r"\frac{\partial u}{\partial t} = 0.01 \nabla^2 u".to_string(),
        domain: Domain {
            dimensions: 1,
            bounds: vec![(0.0, 1.0)],
            resolution: vec![128],
            time_bounds: Some((0.0, 0.5)),
            time_resolution: Some(50),
        },
        boundary_conditions: vec![
            BoundaryCondition {
                boundary: "left".to_string(),
                bc_type: BCType::Dirichlet,
                value: "0".to_string(),
                robin_coefficients: None,
            },
            BoundaryCondition {
                boundary: "right".to_string(),
                bc_type: BCType::Dirichlet,
                value: "0".to_string(),
                robin_coefficients: None,
            },
        ],
        initial_condition: Some(InitialCondition {
            function: "sin(pi*x)".to_string(),
            velocity: None,
        }),
        parameters: PDEParameters {
            diffusion: Some(0.01),
            ..Default::default()
        },
    }
}

/// Create a 1D Burgers' equation problem.
#[must_use]
pub fn create_burgers_1d() -> PDEDefinition {
    PDEDefinition {
        pde_type: PDEType::Burgers,
        equation:
            r"\frac{\partial u}{\partial t} + u\frac{\partial u}{\partial x} = 0.01 \nabla^2 u"
                .to_string(),
        domain: Domain {
            dimensions: 1,
            bounds: vec![(0.0, 2.0 * std::f64::consts::PI)],
            resolution: vec![256],
            time_bounds: Some((0.0, 1.0)),
            time_resolution: Some(100),
        },
        boundary_conditions: vec![BoundaryCondition {
            boundary: "all".to_string(),
            bc_type: BCType::Periodic,
            value: "periodic".to_string(),
            robin_coefficients: None,
        }],
        initial_condition: Some(InitialCondition {
            function: "sin(x)".to_string(),
            velocity: None,
        }),
        parameters: PDEParameters {
            diffusion: Some(0.01),
            ..Default::default()
        },
    }
}

/// Create a 2D Navier-Stokes problem.
#[must_use]
pub fn create_navier_stokes_2d() -> PDEDefinition {
    PDEDefinition {
        pde_type: PDEType::NavierStokes,
        equation: r"\frac{\partial \mathbf{u}}{\partial t} + (\mathbf{u} \cdot \nabla)\mathbf{u} = -\nabla p + \nu \nabla^2 \mathbf{u}".to_string(),
        domain: Domain {
            dimensions: 2,
            bounds: vec![(0.0, 2.0 * std::f64::consts::PI), (0.0, 2.0 * std::f64::consts::PI)],
            resolution: vec![64, 64],
            time_bounds: Some((0.0, 10.0)),
            time_resolution: Some(100),
        },
        boundary_conditions: vec![BoundaryCondition {
            boundary: "all".to_string(),
            bc_type: BCType::Periodic,
            value: "periodic".to_string(),
            robin_coefficients: None,
        }],
        initial_condition: Some(InitialCondition {
            function: "kolmogorov_flow".to_string(),
            velocity: None,
        }),
        parameters: PDEParameters {
            reynolds_number: Some(1000.0),
            diffusion: Some(0.001),
            ..Default::default()
        },
    }
}

/// Create a 1D wave equation problem.
#[must_use]
pub fn create_wave_1d() -> PDEDefinition {
    PDEDefinition {
        pde_type: PDEType::Wave,
        equation: r"\frac{\partial^2 u}{\partial t^2} = c^2 \nabla^2 u".to_string(),
        domain: Domain {
            dimensions: 1,
            bounds: vec![(0.0, 1.0)],
            resolution: vec![128],
            time_bounds: Some((0.0, 2.0)),
            time_resolution: Some(200),
        },
        boundary_conditions: vec![
            BoundaryCondition {
                boundary: "left".to_string(),
                bc_type: BCType::Dirichlet,
                value: "0".to_string(),
                robin_coefficients: None,
            },
            BoundaryCondition {
                boundary: "right".to_string(),
                bc_type: BCType::Dirichlet,
                value: "0".to_string(),
                robin_coefficients: None,
            },
        ],
        initial_condition: Some(InitialCondition {
            function: "gaussian".to_string(),
            velocity: Some("0".to_string()),
        }),
        parameters: PDEParameters {
            wave_speed: Some(1.0),
            ..Default::default()
        },
    }
}

/// Create FNO configuration.
#[must_use]
pub fn create_fno_config() -> OperatorConfig {
    OperatorConfig {
        operator_type: OperatorType::FNO,
        hidden_dim: 64,
        num_layers: 4,
        fourier_modes: Some(12),
        branch_width: None,
        trunk_width: None,
        physics_weight: None,
        activation: ActivationType::GELU,
        residual: true,
        dropout: 0.0,
    }
}

/// Create DeepONet configuration.
#[must_use]
pub fn create_deeponet_config() -> OperatorConfig {
    OperatorConfig {
        operator_type: OperatorType::DeepONet,
        hidden_dim: 100,
        num_layers: 6,
        fourier_modes: None,
        branch_width: Some(100),
        trunk_width: Some(100),
        physics_weight: None,
        activation: ActivationType::Tanh,
        residual: false,
        dropout: 0.0,
    }
}

/// Create PINO configuration.
#[must_use]
pub fn create_pino_config() -> OperatorConfig {
    OperatorConfig {
        operator_type: OperatorType::PINO,
        hidden_dim: 64,
        num_layers: 4,
        fourier_modes: Some(12),
        branch_width: None,
        trunk_width: None,
        physics_weight: Some(0.1),
        activation: ActivationType::GELU,
        residual: true,
        dropout: 0.0,
    }
}

/// Create training configuration.
#[must_use]
pub fn create_training_config() -> TrainingConfig {
    TrainingConfig {
        epochs: 100,
        batch_size: 32,
        learning_rate: 1e-3,
        scheduler: SchedulerType::CosineAnnealing,
        optimizer: neuralop_studio_shared::OptimizerType::AdamW,
        weight_decay: 1e-4,
        num_train_samples: 1000,
        num_val_samples: 100,
        num_test_samples: 100,
        seed: Some(42),
    }
}

/// Create a quick demo configuration (for testing).
#[must_use]
pub fn create_quick_demo_config() -> (PDEDefinition, OperatorConfig, TrainingConfig) {
    let mut pde = create_poisson_2d();
    pde.domain.resolution = vec![32, 32];

    let config = create_fno_config();

    let training = TrainingConfig {
        epochs: 10,
        batch_size: 16,
        learning_rate: 1e-3,
        scheduler: SchedulerType::Constant,
        optimizer: neuralop_studio_shared::OptimizerType::Adam,
        weight_decay: 0.0,
        num_train_samples: 100,
        num_val_samples: 20,
        num_test_samples: 20,
        seed: Some(42),
    };

    (pde, config, training)
}

/// Generate synthetic training data for a PDE.
#[must_use]
pub fn generate_training_data(
    pde: &PDEDefinition,
    num_samples: usize,
    seed: u64,
) -> (Vec<Vec<f64>>, Vec<Vec<f64>>) {
    let grid_size: usize = pde.domain.resolution.iter().product();
    let mut rng_state = seed;

    let random = |state: &mut u64| -> f64 {
        *state = state
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        (*state >> 11) as f64 / (1u64 << 53) as f64
    };

    let random_normal = |state: &mut u64| -> f64 {
        let u1 = random(state) + 1e-10;
        let u2 = random(state);
        (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos()
    };

    let inputs: Vec<Vec<f64>> = (0..num_samples)
        .map(|_| {
            (0..grid_size)
                .map(|j| {
                    let x = (j % pde.domain.resolution[0]) as f64 / pde.domain.resolution[0] as f64;
                    let y = if pde.domain.dimensions > 1 {
                        (j / pde.domain.resolution[0]) as f64
                            / pde.domain.resolution.get(1).copied().unwrap_or(1) as f64
                    } else {
                        0.0
                    };
                    (std::f64::consts::PI * x).sin() * (std::f64::consts::PI * y).sin()
                        + random_normal(&mut rng_state) * 0.1
                })
                .collect()
        })
        .collect();

    let outputs: Vec<Vec<f64>> = inputs
        .iter()
        .map(|input| {
            // Simplified: solution is related to source
            input.iter().map(|&v| v * 0.5).collect()
        })
        .collect();

    (inputs, outputs)
}

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

    #[test]
    fn test_demo_config() {
        let (pde, op, train) = create_demo_config();
        assert_eq!(pde.pde_type, PDEType::Poisson);
        assert_eq!(op.operator_type, OperatorType::FNO);
        assert_eq!(train.epochs, 100);
    }

    #[test]
    fn test_poisson_2d() {
        let pde = create_poisson_2d();
        assert_eq!(pde.domain.dimensions, 2);
        assert_eq!(pde.domain.resolution, vec![64, 64]);
    }

    #[test]
    fn test_heat_1d() {
        let pde = create_heat_1d();
        assert_eq!(pde.domain.dimensions, 1);
        assert!(pde.domain.time_bounds.is_some());
    }

    #[test]
    fn test_burgers_1d() {
        let pde = create_burgers_1d();
        assert_eq!(pde.pde_type, PDEType::Burgers);
        assert!(pde.parameters.diffusion.is_some());
    }

    #[test]
    fn test_navier_stokes_2d() {
        let pde = create_navier_stokes_2d();
        assert_eq!(pde.pde_type, PDEType::NavierStokes);
        assert!(pde.parameters.reynolds_number.is_some());
    }

    #[test]
    fn test_wave_1d() {
        let pde = create_wave_1d();
        assert_eq!(pde.pde_type, PDEType::Wave);
        assert!(pde.parameters.wave_speed.is_some());
    }

    #[test]
    fn test_operator_configs() {
        let fno = create_fno_config();
        assert_eq!(fno.operator_type, OperatorType::FNO);
        assert!(fno.fourier_modes.is_some());

        let deeponet = create_deeponet_config();
        assert_eq!(deeponet.operator_type, OperatorType::DeepONet);
        assert!(deeponet.branch_width.is_some());

        let pino = create_pino_config();
        assert_eq!(pino.operator_type, OperatorType::PINO);
        assert!(pino.physics_weight.is_some());
    }

    #[test]
    fn test_training_config() {
        let config = create_training_config();
        assert_eq!(config.epochs, 100);
        assert_eq!(config.batch_size, 32);
    }

    #[test]
    fn test_quick_demo_config() {
        let (pde, _, train) = create_quick_demo_config();
        assert_eq!(pde.domain.resolution, vec![32, 32]);
        assert_eq!(train.epochs, 10);
    }

    #[test]
    fn test_generate_training_data() {
        let pde = create_poisson_2d();
        let (inputs, outputs) = generate_training_data(&pde, 10, 42);

        assert_eq!(inputs.len(), 10);
        assert_eq!(outputs.len(), 10);
        assert_eq!(inputs[0].len(), 64 * 64);
    }
}