use dynamics::Dynamics;
use cell;
use io::IO;
use std::fs::File;
use std::io::prelude::*;
use std::path::Path;
use image;
use std::error::Error;
use lattice_ops_2d::LatticeOps2d;
use lattice_ops_2d::MultiLatticeOps2d;
use lattice2d::Lattice2d;
use lattice2d::Lattice2dBuilder;
use std::mem;
use std::vec::Vec;
use point2d::Point2d;
use communicator::Communicator2d;
use std::f64;
use std::u8;
use box2d::Box2d;
use ind_hlp;
use rayon::prelude::*;
use sdl2::render::Renderer;
use sdl2::pixels::Color;
use sdl2::rect::Point;
use colormaps;

/// Function that computes the number of elements in
/// in a line of integer length "n" divided in 
/// n_lattice subdomains.
fn elem_per_block(n: usize, n_lattice: usize) -> Vec<usize>{
	let mut n_per_lattice = Vec::with_capacity(n_lattice*mem::size_of::<usize>());
	let first_n = n / n_lattice;
	let mut rem = n % n_lattice;
	for _i in 0..n_lattice {
		if rem > 0 {
			n_per_lattice.push(first_n+1);
			rem -= 1;
		} else {
			n_per_lattice.push(first_n);
		}
	}
	n_per_lattice
}

/// The multi-block structure contains the information
/// about each block for parallemlism.
#[derive(Debug)]
pub struct MultiBlock2d {
	nx: usize,
	ny: usize,
	nx_blocks: usize,
	ny_blocks: usize,
	env: usize,
	boxes_offsets: Vec<(Box2d,Point2d)>,
	comms: Vec<Communicator2d>,
}

impl MultiBlock2d {
	fn new(nx: usize, ny: usize, nx_blocks: usize, ny_blocks: usize, env: usize) -> MultiBlock2d {
		let nx_per_block = elem_per_block(nx, nx_blocks);
		let ny_per_block = elem_per_block(ny, ny_blocks);

		let mut boxes_offsets = Vec::new();
		let mut comms = Vec::new();
		let mut off_x = 0;
		for &x in &nx_per_block {
			let mut off_y = 0;
			for &y in &ny_per_block {
				boxes_offsets.push((Box2d::new(0, (x+2usize*env) as i32,0, (y+2usize*env) as i32),
					                Point2d::new(off_x-env as i32, off_y-env as i32))); // add twice the envelope in each direction
				off_y += y as i32;
			}
			off_x += x as i32;
		}

		for ix in 0..nx_blocks {
			for iy in 0..ny_blocks {
				let i_loc = ind_hlp::get_grid_idx(ix, iy, nx_blocks, ny_blocks);


				let loc_nx = boxes_offsets[i_loc].0.x1 as usize;
				let loc_ny = boxes_offsets[i_loc].0.y1 as usize;

				let mut from_boxes = Vec::new();
				for i in 1..cell::Q {
					let next_i = ((ix as i32 + cell::C[i][0] + nx_blocks as i32) % nx_blocks as i32) as usize;
					let next_j = ((iy as i32 + cell::C[i][1] + ny_blocks as i32) % ny_blocks as i32) as usize;

					let i_next = ind_hlp::get_grid_idx(next_i, next_j, nx_blocks, ny_blocks);
					let next_nx = boxes_offsets[i_next].0.x1;
					let next_ny = boxes_offsets[i_next].0.y1;

					from_boxes.push(Box2d::new(0, next_nx as i32, 0, next_ny as i32))
				}

				comms.push(Communicator2d::new(loc_nx, loc_ny, from_boxes, env));
			}
		}
		MultiBlock2d{nx: nx, ny: ny, nx_blocks: nx_blocks, ny_blocks: ny_blocks, env: env,
		             boxes_offsets: boxes_offsets, comms: comms}
	}

	pub fn get_boxes_offsets<'a>(&'a self) -> &'a Vec<(Box2d,Point2d)> {
		&self.boxes_offsets
	}
}

