use glam::{bool, f32, Mat2, Vec2, vec2};
use crate::riemann;
use Subspace::{Boundary, Inner, Outer};
use metric::Tube;
use coords::{FlatCoordinateSystem, InnerCS, OuterCS};
use crate::tube::coords::FlatRegion;
use crate::types::{FlatTraceResult, Hit, Location, Object, Ray};

pub mod metric;

pub struct Space {
	pub tube: Tube,
	pub objs: Vec<Object>,
}

#[derive(PartialEq, Eq, Debug)]
pub enum Subspace {
	Outer,
	Boundary,
	Inner,
}

impl Space {
	pub fn which_subspace(&self, pt: Vec2) -> Subspace {
		if pt.y.abs() > self.tube.external_halflength {
			Outer
		} else if pt.x.abs() > self.tube.outer_radius {
			Outer
		} else if pt.x.abs() > self.tube.inner_radius {
			Boundary
		} else {
			Inner
		}
	}

	/// Выполняет один шаг трассировки. Работает в любой части пространства, но вне Boundary доступны более эффективные методы.
	/// ray задаётся в основной СК.
	pub fn trace_step(&self, ray: Ray) -> Ray {
		let a: Vec2 = -riemann::contract2(riemann::krist(&self.tube, ray.pos), ray.dir);
		let v = ray.dir + a;
		let p = ray.pos + v;
		Ray { pos: p, dir: v }
	}

	/// Выполняет один шаг перемещения. Работает в любой части пространства.
	/// off задаётся в локальной СК. Рекомендуется считать небольшими шагами.
	pub fn move_step(&self, loc: Location, off: Vec2) -> Location {
		let corr = Mat2::IDENTITY - riemann::contract(riemann::krist(&self.tube, loc.pos), loc.rot * off);
		let p = loc.pos + corr * loc.rot * off;
		Location { pos: p, rot: corr * loc.rot }
	}

	pub fn trace_iter(&self, ray: Ray) -> impl Iterator<Item=Ray> + '_ {
		std::iter::successors(Some(ray), |&ray| Some(self.trace_step(ray)))
	}

	pub fn trace_inner(&self, ray: Ray) -> FlatTraceResult {
		assert_eq!(self.which_subspace(ray.pos), Inner);
		self.trace_flat(InnerCS(self.tube), ray)
	}

	pub fn trace_outer(&self, ray: Ray) -> FlatTraceResult {
		assert_eq!(self.which_subspace(ray.pos), Outer);
		self.trace_flat(OuterCS(self.tube), ray)
	}

	fn trace_flat(&self, cs: impl FlatRegion, ray: Ray) -> FlatTraceResult {
		let ray = cs.global_to_flat(ray);
		let dist = cs.distance_to_boundary(ray);
		let objs = self.list_objects(|loc| cs.global_to_flat(loc));
		FlatTraceResult {
			end: dist.map(|dist| cs.flat_to_global(ray.forward(dist))),
			objects: Self::hit_objects(objs.as_slice(), ray, dist, |pos| cs.flat_to_global(pos)),
		}
	}

	fn trace_boundary(&self, ray: Ray) -> Ray {
		assert_eq!(self.which_subspace(ray.pos), Boundary);
		self.trace_iter(ray)
			.find(|&ray| self.which_subspace(ray.pos) != Boundary)
			.expect("Can't get outta the wall!")
	}

	fn list_objects(&self, tfm: impl Fn(Location) -> Location) -> Vec<Object> {
		self.objs.iter().map(|&Object { id, loc, r }| Object { id, loc: tfm(loc), r }).collect()
	}

	fn hit_objects(objs: &[Object], ray: Ray, limit: Option<f32>, globalize: impl Fn(Vec2) -> Vec2) -> Vec<Hit> {
		let limit = limit.unwrap_or(f32::INFINITY);
		objs.iter()
			.filter_map(|obj| {
				let rel = ray.pos - obj.loc.pos;
				let diff = rel.dot(ray.dir).powi(2) - ray.dir.length_squared() * (rel.length_squared() - obj.r.powi(2));
				if diff > 0.0 {
					let t = (-rel.dot(ray.dir) - diff.sqrt()) / ray.dir.length_squared();
					Some((obj, t))
				} else {
					None
				}
			})
			.filter(|&(_, t)| t >= 0.0 && t < limit)
			.map(|(obj, t)| {
				let pos = ray.forward(t).pos;
				let rel = obj.loc.rot.inverse() * Ray { pos: pos - obj.loc.pos, dir: ray.dir };
				Hit { id: obj.id, distance: t, pos: globalize(pos), rel }
			})
			.collect()
	}

	pub fn line(&self, a: Vec2, b: Vec2, step: f32) -> Vec<Vec2> {
		match self.which_subspace(a) {
			Outer => vec![b],
			Inner => {
				let cs = InnerCS(self.tube);
				let n = ((b - a).length() / step) as usize + 1;
				let a = cs.global_to_flat(a);
				let b = cs.global_to_flat(b);
				(1..=n).map(|k| cs.flat_to_global(a.lerp(b, k as f32 / n as f32))).collect()
			}
			Boundary => panic!("Can't draw a line here!"),
		}
	}
}

