use std::rc::Rc;
use flo_draw::*;
use flo_canvas::*;
use glam::*;
use riemann::{Decomp2, Metric, trace_iter};
use crate::boundary::Loop;

#[cfg(test)]
use approx::assert_abs_diff_eq;

const RAYS_IN_FAN: usize = 101;
const DT: f32 = 0.1;

pub fn main() {
	let space = Coil {
		scale: 3.0,
		r: 300.0,
		w: 50.0,
		m: 10.0,
	};
	let space = Rect {
		inner_radius: 30.0,
		outer_radius: 50.0,
		internal_halflength: 100.0,
		external_halflength: 300.0,
	};
	with_2d_graphics(move || {
		let canvas = create_drawing_window("Refraction");
		canvas.draw(|gc| {
			gc.canvas_height(1000.0);
			space.draw_background(gc);
			space.draw_rays(gc, &space);

			let space = Rc::new(space);
			let parts: Vec<Box<dyn SpacePart>> = vec![
				Box::new(Outside { space: space.clone() }),
				Box::new(Wall { space: space.clone() }),
				Box::new(Inside { space: space.clone() }),
			];
			space.draw_rays(gc, parts.as_slice());
		});
	});
}

trait Interesting {
	fn draw_background(&self, gc: &mut Vec<Draw>);
	fn draw_rays(&self, gc: &mut Vec<Draw>, tracer: &(impl Rayable + ?Sized));
}

trait SpacePart: boundary::Boundary + Metric + SpaceVisual {}

impl<T: boundary::Boundary + Metric + SpaceVisual> SpacePart for T {}

trait SpaceVisual {
	fn color(&self) -> Color;
	fn globalize_loc(&self, pos: Vec2) -> Vec2;
}

impl SpaceVisual for Outside {
	fn color(&self) -> Color {
		Color::Rgba(0.7, 0.7, 0.7, 1.0)
	}

	fn globalize_loc(&self, pos: Vec2) -> Vec2 {
		pos
	}
}

impl SpaceVisual for Wall {
	fn color(&self) -> Color {
		Color::Rgba(1.0, 0.0, 0.0, 1.0)
	}

	fn globalize_loc(&self, pos: Vec2) -> Vec2 {
		pos
	}
}

impl SpaceVisual for Inside {
	fn color(&self) -> Color {
		Color::Rgba(0.0, 0.7, 0.0, 1.0)
	}

	fn globalize_loc(&self, pos: Vec2) -> Vec2 {
		vec2(pos.x, self.space.x(pos.y))
	}
}

impl Rayable for [Box<dyn SpacePart>] {
	fn draw_ray(&self, gc: &mut Vec<Draw>, base: Vec2, dir: Vec2) {
		gc.new_path();
		let dt = DT;
		let mut s = boundary::Id(0);
		let mut p = base;
		let mut v = dir.normalize();
		let part = &*self[s.0 as usize];
		gc.stroke_color(part.color());
		gc.move_to(p.x, p.y);
		for _ in 0..10000 {
			let part = &*self[s.0 as usize];
			let a: Vec2 = -riemann::convolute(riemann::krist(part, p), v);
			v = v + a * dt;
			if let Some((id, base, dir)) = part.next(p, v, dt) {
				gc.stroke();
				gc.new_path();
				let pt = part.globalize_loc(p);
				gc.move_to(pt.x, pt.y);
				s = id;
				p = base;
				v = dir;
				let part = &*self[s.0 as usize];
				let pt = part.globalize_loc(p);
				gc.stroke_color(part.color());
				gc.line_to(pt.x, pt.y);
			} else {
				p = p + v * dt;
				let pt = part.globalize_loc(p);
				gc.line_to(pt.x, pt.y);
			}
		}
		gc.stroke();
	}
}

trait Rayable {
	fn draw_ray(&self, gc: &mut Vec<Draw>, base: Vec2, dir: Vec2);

	fn draw_fan(&self, gc: &mut Vec<Draw>, base: Vec2, dir: Vec2, spread: f32) {
		let dir = dir.normalize();
		let v = vec2(-dir.y, dir.x);
		for y in itertools_num::linspace(-spread, spread, RAYS_IN_FAN) {
			self.draw_ray(gc, base, dir + y * v);
		}
	}
}