pub struct MultiLattice2dBuilder{
	ml: Vec<Lattice2d>,
	multi_block: MultiBlock2d,
	comm_pops: Vec<Vec<Vec<f64>>>,
}

impl MultiLattice2dBuilder {
	pub fn new(nx: usize, ny: usize, nx_lattice: usize, ny_lattice: usize, env: usize) -> Self {
		
		let multi_block = MultiBlock2d::new(nx, ny, nx_lattice, ny_lattice, env);
		let mut comm_pops:Vec<Vec<Vec<f64>>> = Vec::new();
		{
			let comms = &multi_block.comms;
			for comm in comms.iter() {
				let mut pops_next = Vec::new();
				for i_pop in 1..cell::Q {
					let from = &comm.get_from(i_pop-1);
					pops_next.push(vec![0.0; from.size()*cell::Q]);
				}
				comm_pops.push(pops_next);
			}
		}

		let ml = Vec::new();

		MultiLattice2dBuilder{ 
			ml: ml, multi_block: multi_block, comm_pops: comm_pops }
	}

	fn set_gen_pops_eq<Frho,Fu,V>(mut self, c:V, rho: Frho, u: Fu) -> Self 
		where Frho : Fn(i32,i32) -> f64, Fu: Fn(i32,i32) -> [f64; cell::D], V: 'static+Dynamics+Clone
	{
		for b_o in self.multi_block.boxes_offsets.iter() {
			let &(bb,loc) = b_o;
			self.ml.push(Lattice2dBuilder::new(bb.x1 as usize,bb.y1 as usize,loc).set_dyn(c.clone()).set_gen_pops_eq(&rho,&u).finalize());
		}
		self
	}
	fn set_gen_pops_neq<Frho,Fu,Fpi,V>(mut self, c:V, rho: Frho, u: Fu, pi: Fpi) -> Self 
		where Frho : Fn(i32,i32) -> f64, Fu: Fn(i32,i32) -> [f64; cell::D],
			  Fpi: Fn(i32,i32) -> [[f64; cell::D]; cell::D], V: 'static+Dynamics+Clone
	{
		for b_o in self.multi_block.boxes_offsets.iter() {
			let &(bb,loc) = b_o;
			self.ml.push(Lattice2dBuilder::new(bb.x1 as usize,bb.y1 as usize,loc)
				.set_dyn(c.clone()).set_gen_pops_neq(&rho,&u,&pi).finalize());
		}
		self
	}

	fn set_const_pops_eq<V>(mut self, c:V, rho: f64, u: [f64; cell::D]) -> Self where V: 'static+Dynamics+Clone {
		for b_o in self.multi_block.boxes_offsets.iter() {
			let &(bb,loc) = b_o;
			self.ml.push(Lattice2dBuilder::new(bb.x1 as usize,bb.y1 as usize,loc)
				.set_dyn(c.clone()).set_const_pops_eq(rho,u).finalize());
		}
		self
	}
	fn set_const_pops_neq<V>(mut self, c:V, rho: f64, u: [f64; cell::D], pi: [[f64; cell::D]; cell::D]) -> Self 
		where V: 'static+Dynamics+Clone
	{
		for b_o in self.multi_block.boxes_offsets.iter() {
			let &(bb,loc) = b_o;
			self.ml.push(Lattice2dBuilder::new(bb.x1 as usize,bb.y1 as usize,loc)
				.set_dyn(c.clone()).set_const_pops_neq(rho,u,pi).finalize());
		}
		self
	}

	pub fn finalize(self) -> MultiLattice2d {
		MultiLattice2d{ml: self.ml, multi_block: self.multi_block, comm_pops: self.comm_pops,}
	}
}

#[derive(Debug)]
pub struct MultiLattice2d{
	ml: Vec<Lattice2d>,
	multi_block: MultiBlock2d,
	comm_pops: Vec<Vec<Vec<f64>>>,
}