struct Rect {
	pub size: Vec2,
}

impl Rect {
	/// Отражает луч, чтобы все координаты направления были положительны (допустимо благодаря симметрии Rect).
	fn flip_ray(ray: Ray) -> Ray {
		Ray { pos: ray.pos * ray.dir.signum(), dir: ray.dir.abs() }
	}

	fn is_inside(&self, pt: Vec2) -> bool {
		pt.abs().cmplt(self.size).all()
	}

	fn trace_into(&self, ray: Ray) -> Option<f32> {
		let ray = Self::flip_ray(ray);
		// ray.pos.x + t * ray.dir.x = −size.x
		let ts = (-self.size - ray.pos) / ray.dir;
		let t = ts.max_element();
		let pt = ray.pos + t * ray.dir;
		if t < 0.0 { return None; }
		if pt.cmpgt(self.size).any() { return None; }
		Some(t)
	}

	fn trace_out_of(&self, ray: Ray) -> Option<f32> {
		let ray = Self::flip_ray(ray);
		// ray.pos.x + t * ray.dir.x = +size.x
		let ts = (self.size - ray.pos) / ray.dir;
		let t = ts.min_element();
		Some(t)
	}
}

#[test]
fn test_rect() {
	assert_eq!(Rect::flip_ray(Ray { pos: vec2(2.0, 3.0), dir: vec2(4.0, 5.0) }), Ray { pos: vec2(2.0, 3.0), dir: vec2(4.0, 5.0) });
	assert_eq!(Rect::flip_ray(Ray { pos: vec2(2.0, 3.0), dir: vec2(-4.0, 5.0) }), Ray { pos: vec2(-2.0, 3.0), dir: vec2(4.0, 5.0) });
	assert_eq!(Rect::flip_ray(Ray { pos: vec2(2.0, 3.0), dir: vec2(4.0, -5.0) }), Ray { pos: vec2(2.0, -3.0), dir: vec2(4.0, 5.0) });
	assert_eq!(Rect::flip_ray(Ray { pos: vec2(2.0, 3.0), dir: vec2(4.0, 0.0) }), Ray { pos: vec2(2.0, 3.0), dir: vec2(4.0, 0.0) });

	let r = Rect { size: vec2(2.0, 3.0) };

	assert_eq!(r.trace_into(Ray { pos: vec2(3.0, 3.0), dir: vec2(1.0, 1.0) }), None);
	assert_eq!(r.trace_into(Ray { pos: vec2(-3.0, 2.0), dir: vec2(1.0, 0.0) }), Some(1.0));
	assert_eq!(r.trace_into(Ray { pos: vec2(-3.0, 2.0), dir: vec2(-1.0, 0.0) }), None);
	assert_eq!(r.trace_into(Ray { pos: vec2(-3.0, 1.0), dir: vec2(2.0, 2.0) }), Some(0.5));
	assert_eq!(r.trace_into(Ray { pos: vec2(-3.0, 2.1), dir: vec2(2.0, 2.0) }), None);

	assert_eq!(r.trace_into(Ray { pos: vec2(2.0, 3.0), dir: vec2(1.0, 1.0) }), None);
	assert_eq!(r.trace_into(Ray { pos: vec2(-2.0, 3.0), dir: vec2(-1.0, 1.0) }), None);
	assert_eq!(r.trace_into(Ray { pos: vec2(2.0, 3.0), dir: vec2(-1.0, -1.0) }), Some(0.0));
	assert_eq!(r.trace_into(Ray { pos: vec2(2.0, -3.0), dir: vec2(-1.0, 1.0) }), Some(0.0));

	assert_eq!(r.trace_out_of(Ray { pos: vec2(0.0, 0.0), dir: vec2(1.0, 1.0) }), Some(2.0));
	assert_eq!(r.trace_out_of(Ray { pos: vec2(0.0, 0.0), dir: vec2(0.0, 1.0) }), Some(3.0));
	assert_eq!(r.trace_out_of(Ray { pos: vec2(0.0, 1.0), dir: vec2(0.0, -1.0) }), Some(4.0));
	assert_eq!(r.trace_out_of(Ray { pos: vec2(1.0, 1.0), dir: vec2(0.0, -1.0) }), Some(4.0));
	assert_eq!(r.trace_out_of(Ray { pos: vec2(2.0, 3.0), dir: vec2(1.0, 1.0) }), Some(0.0));
}

