use crate::utils::*;

use nalgebra::{base::*, matrix, vector};
use std::marker::PhantomData;

pub trait Settings {}

pub trait Model<T, S: Settings, const D: usize> {
	/// Returns f(x)
	fn f(&self, x: Vect<T, D>) -> Vect<T, D>;
	/// Returns df(x)/dx
	fn df(&self, x: Vect<T, D>) -> Mat<T, D, D>;
	fn get_settings(&self) -> &S;
	fn get_settings_mut(&mut self) -> &mut S;
}

pub trait Coloring<T, P, C, const D: usize> {
	fn prepare<I: Iterator<Item = Vect<T, D>>>(iter: I) -> Option<P>;
	fn color(pre: &P, val: Vect<T, D>) -> C;
}

pub struct MinMaxColoringInfo<T, const D: usize> {
	pub min: [T; D],
	pub max: [T; D],
}

pub trait Constraint<T, const D: usize> {
	fn constrain_mut(&self, val: &mut Vect<T, D>);
	fn constrain(&self, mut val: Vect<T, D>) -> Vect<T, D> {
		self.constrain_mut(&mut val);
		val
	}
}

pub struct MinMaxConstraint<T, const D: usize> {
	pub min: [T; D],
	pub max: [T; D],
}

impl<T: Copy + PartialOrd, const D: usize> Constraint<T, D> for MinMaxConstraint<T, D> {
	fn constrain_mut(&self, val: &mut Vect<T, D>) {
		for ((comp, min), max) in val.iter_mut().zip(self.min.iter()).zip(self.max.iter()) {
			*comp = num_traits::clamp(*comp, *min, *max)
		}
	}
}

// Models

pub mod constrained {
	use super::*;

	#[derive(Clone, Debug)]
	pub struct ConstrainedSettings<T, O, M, S> {
		pub constraint: O,
		pub model: M,
		pub _p: PhantomData<(T, S)>,
	}

	impl<T, O, M, S> Settings for ConstrainedSettings<T, O, M, S> {}

	#[derive(Clone)]
	pub struct Constrained<T, O, M, S> {
		pub s: ConstrainedSettings<T, O, M, S>,
	}

	impl<T, O: Constraint<T, D>, M: Model<T, S, D>, S: Settings, const D: usize>
		Model<T, ConstrainedSettings<T, O, M, S>, D> for Constrained<T, O, M, S>
	{
		fn f(&self, x: Vect<T, D>) -> Vect<T, D> {
			self.s.constraint.constrain(self.s.model.f(x))
		}
		fn df(&self, x: Vect<T, D>) -> Mat<T, D, D> {
			self.s.model.df(x)
		}
		fn get_settings(&self) -> &ConstrainedSettings<T, O, M, S> {
			&self.s
		}
		fn get_settings_mut(&mut self) -> &mut ConstrainedSettings<T, O, M, S> {
			&mut self.s
		}
	}

	impl<T, O, M: Coloring<T, P, C, D>, S, P, C, const D: usize> Coloring<T, P, C, D>
		for Constrained<T, O, M, S>
	{
		fn prepare<I: Iterator<Item = Vect<T, D>>>(iter: I) -> Option<P> {
			M::prepare(iter)
		}
		fn color(pre: &P, val: Vect<T, D>) -> C {
			M::color(pre, val)
		}
	}
}

/// SIR model without vital dynamics
pub mod sir {
	use super::*;

	#[derive(Clone, Debug)]
	pub struct SirSettings<T> {
		/// Transmission probability
		pub beta: T,
		/// Removal probability
		pub gamma: T,
		pub pop: T,
	}

	impl<T> Settings for SirSettings<T> {}

	#[derive(Clone)]
	pub struct Sir<T> {
		pub s: SirSettings<T>,
	}

