use crate::{
	model::{self, Model},
	opti::GradientDescentOptimizer,
	solver::Solver,
	utils::*,
};

use plotters::prelude::*;
use rayon::prelude::*;

const CHART_SIZE: (u32, u32) = (1000, 800); //(500, 400);
const CHART_SIZE_OBJ: (u32, u32) = (960, 960); //(480, 480);

pub fn draw_chart(filename: &str, title: Option<&str>, pop: f64, xlist: &[Vect<f64, 2>], dt: f64) {
	let filepath = format!("target/{}.png", filename);
	let root = BitMapBackend::new(&filepath, CHART_SIZE).into_drawing_area();
	root.fill(&WHITE).unwrap();
	let mut chart = ChartBuilder::on(&root);
	if let Some(title) = title {
		chart.caption(
			title,
			FontDesc::new(FontFamily::Name("cantarell"), 28.0, FontStyle::Normal),
		);
	}
	let mut chart = chart
		.margin_right(12)
		.y_label_area_size(30)
		.x_label_area_size(30)
		.build_cartesian_2d(0.0f64..xlist.len() as f64 * dt, 0.0f64..1.)
		.unwrap();

	chart.configure_mesh().x_desc("Time").draw().unwrap();

	chart
		.draw_series(LineSeries::new(
			xlist.iter().enumerate().map(|(i, x)| (i as f64 * dt, x[0])),
			BLUE,
		))
		.unwrap()
		.label("Susceptible")
		.legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], BLUE));

	chart
		.draw_series(LineSeries::new(
			xlist.iter().enumerate().map(|(i, x)| (i as f64 * dt, x[1])),
			RED,
		))
		.unwrap()
		.label("Infected")
		.legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], RED));

	chart
		.draw_series(LineSeries::new(
			xlist
				.iter()
				.enumerate()
				.map(|(i, x)| (i as f64 * dt, pop - x[0] - x[1])),
			GREEN,
		))
		.unwrap()
		.label("Removed")
		.legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], GREEN));

	chart
		.configure_series_labels()
		.background_style(WHITE.mix(0.8))
		.border_style(BLACK)
		.draw()
		.unwrap();
}

pub fn draw_error_chart(filename: &str, title: Option<&str>, xlist: &[f64]) {
	let filepath = format!("target/{}.png", filename);
	let root = BitMapBackend::new(&filepath, CHART_SIZE).into_drawing_area();
	root.fill(&WHITE).unwrap();
	let mut chart = ChartBuilder::on(&root);
	if let Some(title) = title {
		chart.caption(
			title,
			FontDesc::new(FontFamily::Name("cantarell"), 28.0, FontStyle::Normal),
		);
	}
	let mut chart = chart
		.margin_right(12)
		.y_label_area_size(50)
		.x_label_area_size(30)
		.build_cartesian_2d(0..xlist.len(), (0.0f64..max(xlist)).log_scale())
		.unwrap();

	let printer = plotters::data::float::FloatPrettyPrinter {
		allow_scientific: true,
		min_decimal: 0,
		max_decimal: 2,
	};
	chart
		.configure_mesh()
		.x_desc("Iterations")
		.y_desc("Mean error")
		.y_label_formatter(&|y| printer.print(*y))
		.draw()
		.unwrap();

	chart
		.draw_series(LineSeries::new(xlist.iter().copied().enumerate(), BLACK))
		.unwrap();
}

pub fn draw_error_chart2(
	filename: &str,
	title: Option<&str>,
	xlist_batch: &[f64],
	xlist_sto: &[f64],
) {
	let filepath = format!("target/{}.png", filename);
	let root = BitMapBackend::new(&filepath, CHART_SIZE).into_drawing_area();
	root.fill(&WHITE).unwrap();
	let mut chart = ChartBuilder::on(&root);
	if let Some(title) = title {
		chart.caption(
			title,
			FontDesc::new(FontFamily::Name("cantarell"), 28.0, FontStyle::Normal),
		);
	}
	let mut chart = chart
		.margin_right(12)
		.y_label_area_size(50)
		.x_label_area_size(30)
		.build_cartesian_2d(
			0..xlist_batch.len().max(xlist_sto.len()),
			(0.0f64..max(xlist_batch).max(max(xlist_sto))).log_scale(),
		)
		.unwrap();

	let printer = plotters::data::float::FloatPrettyPrinter {
		allow_scientific: true,
		min_decimal: 0,
		max_decimal: 2,
	};
	chart
		.configure_mesh()
		.x_desc("Iterations")
		.y_desc("Mean error")
		.y_label_formatter(&|y| printer.print(*y))
		.draw()
		.unwrap();

	chart
		.draw_series(LineSeries::new(
			xlist_batch.iter().copied().enumerate(),
			BLACK,
		))
		.unwrap()
		.label("Batch")
		.legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], BLACK));

	chart
		.draw_series(LineSeries::new(xlist_sto.iter().copied().enumerate(), RED))
		.unwrap()
		.label("Stochastic")
		.legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], RED));

	chart
		.configure_series_labels()
		.background_style(WHITE.mix(0.8))
		.border_style(BLACK)
		.position(SeriesLabelPosition::MiddleLeft)
		.draw()
		.unwrap();
}