mod coords {
	use glam::{Mat2, Vec2, vec2};
	use crate::riemann::Metric;
	use crate::types::{Location, Ray};
	use super::{Rect, Tube};

	pub trait FlatCoordinateSystem<T> {
		fn flat_to_global(&self, v: T) -> T;
		fn global_to_flat(&self, v: T) -> T;
	}

	pub trait FlatRegion: FlatCoordinateSystem<Vec2> + FlatCoordinateSystem<Ray> + FlatCoordinateSystem<Location> {
		// Измеряет расстояние до выхода за пределы области вдоль луча ray. Луч задаётся в плоской СК.
		fn distance_to_boundary(&self, _ray: Ray) -> Option<f32> { None }
	}

	trait MetricCS {
		fn global_metric(&self) -> &impl Metric;
	}

	impl<T: FlatCoordinateSystem<Vec2> + MetricCS> FlatCoordinateSystem<Ray> for T {
		fn flat_to_global(&self, ray: Ray) -> Ray {
			let pos = self.flat_to_global(ray.pos);
			Ray {
				pos,
				dir: Mat2::from(self.global_metric().sqrt_at(pos).inverse()) * ray.dir,
			}
		}

		fn global_to_flat(&self, ray: Ray) -> Ray {
			Ray {
				pos: self.global_to_flat(ray.pos),
				dir: Mat2::from(self.global_metric().sqrt_at(ray.pos)) * ray.dir,
			}
		}
	}

	impl<T: FlatCoordinateSystem<Vec2> + MetricCS> FlatCoordinateSystem<Location> for T {
		fn flat_to_global(&self, loc: Location) -> Location {
			let pos = self.flat_to_global(loc.pos);
			Location {
				pos,
				rot: Mat2::from(self.global_metric().sqrt_at(pos).inverse()) * loc.rot,
			}
		}

		fn global_to_flat(&self, loc: Location) -> Location {
			Location {
				pos: self.global_to_flat(loc.pos), // в плоской СК для Inner или её продолжении на Outer
				rot: Mat2::from(self.global_metric().sqrt_at(loc.pos)) * loc.rot,
			}
		}
	}

	pub struct InnerCS(pub Tube);

	impl MetricCS for InnerCS { fn global_metric(&self) -> &impl Metric { &self.0 } }

	impl FlatCoordinateSystem<Vec2> for InnerCS {
		fn flat_to_global(&self, pos: Vec2) -> Vec2 {
			vec2(pos.x, self.0.y(pos.y))
		}

		// Работает только при |pos.x| ≤ inner_radius или |pos.y| ≥ external_halflength.
		fn global_to_flat(&self, pos: Vec2) -> Vec2 {
			vec2(pos.x, self.0.v(pos.y))
		}
	}

	impl FlatRegion for InnerCS {
		fn distance_to_boundary(&self, ray: Ray) -> Option<f32> {
			Rect { size: vec2(self.0.inner_radius, self.0.internal_halflength) }.trace_out_of(ray)
		}
	}