	impl Model<f64, SirSettings<f64>, 2> for Sir<f64> {
		fn f(&self, x: Vector2<f64>) -> Vector2<f64> {
			vector![
				-self.s.beta * x[0] * x[1] / self.s.pop,
				self.s.beta * x[0] * x[1] / self.s.pop - self.s.gamma * x[1]
			]
		}
		fn df(&self, x: Vector2<f64>) -> Matrix2<f64> {
			matrix![
				-self.s.beta*x[1]/self.s.pop, -self.s.beta*x[0]/self.s.pop;
				self.s.beta*x[1]/self.s.pop, self.s.beta*x[0]/self.s.pop-self.s.gamma
			]
		}
		fn get_settings(&self) -> &SirSettings<f64> {
			&self.s
		}
		fn get_settings_mut(&mut self) -> &mut SirSettings<f64> {
			&mut self.s
		}
	}

	impl<T> From<SirSettings<T>> for Vect<T, 3> {
		fn from(s: SirSettings<T>) -> Self {
			vector![s.beta, s.gamma, s.pop]
		}
	}

	impl<T: Scalar + Copy> From<Vect<T, 3>> for SirSettings<T> {
		fn from(v: Vect<T, 3>) -> Self {
			Self {
				beta: v[0],
				gamma: v[1],
				pop: v[2],
			}
		}
	}
}

/// SIR model with vital dynamics, constant population
pub mod sirv {
	use super::*;
	#[derive(Clone, Debug)]
	pub struct SirvSettings<T> {
		/// Transmission probability
		pub beta: T,
		/// Removal probability
		pub gamma: T,
		/// Birth rate
		pub lambda: T,
		/// Death rate
		pub mu: T,
		pub pop: T,
	}

	impl<T> Settings for SirvSettings<T> {}

	#[derive(Clone)]
	pub struct Sirv<T> {
		pub s: SirvSettings<T>,
	}

	impl Model<f64, SirvSettings<f64>, 2> for Sirv<f64> {
		fn f(&self, x: Vector2<f64>) -> Vector2<f64> {
			vector![
				self.s.lambda - self.s.beta * x[0] * x[1] / self.s.pop - self.s.mu * x[0],
				self.s.beta * x[0] * x[1] / self.s.pop - self.s.gamma * x[1] - self.s.mu * x[1]
			]
		}
		fn df(&self, x: Vector2<f64>) -> Matrix2<f64> {
			matrix![
				-self.s.beta*x[1]/self.s.pop - self.s.mu, -self.s.beta*x[0]/self.s.pop;
				self.s.beta*x[1]/self.s.pop, self.s.beta*x[0]/self.s.pop-self.s.gamma - self.s.mu
			]
		}
		fn get_settings(&self) -> &SirvSettings<f64> {
			&self.s
		}
		fn get_settings_mut(&mut self) -> &mut SirvSettings<f64> {
			&mut self.s
		}
	}

	impl<T> From<SirvSettings<T>> for Vect<T, 5> {
		fn from(s: SirvSettings<T>) -> Self {
			vector![s.beta, s.gamma, s.lambda, s.mu, s.pop]
		}
	}

	impl<T: Scalar + Copy> From<Vect<T, 5>> for SirvSettings<T> {
		fn from(v: Vect<T, 5>) -> Self {
			Self {
				beta: v[0],
				gamma: v[1],
				lambda: v[2],
				mu: v[3],
				pop: v[4],
			}
		}
	}
}

/// Giraffe
pub mod giraffe {
	use super::*;

	#[derive(Clone, Debug)]
	pub struct GiraffeSettings<T> {
		pub a_a: T,
		pub a_b: T,
		pub b_a: T,
		pub b_b: T,
	}

	impl<T> Settings for GiraffeSettings<T> {}

	#[derive(Clone)]
	pub struct Giraffe<T> {
		pub s: GiraffeSettings<T>,
	}

