mod charts;
mod live;
mod model;
mod opti;
mod solver;
mod space;
mod utils;

use model::Model;
use opti::GradientDescentOptimizer;
use solver::*;

use nalgebra::vector;
use rand::Rng;
use std::{
	collections::HashMap,
	sync::{Arc, RwLock},
	thread,
	time::Duration,
};

fn main() {
	bike();
	//lyfe();
	//stage();
}

#[allow(dead_code)]
fn lyfe() {
	let size = (800, 800);
	let diffusion = vector![0.2, 0.2];

	let model = model::constrained::Constrained {
		s: model::constrained::ConstrainedSettings {
			model: model::lyfe::Lyfe {
				s: model::lyfe::LyfeSettings {
					da: 0.2,
					db: 0.3,
					f: 0.03,
					r: 0.061,
				},
			},
			constraint: model::MinMaxConstraint {
				min: [0.0, 0.0],
				max: [1.0, 1.0],
			},
			_p: Default::default(),
		},
	};
	/*let solver = ImplicitEulerSolver {
		dt: 0.1,
		tol: 0.000001,
		niters: 100,
	};*/
	let solver = ExplicitEulerSolver { dt: 0.01 };

	let mut space = space::Space {
		model,
		solver,
		old_points: vec![
			space::Point {
				pos: vector![0.0, 0.0],
				diffusion,
			};
			size.0 * size.1
		],
		points: vec![
			space::Point {
				pos: vector![0.0, 0.0],
				diffusion,
			};
			size.0 * size.1
		],
		size,
		sources: HashMap::new(),
		time: 0.0,
		_p: Default::default(),
	};

	let mut rng = rand::thread_rng();
	for _ in 0..100 {
		space.points[rng.gen_range(0..space.old_points.len())].pos[0] = 0.8;
	}
	//space.points[size.0 * size.1 / 2 + size.0 / 2].pos[0] = 0.1;
	//space.points[size.0 * size.1 / 2 + size.0 / 2 + 100].pos[0] = 0.05;

	let space = Arc::new(RwLock::new(space));

	thread::spawn({
		let space = space.clone();
		let interval = Duration::from_millis(1);
		move || loop {
			space.write().unwrap().simulate(0.1);
			std::thread::sleep(interval);
		}
	});

	thread::spawn(move || live::run(space)).join().unwrap();
}

#[allow(dead_code)]
fn giraffe() {
	let size = (800, 800);
	let diffusion = vector![0.1, 0.1];

	let model = model::constrained::Constrained {
		s: model::constrained::ConstrainedSettings {
			model: model::giraffe::Giraffe {
				s: model::giraffe::GiraffeSettings {
					a_a: 0.7,
					a_b: 0.2,
					b_a: -0.5,
					b_b: 0.1,
				},
			},
			constraint: model::MinMaxConstraint {
				min: [0.0, 0.0],
				max: [1.0, 1.0],
			},
			_p: Default::default(),
		},
	};
	/*let solver = ImplicitEulerSolver {
		dt: 0.1,
		tol: 0.000001,
		niters: 100,
	};*/
	let solver = ExplicitEulerSolver { dt: 0.1 };

	let mut space = space::Space {
		model,
		solver,
		old_points: vec![
			space::Point {
				pos: vector![0.0, 0.0],
				diffusion,
			};
			size.0 * size.1
		],
		points: vec![
			space::Point {
				pos: vector![0.0, 0.0],
				diffusion,
			};
			size.0 * size.1
		],
		size,
		sources: HashMap::new(),
		time: 0.0,
		_p: Default::default(),
	};

	let mut rng = rand::thread_rng();
	for _ in 0..100 {
		space.points[rng.gen_range(0..space.old_points.len())].pos[0] = 0.5;
	}
	//space.points[size.0 * size.1 / 2 + size.0 / 2].pos[0] = 0.1;
	//space.points[size.0 * size.1 / 2 + size.0 / 2 + 100].pos[0] = 0.05;

	let space = Arc::new(RwLock::new(space));

	thread::spawn({
		let space = space.clone();
		let interval = Duration::from_millis(1);
		move || loop {
			space.write().unwrap().simulate(0.1);
			std::thread::sleep(interval);
		}
	});

	thread::spawn(move || live::run(space)).join().unwrap();
}

