const MASS_HELIX: f64 = 2.;
const K_SPRING: f64 = 100.;
const FRICTION: f64 = 100.;

use super::{geometry, Design, Helix, Parameters, Point, Vector};

use std::f64::consts::PI;
use std::f64::consts::SQRT_2;

/// Structure that adjust the roll of helices
pub struct RollSystem {
    speed: Vec<f64>,
    acceleration: Vec<f64>,
}

/// Structure that adjust the origin of the helices
pub struct SpringSystem {
    speed: Vec<Vector<f64>>,
    acceleration: Vec<Vector<f64>>,
    intervals: Vec<Option<(isize, isize)>>,
    cross_overs: Vec<(isize, isize, bool, isize, isize, bool)>,
}

// Angle between A and C_ where C is the nucleotide following A and C_ is the projection of
// C on A's circle
fn angle_aoc2(p: &Parameters) -> f64 {
    2. * PI / p.bases_per_turn
}

fn dist_ac(p: &Parameters) -> f64 {
    (dist_ac2(p) * dist_ac2(p) + p.z_step * p.z_step).sqrt()
}

fn dist_ac2(p: &Parameters) -> f64 {
    SQRT_2 * (1. - angle_aoc2(p).cos()).sqrt() * p.helix_radius
}

fn cross_over_force(
    me: &Helix,
    other: &Helix,
    parameters: &Parameters,
    n_self: isize,
    b_self: bool,
    n_other: isize,
    b_other: bool,
) -> (f64, f64) {
    let nucl_self = me.space_pos(parameters, n_self, b_self);
    let nucl_other = other.space_pos(parameters, n_other, b_other);

    let real_dist = nucl_self
        .iter()
        .zip(nucl_other.iter())
        .map(|(a, b)| (a - b) * (a - b))
        .sum::<f64>()
        .sqrt();
    let norm = K_SPRING * (real_dist - dist_ac(parameters));

    let theta_self = me.theta(n_self, b_self, parameters);
    let theta_other = other.theta(n_other, b_other, parameters);

    let vec_self = me.rotate_point([0., -theta_self.sin(), theta_self.cos()]);
    let vec_other = other.rotate_point([0., -theta_other.sin(), theta_other.cos()]);

    (
        (0..3)
            .map(|i| norm * vec_self[i] * (nucl_other[i] - nucl_self[i]) / real_dist)
            .sum(),
        (0..3)
            .map(|i| norm * vec_other[i] * (nucl_self[i] - nucl_other[i]) / real_dist)
            .sum(),
    )
}

fn spring_force(
    me: &Helix,
    other: &Helix,
    parameters: &Parameters,
    n_self: isize,
    b_self: bool,
    n_other: isize,
    b_other: bool,
) -> (Vector<f64>, Vector<f64>) {
    let nucl_self = me.space_pos(parameters, n_self, b_self);
    let nucl_other = other.space_pos(parameters, n_other, b_other);

    let real_dist = nucl_self
        .iter()
        .zip(nucl_other.iter())
        .map(|(a, b)| (a - b) * (a - b))
        .sum::<f64>()
        .sqrt();
    let norm = K_SPRING * (real_dist - dist_ac(parameters)) / real_dist;
    let point_self = Point::from_coord(nucl_self);
    let point_other = Point::from_coord(nucl_other);
    (
        norm * (point_other - point_self),
        norm * (point_self - point_other),
    )
}

impl RollSystem {
    /// Create a system from a design, the system will adjust the helices of the design.
    pub fn from_design<S, T>(design: &Design<S, T>) -> Self {
        let speed = vec![0.; design.helices.len()];
        let acceleration = vec![0.; design.helices.len()];
        RollSystem {
            speed,
            acceleration,
        }
    }

    fn update_acceleration<S, T>(&mut self, design: &Design<S, T>) {
        let cross_overs = design.get_cross_overs();
        for i in 0..self.acceleration.len() {
            self.acceleration[i] = -self.speed[i] * FRICTION / MASS_HELIX;
        }
        for (h1, x1, b1, h2, x2, b2) in cross_overs.iter() {
            if h1 >= h2 {
                continue;
            }
            let me = &design.helices[*h1 as usize];
            let other = &design.helices[*h2 as usize];
            let (delta_1, delta_2) =
                cross_over_force(me, other, &design.parameters.unwrap(), *x1, *b1, *x2, *b2);
            self.acceleration[*h1 as usize] += delta_1 / MASS_HELIX;
            self.acceleration[*h2 as usize] += delta_2 / MASS_HELIX;
        }
    }

    fn update_speed(&mut self, dt: f64) {
        for i in 0..self.speed.len() {
            self.speed[i] += dt * self.acceleration[i];
        }
    }

    fn update_rolls<S, T>(&self, design: &mut Design<S, T>, dt: f64) {
        for i in 0..self.speed.len() {
            design.helices[i].roll += self.speed[i] * dt;
        }
    }