impl<T: Metric> Rayable for T {
	fn draw_ray(&self, gc: &mut Vec<Draw>, base: Vec2, dir: Vec2) {
		let dir = self.globalize(base, dir);
		gc.new_path();
		gc.move_to(base.x, base.y);
		for pt in trace_iter(self, base, dir, DT).take(10000) {
			gc.line_to(pt.x, pt.y);
			if pt.abs().cmpgt(Vec2::splat(1000.0)).any() {
				break;
			}
		}
		gc.stroke();
	}
}

impl Interesting for Coil {
	fn draw_background(&self, gc: &mut Vec<Draw>) {
		gc.new_path();
		gc.circle(0.0, 0.0, self.r + self.w + self.m);
		gc.circle(0.0, 0.0, self.r + self.w);
		gc.circle(0.0, 0.0, self.r - self.w);
		gc.circle(0.0, 0.0, self.r - self.w - self.m);
		gc.winding_rule(WindingRule::EvenOdd);
		gc.fill_color(Color::Rgba(0.8, 0.8, 0.8, 1.0));
		gc.fill();
	}

	fn draw_rays(&self, gc: &mut Vec<Draw>, tracer: &(impl Rayable + ?Sized)) {
		gc.line_width(0.5);
		gc.stroke_color(Color::Rgba(1.0, 0.5, 0.0, 1.0));
		tracer.draw_fan(gc, vec2(-500.0, 0.0), vec2(1.0, 0.0), 1.0);
		gc.stroke_color(Color::Rgba(0.0, 0.5, 1.0, 1.0));
		tracer.draw_fan(gc, vec2(0.0, self.r), vec2(1.0, 0.0), 1.0);
	}
}

impl Interesting for Rect {
	fn draw_background(&self, gc: &mut Vec<Draw>) {
		gc.new_path();
		gc.rect(-self.outer_radius, -self.external_halflength, self.outer_radius, self.external_halflength);
		gc.rect(-self.inner_radius, -self.external_halflength, self.inner_radius, self.external_halflength);
		gc.winding_rule(WindingRule::EvenOdd);
		gc.fill_color(Color::Rgba(0.8, 0.8, 0.8, 1.0));
		gc.fill();
	}

	fn draw_rays(&self, gc: &mut Vec<Draw>, tracer: &(impl Rayable + ?Sized)) {
		gc.line_width(0.5);
		gc.stroke_color(Color::Rgba(1.0, 0.5, 0.0, 1.0));
		tracer.draw_fan(gc, vec2(-500.0, 0.0), vec2(1.0, 0.0), 1.0);
		gc.stroke_color(Color::Rgba(0.0, 0.5, 1.0, 1.0));
		tracer.draw_fan(gc, vec2(0.0, -0.5 * self.external_halflength), vec2(1.0, 1.0), 1.0);
		gc.stroke_color(Color::Rgba(0.2, 0.7, 0.0, 1.0));
		tracer.draw_fan(gc, vec2(-0.5 * self.inner_radius, -1.2 * self.external_halflength), vec2(0.0, 1.0), 1.0);
	}
}

struct Coil {
	m: f32,
	scale: f32,
	r: f32,
	w: f32,
}

impl Metric for Coil {
	fn halfmetric(&self, pos: Vec2) -> Decomp2 {
		let r = pos.length();
		let dir = pos.normalize();

		let s = smoothbox(r, vec2(self.r - self.w, self.r + self.w), self.m);
		let t = 1.0.lerp(self.r / r / self.scale, s);

		Decomp2 {
			ortho: Mat2::from_cols_array(&[
				dir.x, -dir.y,
				dir.y, dir.x,
			]),
			diag: vec2(1.0, t),
		}
	}
}

struct Rect {
	outer_radius: f32,
	inner_radius: f32,
	external_halflength: f32,
	internal_halflength: f32,
}

impl Rect {
	fn γ(&self) -> f32 { self.external_halflength / self.internal_halflength }
	fn ri(&self) -> f32 { self.internal_halflength }
	fn re(&self) -> f32 { self.external_halflength }
	fn a(&self) -> f32 { (1.0 - self.γ()) / self.ri() }
	fn b(&self) -> f32 { 2.0 * self.γ() - 1.0 }