#[allow(dead_code)]
fn stage() {
	let mut rng = rand::thread_rng();

	// ---- Initialization

	let x0 = vector![0.99, 0.01];
	let dt = 0.1;
	let nsamples: usize = 400;
	let nsamples_partial: usize = 40;

	// ---- True data generation

	let settings_true = model::sir::SirSettings {
		beta: 0.6,
		gamma: 0.1,
		pop: 1.0,
	};

	let model = model::sir::Sir {
		s: settings_true.clone(),
	};
	let solver = ImplicitEulerSolver {
		dt: 0.1,
		tol: 0.000001,
		niters: 100,
	};
	let mut xlist_true = Vec::with_capacity(nsamples);
	xlist_true.push(x0);
	let mut x = x0;
	for _ in 0..nsamples - 1 {
		x = solver.f(&model, x);
		xlist_true.push(x);
	}

	// ---- Calibration

	let mut optimizer = GradientDescentOptimizer::new(model, solver);
	let settings_random = model::sir::SirSettings {
		beta: rng.gen(),
		gamma: rng.gen(),
		pop: 1.0,
	};
	*optimizer.model.get_settings_mut() = model::sir::SirSettings {
		beta: 0.38960491052564317,
		gamma: 0.6549130899826807,
		pop: 1.0,
	}; //settings_random.clone();
	let mut optimizer_sto = optimizer.clone();

	let mut xlist_random = Vec::with_capacity(nsamples);
	xlist_random.push(x0);
	let mut x = x0;
	for _ in 0..nsamples - 1 {
		x = optimizer.solver.f(&optimizer.model, x);
		xlist_random.push(x);
	}

	// Batch

	let mut path = Vec::new();
	let mut error = Vec::new();
	for rate in [1.0, 0.1, 0.01, 0.001] {
		let (path_, error_) = &mut optimizer.calibrate_batch_record(
			&xlist_true[..nsamples_partial],
			0.00001,
			rate,
			1000,
			0..2,
		);
		path.append(path_);
		error.append(error_);
	}

	let mut xlist = Vec::with_capacity(nsamples);
	xlist.push(x0);
	let mut x = x0;
	for _ in 0..nsamples - 1 {
		x = optimizer.solver.f(&optimizer.model, x);
		xlist.push(x);
	}

	// Stochastic

	let mut path_sto = Vec::new();
	let mut error_sto = Vec::new();
	for rate in [1.0, 0.1, 0.01, 0.001] {
		let (path_, error_) = &mut optimizer_sto.calibrate_stochastic_record(
			&xlist_true[..nsamples_partial],
			0.00001,
			rate,
			10,
			0..2,
		);
		path_sto.append(path_);
		error_sto.append(error_);
	}

	let mut xlist_sto = Vec::with_capacity(nsamples);
	xlist_sto.push(x0);
	let mut x = x0;
	for _ in 0..nsamples - 1 {
		x = optimizer_sto.solver.f(&optimizer_sto.model, x);
		xlist_sto.push(x);
	}

	// ---- Printing

	println!("Random settings:\n{:?}", settings_random);
	println!("Calibrated settings:\n{:?}", optimizer.model.get_settings());
	println!("True settings:\n{:?}", settings_true);

	// ---- Drawing

	// Main plots

	charts::draw_chart("sir_true", None, settings_true.pop, &xlist_true, dt);
	charts::draw_chart("sir_random", None, settings_random.pop, &xlist_random, dt);
	charts::draw_chart(
		"sir_calibrated",
		None,
		optimizer.model.get_settings().pop,
		&xlist,
		dt,
	);
	charts::draw_error_chart2("error", None, &error, &error_sto);

	charts::plot_objective(
		"obj_partial",
		None,
		optimizer.clone(),
		&xlist_true[..nsamples_partial],
		Some(&path.iter().map(|v| (v[0], v[1])).collect::<Vec<_>>()),
		Some(&path_sto.iter().map(|v| (v[0], v[1])).collect::<Vec<_>>()),
	);
	charts::plot_objective(
		"obj",
		None,
		optimizer.clone(),
		&xlist_true,
		Some(&path.iter().map(|v| (v[0], v[1])).collect::<Vec<_>>()),
		Some(&path_sto.iter().map(|v| (v[0], v[1])).collect::<Vec<_>>()),
	);

	// Implicit/explicit Euler comparison

	{
		let dur = 40f64;
		let settings = model::sir::SirSettings {
			beta: 0.999,
			gamma: 0.5,
			pop: 1.0,
		};
		let model = model::sir::Sir {
			s: settings.clone(),
		};
		let solver_explicit = ExplicitEulerSolver { dt: 1.0 };
		let solver_implicit = ImplicitEulerSolver {
			dt: 1.0,
			tol: 0.000001,
			niters: 100,
		};
		let solver_true = ImplicitEulerSolver {
			dt: 0.001,
			tol: 0.000001,
			niters: 100,
		};
		let nsamples_explicit = (dur / solver_explicit.dt) as usize;
		let nsamples_implicit = (dur / solver_implicit.dt) as usize;
		let nsamples_true = (dur / solver_true.dt) as usize;
		let mut xlist_explicit = Vec::with_capacity(nsamples_explicit);
		xlist_explicit.push(x0);
		let mut xlist_implicit = Vec::with_capacity(nsamples_implicit);
		xlist_implicit.push(x0);
		let mut xlist_true = Vec::with_capacity(nsamples_true);
		xlist_true.push(x0);

		let mut x = x0;
		for _ in 1..nsamples_explicit {
			x = solver_explicit.f(&model, x);
			xlist_explicit.push(x);
		}

		x = x0;
		for _ in 1..nsamples_implicit {
			x = solver_implicit.f(&model, x);
			xlist_implicit.push(x);
		}

		x = x0;
		for _ in 1..nsamples_true {
			x = solver_true.f(&model, x);
			xlist_true.push(x);
		}

		charts::draw_comparison_chart(
			"comp_euler",
			None,
			&settings,
			&xlist_explicit,
			&xlist_implicit,
			&xlist_true,
			solver_explicit.dt,
			solver_implicit.dt,
			solver_true.dt,
		);
	}

	// SIRV charts

	{
		let nsamples = 1000;
		let settings = model::sirv::SirvSettings {
			beta: 0.8,
			gamma: 0.2,
			lambda: 0.025,
			mu: 0.02,
			pop: 1.0,
		};
		let model = model::sirv::Sirv { s: settings };

		let mut xlist = Vec::with_capacity(nsamples);
		xlist.push(x0);

		let mut x = x0;
		for _ in 1..nsamples {
			x = optimizer.solver.f(&model, x);
			xlist.push(x);
		}

		charts::draw_chart("sirv", None, model.get_settings().pop, &xlist, dt);
	}
}

