use glam::{f32, Mat2, Vec2, vec2};

use crate::fns::{self, Limiter};
use crate::riemann::{Decomp2, Metric, Tens2};

#[derive(Copy, Clone, Debug)]
pub struct Tube {
	pub outer_radius: f32,
	pub inner_radius: f32,
	pub external_halflength: f32,
	pub internal_halflength: f32,
}

impl Tube {
	fn fx(&self) -> impl Limiter { fns::SmootherstepLimiter { min: self.inner_radius, max: self.outer_radius } }
	fn fy(&self) -> fns::QuadraticAccelerator { fns::QuadraticAccelerator { internal: self.internal_halflength, external: self.external_halflength } }

	pub fn y(&self, v: f32) -> f32 { self.fy().x(v) }
	pub fn v(&self, y: f32) -> f32 { self.fy().u(y) }
	pub fn dy(&self, v: f32) -> f32 { self.fy().dx(v) }
	pub fn dv(&self, y: f32) -> f32 { self.fy().du(y) }
}

impl Metric for Tube {
	fn sqrt_at(&self, pos: Vec2) -> Decomp2 {
		let sx = self.fx().value(pos.x);
		let sy = self.fy().du(pos.y);
		let s = sx + sy - sx * sy;
		assert!(sx.is_finite());
		assert!(sy.is_finite());
		assert!(sy > 0.0);
		Decomp2 {
			ortho: Mat2::IDENTITY,
			diag: vec2(1.0, s),
		}
	}

	fn part_derivs_at(&self, pos: Vec2) -> Tens2 {
		let sx = self.fx().value(pos.x);
		let sy = self.fy().du(pos.y);
		let s = sx + sy - sx * sy;
		let dsx_dx = self.fx().derivative(pos.x);
		let dsy_dy = self.fy().d2u(pos.y);
		let ds2_dx = 2.0 * s * (1.0 - sy) * dsx_dx;
		let ds2_dy = 2.0 * s * (1.0 - sx) * dsy_dy;
		[
			Mat2::from_cols_array(&[0.0, 0.0, 0.0, ds2_dx]),
			Mat2::from_cols_array(&[0.0, 0.0, 0.0, ds2_dy]),
		]
	}
}

#[cfg(test)]
mod test {
	use approx::assert_abs_diff_eq;
	use glam::{Vec2, vec2};
	use itertools_num::linspace;

	use crate::riemann::{Decomp2, Metric};
	use crate::tube::Space;
	use crate::types::Ray;

	use super::Tube;

	#[test]
	fn test_tube_metric_derivs() {
		struct Approx(Tube);
		impl Metric for Approx {
			fn sqrt_at(&self, pos: Vec2) -> Decomp2 { self.0.sqrt_at(pos) }
		}
		let testee = Tube {
			inner_radius: 30.0,
			outer_radius: 50.0,
			internal_halflength: 100.0,
			external_halflength: 300.0,
		};
		let approx = Approx(testee);
		let epsilon = 1.0e-3;
		let margin = 1.0 / 16.0;
		let mul = 1.0 + margin;
		for x in itertools_num::linspace(-mul * testee.outer_radius, mul * testee.outer_radius, 100) {
			for y in itertools_num::linspace(-mul * testee.external_halflength, mul * testee.external_halflength, 100) {
				let pos = vec2(x, y);
				let computed = testee.part_derivs_at(pos);
				let reference = approx.part_derivs_at(pos);
				let eq = (0..2).all(|coord| computed[coord].abs_diff_eq(reference[coord], epsilon));
				assert!(eq, "Bad derivative computation at {pos}:\n  explicit:  {computed:?}\n  numerical: {reference:?}\n");
			}
		}
	}

	#[test]
	fn test_accelerator() {
		let space = Space {
			tube: Tube {
				inner_radius: 30.0,
				outer_radius: 50.0,
				internal_halflength: 100.0,
				external_halflength: 300.0,
			},
			objs: vec![],
		};
		let ε = 1e-3;
		let off = 10.0;
		let steps = 1024;
		for ax in [-30.0 + ε, -25.0, -3.0, 17.0, 30.0 - ε] {
			for bx in [0.0, ε, 1.0, 7.0, 30.0 - ε] {
				let a = vec2(ax, -(space.tube.external_halflength + off));
				let b = vec2(bx, space.tube.external_halflength + off);
				let Δ = vec2(bx - ax, 2.0 * (space.tube.internal_halflength + off));
				let dir = Δ / (steps as f32);
				let traced = space.trace_iter(Ray { pos: a, dir }).nth(steps).unwrap();
				assert_abs_diff_eq!(traced.pos.x, b.x, epsilon=1.0e-2);
				assert_abs_diff_eq!(traced.pos.y, b.y, epsilon=1.0e1);
				assert_abs_diff_eq!(traced.dir.x, dir.x, epsilon=1.0e-3);
				assert_abs_diff_eq!(traced.dir.y, dir.y, epsilon=1.0e-2);
			}
		}
	}

