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

pub fn main() {
	let space = Coil {
		scale: 3.0,
		r: 300.0,
		w: 50.0,
		m: 10.0,
	};
	let space = Rect {
		scale: 3.0,
		r: vec2(30.0, 300.0),
		m: vec2(20.0, 50.0),
	};
	with_2d_graphics(move || {
		let canvas = create_drawing_window("Refraction");
		canvas.draw(|gc| {
			gc.canvas_height(1000.0);
			space.render(gc);

			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() }),
			];

			parts.draw_fan(gc, vec2(-500.0, 0.0), vec2(1.0, 0.0), 1.0);
		});
	});
}

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, pos.y * self.space.scale)
	}
}

impl Rayable for [Box<dyn SpacePart>] {
	fn draw_ray(&self, gc: &mut Vec<Draw>, base: Vec2, dir: Vec2) {
		gc.new_path();
		let dt = 1.0;
		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, 101) {
			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, 1.0).take(10000) {
			gc.line_to(pt.x, pt.y);
			if pt.abs().cmpgt(Vec2::splat(1000.0)).any() {
				break;
			}
		}
		gc.stroke();
	}
}

trait Renderable {
	fn render(&self, gc: &mut Vec<Draw>);
}

impl Renderable for Coil {
	fn render(&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();
		gc.line_width(0.5);
		gc.stroke_color(Color::Rgba(1.0, 0.5, 0.0, 1.0));
		self.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));
		self.draw_fan(gc, vec2(0.0, self.r), vec2(1.0, 0.0), 1.0);
	}
}

impl Renderable for Rect {
	fn render(&self, gc: &mut Vec<Draw>) {
		gc.new_path();
		gc.rect(-self.r.x - self.m.x, -self.r.y - self.m.y, self.r.x + self.m.x, self.r.y + self.m.y);
		gc.rect(-self.r.x, -self.r.y, self.r.x, self.r.y);
		gc.winding_rule(WindingRule::EvenOdd);
		gc.fill_color(Color::Rgba(0.8, 0.8, 0.8, 1.0));
		gc.fill();
		gc.line_width(0.5);
		gc.stroke_color(Color::Rgba(1.0, 0.5, 0.0, 1.0));
		self.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));
		self.draw_fan(gc, vec2(0.0, -0.5 * self.r.y), vec2(1.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 {
	scale: f32,
	r: Vec2,
	m: Vec2,
}

impl Metric for Rect {
	fn halfmetric(&self, pos: Vec2) -> Decomp2 {
		let s = smoothbox(pos.x, vec2(-self.r.x, self.r.x), self.m.x) * smoothbox(pos.y, vec2(-self.r.y, self.r.y), self.m.y);
		Decomp2 {
			ortho: Mat2::IDENTITY,
			diag: vec2(1.0, 1.0.lerp(1.0 / self.scale, s)),
		}
	}
}

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 size = self.space.r + self.space.m;
		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 (_, dist) = bnd.hit(base, dir)?;
		if dist <= limit {
			return Some((boundary::Id(1), base + dist * dir, 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 osize = self.space.r + self.space.m;
		let isize = self.space.r;
		let obnd = Loop(vec![vec2(-osize.x, -osize.y), vec2(-osize.x, osize.y), vec2(osize.x, osize.y), vec2(osize.x, -osize.y)]);
		let ibnd = Loop(vec![vec2(-isize.x, -isize.y), vec2(isize.x, -isize.y), vec2(isize.x, isize.y), vec2(-isize.x, isize.y)]);
		if let Some((_, dist)) = ibnd.hit(base, dir) {
			if dist <= limit {
				let p = base + dist * dir;
				let v = dir;
				let p = vec2(p.x, p.y / self.space.scale);
				let v = vec2(v.x, v.y / self.space.scale);
				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 = self.space.r;
		let size = vec2(size.x, size.y / self.space.scale);
		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 (_, dist) = bnd.hit(base, dir)?;
		if dist <= limit {
			let p = base + dist * dir;
			let v = dir;
			let p = vec2(p.x, p.y * self.space.scale);
			let v = vec2(v.x, v.y * self.space.scale);
			return Some((boundary::Id(1), p, v));
		}
		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], smoothstep in-between.
fn smoothbox(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);
	smoothstep(lin.clamp(0.0, 1.0))
}