	impl Model<f64, GiraffeSettings<f64>, 2> for Giraffe<f64> {
		fn f(&self, x: Vector2<f64>) -> Vector2<f64> {
			vector![
				self.s.a_a * x[0] + self.s.b_a * x[1],
				self.s.a_b * x[0] + self.s.b_b * x[1]
			]
		}
		fn df(&self, _x: Vector2<f64>) -> Matrix2<f64> {
			matrix![
				self.s.a_a, self.s.b_a;
				self.s.a_b, self.s.b_b
			]
		}
		fn get_settings(&self) -> &GiraffeSettings<f64> {
			&self.s
		}
		fn get_settings_mut(&mut self) -> &mut GiraffeSettings<f64> {
			&mut self.s
		}
	}

	impl<T> From<GiraffeSettings<T>> for Vect<T, 4> {
		fn from(s: GiraffeSettings<T>) -> Self {
			vector![s.a_a, s.a_b, s.b_a, s.b_b]
		}
	}

	impl<T: Scalar + Copy> From<Vect<T, 4>> for GiraffeSettings<T> {
		fn from(v: Vect<T, 4>) -> Self {
			Self {
				a_a: v[0],
				a_b: v[1],
				b_a: v[2],
				b_b: v[3],
			}
		}
	}

	/*impl Coloring<f64, MinMaxColoringInfo<f64, 2>, [u8; 3], 2> for Giraffe<f64> {
	fn prepare<I: Iterator<Item = Vect<f64, 2>>>(
		iter: I,
	) -> Option<MinMaxColoringInfo<f64, 2>> {
		let mut r = MinMaxColoringInfo {
			min: [1.0, 1.0],
			max: [0.0, 0.0],
		};
		for val in iter {
			if val[0] < r.min[0] {
				r.min[0] = val[0];
			}
			if val[0] > r.max[0] {
				r.max[0] = val[0];
			}
			if val[1] < r.min[1] {
				r.min[1] = val[1];
			}
			if val[1] > r.max[1] {
				r.max[1] = val[1];
			}
		}
		if r.min[0] == r.max[0] || r.min[1] == r.max[1] {
			None
		} else {
			Some(r)
		}
	}
	fn color(pre: &MinMaxColoringInfo<f64, 2>, val: Vect<f64, 2>) -> [u8; 3] {
		[
			((val[1] - pre.min[1]).abs() / (pre.max[1] - pre.min[1]).abs() * 255.0) as u8,
			0,
			((val[0] - pre.min[0]).abs() / (pre.max[0] - pre.min[0]).abs() * 255.0) as u8,
		]
	}*/

	impl Coloring<f64, (), [u8; 3], 2> for Giraffe<f64> {
		fn prepare<I: Iterator<Item = Vect<f64, 2>>>(_iter: I) -> Option<()> {
			Some(())
		}
		fn color(_pre: &(), val: Vect<f64, 2>) -> [u8; 3] {
			[
				(val[0] * 255.0) as u8,
				((val[0] + val[1]) * 127.5) as u8,
				255 - (val[1] * 255.0) as u8,
			]
		}
	}
}

/// https://arxiv.org/abs/2210.05227
pub mod lyfe {
	use super::*;

	#[derive(Clone, Debug)]
	pub struct LyfeSettings<T> {
		pub da: T,
		pub db: T,
		pub f: T,
		pub r: T,
	}

	impl<T> Settings for LyfeSettings<T> {}

	#[derive(Clone)]
	pub struct Lyfe<T> {
		pub s: LyfeSettings<T>,
	}

	impl Model<f64, LyfeSettings<f64>, 2> for Lyfe<f64> {
		fn f(&self, x: Vector2<f64>) -> Vector2<f64> {
			vector![
				x[0] * x[1].powi(2) + self.s.f * (1.0 - x[0]),
				x[0] * x[1].powi(2) - (self.s.f + self.s.r) * x[1]
			]
		}
		fn df(&self, x: Vector2<f64>) -> Matrix2<f64> {
			matrix![
				x[1].powi(2)-self.s.f, 2.0*x[0]*x[1];
				x[1].powi(2), 2.0*x[0]*x[1] - (self.s.f+self.s.r)
			]
		}
		fn get_settings(&self) -> &LyfeSettings<f64> {
			&self.s
		}
		fn get_settings_mut(&mut self) -> &mut LyfeSettings<f64> {
			&mut self.s
		}
	}