	#[test]
	#[ignore]
	fn test_accelerator_slow() {
		let space = Space {
			tube: Tube {
				inner_radius: 30.0,
				outer_radius: 50.0,
				internal_halflength: 100.0,
				external_halflength: 300.0,
			},
			objs: vec![],
		};
		let ε = 1e-3;
		let off = 10.0;
		let steps = 4096;
		for ax in linspace(-space.tube.inner_radius + ε, space.tube.inner_radius - ε, 20) {
			for bx in linspace(-space.tube.inner_radius + ε, space.tube.inner_radius - ε, 20) {
				let a = vec2(ax, -(space.tube.external_halflength + off));
				let b = vec2(bx, space.tube.external_halflength + off);
				let Δ = vec2(bx - ax, 2.0 * (space.tube.internal_halflength + off));
				let dir = Δ / (steps as f32);
				let traced = space.trace_iter(Ray { pos: a, dir }).nth(steps).unwrap();
				assert_abs_diff_eq!(traced.pos.x, b.x, epsilon=1.0e-2);
				assert_abs_diff_eq!(traced.pos.y, b.y, epsilon=1.0e0);
				assert_abs_diff_eq!(traced.dir.x, dir.x, epsilon=1.0e-3);
				assert_abs_diff_eq!(traced.dir.y, dir.y, epsilon=1.0e-3);
			}
		}
	}

	#[test]
	fn test_accelerator_inner_edge() {
		let space = Space {
			tube: Tube {
				inner_radius: 30.0,
				outer_radius: 50.0,
				internal_halflength: 100.0,
				external_halflength: 300.0,
			},
			objs: vec![],
		};
		let ε = 1e-3;
		let off = 10.0;
		let steps = 10000;
		for x in [space.tube.inner_radius - ε, space.tube.inner_radius + ε] {
			let a = vec2(x, -(space.tube.external_halflength + off));
			let b = vec2(x, space.tube.external_halflength + off);
			let Δ = vec2(0.0, 2.0 * (space.tube.internal_halflength + off));
			let dir = Δ / (steps as f32);
			let traced = space.trace_iter(Ray { pos: a, dir }).nth(steps).unwrap();
			assert_abs_diff_eq!(traced.pos.x, b.x, epsilon=1.0e-1);
			assert_abs_diff_eq!(traced.pos.y, b.y, epsilon=1.0e0);
			assert_abs_diff_eq!(traced.dir.x, dir.x, epsilon=1.0e-2);
			assert_abs_diff_eq!(traced.dir.y, dir.y, epsilon=1.0e-2);
		}
	}

	#[test]
	fn test_accelerator_outer_edge() {
		let space = Space {
			tube: Tube {
				inner_radius: 30.0,
				outer_radius: 50.0,
				internal_halflength: 100.0,
				external_halflength: 300.0,
			},
			objs: vec![],
		};
		let ε = 1e-3;
		let off = 10.0;
		let steps = 4096;
		for x in [space.tube.outer_radius + ε, space.tube.outer_radius - ε] {
			let a = vec2(x, -(space.tube.external_halflength + off));
			let b = vec2(x, space.tube.external_halflength + off);
			let Δ = vec2(0.0, 2.0 * (space.tube.external_halflength + off));
			let dir = Δ / (steps as f32);
			let traced = space.trace_iter(Ray { pos: a, dir }).nth(steps).unwrap();
			assert_abs_diff_eq!(traced.pos.x, b.x, epsilon=2.0e0);
			assert_abs_diff_eq!(traced.pos.y, b.y, epsilon=1.0e0);
			assert_abs_diff_eq!(traced.dir.x, dir.x, epsilon=1.0e-2);
			assert_abs_diff_eq!(traced.dir.y, dir.y, epsilon=1.0e-2);
		}
	}
}