#![feature(concat_idents)]
mod rust_cfd;
mod setup;

use clap::Parser;

pub use crate::setup::*;

use hercules_rt::{runner, HerculesImmBox, HerculesImmBoxTo, HerculesMutBox, HerculesMutBoxTo};

juno_build::juno!("euler");
juno_build::juno!("pre_euler");

#[derive(Parser)]
#[clap(author, version, about, long_about = None)]
pub struct CFDInputs {
    pub data_file: String,
    pub iterations: usize,
    pub block_size: usize,
    #[clap(short = None, long = Some("pre-euler"))]
    pub pre_euler: bool,
}

fn run_euler(
    nelr: usize,
    iterations: usize,
    mut variables: Variables,
    areas: &[f32],
    elements_surrounding_elements: &[i32],
    normals: &Normals,
    ff_variable: &Variable,
    ff_fc_density_energy: &Float3,
    ff_fc_momentum_x: &Float3,
    ff_fc_momentum_y: &Float3,
    ff_fc_momentum_z: &Float3,
) -> Variables {
    let mut v_density = HerculesMutBox::from(variables.density.as_mut_slice());
    let mut v_momentum_x = HerculesMutBox::from(variables.momentum.x.as_mut_slice());
    let mut v_momentum_y = HerculesMutBox::from(variables.momentum.y.as_mut_slice());
    let mut v_momentum_z = HerculesMutBox::from(variables.momentum.z.as_mut_slice());
    let mut v_energy = HerculesMutBox::from(variables.energy.as_mut_slice());

    let areas = HerculesImmBox::from(areas);
    let elements_surrounding_elements = HerculesImmBox::from(elements_surrounding_elements);

    let normals_x = HerculesImmBox::from(normals.x.as_slice());
    let normals_y = HerculesImmBox::from(normals.y.as_slice());
    let normals_z = HerculesImmBox::from(normals.z.as_slice());

    let mut runner = runner!(euler);
    let (density, momentum_x, momentum_y, momentum_z, energy) = async_std::task::block_on(async {
        runner
            .run(
                nelr as u64,
                iterations as u64,
                v_density.to(),
                v_momentum_x.to(),
                v_momentum_y.to(),
                v_momentum_z.to(),
                v_energy.to(),
                areas.to(),
                elements_surrounding_elements.to(),
                normals_x.to(),
                normals_y.to(),
                normals_z.to(),
                ff_variable.density,
                ff_variable.momentum.x,
                ff_variable.momentum.y,
                ff_variable.momentum.z,
                ff_variable.energy,
                ff_fc_density_energy.x,
                ff_fc_density_energy.y,
                ff_fc_density_energy.z,
                ff_fc_momentum_x.x,
                ff_fc_momentum_x.y,
                ff_fc_momentum_x.z,
                ff_fc_momentum_y.x,
                ff_fc_momentum_y.y,
                ff_fc_momentum_y.z,
                ff_fc_momentum_z.x,
                ff_fc_momentum_z.y,
                ff_fc_momentum_z.z,
            )
            .await
    });

    Variables {
        density: AlignedSlice::from_slice(HerculesMutBox::from(density).as_slice()),
        momentum: Momentum {
            x: AlignedSlice::from_slice(HerculesMutBox::from(momentum_x).as_slice()),
            y: AlignedSlice::from_slice(HerculesMutBox::from(momentum_y).as_slice()),
            z: AlignedSlice::from_slice(HerculesMutBox::from(momentum_z).as_slice()),
        },
        energy: AlignedSlice::from_slice(HerculesMutBox::from(energy).as_slice()),
    }
}

fn run_pre_euler(
    nelr: usize,
    iterations: usize,
    mut variables: Variables,
    areas: &[f32],
    elements_surrounding_elements: &[i32],
    normals: &Normals,
    ff_variable: &Variable,
    ff_fc_density_energy: &Float3,
    ff_fc_momentum_x: &Float3,
    ff_fc_momentum_y: &Float3,
    ff_fc_momentum_z: &Float3,
) -> Variables {
    let mut v_density = HerculesMutBox::from(variables.density.as_mut_slice());
    let mut v_momentum_x = HerculesMutBox::from(variables.momentum.x.as_mut_slice());
    let mut v_momentum_y = HerculesMutBox::from(variables.momentum.y.as_mut_slice());
    let mut v_momentum_z = HerculesMutBox::from(variables.momentum.z.as_mut_slice());
    let mut v_energy = HerculesMutBox::from(variables.energy.as_mut_slice());

    let areas = HerculesImmBox::from(areas);
    let elements_surrounding_elements = HerculesImmBox::from(elements_surrounding_elements);

    let normals_x = HerculesImmBox::from(normals.x.as_slice());
    let normals_y = HerculesImmBox::from(normals.y.as_slice());
    let normals_z = HerculesImmBox::from(normals.z.as_slice());

    let mut runner = runner!(pre_euler);
    let (density, momentum_x, momentum_y, momentum_z, energy) = async_std::task::block_on(async {
        runner
            .run(
                nelr as u64,
                iterations as u64,
                v_density.to(),
                v_momentum_x.to(),
                v_momentum_y.to(),
                v_momentum_z.to(),
                v_energy.to(),
                areas.to(),
                elements_surrounding_elements.to(),
                normals_x.to(),
                normals_y.to(),
                normals_z.to(),
                ff_variable.density,
                ff_variable.momentum.x,
                ff_variable.momentum.y,
                ff_variable.momentum.z,
                ff_variable.energy,
                ff_fc_density_energy.x,
                ff_fc_density_energy.y,
                ff_fc_density_energy.z,
                ff_fc_momentum_x.x,
                ff_fc_momentum_x.y,
                ff_fc_momentum_x.z,
                ff_fc_momentum_y.x,
                ff_fc_momentum_y.y,
                ff_fc_momentum_y.z,
                ff_fc_momentum_z.x,
                ff_fc_momentum_z.y,
                ff_fc_momentum_z.z,
            )
            .await
    });

    Variables {
        density: AlignedSlice::from_slice(HerculesMutBox::from(density).as_slice()),
        momentum: Momentum {
            x: AlignedSlice::from_slice(HerculesMutBox::from(momentum_x).as_slice()),
            y: AlignedSlice::from_slice(HerculesMutBox::from(momentum_y).as_slice()),
            z: AlignedSlice::from_slice(HerculesMutBox::from(momentum_z).as_slice()),
        },
        energy: AlignedSlice::from_slice(HerculesMutBox::from(energy).as_slice()),
    }
}