pub fn plot_objective<
	R: Solver<f64, model::sir::SirSettings<f64>, model::sir::Sir<f64>, 2> + Clone + Sync,
>(
	filename: &str,
	title: Option<&str>,
	optimizer: GradientDescentOptimizer<
		f64,
		model::sir::SirSettings<f64>,
		model::sir::Sir<f64>,
		R,
		2,
		3,
	>,
	ylist_true: &[Vect<f64, 2>],
	path_batch: Option<&[(f64, f64)]>,
	path_sto: Option<&[(f64, f64)]>,
) {
	let filepath = format!("target/{}.png", filename);
	let root = BitMapBackend::new(&filepath, CHART_SIZE_OBJ).into_drawing_area();
	root.fill(&WHITE).unwrap();
	let mut chart = ChartBuilder::on(&root);
	if let Some(title) = title {
		chart.caption(
			title,
			FontDesc::new(FontFamily::Name("cantarell"), 28.0, FontStyle::Normal),
		);
	}
	let mut chart = chart
		.margin_right(12)
		.x_label_area_size(30)
		.y_label_area_size(40)
		.build_cartesian_2d(0.0f64..1., 0.0f64..1.)
		.unwrap();

	chart
		.configure_mesh()
		.x_desc("beta")
		.y_desc("gamma")
		.draw()
		.unwrap();

	let area = chart.plotting_area();

	let range = area.get_pixel_range();
	let (pw, ph) = (range.0.end - range.0.start, range.1.end - range.1.start);
	let (xr, yr) = (chart.x_range(), chart.y_range());
	let step = (
		(xr.end - xr.start) / pw as f64,
		(yr.end - yr.start) / ph as f64,
	);

	let mut min = f64::MAX;
	let mut max = f64::MIN;
	let vals: Vec<(f64, f64, f64)> = (0..pw * ph)
		.into_par_iter()
		.map(|i| {
			let (x, y) = (
				xr.start + step.0 * (i % pw) as f64,
				yr.start + step.1 * (i / pw) as f64,
			);
			let mut optimizer = optimizer.clone();
			let s = optimizer.model.get_settings_mut();
			s.beta = x;
			s.gamma = y;
			let val = optimizer.objective_batch(&optimizer.model, ylist_true);
			(x, y, val)
		})
		.collect();
	vals.iter().for_each(|(_, _, val)| {
		if *val > max {
			max = *val;
		}
		if *val < min {
			min = *val;
		}
	});
	let ampl = 0.825 / (max - min);

	for (x, y, c) in vals {
		area.draw_pixel((x, y), &HSLColor((c - min) * ampl, 1.0, 0.5))
			.unwrap();
	}

	if let Some(path_sto) = path_sto {
		chart
			.draw_series(std::iter::once(PathElement::new(
				path_sto,
				RGBColor(128, 128, 128),
			)))
			.unwrap()
			.label("Stochastic")
			.legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], RGBColor(128, 128, 128)));
	}

	if let Some(path_batch) = path_batch {
		chart
			.draw_series(std::iter::once(PathElement::new(path_batch, BLACK)))
			.unwrap()
			.label("Batch")
			.legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], BLACK));
	}

	chart
		.configure_series_labels()
		.background_style(WHITE.mix(0.8))
		.border_style(BLACK)
		.position(SeriesLabelPosition::UpperRight)
		.draw()
		.unwrap();
}