	fn root(&self, x: f32) -> f32 { ((2.0 * self.γ() - 1.0).powi(2) + 4.0 * (1.0 - self.γ()) * x / self.ri()).sqrt() }
	fn d(&self, u: f32) -> f32 { 2.0 * self.a() * u + self.b() }
	pub fn x(&self, u: f32) -> f32 { (self.a() * u.abs() + self.b()) * u }
	pub fn u(&self, x: f32) -> f32 { 0.5 * self.ri() * (1.0 - 2.0 * self.γ() + self.root(x.abs())) / (1.0 - self.γ()) * x.signum() }
	pub fn dx(&self, u: f32, du: f32) -> f32 { du * self.d(u.abs()) }
	pub fn du(&self, x: f32, dx: f32) -> f32 { dx / self.root(x.abs()) }
}

#[test]
fn test_rect() {
	let r = Rect {
		outer_radius: 50.0,
		inner_radius: 20.0,
		external_halflength: 100.0,
		internal_halflength: 10.0,
	};
	assert_abs_diff_eq!(r.x(r.internal_halflength), r.external_halflength, epsilon = 1.0e-5);
	assert_abs_diff_eq!(r.x(-r.internal_halflength), -r.external_halflength, epsilon = 1.0e-5);
	assert_abs_diff_eq!(r.u(r.external_halflength), r.internal_halflength, epsilon = 1.0e-5);
	assert_abs_diff_eq!(r.u(-r.external_halflength), -r.internal_halflength, epsilon = 1.0e-5);
	assert_abs_diff_eq!(r.u(r.x(1.0)), 1.0, epsilon = 1.0e-5);
	assert_abs_diff_eq!(r.u(r.x(5.0)), 5.0, epsilon = 1.0e-5);
	assert_abs_diff_eq!(r.u(r.x(-5.0)), -5.0, epsilon = 1.0e-5);
	assert_abs_diff_eq!(r.u(r.x(10.0)), 10.0, epsilon = 1.0e-5);
	assert_abs_diff_eq!(r.x(r.u(10.0)), 10.0, epsilon = 1.0e-5);
	assert_abs_diff_eq!(r.x(r.u(50.0)), 50.0, epsilon = 1.0e-5);
	assert_abs_diff_eq!(r.x(r.u(-50.0)), -50.0, epsilon = 1.0e-5);
	assert_abs_diff_eq!(r.x(r.u(100.0)), 100.0, epsilon = 1.0e-5);
	assert_abs_diff_eq!(r.d(r.u(10.0)), r.root(10.0), epsilon = 1.0e-5);
	assert_abs_diff_eq!(r.d(r.u(50.0)), r.root(50.0), epsilon = 1.0e-5);
	assert_abs_diff_eq!(r.d(r.u(100.0)), r.root(100.0), epsilon = 1.0e-5);
	assert_abs_diff_eq!(r.du(10.0, r.dx(r.u(10.0), 3.0)), 3.0, epsilon = 1.0e-5);
	assert_abs_diff_eq!(r.du(50.0, r.dx(r.u(50.0), 3.0)), 3.0, epsilon = 1.0e-5);
	assert_abs_diff_eq!(r.du(-50.0, r.dx(r.u(-50.0), 3.0)), 3.0, epsilon = 1.0e-5);
	assert_abs_diff_eq!(r.du(100.0, r.dx(r.u(100.0), 3.0)), 3.0, epsilon = 1.0e-5);
	assert_abs_diff_eq!(r.dx(1.0, r.du(r.x(1.0), 3.0)), 3.0, epsilon = 1.0e-5);
	assert_abs_diff_eq!(r.dx(5.0, r.du(r.x(5.0), 3.0)), 3.0, epsilon = 1.0e-5);
	assert_abs_diff_eq!(r.dx(-5.0, r.du(r.x(-5.0), 3.0)), 3.0, epsilon = 1.0e-5);
	assert_abs_diff_eq!(r.dx(10.0, r.du(r.x(10.0), 3.0)), 3.0, epsilon = 1.0e-5);
}