    /// Adjuste the helices of the design, do not show intermediate steps
    pub fn solve<S, T>(&mut self, design: &mut Design<S, T>, dt: f64) {
        let mut nb_step = 0;
        let mut done = false;
        while !done && nb_step < 10000 {
            self.update_rolls(design, dt);
            self.update_speed(dt);
            self.update_acceleration(design);
            println!("acceleration {:?}", self.acceleration);
            done = self.acceleration.iter().map(|x| x.abs()).sum::<f64>() < 1e-8;
            nb_step += 1;
        }
    }

    /// Do one step of simulation with time step dt
    pub fn solve_one_step<S, T>(&mut self, design: &mut Design<S, T>, dt: f64) -> f64 {
        self.update_rolls(design, dt);
        self.update_speed(dt);
        self.update_acceleration(design);
        self.acceleration.iter().map(|x| x.abs()).sum::<f64>()
    }
}

impl SpringSystem {
    /// Create a system from a design, the system will adjust the helices of the design.
    pub fn from_design<S, T>(design: &Design<S, T>) -> Self {
        let speed = vec![Vector::zero(); design.helices.len()];
        let acceleration = vec![Vector::zero(); design.helices.len()];
        let intervals = {
            let n = design.helices.len();
            (0..n).map(|i| design.get_interval_helix(i)).collect()
        };
        let cross_overs = design.get_cross_overs();
        SpringSystem {
            speed,
            acceleration,
            intervals,
            cross_overs,
        }
    }

    fn update_acceleration<S, T>(&mut self, design: &Design<S, T>) {
        for i in 0..self.acceleration.len() {
            self.acceleration[i] = -self.speed[i] * FRICTION / MASS_HELIX;
        }
        for (h1, x1, b1, h2, x2, b2) in self.cross_overs.iter() {
            if h1 >= h2 {
                continue;
            }
            let me = &design.helices[*h1 as usize];
            let other = &design.helices[*h2 as usize];
            let (delta_1, delta_2) =
                spring_force(me, other, &design.parameters.unwrap(), *x1, *b1, *x2, *b2);
            self.acceleration[*h1 as usize] += delta_1 / MASS_HELIX;
            self.acceleration[*h2 as usize] += delta_2 / MASS_HELIX;
        }
        let nb_helices = design.helices.len();
        let param = &design.parameters.expect("parameters");

        let r = 2. * param.helix_radius + param.inter_helix_gap;

        for i in 0..(nb_helices - 1) {
            if self.intervals[i].is_none() {
                continue;
            }
            let a = Point::from_coord(
                design.helices[i].axis_pos(param, self.intervals[i].expect("interval").0),
            );
            let b = Point::from_coord(
                design.helices[i].axis_pos(param, self.intervals[i].expect("interval").1),
            );
            for j in (i + 1)..nb_helices {
                if self.intervals[j].is_none() {
                    continue;
                }
                let c = Point::from_coord(
                    design.helices[j].axis_pos(param, self.intervals[j].expect("interval").0),
                );
                let d = Point::from_coord(
                    design.helices[j].axis_pos(param, self.intervals[j].expect("interval").1),
                );
                let (dist, vec) = geometry::distance_segment(a, b, c, d);
                if dist < r {
                    let norm = ((dist - r) / dist).powi(4) / MASS_HELIX * 100000.;
                    self.acceleration[i] += norm * vec;
                    self.acceleration[j] += -norm * vec;
                }
            }
        }
    }

    fn update_speed(&mut self, dt: f64) {
        for i in 0..self.speed.len() {
            self.speed[i] += dt * self.acceleration[i];
        }
    }

    fn update_origin<S, T>(&self, design: &mut Design<S, T>, dt: f64) {
        for i in 0..self.speed.len() {
            let delta = self.speed[i] * dt;
            design.helices[i].origin.x += delta.x;
            design.helices[i].origin.y += delta.y;
            design.helices[i].origin.z += delta.z;
        }
    }

    /// Adjuste the helices of the design, do not show intermediate steps
    pub fn solve<S, T>(&mut self, design: &mut Design<S, T>, dt: f64) {
        let mut nb_step = 0;
        let mut done = false;
        while !done && nb_step < 10000 {
            self.update_origin(design, dt);
            self.update_speed(dt);
            self.update_acceleration(design);
            println!("acceleration {:?}", self.acceleration);
            done = self.acceleration.iter().map(|x| x.norm()).sum::<f64>() < 1e-8;
            nb_step += 1;
        }
    }

    /// Do one step of simulation with time step dt
    pub fn solve_one_step<S, T>(&mut self, design: &mut Design<S, T>, dt: f64) -> f64 {
        self.update_origin(design, dt);
        self.update_speed(dt);
        self.update_acceleration(design);
        self.acceleration.iter().map(|x| x.norm()).sum::<f64>()
    }
}