	impl<T> From<LyfeSettings<T>> for Vect<T, 4> {
		fn from(s: LyfeSettings<T>) -> Self {
			vector![s.da, s.db, s.f, s.r]
		}
	}

	impl<T: Scalar + Copy> From<Vect<T, 4>> for LyfeSettings<T> {
		fn from(v: Vect<T, 4>) -> Self {
			Self {
				da: v[0],
				db: v[1],
				f: v[2],
				r: v[3],
			}
		}
	}

	/*impl Coloring<f64, MinMaxColoringInfo<f64, 2>, [u8; 3], 2> for Lyfe<f64> {
	fn prepare<I: Iterator<Item = Vect<f64, 2>>>(
		iter: I,
	) -> Option<MinMaxColoringInfo<f64, 2>> {
		let mut r = MinMaxColoringInfo {
			min: [1.0, 1.0],
			max: [0.0, 0.0],
		};
		for val in iter {
			if val[0] < r.min[0] {
				r.min[0] = val[0];
			}
			if val[0] > r.max[0] {
				r.max[0] = val[0];
			}
			if val[1] < r.min[1] {
				r.min[1] = val[1];
			}
			if val[1] > r.max[1] {
				r.max[1] = val[1];
			}
		}
		if r.min[0] == r.max[0] || r.min[1] == r.max[1] {
			None
		} else {
			Some(r)
		}
	}
	fn color(pre: &MinMaxColoringInfo<f64, 2>, val: Vect<f64, 2>) -> [u8; 3] {
		[
			((val[1] - pre.min[1]).abs() / (pre.max[1] - pre.min[1]).abs() * 255.0) as u8,
			0,
			((val[0] - pre.min[0]).abs() / (pre.max[0] - pre.min[0]).abs() * 255.0) as u8,
		]
	}*/

	impl Coloring<f64, (), [u8; 3], 2> for Lyfe<f64> {
		fn prepare<I: Iterator<Item = Vect<f64, 2>>>(_iter: I) -> Option<()> {
			Some(())
		}
		fn color(_pre: &(), val: Vect<f64, 2>) -> [u8; 3] {
			[
				(val[0] * 255.0) as u8,
				((val[0] + val[1]) * 127.5) as u8,
				255 - (val[1] * 255.0) as u8,
			]
		}
	}
}

pub mod bike {
	use super::*;

	#[derive(Clone, Debug)]
	pub struct BikeSettings<T> {
		/// Pénétration dans l'air
		pub cx: T,
		/// Gravité
		pub g: T,
		/// Masse
		pub m: T,
		/// Pente
		pub th: T,
	}

	impl<T> Settings for BikeSettings<T> {}

	fn b(x: f64, v: f64) -> f64 {
		80.
	}

	fn db(x: f64, v: f64) -> Vector2<f64> {
		vector![0., 0.]
	}

	#[derive(Clone)]
	pub struct Bike<T> {
		pub s: BikeSettings<T>,
	}

	impl Model<f64, BikeSettings<f64>, 2> for Bike<f64> {
		fn f(&self, x: Vector2<f64>) -> Vector2<f64> {
			vector![
				x[1],
				self.s.g * self.s.th.sin() - (self.s.cx * x[1].powi(2) + b(x[0], x[1])) / self.s.m
			]
		}
		fn df(&self, x: Vector2<f64>) -> Matrix2<f64> {
			let dbx = db(x[0], x[1]);
			matrix![
				0., 1.;
				-dbx[0]/self.s.m, -(2.*self.s.cx*x[1]+dbx[1])/self.s.m
			]
		}
		fn get_settings(&self) -> &BikeSettings<f64> {
			&self.s
		}
		fn get_settings_mut(&mut self) -> &mut BikeSettings<f64> {
			&mut self.s
		}
	}

