use std::fs::OpenOptions;
use std::io::BufWriter;

use serde::Serialize;
use log::{debug, trace};
use glam::DVec3;

use crate::config::Body;
use crate::types::{norm_mass, norm_pos, norm_vel, norm_time, real_pos, real_vel, real_time, real_body};

pub struct Simulation {
    pub bodies: Vec<Body>,
    step_size: f64,
    steps_per_save: usize,
    save_file: String 
}

#[derive(Serialize)]
struct Snapshot<'a> {
    time: f64,
    bodies: &'a [Body],
}

impl Simulation {
    pub fn new(
        mut bodies: Vec<Body>, 
        mut step_size: f64, 
        steps_per_save: usize,
        save_file: String,
    ) -> Self {
        let n = bodies.len();
        for i in 0..n {
            bodies[i].mass = norm_mass(bodies[i].mass);
            bodies[i].position = norm_pos(bodies[i].position);
            bodies[i].velocity = norm_vel(bodies[i].velocity);
        }
        step_size = norm_time(step_size);
        Self {bodies, step_size, steps_per_save, save_file}
    }

    pub fn run<F>(&mut self, steps: usize, mut on_step: Option<F>)
    where F: FnMut() {
        let file = OpenOptions::new()
            .create(true)
            .append(true)
            .open(&self.save_file);
        let mut writer = BufWriter::new(file.unwrap());

        for i in 0..steps {
            debug!("Step: {}", i);
            self.reset_accelerations();
            self.caclulate_accelerations();
            self.step_bodies();
            if i % self.steps_per_save == 0 {
                //save the state
                let real_bodies: Vec<Body> = self.bodies.iter().map(real_body).collect();
                let snapshot = Snapshot {
                    time: real_time(i as f64*self.step_size),
                    bodies: &real_bodies,
                };

                bincode::serialize_into(&mut writer, &snapshot).expect("Couldn't write to trajectory. ");  
                //Do the progress bar
                if let Some(f) = &mut on_step {
                    f();
                }
            }
        }
    }

    fn caclulate_accelerations(&mut self) {
        let r_hat_over_r3 = self.get_rhat_over_r_three();
        let n = self.bodies.len();
        for i in 0..(n-1) { 
            let mass_i = self.bodies[i].mass;
            for j in (i+1)..n {
                let mass_j = self.bodies[j].mass;
                self.bodies[i].acceleration += r_hat_over_r3[i][j]*mass_j;
                self.bodies[j].acceleration -= r_hat_over_r3[i][j]*mass_i;
            }
        }
    }

    fn reset_accelerations(&mut self) {
        for i in 0..self.bodies.len() {
            self.bodies[i].acceleration = DVec3::ZERO;
        }
    }

    fn get_rhat_over_r_three (&self) -> Vec<Vec<DVec3>>{
        let n = self.bodies.len();
        let mut r_hat_over_r3 = vec![vec![DVec3::ZERO; n]; n];
        for i in 0..(n-1) {
            for j in (i+1)..n {
                let r_ij = self.bodies[j].position - self.bodies[i].position;
                let dist = r_ij.length();
                r_hat_over_r3[i][j] = r_ij / dist.powf(3.0);
            }
        }
        trace!("rhat over r3: {:?}", r_hat_over_r3);
        return r_hat_over_r3;
    }

    fn step_bodies(&mut self) {
        for body in &mut self.bodies {
            //do for each of three (x,y,z) dimensions
            body.position += self.step_size*body.velocity;
            body.velocity += self.step_size*body.acceleration;
            trace!("{} now at {:?}", body.name, real_pos(body.position));
            trace!("{} moving at {:?}", body.name, real_vel(body.velocity));
        }
    }
}