use std::{
	collections::{BTreeMap, HashMap},
	time::Instant,
};

fn main () {
	// PUC Lua is 250 ms
	// I'm using AOT Rust here so anything more than that 
	// is simply embarrassing
	
	fn timer <F: Fn ()> (f: F, label: &str) {
		let start = Instant::now ();
		f ();
		let stop = Instant::now ();
		println! ("{}: {} ms", label, (stop - start).as_millis ());
	}
	
	// timer (|| n_body::<BTreeTable> (100000), "BTreeTable");
	// timer (|| n_body::<BTreeTable2> (100000), "BTreeTable2");
	// timer (|| n_body::<HashTable> (100000), "HashTable");
	// timer (|| n_body_2::<HashTable2> (100000), "HashTable2");
	
	// Fastest on my system, at 83 ms
	timer (|| n_body_2::<BTreeTable3> (100000), "BTreeTable3");
	timer (|| n_body_2::<BTreeTable3> (500000), "BTreeTable3");
	timer (|| n_body_3::<BTreeTable4> (500000), "BTreeTable4");
}

trait Table {
	fn get (&self, key: &str) -> f64;
	fn set (&mut self, key: &'static str, value: f64);
}

#[derive (Default)]
struct BTreeTable (BTreeMap <String, f64>);

impl Table for BTreeTable {
	fn get (&self, key: &str) -> f64 {
		*self.0.get (key).unwrap ()
	}
	