impl Metric for Rect {
	fn halfmetric(&self, pos: Vec2) -> Decomp2 {
		let x = pos.y.abs();
		let sx = ((pos.x.abs() - self.inner_radius) / (self.outer_radius - self.inner_radius)).clamp(0.0, 1.0);
		let sy = if x <= self.external_halflength { 1.0 / self.root(x) } else { 1.0 };
		assert!(sx.is_finite());
		assert!(sy.is_finite());
		assert!(sy > 0.0);
		Decomp2 {
			ortho: Mat2::IDENTITY,
			diag: vec2(1.0, sy.lerp(1.0, sx)),
		}
	}
}

struct Flat;

impl Metric for Flat {
	fn halfmetric(&self, pos: Vec2) -> Decomp2 {
		Decomp2 {
			ortho: Mat2::IDENTITY,
			diag: Vec2::splat(1.0),
		}
	}
}

struct Outside {
	space: Rc<Rect>,
}

struct Wall {
	space: Rc<Rect>,
}

struct Inside {
	space: Rc<Rect>,
}

impl boundary::Boundary for Outside {
	fn next(&self, base: Vec2, dir: Vec2, limit: f32) -> Option<(boundary::Id, Vec2, Vec2)> {
		let or = self.space.outer_radius;
		let ir = self.space.inner_radius;
		let hl = self.space.external_halflength;
		let bnd = Loop(vec![
			vec2(-or, -hl), vec2(-ir, -hl), vec2(ir, -hl), vec2(or, -hl),
			vec2(or, hl), vec2(ir, hl), vec2(-ir, hl), vec2(-or, hl),
		]);
		let (side, dist) = bnd.hit(base, dir)?;
		if dist <= limit {
			let pt = base + dist * dir;
			return match side {
				1 | 5 => Some((boundary::Id(2), vec2(pt.x, self.space.u(pt.y)), vec2(dir.x, self.space.du(pt.y, dir.y)))),
				_ => Some((boundary::Id(1), pt, dir)),
			};
		}
		None
	}
}

impl Metric for Outside {
	fn halfmetric(&self, pos: Vec2) -> Decomp2 {
		Flat {}.halfmetric(pos)
	}
}

impl boundary::Boundary for Wall {
	fn next(&self, base: Vec2, dir: Vec2, limit: f32) -> Option<(boundary::Id, Vec2, Vec2)> {
		let or = self.space.outer_radius;
		let ir = self.space.inner_radius;
		let hl = self.space.external_halflength;
		let obnd = Loop(vec![vec2(-or, -hl), vec2(-or, hl), vec2(or, hl), vec2(or, -hl)]);
		let ibnd = Loop(vec![vec2(-ir, -hl), vec2(ir, -hl), vec2(ir, hl), vec2(-ir, hl)]);
		if let Some((_, dist)) = ibnd.hit(base, dir) {
			if dist <= limit {
				let p = base + dist * dir;
				let v = dir;
				let v = vec2(v.x, self.space.du(p.y, v.y));
				let p = vec2(p.x, self.space.u(p.y));
				return Some((boundary::Id(2), p, v));
			}
		}
		if let Some((_, dist)) = obnd.hit(base, dir) {
			if dist <= limit {
				return Some((boundary::Id(0), base + dist * dir, dir));
			}
		}
		None
	}
}

impl Metric for Wall {
	fn halfmetric(&self, pos: Vec2) -> Decomp2 {
		self.space.halfmetric(pos)
	}
}

impl boundary::Boundary for Inside {
	fn next(&self, base: Vec2, dir: Vec2, limit: f32) -> Option<(boundary::Id, Vec2, Vec2)> {
		let size = vec2(self.space.inner_radius, self.space.internal_halflength);
		let bnd = Loop(vec![vec2(-size.x, -size.y), vec2(-size.x, size.y), vec2(size.x, size.y), vec2(size.x, -size.y)]);
		let (side, dist) = bnd.hit(base, dir)?;
		if dist <= limit {
			let p = base + dist * dir;
			let v = dir;
			let v = vec2(v.x, self.space.dx(p.y, v.y));
			let p = vec2(p.x, self.space.x(p.y));
			return match side {
				0 | 2 => Some((boundary::Id(1), p, v)),
				1 | 3 => Some((boundary::Id(0), p, v)),
				_ => panic!()
			};
		}
		None
	}
}