impl MultiLattice2d {
	pub fn set_gen_dyn<V, Fmask>(&mut self, c: V, foo: Fmask)  
		where V: 'static+Dynamics+Clone, Fmask : Fn(i32,i32) -> bool
	{
		for i in 0..self.ml.len() {
			self.ml[i].set_gen_dyn(c.clone(), &foo);
		}
	}

	pub fn new_eq_const<V>(nx: usize, ny: usize, nx_lattice: usize, ny_lattice: usize, env: usize, 
						rho: f64, u: [f64; cell::D], c: V) -> Self where V: 'static+Dynamics+Clone
	{
		MultiLattice2dBuilder::new(nx,ny,nx_lattice,ny_lattice,env).set_const_pops_eq(c,rho,u).finalize()
	}

	pub fn new_eq_neq_const<V>(nx: usize, ny: usize, nx_lattice: usize, ny_lattice: usize, env: usize, 
		rho: f64, u: [f64; cell::D], pi: [[f64; cell::D]; cell::D], c: V) -> Self where V: 'static+Dynamics+Clone
	{
		MultiLattice2dBuilder::new(nx,ny,nx_lattice,ny_lattice,env).set_const_pops_neq(c,rho,u,pi).finalize()
	}

	pub fn new_eq<Frho,Fu,V>(nx: usize, ny: usize, nx_lattice: usize, ny_lattice: usize, env: usize, 
			rho: Frho, u: Fu, c: V) -> Self
		where Frho : Fn(i32,i32) -> f64, Fu: Fn(i32,i32) -> [f64; cell::D],  V: 'static+Dynamics+Clone
	{
		MultiLattice2dBuilder::new(nx,ny,nx_lattice,ny_lattice,env).set_gen_pops_eq(c,rho,u).finalize()
	}

	pub fn new_eq_neq<Frho,Fu,Fpi,V>(nx: usize, ny: usize, nx_lattice: usize, ny_lattice: usize, env: usize, 
			rho: Frho, u: Fu, pi: Fpi, c: V) -> Self
		where Frho : Fn(i32,i32) -> f64, Fu: Fn(i32,i32) -> [f64; cell::D], Fpi: Fn(i32,i32) -> [[f64; cell::D]; cell::D],
		V: 'static+Dynamics+Clone
	{
		MultiLattice2dBuilder::new(nx,ny,nx_lattice,ny_lattice,env).set_gen_pops_neq(c,rho,u,pi).finalize()
	}

	pub fn get_nx(&self) -> usize {
		self.multi_block.nx
	}

	pub fn get_ny(&self) -> usize {
		self.multi_block.ny
	}

	pub fn get_nx_lattice(&self) -> usize {
		self.multi_block.nx_blocks
	}

	pub fn get_ny_lattice(&self) -> usize {
		self.multi_block.ny_blocks
	}

	pub fn get_env(&self) -> usize {
		self.multi_block.env
	}

	pub fn get_multi_block<'a>(&'a self) -> &'a MultiBlock2d {
		&self.multi_block
	}

	pub fn get_lattice<'a>(&'a self, i: usize) -> &'a Lattice2d {
		&self.ml[i]
	}

	pub fn get_lattices<'a>(&'a self) -> &'a Vec<Lattice2d> {
		&self.ml
	}

	pub fn get_mut_lattices<'a>(&'a mut self) -> &'a mut Vec<Lattice2d> {
		&mut self.ml
	}

	pub fn get_mut_lattice<'a>(&'a mut self, i: usize) -> &'a mut Lattice2d {
		&mut self.ml[i]
	}
}