	fn set (&mut self, key: &str, value: f64) {
		self.0.insert (key.to_string (), value);
	}
}

#[derive (Default)]
struct BTreeTable2 (BTreeMap <&'static str, f64>);

impl Table for BTreeTable2 {
	fn get (&self, key: &str) -> f64 {
		*self.0.get (key).unwrap ()
	}
	
	fn set (&mut self, key: &'static str, value: f64) {
		self.0.insert (key, value);
	}
}

#[derive (Default)]
struct HashTable (HashMap <String, f64>);

impl Table for HashTable {
	fn get (&self, key: &str) -> f64 {
		*self.0.get (key).unwrap ()
	}
	
	fn set (&mut self, key: &str, value: f64) {
		self.0.insert (key.to_string (), value);
	}
}

const PI: f64 = 3.141592653589793f64;
const SOLAR_MASS: f64 = 4.0 * PI * PI;

fn n_body <T: Default + Table> (num_iters: usize) {
	const X: &str = "x";
	const Y: &str = "y";
	const Z: &str = "z";
	const VX: &str = "vx";
	const VY: &str = "vy";
	const VZ: &str = "vz";
	const MASS: &str = "mass";
	
	let days_per_year = 365.24;
	
	let mut bodies = vec! [];
	{
		let mut t = T::default ();
		t.set (X, 0.0);
		t.set (Y, 0.0);
		t.set (Z, 0.0);
		t.set (VX, 0.0);
		t.set (VY, 0.0);
		t.set (VZ, 0.0);
		t.set (MASS, SOLAR_MASS);
		
		bodies.push (t);
	}
	{
		let mut t = T::default ();
		t.set (X, 4.84143144246472090e+00);
		t.set (Y, -1.16032004402742839e+00);
		t.set (Z, -1.03622044471123109e-01);
		t.set (VX, 1.66007664274403694e-03 * days_per_year);
		t.set (VY, 7.69901118419740425e-03 * days_per_year);
		t.set (VZ, -6.90460016972063023e-05 * days_per_year);
		t.set (MASS, 9.54791938424326609e-04 * SOLAR_MASS);
		
		bodies.push (t);
	}
	{
		let mut t = T::default ();
		t.set (X, 8.34336671824457987e+00);
		t.set (Y, 4.12479856412430479e+00);
		t.set (Z, -4.03523417114321381e-01);
		t.set (VX, -2.76742510726862411e-03 * days_per_year);
		t.set (VY, 4.99852801234917238e-03 * days_per_year);
		t.set (VZ, 2.30417297573763929e-05 * days_per_year);
		t.set (MASS, 2.85885980666130812e-04 * SOLAR_MASS);
		
		bodies.push (t);
	}
	{
		let mut t = T::default ();
		t.set (X, 1.28943695621391310e+01);
		t.set (Y, -1.51111514016986312e+01);
		t.set (Z, -2.23307578892655734e-01);
		t.set (VX, 2.96460137564761618e-03 * days_per_year);
		t.set (VY, 2.37847173959480950e-03 * days_per_year);
		t.set (VZ, -2.96589568540237556e-05 * days_per_year);
		t.set (MASS, 4.36624404335156298e-05 * SOLAR_MASS);
		
		bodies.push (t);
	}
	{
		let mut t = T::default ();
		t.set (X, 1.53796971148509165e+01);
		t.set (Y, -2.59193146099879641e+01);
		t.set (Z, 1.79258772950371181e-01);
		t.set (VX, 2.68067772490389322e-03 * days_per_year);
		t.set (VY, 1.62824170038242295e-03 * days_per_year);
		t.set (VZ, -9.51592254519715870e-05 * days_per_year);
		t.set (MASS, 5.15138902046611451e-05 * SOLAR_MASS);
		
		bodies.push (t);
	}
	
	fn advance <T: Table> (bodies: &mut [T], dt: f64) {
		let nbody = bodies.len ();
		
		for i in 0..nbody {
			let bix;
			let biy;
			let biz;
			let bimass;
			let mut bivx;
			let mut bivy;
			let mut bivz;
			{
				let bi = &mut bodies [i];
				bix = bi.get (X);
				biy = bi.get (Y);
				biz = bi.get (Z);
				bimass = bi.get (MASS);
				bivx = bi.get (VX);
				bivy = bi.get (VY);
				bivz = bi.get (VZ);
			}
			
			for j in i + 1..nbody {
				let bj = &mut bodies [j];
				let dx = bix - bj.get (X);
				let dy = biy - bj.get (Y);
				let dz = biz - bj.get (Z);
				let mut mag = (dx * dx + dy * dy + dz * dz).sqrt ();
				mag = dt / (mag * mag * mag);
				let mut bm = bj.get (MASS) * mag;
				bivx -= dx * bm;
				bivy -= dy * bm;
				bivz -= dz * bm;
				bm = bimass * mag;
				bj.set (VX, bj.get (VX) + dx * bm);
				bj.set (VY, bj.get (VY) + dy * bm);
				bj.set (VZ, bj.get (VZ) + dz * bm);
			}
			
			let bi = &mut bodies [i];
			bi.set (VX, bivx);
			bi.set (VY, bivy);
			bi.set (VZ, bivz);
			bi.set (X, bix + dt * bivx);
			bi.set (Y, biy + dt * bivy);
			bi.set (Z, biz + dt * bivz);
		}
	}
	
	fn energy <T: Table> (bodies: &[T]) -> f64 {
		let mut e = 0.0;
		for (i, bi) in bodies.iter ().enumerate () {
			let vx = bi.get (VX);
			let vy = bi.get (VY);
			let vz = bi.get (VZ);
			let bim = bi.get (MASS);
			
			e += 0.5 * bim * (vx * vx + vy * vy + vz * vz);
			
			for j in i + 1..bodies.len () {
				let bj = &bodies [j];
				let dx = bi.get (X) - bj.get (X);
				let dy = bi.get (Y) - bj.get (Y);
				let dz = bi.get (Z) - bj.get (Z);
				let distance = (dx * dx + dy * dy + dz * dz).sqrt ();
				e -= (bim * bj.get (MASS)) / distance;
			}
		}
		
		e
	}
	
	fn offset_momentum <T: Table> (b: &mut [T]) {
		let mut px = 0.0;
		let mut py = 0.0;
		let mut pz = 0.0;
		
		for bi in b.iter () {
			let bim = bi.get (MASS);
			px += bi.get (VX) * bim;
			py += bi.get (VY) * bim;
			pz += bi.get (VZ) * bim;
		}
		
		let bimass = b [0].get (MASS);
		b [0].set (VX, -px / bimass);
		b [0].set (VY, -py / bimass);
		b [0].set (VZ, -pz / bimass);
	}
	
	offset_momentum (&mut bodies);
	println! ("{:0.9}", energy (&bodies));
	for _ in 0..num_iters {
		advance (&mut bodies, 0.01);
	}
	println! ("{:0.9}", energy (&bodies));
}

trait Table2 {
	fn get (&self, key: i64) -> f64;
	fn set (&mut self, key: i64, value: f64);
}

#[derive (Default)]
struct BTreeTable3 (BTreeMap <i64, f64>);

impl Table2 for BTreeTable3 {
	fn get (&self, key: i64) -> f64 {
		*self.0.get (&key).unwrap ()
	}
	
	fn set (&mut self, key: i64, value: f64) {
		self.0.insert (key, value);
	}
}

#[derive (Default)]
struct HashTable2 (HashMap <i64, f64>);

impl Table2 for HashTable2 {
	fn get (&self, key: i64) -> f64 {
		*self.0.get (&key).unwrap ()
	}
	
	fn set (&mut self, key: i64, value: f64) {
		self.0.insert (key, value);
	}
}

fn n_body_2 <T: Default + Table2> (num_iters: usize) {
	const X: i64 = 1;
	const Y: i64 = 2;
	const Z: i64 = 3;
	const VX: i64 = 4;
	const VY: i64 = 5;
	const VZ: i64 = 6;
	const MASS: i64 = 7;
	
	let days_per_year = 365.24;
	
	let mut bodies = vec! [];
	{
		let mut t = T::default ();
		t.set (X, 0.0);
		t.set (Y, 0.0);
		t.set (Z, 0.0);
		t.set (VX, 0.0);
		t.set (VY, 0.0);
		t.set (VZ, 0.0);
		t.set (MASS, SOLAR_MASS);
		
		bodies.push (t);
	}
	{
		let mut t = T::default ();
		t.set (X, 4.84143144246472090e+00);
		t.set (Y, -1.16032004402742839e+00);
		t.set (Z, -1.03622044471123109e-01);
		t.set (VX, 1.66007664274403694e-03 * days_per_year);
		t.set (VY, 7.69901118419740425e-03 * days_per_year);
		t.set (VZ, -6.90460016972063023e-05 * days_per_year);
		t.set (MASS, 9.54791938424326609e-04 * SOLAR_MASS);
		
		bodies.push (t);
	}
	{
		let mut t = T::default ();
		t.set (X, 8.34336671824457987e+00);
		t.set (Y, 4.12479856412430479e+00);
		t.set (Z, -4.03523417114321381e-01);
		t.set (VX, -2.76742510726862411e-03 * days_per_year);
		t.set (VY, 4.99852801234917238e-03 * days_per_year);
		t.set (VZ, 2.30417297573763929e-05 * days_per_year);
		t.set (MASS, 2.85885980666130812e-04 * SOLAR_MASS);
		
		bodies.push (t);
	}
	{
		let mut t = T::default ();
		t.set (X, 1.28943695621391310e+01);
		t.set (Y, -1.51111514016986312e+01);
		t.set (Z, -2.23307578892655734e-01);
		t.set (VX, 2.96460137564761618e-03 * days_per_year);
		t.set (VY, 2.37847173959480950e-03 * days_per_year);
		t.set (VZ, -2.96589568540237556e-05 * days_per_year);
		t.set (MASS, 4.36624404335156298e-05 * SOLAR_MASS);
		
		bodies.push (t);
	}
	{
		let mut t = T::default ();
		t.set (X, 1.53796971148509165e+01);
		t.set (Y, -2.59193146099879641e+01);
		t.set (Z, 1.79258772950371181e-01);
		t.set (VX, 2.68067772490389322e-03 * days_per_year);
		t.set (VY, 1.62824170038242295e-03 * days_per_year);
		t.set (VZ, -9.51592254519715870e-05 * days_per_year);
		t.set (MASS, 5.15138902046611451e-05 * SOLAR_MASS);
		
		bodies.push (t);
	}
	
	fn advance <T: Table2> (bodies: &mut [T], dt: f64) {
		let nbody = bodies.len ();
		
		for i in 0..nbody {
			let bix;
			let biy;
			let biz;
			let bimass;
			let mut bivx;
			let mut bivy;
			let mut bivz;
			{
				let bi = &mut bodies [i];
				bix = bi.get (X);
				biy = bi.get (Y);
				biz = bi.get (Z);
				bimass = bi.get (MASS);
				bivx = bi.get (VX);
				bivy = bi.get (VY);
				bivz = bi.get (VZ);
			}
			
			for j in i + 1..nbody {
				let bj = &mut bodies [j];
				let dx = bix - bj.get (X);
				let dy = biy - bj.get (Y);
				let dz = biz - bj.get (Z);
				let mut mag = (dx * dx + dy * dy + dz * dz).sqrt ();
				mag = dt / (mag * mag * mag);
				let mut bm = bj.get (MASS) * mag;
				bivx -= dx * bm;
				bivy -= dy * bm;
				bivz -= dz * bm;
				bm = bimass * mag;
				bj.set (VX, bj.get (VX) + dx * bm);
				bj.set (VY, bj.get (VY) + dy * bm);
				bj.set (VZ, bj.get (VZ) + dz * bm);
			}
			
			let bi = &mut bodies [i];
			bi.set (VX, bivx);
			bi.set (VY, bivy);
			bi.set (VZ, bivz);
			bi.set (X, bix + dt * bivx);
			bi.set (Y, biy + dt * bivy);
			bi.set (Z, biz + dt * bivz);
		}
	}
	
	fn energy <T: Table2> (bodies: &[T]) -> f64 {
		let mut e = 0.0;
		for (i, bi) in bodies.iter ().enumerate () {
			let vx = bi.get (VX);
			let vy = bi.get (VY);
			let vz = bi.get (VZ);
			let bim = bi.get (MASS);
			
			e += 0.5 * bim * (vx * vx + vy * vy + vz * vz);
			
			for j in i + 1..bodies.len () {
				let bj = &bodies [j];
				let dx = bi.get (X) - bj.get (X);
				let dy = bi.get (Y) - bj.get (Y);
				let dz = bi.get (Z) - bj.get (Z);
				let distance = (dx * dx + dy * dy + dz * dz).sqrt ();
				e -= (bim * bj.get (MASS)) / distance;
			}
		}
		
		e
	}
	
	fn offset_momentum <T: Table2> (b: &mut [T]) {
		let mut px = 0.0;
		let mut py = 0.0;
		let mut pz = 0.0;
		
		for bi in b.iter () {
			let bim = bi.get (MASS);
			px += bi.get (VX) * bim;
			py += bi.get (VY) * bim;
			pz += bi.get (VZ) * bim;
		}
		
		let bimass = b [0].get (MASS);
		b [0].set (VX, -px / bimass);
		b [0].set (VY, -py / bimass);
		b [0].set (VZ, -pz / bimass);
	}
	
	offset_momentum (&mut bodies);
	println! ("{:0.9}", energy (&bodies));
	for _ in 0..num_iters {
		advance (&mut bodies, 0.01);
	}
	println! ("{:0.9}", energy (&bodies));
}

trait Table3 {
	fn get (&self, key: u64) -> f64;
	fn set (&mut self, key: u64, value: f64);
}

#[derive (Default)]
struct BTreeTable4 (BTreeMap <u64, f64>);

impl Table3 for BTreeTable4 {
	fn get (&self, key: u64) -> f64 {
		*self.0.get (&key).unwrap ()
	}
	
	fn set (&mut self, key: u64, value: f64) {
		self.0.insert (key, value);
	}
}

fn n_body_3 <T: Default + Table3> (num_iters: usize) {
	const X: u64 = 1;
	const Y: u64 = 2;
	const Z: u64 = 3;
	const VX: u64 = 4;
	const VY: u64 = 5;
	const VZ: u64 = 6;
	const MASS: u64 = 7;
	
	let days_per_year = 365.24;
	
	let mut bodies = vec! [];
	{
		let mut t = T::default ();
		t.set (X, 0.0);
		t.set (Y, 0.0);
		t.set (Z, 0.0);
		t.set (VX, 0.0);
		t.set (VY, 0.0);
		t.set (VZ, 0.0);
		t.set (MASS, SOLAR_MASS);
		
		bodies.push (t);
	}
	{
		let mut t = T::default ();
		t.set (X, 4.84143144246472090e+00);
		t.set (Y, -1.16032004402742839e+00);
		t.set (Z, -1.03622044471123109e-01);
		t.set (VX, 1.66007664274403694e-03 * days_per_year);
		t.set (VY, 7.69901118419740425e-03 * days_per_year);
		t.set (VZ, -6.90460016972063023e-05 * days_per_year);
		t.set (MASS, 9.54791938424326609e-04 * SOLAR_MASS);
		
		bodies.push (t);
	}
	{
		let mut t = T::default ();
		t.set (X, 8.34336671824457987e+00);
		t.set (Y, 4.12479856412430479e+00);
		t.set (Z, -4.03523417114321381e-01);
		t.set (VX, -2.76742510726862411e-03 * days_per_year);
		t.set (VY, 4.99852801234917238e-03 * days_per_year);
		t.set (VZ, 2.30417297573763929e-05 * days_per_year);
		t.set (MASS, 2.85885980666130812e-04 * SOLAR_MASS);
		
		bodies.push (t);
	}
	{
		let mut t = T::default ();
		t.set (X, 1.28943695621391310e+01);
		t.set (Y, -1.51111514016986312e+01);
		t.set (Z, -2.23307578892655734e-01);
		t.set (VX, 2.96460137564761618e-03 * days_per_year);
		t.set (VY, 2.37847173959480950e-03 * days_per_year);
		t.set (VZ, -2.96589568540237556e-05 * days_per_year);
		t.set (MASS, 4.36624404335156298e-05 * SOLAR_MASS);
		
		bodies.push (t);
	}
	{
		let mut t = T::default ();
		t.set (X, 1.53796971148509165e+01);
		t.set (Y, -2.59193146099879641e+01);
		t.set (Z, 1.79258772950371181e-01);
		t.set (VX, 2.68067772490389322e-03 * days_per_year);
		t.set (VY, 1.62824170038242295e-03 * days_per_year);
		t.set (VZ, -9.51592254519715870e-05 * days_per_year);
		t.set (MASS, 5.15138902046611451e-05 * SOLAR_MASS);
		
		bodies.push (t);
	}
	
	fn advance <T: Table3> (bodies: &mut [T], dt: f64) {
		let nbody = bodies.len ();
		
		for i in 0..nbody {
			let bix;
			let biy;
			let biz;
			let bimass;
			let mut bivx;
			let mut bivy;
			let mut bivz;
			{
				let bi = &mut bodies [i];
				bix = bi.get (X);
				biy = bi.get (Y);
				biz = bi.get (Z);
				bimass = bi.get (MASS);
				bivx = bi.get (VX);
				bivy = bi.get (VY);
				bivz = bi.get (VZ);
			}
			
			for j in i + 1..nbody {
				let bj = &mut bodies [j];
				let dx = bix - bj.get (X);
				let dy = biy - bj.get (Y);
				let dz = biz - bj.get (Z);
				let mut mag = (dx * dx + dy * dy + dz * dz).sqrt ();
				mag = dt / (mag * mag * mag);
				let mut bm = bj.get (MASS) * mag;
				bivx -= dx * bm;
				bivy -= dy * bm;
				bivz -= dz * bm;
				bm = bimass * mag;
				bj.set (VX, bj.get (VX) + dx * bm);
				bj.set (VY, bj.get (VY) + dy * bm);
				bj.set (VZ, bj.get (VZ) + dz * bm);
			}
			
			let bi = &mut bodies [i];
			bi.set (VX, bivx);
			bi.set (VY, bivy);
			bi.set (VZ, bivz);
			bi.set (X, bix + dt * bivx);
			bi.set (Y, biy + dt * bivy);
			bi.set (Z, biz + dt * bivz);
		}
	}
	
	fn energy <T: Table3> (bodies: &[T]) -> f64 {
		let mut e = 0.0;
		for (i, bi) in bodies.iter ().enumerate () {
			let vx = bi.get (VX);
			let vy = bi.get (VY);
			let vz = bi.get (VZ);
			let bim = bi.get (MASS);
			
			e += 0.5 * bim * (vx * vx + vy * vy + vz * vz);
			
			for j in i + 1..bodies.len () {
				let bj = &bodies [j];
				let dx = bi.get (X) - bj.get (X);
				let dy = bi.get (Y) - bj.get (Y);
				let dz = bi.get (Z) - bj.get (Z);
				let distance = (dx * dx + dy * dy + dz * dz).sqrt ();
				e -= (bim * bj.get (MASS)) / distance;
			}
		}
		
		e
	}
	
	fn offset_momentum <T: Table3> (b: &mut [T]) {
		let mut px = 0.0;
		let mut py = 0.0;
		let mut pz = 0.0;
		
		for bi in b.iter () {
			let bim = bi.get (MASS);
			px += bi.get (VX) * bim;
			py += bi.get (VY) * bim;
			pz += bi.get (VZ) * bim;
		}
		
		let bimass = b [0].get (MASS);
		b [0].set (VX, -px / bimass);
		b [0].set (VY, -py / bimass);
		b [0].set (VZ, -pz / bimass);
	}
	
	offset_momentum (&mut bodies);
	println! ("{:0.9}", energy (&bodies));
	for _ in 0..num_iters {
		advance (&mut bodies, 0.01);
	}
	println! ("{:0.9}", energy (&bodies));
}