	impl<T> From<BikeSettings<T>> for Vect<T, 4> {
		fn from(s: BikeSettings<T>) -> Self {
			vector![s.cx, s.g, s.m, s.th]
		}
	}

	impl<T: Scalar + Copy> From<Vect<T, 4>> for BikeSettings<T> {
		fn from(v: Vect<T, 4>) -> Self {
			Self {
				cx: v[0],
				g: v[1],
				m: v[2],
				th: v[3],
			}
		}
	}

	/*impl Coloring<f64, MinMaxColoringInfo<f64, 2>, [u8; 3], 2> for Lyfe<f64> {
	fn prepare<I: Iterator<Item = Vect<f64, 2>>>(
		iter: I,
	) -> Option<MinMaxColoringInfo<f64, 2>> {
		let mut r = MinMaxColoringInfo {
			min: [1.0, 1.0],
			max: [0.0, 0.0],
		};
		for val in iter {
			if val[0] < r.min[0] {
				r.min[0] = val[0];
			}
			if val[0] > r.max[0] {
				r.max[0] = val[0];
			}
			if val[1] < r.min[1] {
				r.min[1] = val[1];
			}
			if val[1] > r.max[1] {
				r.max[1] = val[1];
			}
		}
		if r.min[0] == r.max[0] || r.min[1] == r.max[1] {
			None
		} else {
			Some(r)
		}
	}
	fn color(pre: &MinMaxColoringInfo<f64, 2>, val: Vect<f64, 2>) -> [u8; 3] {
		[
			((val[1] - pre.min[1]).abs() / (pre.max[1] - pre.min[1]).abs() * 255.0) as u8,
			0,
			((val[0] - pre.min[0]).abs() / (pre.max[0] - pre.min[0]).abs() * 255.0) as u8,
		]
	}*/

	impl Coloring<f64, (), [u8; 3], 2> for Bike<f64> {
		fn prepare<I: Iterator<Item = Vect<f64, 2>>>(_iter: I) -> Option<()> {
			Some(())
		}
		fn color(_pre: &(), val: Vect<f64, 2>) -> [u8; 3] {
			[
				(val[0] * 255.0) as u8,
				((val[0] + val[1]) * 127.5) as u8,
				255 - (val[1] * 255.0) as u8,
			]
		}
	}
}

pub mod bike2 {
	use super::*;

	#[derive(Clone, Debug)]
	pub struct BikeSettings<T> {
		/// Pénétration dans l'air
		pub cx: T,
		/// Gravité
		pub g: T,
		/// Masse
		pub m: T,
		/// Pente
		pub th: T,
		/// (x, v, settings) -> ((b1, b2), (b1', b2'))
		pub b: fn(T, T, &BikeSettings<T>) -> ((T, T), (T, T)),
	}
	impl<T> Settings for BikeSettings<T> {}

	#[derive(Clone)]
	pub struct Bike<T> {
		pub s: BikeSettings<T>,
	}

	impl Model<f64, BikeSettings<f64>, 2> for Bike<f64> {
		fn f(&self, x: Vector2<f64>) -> Vector2<f64> {
			let ((b1, b2), (_db1, _db2)) = (self.s.b)(x[0], x[1], &self.s);
			vector![
				x[1],
				self.s.g * self.s.th.sin() - (self.s.cx * x[1].powi(2) + b1 + b2) / self.s.m
			]
		}
		fn df(&self, x: Vector2<f64>) -> Matrix2<f64> {
			let ((_b1, _b2), (db1, db2)) = (self.s.b)(x[0], x[1], &self.s);
			matrix![
				0., 1.;
				0., -(2.*self.s.cx*x[1]+db1+db2)/self.s.m
			]
		}
		fn get_settings(&self) -> &BikeSettings<f64> {
			&self.s
		}
		fn get_settings_mut(&mut self) -> &mut BikeSettings<f64> {
			&mut self.s
		}
	}
}