use core::f64::consts::PI;
use std::time::Instant;
use gauss_quad::GaussLegendre;
use num::Float;
use std::collections::HashMap;
use std::io;
use std::rc::Rc;

use std::{
    fs::File,
    io::BufReader,
};

use crate::gq_storage; // This should be moved to main.rs after new archtitecture (see main.rs TODO)

pub type Complex = num::complex::Complex<f64>;

const M_ROH: f64 = 0.7736; 
const M_PI: f64 = 0.13957;
pub const S0: f64 = 4.0 * M_PI * M_PI;
const M_K: f64 = 0.496;
const B0: f64 = 1.055;
const B1: f64 = 0.15;
const LAMBDA1: f64 = 1.57;
const LAMBDA2: f64 = -1.96;
const LAMBDA_HIGH: f64 = 10.0;
pub const EPSILON: f64 = 1e-15;
const S_MID_CUTOFF: f64 = 1.42 * 1.42;
const INF: f64 = 1e5;

#[derive(Default)]
pub struct Cache {
    pub gauss_quad_lut: HashMap<u32, (Rc<[f64]>, Rc<[f64]>), nohash_hasher::BuildNoHashHasher<u32>>, 
    pub delta_lut: HashMap<u64, f64, nohash_hasher::BuildNoHashHasher<u64>>, 
    pub omnes_lut: HashMap<u32, HashMap<u64, Complex, nohash_hasher::BuildNoHashHasher<u64>>, nohash_hasher::BuildNoHashHasher<u32>> // First key is n, second is s
}

impl Cache {
    pub fn from_file(filepath: &str) -> io::Result<Cache> {
        let file = File::open(filepath)?;
        let mut file = BufReader::new(file);
    
        let me = Cache{
            gauss_quad_lut: gq_storage::deserialize(&mut file, |_| true)?.into_iter()
            .map(|(roots, weights)| (u32::try_from(roots.len()).unwrap(), (Rc::from(roots), Rc::from(weights)))).collect(),
            delta_lut: HashMap::with_hasher(Default::default()),
            omnes_lut: HashMap::with_hasher(Default::default())
        };
        Ok(me)
    }
    pub fn from_slice(data: &[u8]) -> io::Result<Cache> {
        let mut cursor = std::io::Cursor::new(data);
        let me = Cache{
            gauss_quad_lut: gq_storage::deserialize(&mut cursor, |_| true)?.into_iter()
            .map(|(roots, weights)| (u32::try_from(roots.len()).unwrap(), (Rc::from(roots), Rc::from(weights)))).collect(),
            delta_lut: HashMap::with_hasher(Default::default()),
            omnes_lut: HashMap::with_hasher(Default::default())
        };
        Ok(me)
    }
}

// derivated values

fn lambda_0() -> f64 {
    delta((2.0*M_K).powi(2))
}

fn delta_mid_cutoff() -> f64 {
    delta(S_MID_CUTOFF)
}

fn atan_shift(x: f64) -> f64 {
    let atan_std: f64 = x.atan();
    if atan_std < 0.0 {
        atan_std + PI
    } else {
        atan_std
    }
}

fn integrate<G:FnMut(f64) -> f64>(roots: &[f64], weights: &[f64], a:f64, b:f64, mut f: G) -> f64 {
    let mut sum: f64 = 0.0;
    let mut i: usize = 0;
    if roots.len() != weights.len() {
        panic!("Error: roots and weights are of different length");
    }

    while i < roots.len() {
        sum += weights[i] * f(roots[i]*(b-a)/2.0 + (a+b)/2.0);
        i += 1;
    }
    sum * (b-a)/2.0
}

fn integrate_complex<G:FnMut(f64) -> Complex>(roots: &[f64], weights: &[f64], a:f64, b:f64, mut f: G) -> Complex {
    let mut sum: Complex = Complex::new(0.0,0.0);
    let mut i: usize = 0;
    if roots.len() != weights.len() {
        panic!("Error: roots and weights are of different length");
    }
    while i < roots.len() {
        sum += weights[i] * f(roots[i]*(b-a)/2.0 + (a+b)/2.0);
        i += 1;
    }
    sum * (b-a)/2.0
}

fn omega(s: f64) -> f64 {
    let a = s.sqrt();
    let b = (1.05f64.powi(2) - s).sqrt();
    (a - b) / (a + b)
}

fn k(s: f64) -> f64 {
    (s/4.0 - M_PI.powi(2)).sqrt()
}