	pub struct OuterCS(pub Tube);

	impl MetricCS for OuterCS { fn global_metric(&self) -> &impl Metric { &self.0 } }

	impl FlatCoordinateSystem<Vec2> for OuterCS {
		fn flat_to_global(&self, pos: Vec2) -> Vec2 {
			let inner = Rect { size: vec2(self.0.inner_radius + 1.0, self.0.external_halflength) };
			if inner.is_inside(pos) {
				let Vec2 { x, y: v } = pos;
				let y = self.0.y(v - v.signum() * (self.0.external_halflength - self.0.internal_halflength));
				vec2(x, y)
			} else {
				pos
			}
		}

		fn global_to_flat(&self, pos: Vec2) -> Vec2 {
			let inner = Rect { size: vec2(self.0.inner_radius + 1.0, self.0.external_halflength) };
			if inner.is_inside(pos) {
				let Vec2 { x: u, y } = pos; // в основной СК
				let v = self.0.v(y) + y.signum() * (self.0.external_halflength - self.0.internal_halflength);
				vec2(u, v) // в плоском продолжении СК Outer на область Inner
			} else {
				pos
			}
		}
	}

	impl FlatRegion for OuterCS {
		fn distance_to_boundary(&self, ray: Ray) -> Option<f32> {
			Rect { size: vec2(self.0.outer_radius, self.0.external_halflength) }.trace_into(ray)
		}
	}

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

		fn test_flat_region(region: &impl FlatRegion, range_global: (Vec2, Vec2), range_flat: (Vec2, Vec2)) {
			const ε: f32 = 1e-3;
			macro_rules! assert_eq_at {
				($at: expr, $left: expr, $right: expr) => {
					let at = $at;
					let left = $left;
					let right = $right;
					assert!(left.abs_diff_eq(right, ε), "Assertion failed at {at}:\n  left: {left} = {}\n right: {right} = {}", stringify!($left), stringify!($right));
				};
			}
			fn check_range(name_a: &str, a: Vec2, range_a: (Vec2, Vec2), name_b: &str, b: Vec2, range_b: (Vec2, Vec2)) {
				assert!(b.cmpge(range_b.0 - ε).all() && b.cmple(range_b.1 + ε).all(), "Assertion failed:\nAt {name_a}: {a}, from range: {range_a:?}\nGot {name_b}: {b}, which is out of range {range_b:?}");
				// TODO sort out when to check these conditions:
				if a.x.abs_diff_eq(&range_a.0.x, ε) { assert_abs_diff_eq!(b.x, range_b.0.x, epsilon=ε); }
				if a.y.abs_diff_eq(&range_a.0.y, ε) { assert_abs_diff_eq!(b.y, range_b.0.y, epsilon=ε); }
				if a.x.abs_diff_eq(&range_a.1.x, ε) { assert_abs_diff_eq!(b.x, range_b.1.x, epsilon=ε); }
				if a.y.abs_diff_eq(&range_a.1.y, ε) { assert_abs_diff_eq!(b.y, range_b.1.y, epsilon=ε); }
			}
			for x in linspace(range_global.0.x, range_global.1.x, 20) {
				for y in linspace(range_global.0.y, range_global.1.y, 20) {
					let pos_global = vec2(x, y);
					let pos_flat = region.global_to_flat(pos_global);
					check_range("global", pos_global, range_global, "flat", pos_flat, range_flat);
					assert_eq_at!(pos_global, region.global_to_flat(Location { pos: pos_global, rot: Mat2::IDENTITY }).pos, pos_flat);
					assert_eq_at!(pos_global, region.flat_to_global(pos_flat), pos_global);
					assert_eq_at!(pos_global, region.flat_to_global(region.global_to_flat(Location { pos: pos_global, rot: Mat2::IDENTITY })).rot, Mat2::IDENTITY);
				}
			}
			for x in linspace(range_flat.0.x, range_flat.1.x, 20) {
				for y in linspace(range_flat.0.y, range_flat.1.y, 20) {
					let pos_flat = vec2(x, y);
					let pos_global = region.flat_to_global(pos_flat);
					check_range("flat", pos_flat, range_flat, "global", pos_global, range_global);
					assert_eq_at!(pos_flat, region.flat_to_global(Location { pos: pos_flat, rot: Mat2::IDENTITY }).pos, pos_global);
					assert_eq_at!(pos_flat, region.global_to_flat(pos_global), pos_flat);
					assert_eq_at!(pos_global, region.global_to_flat(region.flat_to_global(Location { pos: pos_global, rot: Mat2::IDENTITY })).rot, Mat2::IDENTITY);
				}
			}
		}

