use std::f32::consts::{FRAC_PI_2, PI};

use flo_canvas::*;
use flo_draw::*;
use glam::*;

use riemann::{Metric, trace_iter};
use tube::metric::Tube;
use tube::Space;
use tube::Subspace::{Boundary, Inner, Outer};
use types::{Location, Object, Ray};

mod riemann;
mod fns;
mod float_fun;
mod tube;
mod types;

const DT: f32 = 0.1;

fn draw_loop(gc: &mut Vec<Draw>, mut pts: impl Iterator<Item=Vec2>) {
	gc.new_path();
	let Some(first) = pts.next() else { return; };
	gc.move_to(first.x, first.y);
	for pt in pts {
		gc.line_to(pt.x, pt.y);
	}
	gc.close_path();
	gc.stroke();
}

pub fn main() {
	with_2d_graphics(move || {
		let canvas = create_drawing_window("Refraction");
		canvas.draw(|gc| {
			let tube = Tube {
				inner_radius: 30.0,
				outer_radius: 50.0,
				internal_halflength: 100.0,
				external_halflength: 300.0,
			};

			let objs: Vec<_> = [-1.25, -1.00, -0.85, -0.50, 0.00, 0.40, 0.70, 0.95, 1.05]
				.iter()
				.enumerate()
				.map(|(k, &y)| Object {
					id: k as i32,
					loc: {
						let pos = vec2(0.0, y * tube.external_halflength);
						let adj: Mat2 = tube.sqrt_at(pos).inverse().into();
						let rot = Mat2::from_angle(y);
						Location {
							pos,
							rot: adj * rot,
						}
					},
					r: 20.0,
				})
				.collect();
			let space = Space { tube, objs };

			gc.canvas_height(500.0);
			gc.transform(Transform2D::rotate(FRAC_PI_2));
			tube.render(gc);
			gc.line_width(0.5);
			// gc.stroke_color(Color::Rgba(1.0, 0.5, 0.0, 0.5));
			// draw_fan(gc, &tube, vec2(-500.0, 0.0), vec2(1.0, 0.0), 1.0);
			gc.stroke_color(Color::Rgba(1.0, 0.5, 0.0, 1.0));
			draw_fan_2(gc, &space, vec2(-500.0, 0.0), vec2(1.0, 0.0), 1.0);
			gc.stroke_color(Color::Rgba(0.5, 1.0, 0.0, 1.0));
			draw_fan_2(gc, &space, vec2(-2.5 * tube.outer_radius, 1.25 * tube.external_halflength), vec2(1.0, -1.0), 1.0);
			draw_track(gc, &space, vec2(-500.0, 0.0), vec2(1.0, 0.2));
			draw_track(gc, &space, vec2(-500.0, 0.0), vec2(1.0, 0.5));
			draw_track(gc, &space, vec2(-0.5 * tube.inner_radius, -1.25 * tube.external_halflength), vec2(0.1, 1.0));

			let circle_segments = 47;
			for obj in &space.objs {
				let pos = obj.loc.pos;
				gc.new_path();
				gc.circle(pos.x, pos.y, 5.0);
				gc.fill_color(Color::Rgba(0.0, 0.5, 1.0, 1.0));
				gc.fill();

				gc.stroke_color(Color::Rgba(0.0, 0.0, 0.0, 0.5));
				draw_loop(gc, itertools_num::linspace(0.0, 2.0 * PI, circle_segments).skip(1).map(|φ| {
					let dir = Vec2::from_angle(φ) * obj.r;
					let dir = obj.loc.rot * dir;
					pos + dir
				}));
				gc.stroke_color(Color::Rgba(0.0, 0.5, 1.0, 0.5));
				draw_loop(gc, itertools_num::linspace(0.0, 2.0 * PI, circle_segments).skip(1).map(|φ| {
					let dir = Vec2::from_angle(φ) * obj.r;
					let dir = obj.loc.rot * dir;
					space.trace_step(Ray { pos, dir }).pos
				}));
				gc.stroke_color(Color::Rgba(0.5, 0.0, 1.0, 1.0));
				draw_loop(gc, itertools_num::linspace(0.0, 2.0 * PI, circle_segments).skip(1).map(|φ| {
					let n = obj.r.floor();
					let d = obj.r / n;
					let dir = Vec2::from_angle(φ);
					let dir = obj.loc.rot * dir * d;
					space.trace_iter(Ray { pos, dir }).nth(n as usize).unwrap().pos
				}));
			}
		});
	});
}