pub fn delta_with_lut(cache:&mut Cache, s: f64) -> f64 {
    match cache.delta_lut.get(&s.to_bits()) {
        Some(val) => {
            *val
        }
        None => {
            let val = delta(s);
            cache.delta_lut.insert(s.to_bits(), val);
            val
        }
    }
}

pub fn delta(s: f64) -> f64 {
    if s <= 2.0 * M_K {
        atan_shift(
            (s.sqrt() / (2.0*k(s).powi(3))
             * (M_ROH.powi(2) - s) 
             * ( (2.0*M_PI.powi(3)) / (M_ROH.powi(2) * s.sqrt()) + B0 + B1 * omega(s))
            ).powi(-1)
        )
    } else if s <= S_MID_CUTOFF {
        let par = 0.5 * s.sqrt() / M_K - 1.0;
        lambda_0() + LAMBDA1 * par + LAMBDA2 * par.powi(2)
    } else {
        PI + 
        (delta_mid_cutoff() - PI)
         * (LAMBDA_HIGH * LAMBDA_HIGH + S_MID_CUTOFF)
         / (LAMBDA_HIGH * LAMBDA_HIGH + s)
    }
}

pub fn omnes_integrand_tan(cache:&mut Cache, s_tick:f64, s:f64) -> Complex {
    let sub = s_tick.tan();
    (delta_with_lut(cache, sub) - delta_with_lut(cache, s)) 
    / (
        sub
        * (sub - s + Complex::new(0.0, EPSILON)) 
        * s_tick.cos().powi(2)
    )
}

pub fn omnes_integrand_notan(cache:&mut Cache, s_tick:f64, s:f64) -> Complex {
    (delta_with_lut(cache, s_tick) - delta_with_lut(cache, s)) 
    / (s_tick * (s_tick - s + Complex::new(0.0, EPSILON)))
}

pub fn omnes(cache:&mut Cache, s:f64, n:u32, use_tan: bool) -> Complex {
    let roots: Rc<[f64]>;
    let weights: Rc<[f64]>;

    if let Some(inner_lut) = cache.omnes_lut.get(&n) {
        if let Some(res) = inner_lut.get(&s.to_bits()) {
            return *res
        }
    }
    
    if !cache.omnes_lut.contains_key(&n) {
        cache.omnes_lut.insert(n, HashMap::with_hasher(Default::default()));
    }   
    
    match cache.gauss_quad_lut.get_mut(&n) {
        Some(gq_values) => {
            roots = Rc::clone(&gq_values.0);
            weights = Rc::clone(&gq_values.1);
        }
        None => {
            let gq_values = GaussLegendre::nodes_and_weights(n.try_into().unwrap());    
            roots = Rc::from(gq_values.0.as_slice());
            weights = Rc::from(gq_values.1.as_slice());
        }
    }
    
    if !cache.gauss_quad_lut.contains_key(&n) {
        let gq_values = GaussLegendre::nodes_and_weights(n.try_into().unwrap());
        cache.gauss_quad_lut.insert(n, (Rc::from(gq_values.0), Rc::from(gq_values.1)));
        println!("n_omnes={} is not included in the lookup file. Consider generating a new lookup file via the appropriate function in gq_storage.rs", n); 
    }

    let intgrl = if use_tan {
        integrate_complex(&roots, &weights, S0.atan(), PI / 2.0, |s_tick| omnes_integrand_tan(cache, s_tick, s))
    } else {
        integrate_complex(&roots, &weights, S0, INF, |s_tick| omnes_integrand_notan(cache, s_tick, s))
    };
    let omnes_res = (s/PI*(intgrl)).exp() * (S0 / Complex::new(S0 - s, -EPSILON)).powf(delta_with_lut(cache, s) / PI);
    cache.omnes_lut.get_mut(&n).unwrap().insert(s.to_bits(), omnes_res);

    omnes_res
}

pub fn phi0_integrand_tan_xsub(cache:&mut Cache, s_tick: f64, s: f64, n_omnes:u32, use_tan_omnes: bool) -> f64 {
    let x_sub = s_tick.tan().powi(2) + S0;
    let omnes_s = omnes(cache, s, n_omnes, use_tan_omnes);
    let omnes_x = omnes(cache, x_sub, n_omnes, use_tan_omnes);
    (omnes_x / omnes_s).norm_sqr().ln()
    / ( (x_sub)* (x_sub-s) )
    / s_tick.cos().powi(2) // Don't use x_sub here, since this term is from tan substitution itself
}