impl MultiLatticeOps2d for MultiLattice2d {
	fn communicate(&mut self) {
		let mut comm_pops = vec![]; 
		self.multi_block.comms.par_iter().enumerate().map( |(i,comm)| 
			{
				let (x,y) = ind_hlp::get_grid_indices(i, self.get_nx_lattice(), self.get_ny_lattice());
				let mut pops_next = Vec::new();
				for i_pop in 1..cell::Q {
					let next_x = ((x as i32 + cell::C[i_pop][0] + self.get_nx_lattice() as i32) % self.get_nx_lattice() as i32) as usize;
					let next_y = ((y as i32 + cell::C[i_pop][1] + self.get_ny_lattice() as i32) % self.get_ny_lattice() as i32) as usize;

					let from = &comm.get_from(i_pop-1);
					let i_next = ind_hlp::get_grid_idx(next_x, next_y, self.get_nx_lattice(), self.get_ny_lattice());

					pops_next.push(self.ml[i_next].get_pops_from_box(from));
				}
				pops_next

			}).collect_into(&mut comm_pops);

		self.multi_block.comms.par_iter()
			.zip(self.ml.par_iter_mut())
			.zip(comm_pops.par_iter())
			.for_each( move |((comm, l), pops)| {
			l.copy_pops(comm.get_to(), &pops);
		});
	}

}

impl LatticeOps2d for MultiLattice2d {
	fn collide(&mut self) {
		self.ml[..].par_iter_mut().for_each(|l| l.collide());
	}
	
	fn periodic_stream(&mut self) {
		self.ml[..].par_iter_mut().for_each(|l| l.periodic_stream());

	}

	fn bulk_stream(&mut self) {
		self.ml[..].par_iter_mut().for_each(|l| l.bulk_stream());
	}

	fn bd_stream(&mut self) {
		self.ml[..].par_iter_mut().for_each(|l| l.bd_stream());
	}

	fn collide_and_stream(&mut self) {
		self.ml[..].par_iter_mut().for_each(|l| {
			l.collide_and_stream();
		});
		self.communicate();
	}
}


// ========== Implementation of IO trait for MultiLattice2d =========== //

fn float_max(v: &Vec<f64>) -> f64 {
	v.iter().cloned().fold(f64::NAN, f64::max)
}

fn float_min(v: &Vec<f64>) -> f64 {
	v.iter().cloned().fold(f64::NAN, f64::min)
}


impl IO for MultiLattice2d {
	fn write_scalar<Fun>(&self, file: &mut File, foo: Fun) where Fun: Fn(&[f64; cell::Q]) -> f64 {
		let mut buf : String = String::new();

		let mut l_ny:usize  = (self.get_ny()+1) as usize;
		for y in 0..self.get_ny_lattice() {

			let mut y_l = self.get_env();
			while y_l < l_ny {
				
				for x in 0..self.get_nx_lattice() {
					let i_l = ind_hlp::get_grid_idx(x as usize,y as usize, self.get_nx_lattice(), self.get_ny_lattice());
					let lattice = &self.ml[i_l];
					l_ny = lattice.get_ny()-self.get_env();

					for x_l in self.get_env()..lattice.get_nx()-self.get_env() {
						let i = ind_hlp::get_grid_idx(x_l, y_l, lattice.get_nx(), lattice.get_ny());
						buf = buf + &(foo(lattice.get_pops(i)).to_string())+" ";
					}
				}
				y_l += 1;
			}
		}

		match file.write_all(buf.as_bytes()) {
	        Err(why) => {
	            panic!("couldn't write to file : {}", why.description())
	        },
	        Ok(_) => println!("successfully wrote to file."),
	    }
	}