#[allow(dead_code)]
fn bike() {
	let mut rng = rand::thread_rng();

	// ---- Initialization

	let x0 = vector![0.0, 60. / 3.6];
	let dt = 0.1;
	let nsamples: usize = 2000;

	// ---- Data generation

	let settings_true = model::bike::BikeSettings::<f64> {
		cx: 0.25,
		g: 9.81,
		m: 70.0,
		th: 0.11,
	};
	println!(
		"true: A={} ; B={}",
		-settings_true.cx / settings_true.m,
		settings_true.g * settings_true.th.sin() - 80. / settings_true.m
	);

	let model = model::bike::Bike {
		s: settings_true.clone(),
	};
	/*let solver = ImplicitEulerSolver {
		dt,
		tol: 0.000001,
		niters: 100,
	};*/
	let solver = ExplicitEulerSolver { dt };
	let mut xlist_true = Vec::with_capacity(nsamples);
	xlist_true.push(x0);
	let mut x = x0;
	for _ in 0..nsamples - 1 {
		x = solver.f(&model, x);
		if x[1] < 0. {
			x[1] = 0.;
		}
		xlist_true.push(x);
	}

	// -- Alternative settings

	// Greater theta

	let mut settings_greater_th = settings_true.clone();
	settings_greater_th.th = 0.12;
	println!(
		"gtth: A={} ; B={}",
		-settings_greater_th.cx / settings_greater_th.m,
		settings_greater_th.g * settings_greater_th.th.sin() - 80. / settings_greater_th.m
	);

	let xlist_greater_th = {
		let model = model::bike::Bike {
			s: settings_greater_th.clone(),
		};
		let mut xlist = Vec::with_capacity(nsamples);
		xlist.push(x0);
		let mut x = x0;
		for _ in 0..nsamples - 1 {
			x = solver.f(&model, x);
			if x[1] < 0. {
				x[1] = 0.;
			}
			xlist.push(x);
		}
		xlist
	};

	// Optimal braking

	let mut settings_opti1 = model::bike2::BikeSettings {
		cx: settings_true.cx,
		g: settings_true.g,
		m: settings_true.m,
		th: std::f64::consts::PI / 180. * 15.,
		b: |x, v, s| {
			let mu = 0.1;
			let gx = 0.65;
			let gy = 1.05;
			let magic = 1.0;
			(
				(
					s.m * s.g * (mu * s.th.cos() + s.th.sin()) - s.cx * v * v,
					0.,
				),
				(-2. * s.cx * v, 0.),
			)
		},
	};

	let (xlist_opti1, blist_opti1) = {
		let model = model::bike2::Bike {
			s: settings_opti1.clone(),
		};
		let mut xlist = Vec::with_capacity(nsamples);
		xlist.push(x0);
		let mut blist = Vec::with_capacity(nsamples);
		let mut x = x0;
		for _ in 0..nsamples - 1 {
			blist.push((settings_opti1.b)(x[0], x[1], &settings_opti1).0 .0);
			x = solver.f(&model, x);
			if x[1] < 0. {
				x[1] = 0.;
				break;
			}
			xlist.push(x);
		}
		(xlist, blist)
	};

	let mut settings_opti2 = model::bike2::BikeSettings {
		cx: settings_true.cx,
		g: settings_true.g,
		m: settings_true.m,
		th: std::f64::consts::PI / 180. * 15.,
		b: |x, v, s| {
			let mu = 0.1;
			let gx = 0.65;
			let gy = 1.05;
			let magic = 1.0;
			(
				(
					s.m * s.g * (mu * s.th.cos() + s.th.sin()) - s.cx * v * v,
					s.m * s.g * s.th.cos() * (mu + gx / gy),
				),
				(-2. * s.cx * v, 0.),
			)
		},
	};

	let xlist_opti2 = {
		let model = model::bike2::Bike {
			s: settings_opti2.clone(),
		};
		let mut xlist = Vec::with_capacity(nsamples);
		xlist.push(x0);
		let mut x = x0;
		for _ in 0..nsamples - 1 {
			x = solver.f(&model, x);
			if x[1] < 0. {
				x[1] = 0.;
				break;
			}
			xlist.push(x);
		}
		xlist
	};

	// -- ODE solution

	let xlist_ode = {
		let settings = settings_true.clone();
		let a = -settings.cx / settings.m;
		let b = settings.g * settings.th.sin() - 80. / settings.m;
		/*let r = (-a*b).sqrt();
		assert!(!r.is_nan());
		let alpha = (1.+x0[1]*a/r)/(1.-x0[1]*a/r);*/
		let r = (a * b).sqrt();
		assert!(!r.is_nan());
		let alpha = -r / (a * x0[1]);
		println!("alpha: {alpha}");
		let stop = (1. / alpha).atan() / r;
		println!("Stop: {stop}");

		let bmax =
			settings.m * settings.g * (settings.th.cos() + settings.th.sin()) - settings.cx * 27.;
		println!("bmax: {bmax}");

		let mut xlist = Vec::with_capacity(nsamples);
		xlist.push(x0);
		let mut x = x0;
		for t in 0..nsamples - 1 {
			let t = t as f64 * dt;
			//dbg!((beta*c*(-t*c).exp()-alpha*c*(t*c).exp()));
			//dbg!((alpha*(t*c).exp()+beta*(-t*c).exp()));
			//let v = (beta*c*(-t*c).exp()-alpha*c*(t*c).exp())/a/(alpha*(t*c).exp()+beta*(-t*c).exp());
			//let v = r/a*(alpha*(-t*r).exp()-(t*r).exp())/(alpha*(-t*r).exp()+(t*r).exp());
			let mut v = r / a * (alpha * (t * r).sin() - (t * r).cos())
				/ (alpha * (t * r).cos() + (t * r).sin());
			if v.is_nan() {
				panic!("NaN");
			}
			v = v.max(0.).min(100.);
			x = vector![0., v];
			xlist.push(x);
			if t > stop {
				break;
			}
		}
		xlist
	};

	// ---- Drawing

	// Main plots

	charts::draw_bike_chart(
		"bike_x",
		None,
		&[(
			"x (m)",
			&xlist_true.iter().map(|x| x[0]).collect::<Vec<_>>(),
		)],
		dt,
	);
	charts::draw_bike_chart(
		"bike_v",
		None,
		&[(
			"v (m/s)",
			&xlist_true.iter().map(|x| x[1]).collect::<Vec<_>>(),
		)],
		dt,
	);
	charts::draw_bike_chart2(
		"bike_opti1",
		None,
		&[(
			"v (m/s)",
			&xlist_opti1.iter().map(|x| x[1]).collect::<Vec<_>>(),
		)],
		&[("b", &blist_opti1)],
		dt,
	);
	charts::draw_bike_chart(
		"bike_opti2",
		None,
		&[
			(
				"v (m/s) (arrière)",
				&xlist_opti1.iter().map(|x| x[1]).collect::<Vec<_>>(),
			),
			(
				"v (m/s) (arrière+avant)",
				&xlist_opti2.iter().map(|x| x[1]).collect::<Vec<_>>(),
			),
		],
		dt,
	);
	charts::draw_bike_chart(
		"bike_var_theta",
		None,
		&[
			(
				&format!("v (θ={}, B<0)", settings_true.th),
				&xlist_true.iter().map(|x| x[1]).collect::<Vec<_>>(),
			),
			(
				&format!("v (θ={}, B>0)", settings_greater_th.th),
				&xlist_greater_th.iter().map(|x| x[1]).collect::<Vec<_>>(),
			),
		],
		dt,
	);
	charts::draw_bike_chart(
		"bike_ode_v",
		None,
		&[(
			"v (m/s)",
			&xlist_ode.iter().map(|x| x[1]).collect::<Vec<_>>(),
		)],
		dt,
	);
}