pub fn phi0_integrand_notan_xsub(cache:&mut Cache, s_tick: f64, s: f64, n_omnes:u32, use_tan_omnes: bool) -> f64 {
    let x_sub = s_tick.powi(2) + S0;
    let omnes_s = omnes(cache, s, n_omnes, use_tan_omnes);
    let omnes_x = omnes(cache, x_sub, n_omnes, use_tan_omnes);
    (omnes_x / omnes_s).norm_sqr().ln()
    / ( (x_sub)* (x_sub-s) )
}

pub fn phi0_integrand_tan_noxsub(cache:&mut Cache, s_tick: f64, s: f64, n_omnes:u32, use_tan_omnes: bool) -> f64 {
    let tansub = s_tick.tan();
    println!("lol {:?}", omnes(cache, tansub, n_omnes, use_tan_omnes));
    let omnes_s = omnes(cache, s, n_omnes, use_tan_omnes);
    let omnes_s_tick = omnes(cache, tansub, n_omnes, use_tan_omnes);
    let res = (omnes_s_tick / omnes_s).norm_sqr().ln()
    / ( (tansub - S0).sqrt() * tansub * (tansub - s) ) 
    / s_tick.cos().powi(2);
    println!("{:?} {:?} {:?} {:?} {:?}", s_tick, tansub, omnes_s, omnes_s_tick, res);
    println!("LOL {:?}", omnes(cache, tansub, n_omnes, use_tan_omnes));
    res
}

pub fn phi0_integrand_notan_noxsub(cache:&mut Cache, s: f64, s_tick: f64, n_omnes:u32, use_tan_omnes: bool) -> f64 {
    
    let omnes_s = omnes(cache, s, n_omnes, use_tan_omnes);
    let omnes_s_tick = omnes(cache, s_tick, n_omnes, use_tan_omnes);
    let res = (omnes_s_tick / omnes_s).norm_sqr().ln()
    / ( (s_tick - S0).sqrt() * s_tick * (s_tick - s) );
    println!("{:?} {:?} {:?} {:?} {:?}", s, s_tick, omnes_s, omnes_s_tick, res);
    
    res

}

pub fn phi0(cache:&mut Cache, s: f64, n_omnes: u32, n_phi0: u32, use_tan_omnes: bool, use_tan_phi0: bool, use_xsub: bool) -> f64 {
    //let t0 = Instant::now();
    if !cache.gauss_quad_lut.contains_key(&n_phi0) {
        let gq_values = GaussLegendre::nodes_and_weights(n_phi0.try_into().unwrap());
        cache.gauss_quad_lut.insert(n_phi0, (Rc::from(gq_values.0), Rc::from(gq_values.1)));
        println!("n_omnes={} is not included in the lookup file. Consider generating a new lookup file via the appropriate function in gq_storage.rs", n_phi0);
    }
    let (ref roots, ref weights) = cache.gauss_quad_lut.get(&n_phi0).unwrap();
    let (roots, weights) = (Rc::clone(roots), Rc::clone(weights));

    let analyt = omnes(cache, s, n_omnes, use_tan_omnes).norm_sqr().ln() / S0.sqrt() / 2.0;
    let intgrl = if use_tan_phi0 {
        if use_xsub {
            s / PI * integrate(&*roots, &*weights, 0.0, PI / 2.0, |x| phi0_integrand_tan_xsub(cache, x, s, n_omnes, use_tan_omnes))
        } else {
            s / PI * integrate(&*roots, &*weights, S0.atan(), PI / 2.0, |x| phi0_integrand_tan_noxsub(cache, x, s, n_omnes, use_tan_omnes))
        }
    } else {
        if use_xsub {
            s / PI * integrate(&*roots, &*weights, 0.0, INF, |x| phi0_integrand_notan_xsub(cache, x, s, n_omnes, use_tan_omnes))
        } else {
            s / PI * integrate(&*roots, &*weights, S0, INF, |x| phi0_integrand_notan_noxsub(cache, x, s, n_omnes, use_tan_omnes))
        }
    };
    //println!("{:?}", t0.elapsed());
    -(s-S0).sqrt() * (intgrl - analyt)
}

#[allow(unused)]
pub fn bit_pattern_randomness_tester(a:f64, b:f64, n_points: u64) -> Vec<f64>{
    let mut modulos = vec![0.0; n_points as usize];
    let x_values: Vec<u64> = (0..modulos.len() as u32).map(|x| (f64::from(x) * ((b - a) / f64::from(modulos.len() as u32)) + a).to_bits()).collect();
    for i in 0..x_values.len() {
        let val = (x_values[i] % n_points) as usize;
        modulos[val] += 1.0;
    }
    modulos
}