pub fn draw_comparison_chart(
	filename: &str,
	title: Option<&str>,
	s: &model::sir::SirSettings<f64>,
	xlist_explicit: &[Vect<f64, 2>],
	xlist_implicit: &[Vect<f64, 2>],
	xlist_true: &[Vect<f64, 2>],
	dt_explicit: f64,
	dt_implicit: f64,
	dt_true: f64,
) {
	let filepath = format!("target/{}.png", filename);
	let root = BitMapBackend::new(&filepath, CHART_SIZE).into_drawing_area();
	root.fill(&WHITE).unwrap();
	let mut chart = ChartBuilder::on(&root);
	if let Some(title) = title {
		chart.caption(
			title,
			FontDesc::new(FontFamily::Name("cantarell"), 28.0, FontStyle::Normal),
		);
	}
	let mut chart = chart
		.margin_right(12)
		.y_label_area_size(30)
		.x_label_area_size(30)
		.build_cartesian_2d(0.0f64..xlist_true.len() as f64 * dt_true, 0.0f64..1.)
		.unwrap();

	chart.configure_mesh().x_desc("Time").draw().unwrap();

	chart
		.draw_series(LineSeries::new(
			xlist_explicit
				.iter()
				.enumerate()
				.map(|(i, x)| (i as f64 * dt_explicit, x[0])),
			ShapeStyle::from(BLUE).stroke_width(3),
		))
		.unwrap()
		.label("Susceptible (explicit)")
		.legend(|(x, y)| {
			PathElement::new(
				vec![(x, y), (x + 20, y)],
				ShapeStyle::from(BLUE).stroke_width(3),
			)
		});

	chart
		.draw_series(LineSeries::new(
			xlist_explicit
				.iter()
				.enumerate()
				.map(|(i, x)| (i as f64 * dt_explicit, x[1])),
			ShapeStyle::from(RED).stroke_width(3),
		))
		.unwrap()
		.label("Infected (explicit)")
		.legend(|(x, y)| {
			PathElement::new(
				vec![(x, y), (x + 20, y)],
				ShapeStyle::from(RED).stroke_width(3),
			)
		});

	chart
		.draw_series(LineSeries::new(
			xlist_explicit
				.iter()
				.enumerate()
				.map(|(i, x)| (i as f64 * dt_explicit, s.pop - x[0] - x[1])),
			ShapeStyle::from(GREEN).stroke_width(3),
		))
		.unwrap()
		.label("Removed (explicit)")
		.legend(|(x, y)| {
			PathElement::new(
				vec![(x, y), (x + 20, y)],
				ShapeStyle::from(GREEN).stroke_width(3),
			)
		});

	chart
		.draw_series(LineSeries::new(
			xlist_implicit
				.iter()
				.enumerate()
				.map(|(i, x)| (i as f64 * dt_implicit, x[0])),
			BLUE,
		))
		.unwrap()
		.label("Susceptible (implicit)")
		.legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], BLUE));

	chart
		.draw_series(LineSeries::new(
			xlist_implicit
				.iter()
				.enumerate()
				.map(|(i, x)| (i as f64 * dt_implicit, x[1])),
			RED,
		))
		.unwrap()
		.label("Infected (implicit)")
		.legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], RED));

	chart
		.draw_series(LineSeries::new(
			xlist_implicit
				.iter()
				.enumerate()
				.map(|(i, x)| (i as f64 * dt_implicit, s.pop - x[0] - x[1])),
			GREEN,
		))
		.unwrap()
		.label("Removed (implicit)")
		.legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], GREEN));

	chart
		.draw_series(LineSeries::new(
			xlist_true
				.iter()
				.enumerate()
				.map(|(i, x)| (i as f64 * dt_true, x[0])),
			RGBColor(0, 0, 128),
		))
		.unwrap()
		.label("Susceptible (true)")
		.legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], RGBColor(0, 0, 128)));

	chart
		.draw_series(LineSeries::new(
			xlist_true
				.iter()
				.enumerate()
				.map(|(i, x)| (i as f64 * dt_true, x[1])),
			RGBColor(128, 0, 0),
		))
		.unwrap()
		.label("Infected (true)")
		.legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], RGBColor(128, 0, 0)));

	chart
		.draw_series(LineSeries::new(
			xlist_true
				.iter()
				.enumerate()
				.map(|(i, x)| (i as f64 * dt_true, s.pop - x[0] - x[1])),
			RGBColor(0, 128, 0),
		))
		.unwrap()
		.label("Removed (true)")
		.legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], RGBColor(0, 128, 0)));

	chart
		.configure_series_labels()
		.background_style(WHITE.mix(0.8))
		.border_style(BLACK)
		.draw()
		.unwrap();
}