fn draw_ray_2(gc: &mut Vec<Draw>, space: &Space, base: Vec2, dir: Vec2) {
	let mut hits = Vec::<Draw>::new();
	let dir = space.tube.globalize(base, dir);
	gc.new_path();
	gc.move_to(base.x, base.y);
	let mut ray = Ray { pos: base, dir: space.tube.normalize_vec_at(base, dir) * DT };
	for _ in 0..10000 {
		ray = space.trace_step(ray);
		gc.line_to(ray.pos.x, ray.pos.y);
		if ray.pos.abs().cmpgt(Vec2::splat(1000.0)).any() {
			break;
		}
		let sub = space.which_subspace(ray.pos);
		if sub == Boundary {
			continue;
		}
		gc.stroke();
		gc.new_dash_pattern();
		// gc.dash_length(6.0);
		gc.new_path();
		gc.move_to(ray.pos.x, ray.pos.y);
		let ret = match sub {
			Inner => space.trace_inner(ray),
			Outer => space.trace_outer(ray),
			Boundary => panic!(),
		};
		for hit in ret.objects {
			let obj = space.objs[hit.id as usize];
			hits.move_to(obj.loc.pos.x, obj.loc.pos.y);
			for pt in trace_iter(&space.tube, obj.loc.pos, obj.loc.rot * hit.rel.pos, hit.rel.pos.length() / 100.0).take(100) {
				hits.line_to(pt.x, pt.y);
			}
			hits.circle(hit.pos.x, hit.pos.y, 1.5);
			let Ray { pos: rel, dir } = hit.rel;
			let diff = rel.dot(dir).powi(2) - dir.length_squared() * (rel.length_squared() - obj.r.powi(2));
			assert!(diff >= 0.0);
			let t = (-rel.dot(dir) + diff.sqrt()) / dir.length_squared();
			let rel2 = hit.rel.forward(t).pos;
			let pos2 = trace_iter(&space.tube, obj.loc.pos, obj.loc.rot * rel2, rel2.length() / 100.0).nth(100).unwrap();
			hits.move_to(pos2.x - 1.0, pos2.y - 1.0);
			hits.line_to(pos2.x + 1.0, pos2.y + 1.0);
			hits.move_to(pos2.x - 1.0, pos2.y + 1.0);
			hits.line_to(pos2.x + 1.0, pos2.y - 1.0);
		}
		let a = ray.pos;
		ray = match ret.end {
			Some(r) => r,
			None => {
				ray = ray.forward(1000.0 / DT);
				gc.line_to(ray.pos.x, ray.pos.y);
				break;
			}
		};
		for p in space.line(a, ray.pos, 10.0) {
			gc.line_to(p.x, p.y);
		}
		gc.stroke();
		gc.new_dash_pattern();
		gc.new_path();
		gc.move_to(ray.pos.x, ray.pos.y);
	}
	gc.stroke();
	gc.new_path();
	gc.new_dash_pattern();
	gc.append(&mut hits);
	gc.stroke();
}

fn draw_fan_2(gc: &mut Vec<Draw>, space: &Space, 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) {
		draw_ray_2(gc, space, base, dir + y * v);
	}
}

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

fn draw_track(gc: &mut Vec<Draw>, space: &Space, start: Vec2, dir: Vec2) {
	const SCALE: f32 = 5.0;
	const STEP: f32 = 2.0 * SCALE;
	// let mut loc = Location { pos: start, rot: Mat2::IDENTITY };
	// let dir = space.tube.globalize(start, dir);
	// let v = space.tube.normalize(start, dir);
	let mut loc = Location { pos: start, rot: mat2(dir, vec2(-dir.y, dir.x)) };
	let v = vec2(1.0, 0.0);
	let mut draw = |loc: &Location| {
		let p = loc.pos;
		let ax = p + loc.rot.x_axis * SCALE;
		let ay = p + loc.rot.y_axis * SCALE;
		gc.new_path();
		gc.stroke_color(Color::Rgba(0.7, 0.0, 0.0, 1.0));
		gc.move_to(p.x, p.y);
		gc.line_to(ax.x, ax.y);
		gc.stroke();
		gc.new_path();
		gc.stroke_color(Color::Rgba(0.0, 0.7, 0.0, 1.0));
		gc.move_to(p.x, p.y);
		gc.line_to(ay.x, ay.y);
		gc.stroke();
	};
	draw(&loc);
	for _ in 0..1000 {
		let N = (STEP / DT).floor() as i32;
		for _ in 0..N {
			loc = space.move_step(loc, v * DT);
		}
		draw(&loc);
	}
}

fn draw_fan(gc: &mut Vec<Draw>, space: &impl Metric, 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) {
		draw_ray(gc, space, base, dir + y * v);
	}
}

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

impl Renderable for Tube {
	fn render(&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();
	}
}