fn compare_float(x: f32, y: f32) -> bool {
    (x - y).abs() < 1e-5
}

fn compare_floats(xs: &Variables, ys: &Variables) -> bool {
    let xs_density = xs.density.as_slice();
    let xs_momentum_x = xs.momentum.x.as_slice();
    let xs_momentum_y = xs.momentum.y.as_slice();
    let xs_momentum_z = xs.momentum.z.as_slice();
    let xs_energy = xs.energy.as_slice();

    let ys_density = ys.density.as_slice();
    let ys_momentum_x = ys.momentum.x.as_slice();
    let ys_momentum_y = ys.momentum.y.as_slice();
    let ys_momentum_z = ys.momentum.z.as_slice();
    let ys_energy = ys.energy.as_slice();

    xs_density.len() == ys_density.len()
        && xs_density
            .iter()
            .zip(ys_density.iter())
            .all(|(x, y)| compare_float(*x, *y))
        && xs_momentum_x.len() == ys_momentum_x.len()
        && xs_momentum_x
            .iter()
            .zip(ys_momentum_x.iter())
            .all(|(x, y)| compare_float(*x, *y))
        && xs_momentum_y.len() == ys_momentum_y.len()
        && xs_momentum_y
            .iter()
            .zip(ys_momentum_y.iter())
            .all(|(x, y)| compare_float(*x, *y))
        && xs_momentum_z.len() == ys_momentum_z.len()
        && xs_momentum_z
            .iter()
            .zip(ys_momentum_z.iter())
            .all(|(x, y)| compare_float(*x, *y))
        && xs_energy.len() == ys_energy.len()
        && xs_energy
            .iter()
            .zip(ys_energy.iter())
            .all(|(x, y)| compare_float(*x, *y))
}

pub fn cfd_harness(args: CFDInputs) {
    let CFDInputs {
        data_file,
        iterations,
        block_size,
        pre_euler,
    } = args;

    let FarFieldConditions {
        ff_variable,
        ff_fc_momentum_x,
        ff_fc_momentum_y,
        ff_fc_momentum_z,
        ff_fc_density_energy,
    } = set_far_field_conditions();

    let GeometryData {
        nelr,
        areas,
        elements_surrounding_elements,
        normals,
    } = read_domain_geometry(data_file, block_size);

    let variables = initialize_variables(nelr, &ff_variable);
    println!("Running CFD with nelr = {}.", nelr);

    let res_juno = if pre_euler {
        run_pre_euler(
            nelr,
            iterations,
            variables.clone(),
            areas.as_slice(),
            elements_surrounding_elements.as_slice(),
            &normals,
            &ff_variable,
            &ff_fc_density_energy,
            &ff_fc_momentum_x,
            &ff_fc_momentum_y,
            &ff_fc_momentum_z,
        )
    } else {
        run_euler(
            nelr,
            iterations,
            variables.clone(),
            areas.as_slice(),
            elements_surrounding_elements.as_slice(),
            &normals,
            &ff_variable,
            &ff_fc_density_energy,
            &ff_fc_momentum_x,
            &ff_fc_momentum_y,
            &ff_fc_momentum_z,
        )
    };
    let res_rust = if pre_euler {
        rust_cfd::pre_euler(
            nelr,
            iterations,
            variables,
            areas.as_slice(),
            elements_surrounding_elements.as_slice(),
            &normals,
            &ff_variable,
            &ff_fc_density_energy,
            &ff_fc_momentum_x,
            &ff_fc_momentum_y,
            &ff_fc_momentum_z,
        )
    } else {
        rust_cfd::euler(
            nelr,
            iterations,
            variables,
            areas.as_slice(),
            elements_surrounding_elements.as_slice(),
            &normals,
            &ff_variable,
            &ff_fc_density_energy,
            &ff_fc_momentum_x,
            &ff_fc_momentum_y,
            &ff_fc_momentum_z,
        )
    };

    if !compare_floats(&res_juno, &res_rust) {
        panic!("Mismatch in results");
    }
}