use itertools::multizip;
use kalast::{
    compute_conductivity, compute_diffusivity, compute_ground_step_from_basic,
    compute_ground_step_from_time_and_stability, linspace, matrix_flat_to_vector,
    EnvironmentSystem, ExportFile, Properties, SingleBody, Time, World, DAY, NUMERICAL_STABILITY,
    TAU, V3, VX,
};

fn numerical(properties: &MyProperties) -> VX<f64> {
    // Define a 1D rod, initialized from a single facet.
    // Ground temperatures cannot be initialized before the time step is share with the body, which
    // happens after the world definition.
    let mut rod = SingleBody::new(
        "1D rod",
        "rsc/obj/facet.obj",
        V3::new(0.0, 0.0, 0.0),
        Properties::new(
            0.0,                        // rotation period
            0.0,                        // revolution period
            0.0,                        // obliquity
            properties.thermal_inertia, // thermal inertia
            properties.density,         // density
            properties.heat_capacity,   // heat capacity
            0.0,                        // albedo
            0.0,                        // emissivity
        ),
    );

    // We manually set the ground depth as the revolution period is zero (causing the thermal skin
    // depth to be zero too).
    rod.body_mut()
        .properties
        .set_ground_depth(properties.ground_depth);

    // The Dirichlet problem is defined such as the temperature are fixed to zero at both boundary
    // conditions of the 1D rod during the whole simulation.
    rod.body_mut().set_fixed_surface_temperatures(true);
    rod.body_mut().set_fixed_ground_limit_temperatures(true);

    // Time
    let time = Time::new(properties.duration, properties.time_step);

    // World
    let mut world = World::new(time, rod);
    world.no_sun(true);

    // Now that the world is defined, the time step is shared with the rod. Thus the ground
    // temperatures matrix is allocated with correct size, and can now be initialized.
    // Here we set the rod to 300K but boundaries to 0K.
    let rod_size = world
        .environment_system
        .body_mut()
        .ground_temperatures_mut()
        .ncols();
    world
        .environment_system
        .body_mut()
        .ground_temperatures_mut()
        .fill(properties.temperatures_initial);
    world
        .environment_system
        .body_mut()
        .ground_temperatures_mut()[(0, 0)] = 0.;
    world
        .environment_system
        .body_mut()
        .ground_temperatures_mut()[(0, rod_size - 1)] = 0.0;

    world.start();

    let temperatures =
        matrix_flat_to_vector(&world.environment_system.body().ground_temperatures().row(0));

    temperatures
}

fn analytical(properties: &MyProperties, number_combination: usize) -> VX<f64> {
    // Complete necessary properties parameters.
    let conductivity = compute_conductivity(
        properties.thermal_inertia,
        properties.density,
        properties.heat_capacity,
    );
    let diffusivity =
        compute_diffusivity(conductivity, properties.density, properties.heat_capacity);
    let ground_step = compute_ground_step_from_time_and_stability(
        properties.time_step,
        diffusivity,
        NUMERICAL_STABILITY,
    );
    let mut ground_vector: VX<f64> = linspace(0.0, properties.ground_depth, ground_step);
    let vector_size = ground_vector.len();
    ground_vector[vector_size - 1] = 0.0;

    // Initialize temperatures
    let mut temperatures = VX::zeros(vector_size);

    // Loop sinus combinations.
    for combination in 1..number_combination + 1 {
        let combination_temperature = -4.0 * properties.temperatures_initial
            / (TAU * combination as f64)
            * ((TAU * combination as f64 / 2.0).cos() - 1.0);
        let sinus = VX::from_iterator(
            vector_size,
            (combination as f64 * TAU * ground_vector.clone() / (2.0 * properties.ground_depth))
                .iter()
                .map(|e| e.sin()),
        );
        let exponential = (-(combination as f64 * TAU / (2.0 * properties.ground_depth)).powi(2)
            * diffusivity
            * properties.duration)
            .exp();
        temperatures += combination_temperature * sinus * exponential;
    }

    temperatures
}

struct MyProperties {
    time_step: f64,
    ground_depth: f64,
    duration: f64,
    thermal_inertia: f64,
    density: f64,
    heat_capacity: f64,
    temperatures_initial: f64,
}

fn main() {
    // Define properties.
    let properties = MyProperties {
        duration: 1.0 * DAY,
        time_step: 1.0,
        ground_depth: 1.0,
        thermal_inertia: 500.0,
        density: 2146.0,
        heat_capacity: 600.0,
        temperatures_initial: 300.0,
    };
    let number_combination = 1e1 as usize;

    // Compute temperatures.
    let ground_vector: VX<f64> = linspace(
        0.0,
        properties.ground_depth,
        compute_ground_step_from_basic(
            properties.time_step,
            properties.thermal_inertia,
            properties.density,
            properties.heat_capacity,
            NUMERICAL_STABILITY,
        ),
    );
    let temperatures_numerical = numerical(&properties);
    let temperatures_analytical = analytical(&properties, number_combination);

    // Some prints to check.
    println!("NUMERICAL\n---------");
    println!(
        "max: {:.10?}, mean: {:?}",
        temperatures_numerical.argmax(),
        temperatures_numerical.mean(),
    );
    println!("\nANALYTICAL\n----------");
    println!(
        "max: {:.10?}, mean: {:?}",
        temperatures_analytical.argmax(),
        temperatures_analytical.mean(),
    );
    println!(
        "{} {} {}",
        ground_vector.len(),
        temperatures_numerical.len(),
        temperatures_analytical.len()
    );

    // Export data.
    let mut file = ExportFile::new("rsc/data/tmp.txt");
    file.add_to_header(format!("Dirichlet problem",));
    file.add_to_header(format!("Duration: {} s", properties.duration));
    file.add_to_header(format!("Time step: {} s", properties.time_step));
    file.add_to_header(format!("Ground depth: {} m", properties.ground_depth));
    file.add_to_header(format!(
        "Thermal intertia: {:0} J.K^{{-1}}.m^{{-2}}.s^{{-1/2}}",
        properties.thermal_inertia
    ));
    file.add_to_header(format!("Density: {:0} kg.m^{{-3}}", properties.density));
    file.add_to_header(format!(
        "Heat capacity: {:0} m^{{2}}.s^{{-2}}.K^{{-1}}",
        properties.heat_capacity
    ));
    file.add_to_header(format!(
        "Initial temperature: {:0} K",
        properties.temperatures_initial
    ));
    file.add_to_header(format!(
        "Number combination sinus: {:0}",
        number_combination
    ));
    file.add_to_header(format!(
        "Depth [m] | Temperatures numeric [K] | Temperatures analytic [K]",
    ));
    for (depth, numeric, analytic) in multizip((
        ground_vector.iter(),
        temperatures_numerical.iter(),
        temperatures_analytical.iter(),
    )) {
        file.add_to_contents(format!("{:.6e} {:.6e} {:.6e}", depth, numeric, analytic));
    }
    file.write();
}