		#[test]
		fn test_mapper_inner() {
			let mapper = InnerCS(Tube {
				inner_radius: 30.0,
				outer_radius: 50.0,
				internal_halflength: 100.0,
				external_halflength: 300.0,
			});
			test_flat_region(&mapper, (vec2(-30.0, -300.0), vec2(30.0, 300.0)), (vec2(-30.0, -100.0), vec2(30.0, 100.0)));
			test_flat_region(&mapper, (vec2(-60.0, -400.0), vec2(60.0, -300.0)), (vec2(-60.0, -200.0), vec2(60.0, -100.0)));
			test_flat_region(&mapper, (vec2(-60.0, 300.0), vec2(60.0, 400.0)), (vec2(-60.0, 100.0), vec2(60.0, 200.0)));
		}

		#[test]
		fn test_mapper_outer() {
			let mapper = OuterCS(Tube {
				inner_radius: 30.0,
				outer_radius: 50.0,
				internal_halflength: 100.0,
				external_halflength: 300.0,
			});
			// TODO replace 200.20016 with something sane
			test_flat_region(&mapper, (vec2(-30.0, -300.0), vec2(30.0, -1.0)), (vec2(-30.0, -300.0), vec2(30.0, -200.20016)));
			test_flat_region(&mapper, (vec2(-30.0, 1.0), vec2(30.0, 300.0)), (vec2(-30.0, 200.20016), vec2(30.0, 300.0)));
			test_flat_region(&mapper, (vec2(-60.0, -400.0), vec2(60.0, -300.0)), (vec2(-60.0, -400.0), vec2(60.0, -300.0)));
			test_flat_region(&mapper, (vec2(-60.0, 300.0), vec2(60.0, 400.0)), (vec2(-60.0, 300.0), vec2(60.0, 400.0)));
			// straight
			for x in linspace(-60., 60., 20) {
				for y in linspace(-320., 320., 20) {
					assert_eq!(mapper.global_to_flat(Location { pos: vec2(x, y), rot: Mat2::IDENTITY }).pos.x, x);
				}
			}
			// symmetrical
			for x in linspace(0., 60., 20) {
				for y in linspace(0., 320., 20) {
					let pp = mapper.global_to_flat(Location { pos: vec2(x, y), rot: Mat2::IDENTITY }).pos;
					let np = mapper.global_to_flat(Location { pos: vec2(-x, y), rot: Mat2::IDENTITY }).pos;
					let pn = mapper.global_to_flat(Location { pos: vec2(x, -y), rot: Mat2::IDENTITY }).pos;
					let nn = mapper.global_to_flat(Location { pos: vec2(-x, -y), rot: Mat2::IDENTITY }).pos;
					assert_eq!(np, vec2(-pp.x, pp.y));
					assert_eq!(pn, vec2(pp.x, -pp.y));
					assert_eq!(nn, vec2(-pp.x, -pp.y));
				}
			}
			// clean boundary
			for x in linspace(50., 60., 20) {
				for y in linspace(0., 320., 20) {
					assert_eq!(mapper.global_to_flat(Location { pos: vec2(x, y), rot: Mat2::IDENTITY }).pos.y, y);
				}
			}
			for x in linspace(0., 60., 20) {
				for y in linspace(300., 320., 20) {
					assert_eq!(mapper.global_to_flat(Location { pos: vec2(x, y), rot: Mat2::IDENTITY }).pos.y, y);
				}
			}
			// accelerating
			for x in linspace(-29., 29., 20) {
				for y in linspace(1., 299., 20) {
					let v = mapper.global_to_flat(Location { pos: vec2(x, y), rot: Mat2::IDENTITY }).pos.y;
					assert!(v > 200.0);
					assert!(v > y);
				}
			}
		}
	}
}