impl Metric for Inside {
	fn halfmetric(&self, pos: Vec2) -> Decomp2 {
		Flat {}.halfmetric(pos)
	}
}

mod boundary {
	use glam::*;

	pub struct Id(pub u8);

	pub trait Boundary {
		fn next(&self, base: Vec2, dir: Vec2, limit: f32) -> Option<(Id, Vec2, Vec2)>;
	}

	pub struct Loop(pub Vec<Vec2>);

	impl Loop {
		pub fn hit(&self, base: Vec2, dir: Vec2) -> Option<(usize, f32)> {
			self.0.iter().enumerate().filter_map(|(k, &a)| {
				let b = self.0[(k + 1) % self.0.len()];
				let u = mat2(a - base, dir).determinant();
				let v = mat2(b - base, dir).determinant();
				if u < 0.0 && v > 0.0 {
					let dist = mat2(a - base, b - a).determinant() / mat2(dir, b - a).determinant();
					if dist >= 0.0 { Some((k, dist)) } else { None }
				} else {
					None
				}
			}).min_by(|(k1, dist1), (k2, dist2)| dist1.total_cmp(dist2))
		}
	}

	#[test]
	fn test_loop() {
		let tri = Loop(vec![vec2(-1.0, -1.0), vec2(1.0, -1.0), vec2(0.0, 1.0)]);
		assert_eq!(tri.hit(vec2(0.0, -2.0), vec2(0.0, 1.0)), Some((0, 1.0)));
		assert_eq!(tri.hit(vec2(0.0, -2.0), vec2(0.0, 0.5)), Some((0, 2.0)));
		assert_eq!(tri.hit(vec2(0.0, -2.0), vec2(1.0, 1.0)), None);
		assert_eq!(tri.hit(vec2(0.0, 0.0), vec2(0.0, 1.0)), None);
		assert_eq!(tri.hit(vec2(0.0, 0.0), vec2(0.0, -1.0)), None);
		assert_eq!(tri.hit(vec2(-1.5, 0.5), vec2(2.0, -1.0)), Some((2, 0.5)));
		assert_eq!(tri.hit(vec2(-1.5, 0.5), vec2(-2.0, 1.0)), None);
	}
}

mod riemann {
	use glam::*;

	pub struct Decomp2 {
		pub ortho: Mat2,
		pub diag: Vec2,
	}

	impl Decomp2 {
		fn square(&self) -> Self {
			Self {
				ortho: self.ortho,
				diag: self.diag * self.diag,
			}
		}

		fn inverse(&self) -> Self {
			Self {
				ortho: self.ortho,
				diag: Vec2::splat(1.0) / self.diag,
			}
		}
	}

	impl From<Decomp2> for Mat2 {
		fn from(value: Decomp2) -> Self {
			value.ortho.transpose() * Mat2::from_diagonal(value.diag) * value.ortho
		}
	}

	type Tens2 = [Mat2; 2];

	pub trait Metric {
		fn halfmetric(&self, pos: Vec2) -> Decomp2;

		fn metric(&self, pos: Vec2) -> Mat2 {
			self.halfmetric(pos).square().into()
		}

		fn invmetric(&self, pos: Vec2) -> Mat2 {
			self.halfmetric(pos).square().inverse().into()
		}

		fn dmetric(&self, pos: Vec2) -> Tens2 {
			part_deriv(|p| self.metric(p), pos, 1.0e-3)
		}

		fn length(&self, at: Vec2, v: Vec2) -> f32 {
			v.dot(self.metric(at) * v).sqrt()
		}

		fn normalize(&self, at: Vec2, v: Vec2) -> Vec2 {
			v / self.length(at, v)
		}

		fn globalize(&self, at: Vec2, v: Vec2) -> Vec2 {
			Mat2::from(self.halfmetric(at).inverse()) * v
		}
	}