	fn write_scalar_to_png<Fun>(&self, fname: String, foo: Fun) where Fun: Fn(&[f64; cell::Q]) -> f64 {
		let mut buf : String = String::new();

		let mut l_ny:usize  = (self.get_ny()+1) as usize;
		for y in 0..self.get_ny_lattice() {

			let mut y_l = self.get_env();
			while y_l < l_ny {
				
				for x in 0..self.get_nx_lattice() {
					let i_l = ind_hlp::get_grid_idx(x as usize,y as usize, self.get_nx_lattice(), self.get_ny_lattice());
					let lattice = &self.ml[i_l];
					l_ny = lattice.get_ny()-self.get_env();

					for x_l in self.get_env()..lattice.get_nx()-self.get_env() {
						let i = ind_hlp::get_grid_idx(x_l, y_l, lattice.get_nx(), lattice.get_ny());
						buf = buf + &(foo(lattice.get_pops(i)).to_string())+" ";
					}
				}
				y_l += 1;
			}
		}

		let buf_str: &str = &buf;
		let v_f64:Vec<f64> = buf_str.split_whitespace().map(|x| x.parse::<f64>().unwrap() as f64).collect();

		let max = float_max(&v_f64);
		let min = float_min(&v_f64);
		let v_u8:Vec<u8> = v_f64.iter().map(|x| ((x-min)/(max-min)* (u8::MAX as f64)) as u8 ).collect();

	    // Create a new ImgBuf with width: imgx and height: imgy
	    let mut imgbuf = image::ImageBuffer::new(self.get_nx() as u32, self.get_ny() as u32);

	    // Iterate over the coordiantes and pixels of the image
	    for (x, y, pixel) in imgbuf.enumerate_pixels_mut() {
	    	// let i:usize = (y + x * self.get_ny() as u32) as usize;
	    	let i:usize = ind_hlp::get_grid_idx(y as usize,x as usize,self.get_ny(), self.get_nx()) as usize;

	        // Create an 8bit pixel of type Luma and value i
	        // and assign in to the pixel at position (x, y)
	        *pixel = image::Luma([v_u8[i] as u8]);
	    }

	    // Save the image as “fractal.png”
	    let ref mut fout = File::create(&Path::new(&fname)).unwrap();

	    // We must indicate the image’s color type and what format to save as
	    let _ = image::ImageLuma8(imgbuf).save(fout, image::PNG);
	}

	#[allow(unused_must_use)]
	fn show_scalar<Fun>(&self, r: &mut Renderer, foo: Fun, cm: &colormaps::ColorMap) where Fun: Fn(&[f64; cell::Q]) -> f64 {
		// clear window
	    let white = Color::RGB(255, 255, 255);
		r.set_draw_color(white);
	    r.clear();

		let mut v_f64 = Vec::new();

		let mut l_ny:usize  = (self.get_ny()+1) as usize;
		for y in 0..self.get_ny_lattice() {

			let mut y_l = self.get_env();
			while y_l < l_ny {
				
				for x in 0..self.get_nx_lattice() {
					let i_l = ind_hlp::get_grid_idx(x as usize,y as usize, self.get_nx_lattice(), self.get_ny_lattice());
					let lattice = &self.ml[i_l];
					l_ny = lattice.get_ny()-self.get_env();

					for x_l in self.get_env()..lattice.get_nx()-self.get_env() {
						let i = ind_hlp::get_grid_idx(x_l, y_l, lattice.get_nx(), lattice.get_ny());
						v_f64.push(foo(lattice.get_pops(i)));
					}
				}
				y_l += 1;
			}
		}

		let max = float_max(&v_f64);
		let min = float_min(&v_f64);
		// let v_u32:Vec<u32> = v_f64.iter().map(|x| ((x-min)/(max-min)* (u32::MAX as f64)) as u32 ).collect();
		let v:Vec<f64> = v_f64.iter().map(|x| ( (x-min)/(max-min) )).collect();

		for x in 0..self.get_nx() {
			for y in 0..self.get_ny() {
					let i = ind_hlp::get_grid_idx(y as usize,x as usize, self.get_ny(), self.get_nx());

					let get_color = cm.get_color(&v[i]);
					match get_color {
					    Ok(color) => {
					    	r.set_draw_color(color);
	        				r.draw_point(Point::new(x as i32, y as i32));
	        			},
					    Err(error) => println!("{:?}", error),
					}
					
			}
		}
		r.present();
	}
}