use std::{error::Error, ops::Div, time::Instant};

use openshmem_rs::ShmemCtx;

use rand::prelude::*;

const BODIES_PER_PE: usize = 200000;
const RANGE: f32 = 2000000.;
const ITERATIONS: usize = 100;
const TIMESTEP: f32 = 0.01;

fn random_in_range(lo: f32, hi: f32, rng: &mut ThreadRng) -> f32 {
    lo + (rng.gen::<f32>() * (hi - lo))
}

fn which_pe(body_idx: usize) -> usize {
    body_idx / BODIES_PER_PE
}

fn calculate_gravity_1d(x_on: f32, x_by: f32, mass_on: f32, mass_by: f32) -> f32 {
    const G: f32 = 6.674e-8;
    let f = (G * mass_on * mass_by).div((x_on - x_by).powi(2));
    if f.is_finite() {
        f
    } else {
        0.
    }
}

fn main() -> Result<(), Box<dyn Error>> {
    let ctx = ShmemCtx::init()?;
    let is_root = ctx.my_pe().raw() == 0;
    let shmalloc = ctx.shmallocator();
    let mut rng = thread_rng();

    if is_root {
        println!(
            "starting generation of {} bodies... ({BODIES_PER_PE} per pe)",
            BODIES_PER_PE * ctx.n_pes()
        );
    }
    ctx.barrier_all();
    let mut start_of_iter = Instant::now();
    let mut time_reduce = shmalloc.shbox(0.0f32);
    let masses = shmalloc.array_gen(|_| random_in_range(1000., 10000., &mut rng), BODIES_PER_PE);
    let mut xs = shmalloc.array_gen(|_| random_in_range(-RANGE, RANGE, &mut rng), BODIES_PER_PE);
    let mut ys = shmalloc.array_gen(|_| random_in_range(-RANGE, RANGE, &mut rng), BODIES_PER_PE);
    let mut vxs = shmalloc.array_default(BODIES_PER_PE);
    let mut vys = shmalloc.array_default(BODIES_PER_PE);
    *time_reduce = start_of_iter.elapsed().as_secs_f32();
    ctx.barrier_all();
    time_reduce.reduce_max(&ctx);
    if is_root {
        println!(
            "generated {} bodies in {:.3}s ({}/s)",
            BODIES_PER_PE * ctx.n_pes(),
            *time_reduce,
            (BODIES_PER_PE * ctx.n_pes()) as f32 / *time_reduce
        );
    }
    ctx.barrier_all();

    // at any time, RAM usage is roughly 2 * BODIES_PER_PE * 3 * size_of::<f32>(),
    // as opposed to PEs * BODIES_PER_PE * 3 * size_of::<f32>().

    // hint to the compiler that we can elide some bounds checks.
    assert!(masses.get(BODIES_PER_PE - 1).is_some());
    assert!(xs.get(BODIES_PER_PE - 1).is_some());
    assert!(ys.get(BODIES_PER_PE - 1).is_some());
    assert!(vys.get(BODIES_PER_PE - 1).is_some());
    assert!(vxs.get(BODIES_PER_PE - 1).is_some());
    for iter in 0..ITERATIONS {
        start_of_iter = Instant::now();
        if is_root {
            println!("starting iteration {iter:>3}");
        }
        // calculate dv
        for pe in 0..ctx.n_pes() {
            let their_masses = ctx.pe(pe).read(&masses, ..);
            let their_xs = ctx.pe(pe).read(&xs, ..);
            let their_ys = ctx.pe(pe).read(&ys, ..);

            for (their_idx, &their_mass) in their_masses.into_iter().enumerate()
            {
                for (idx, mass) in masses.iter().enumerate() {
                    let fx = calculate_gravity_1d(xs[idx], their_xs[their_idx], masses[idx], their_mass);
                    let fy = calculate_gravity_1d(ys[idx], their_ys[their_idx], masses[idx], their_mass);
                    vxs[idx] += fx / mass;
                    vys[idx] += fy / mass;
                }
            }
        }
        ctx.barrier_all();
        // calculate dx
        for idx in 0..BODIES_PER_PE {

        }
        *time_reduce = start_of_iter.elapsed().as_secs_f32();
        ctx.barrier_all();
    }

    Ok(())
}