	pub struct TraceIter<'a, M: Metric> {
		space: &'a M,
		p: Vec2,
		v: Vec2,
		dt: f32,
	}

	impl<'a, M: Metric> Iterator for TraceIter<'a, M> {
		type Item = Vec2;

		fn next(&mut self) -> Option<Self::Item> {
			let a: Vec2 = -convolute(krist(self.space, self.p), self.v);
			self.v = self.v + a * self.dt;
			self.p = self.p + self.v * self.dt;
			Some(self.p)
		}
	}

	pub fn trace_iter<M: Metric>(space: &M, base: Vec2, dir: Vec2, dt: f32) -> TraceIter<M> {
		TraceIter {
			space,
			p: base,
			v: space.normalize(base, dir),
			dt,
		}
	}

	pub fn krist(space: &(impl Metric + ?Sized), pos: Vec2) -> Tens2 {
		// Γ^i_k_l = .5 * g^i^m * (g_m_k,l + g_m_l,k - g_k_l,m)
		let g = space.invmetric(pos); // с верхними индексами
		let d = space.dmetric(pos);
		// ret[i][l][k] = sum((m) => .5f * g[m][i] * (d[k][l][m] + d[l][k][m] - d[m][k][l]))
		make_tens2(|i, l, k| 0.5 * (0..2).map(|m| g.col(m)[i] * (d[l].col(k)[m] + d[k].col(m)[l] - d[m].col(k)[l])).sum::<f32>())
	}

	fn dir_deriv(f: impl Fn(Vec2) -> Mat2, pos: Vec2, delta: Vec2) -> Mat2 {
		(f(pos + delta) - f(pos - delta)) / (2.0 * delta.length())
	}

	fn part_deriv(f: impl Fn(Vec2) -> Mat2, pos: Vec2, eps: f32) -> Tens2 {
		[
			dir_deriv(&f, pos, vec2(eps, 0.0)),
			dir_deriv(&f, pos, vec2(0.0, eps)),
		]
	}

	pub fn convolute(t: Tens2, v: Vec2) -> Vec2 {
		vec2(
			v.dot(t[0] * v),
			v.dot(t[1] * v),
		)
	}

	fn make_vec2(f: impl Fn(usize) -> f32) -> Vec2 {
		Vec2::from_array(std::array::from_fn(|i| f(i)))
	}

	fn make_mat2(f: impl Fn(usize, usize) -> f32) -> Mat2 {
		Mat2::from_cols_array_2d(&std::array::from_fn(|i| std::array::from_fn(|j| f(i, j))))
	}

	fn make_tens2(f: impl Fn(usize, usize, usize) -> f32) -> Tens2 {
		std::array::from_fn(|i| make_mat2(|j, k| f(i, j, k)))
	}

	#[test]
	fn m2() {
		let m = make_mat2(|i, j| (i + 2 * j) as f32);
		assert_eq!(m.col(0)[0], 0.0);
		assert_eq!(m.col(1)[0], 1.0);
		assert_eq!(m.col(0)[1], 2.0);
		assert_eq!(m.col(1)[1], 3.0);
	}

	#[test]
	fn t2() {
		let t = make_tens2(|i, j, k| (i + 2 * j + 4 * k) as f32);
		assert_eq!(t[0].col(0)[0], 0.0);
		assert_eq!(t[1].col(0)[0], 1.0);
		assert_eq!(t[0].col(1)[0], 2.0);
		assert_eq!(t[1].col(1)[0], 3.0);
		assert_eq!(t[0].col(0)[1], 4.0);
		assert_eq!(t[1].col(0)[1], 5.0);
		assert_eq!(t[0].col(1)[1], 6.0);
		assert_eq!(t[1].col(1)[1], 7.0);
	}
}

fn smoothstep(x: f32) -> f32 {
	3.0 * x * x - 2.0 * x * x * x
}

/// 1.0 for val∈[range.x, range.y], 0.0 for val∉[range.x−pad, range.y+pad], linear in-between.
fn trapezoid(val: f32, range: Vec2, pad: f32) -> f32 {
	let slope1 = 1.0 + (val - range.x) / pad;
	let slope2 = 1.0 - (val - range.y) / pad;
	let lin = slope1.min(slope2);
	lin.clamp(0.0, 1.0)
}

/// 1.0 for val∈[range.x, range.y], 0.0 for val∉[range.x−pad, range.y+pad], smoothstep in-between.
fn smoothbox(val: f32, range: Vec2, pad: f32) -> f32 {
	smoothstep(trapezoid(val, range, pad))
}