pub fn draw_bike_chart(filename: &str, title: Option<&str>, xlists: &[(&str, &[f64])], dt: f64) {
	let max_x = xlists
		.iter()
		.map(|(_, xlist)| xlist.len())
		.max()
		.expect("at least one series expected");
	let max_y = *xlists
		.iter()
		.map(|(_, xlist)| {
			xlist
				.iter()
				.max_by(|a, b| a.total_cmp(b))
				.expect("at least one sample per series expected")
		})
		.max_by(|a, b| a.total_cmp(b))
		.unwrap();
	let filepath = format!("target/{}.png", filename);
	let root = BitMapBackend::new(&filepath, (640, 480)).into_drawing_area();
	root.fill(&WHITE).unwrap();
	let mut chart = ChartBuilder::on(&root);
	if let Some(title) = title {
		chart.caption(
			title,
			FontDesc::new(FontFamily::Name("cantarell"), 28.0, FontStyle::Normal),
		);
	}
	let mut chart = chart
		.margin_right(12)
		.margin_top(12)
		.y_label_area_size(30)
		.x_label_area_size(30)
		.build_cartesian_2d(0.0f64..max_x as f64 * dt, 0.0f64..max_y)
		.unwrap();

	chart.configure_mesh().x_desc("Temps (s)").draw().unwrap();

	for (list_i, (label, xlist)) in xlists.into_iter().enumerate() {
		chart
			.draw_series(LineSeries::new(
				xlist.iter().enumerate().map(|(i, x)| (i as f64 * dt, *x)),
				Palette100::pick(list_i + 1).stroke_width(2),
			))
			.unwrap()
			.label(*label)
			.legend(move |(x, y)| {
				PathElement::new(
					vec![(x, y), (x + 20, y)],
					Palette100::pick(list_i + 1).stroke_width(2),
				)
			});
	}

	chart
		.configure_series_labels()
		.border_style(BLACK)
		.background_style(WHITE.mix(0.8))
		.label_font(("Libertinus Serif", 20))
		.draw()
		.unwrap();
}

pub fn draw_bike_chart2(
	filename: &str,
	title: Option<&str>,
	xlists1: &[(&str, &[f64])],
	xlists2: &[(&str, &[f64])],
	dt: f64,
) {
	let max_x = xlists1
		.iter()
		.chain(xlists2.iter())
		.map(|(_, xlist)| xlist.len())
		.max()
		.expect("at least one series expected");
	let max_y1 = *xlists1
		.iter()
		.map(|(_, xlist)| {
			xlist
				.iter()
				.max_by(|a, b| a.total_cmp(b))
				.expect("at least one sample per series expected")
		})
		.max_by(|a, b| a.total_cmp(b))
		.unwrap();
	let max_y2 = *xlists2
		.iter()
		.map(|(_, xlist)| {
			xlist
				.iter()
				.max_by(|a, b| a.total_cmp(b))
				.expect("at least one sample per series expected")
		})
		.max_by(|a, b| a.total_cmp(b))
		.unwrap();
	let filepath = format!("target/{}.png", filename);
	let root = BitMapBackend::new(&filepath, (640, 480)).into_drawing_area();
	root.fill(&WHITE).unwrap();
	let mut chart = ChartBuilder::on(&root);
	if let Some(title) = title {
		chart.caption(
			title,
			FontDesc::new(FontFamily::Name("cantarell"), 28.0, FontStyle::Normal),
		);
	}
	let mut chart = chart
		.margin_right(12)
		.margin_top(12)
		.y_label_area_size(50)
		.x_label_area_size(30)
		.right_y_label_area_size(50)
		.build_cartesian_2d(0.0f64..max_x as f64 * dt, 0.0f64..max_y1)
		.unwrap()
		.set_secondary_coord(0.0f64..max_x as f64 * dt, 0.0f64..max_y2);

	chart
		.configure_mesh()
		.x_desc("Temps (s)")
		.y_desc("Vitesse (m/s)")
		.draw()
		.unwrap();
	chart
		.configure_secondary_axes()
		.y_desc("Freinage")
		.draw()
		.unwrap();

	for (list_i, (label, xlist)) in xlists1.into_iter().enumerate() {
		chart
			.draw_series(LineSeries::new(
				xlist.iter().enumerate().map(|(i, x)| (i as f64 * dt, *x)),
				Palette100::pick(list_i + 1).stroke_width(2),
			))
			.unwrap()
			.label(*label)
			.legend(move |(x, y)| {
				PathElement::new(
					vec![(x, y), (x + 20, y)],
					Palette100::pick(list_i + 1).stroke_width(2),
				)
			});
	}

	for (list_i, (label, xlist)) in (xlists1.len()..).zip(xlists2.into_iter()) {
		chart
			.draw_secondary_series(LineSeries::new(
				xlist.iter().enumerate().map(|(i, x)| (i as f64 * dt, *x)),
				Palette100::pick(list_i + 1).stroke_width(2),
			))
			.unwrap()
			.label(*label)
			.legend(move |(x, y)| {
				PathElement::new(
					vec![(x, y), (x + 20, y)],
					Palette100::pick(list_i + 1).stroke_width(2),
				)
			});
	}

	chart
		.configure_series_labels()
		.border_style(BLACK)
		.background_style(WHITE.mix(0.8))
		.label_font(("Libertinus Serif", 20))
		.draw()
		.unwrap();
}