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

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

use refraction::ifaces::{DebugTraceable, Traceable};
use refraction::mathx::MatExt;
use refraction::riemann::{trace_iter, Metric};
use refraction::tube::metric::Tube;
use refraction::tube::Space;
use refraction::types::{Location, Object, Ray};
use refraction::DT;

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: put_object(
						&tube,
						vec2(0.0, y * tube.external_halflength),
						Mat2::from_angle(y),
					),
					r: 20.0,
				})
				.collect();
			let space = Space { tube, objs };
			let cam1 = put_object(&space.tube, vec2(-500., 0.), Mat2::IDENTITY);
			let cam2 = put_object(
				&space.tube,
				vec2(-2.5 * tube.outer_radius, 1.25 * tube.external_halflength),
				mat2(vec2(1., -1.), vec2(1., 1.)),
			);
			let cam3 = put_object(
				&space.tube,
				vec2(0.25 * tube.inner_radius, 0.25 * tube.external_halflength),
				mat2(vec2(0., -1.), vec2(1., 0.)),
			);

			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(0.0, 0.8, 1.0, 1.0));
			draw_fan_2(gc, &space, cam3, 1.0);
			gc.stroke_color(Color::Rgba(0.5, 1.0, 0.0, 1.0));
			draw_fan_2(gc, &space, cam2, 1.0);
			gc.stroke_color(Color::Rgba(1.0, 0.5, 0.0, 1.0));
			draw_fan_2(gc, &space, cam1, 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 rel_to_abs(space: &impl Metric, base: &Location, rel: Vec2, steps: usize) -> Vec2 {
	let c = 1.0 / (steps as f32);
	trace_iter(space, base.pos, base.rot * rel, c * rel.length())
		.nth(steps - 1)
		.unwrap()
}

/// Converts a position and a rotation to a [Location]. Only the X direction is preserved from `rot` to ensure the resulting Location describes an orthonormal coordinate system.
fn put_object(space: &impl Metric, pos: Vec2, rot: Mat2) -> Location {
	let metric_sqrt = space.sqrt_at(pos);
	let metric_inv_sqrt = space.sqrt_at(pos).inverse();
	let rot = metric_inv_sqrt * (metric_sqrt * rot).orthonormalize();
	Location { pos, rot }
}

#[test]
fn test_put_object() {
	use approx::assert_abs_diff_eq;

	let ε = 1e-5;
	let m = refraction::riemann::samples::ScaledMetric {
		scale: vec2(3., 4.),
	};

	let loc = put_object(&m, vec2(1., 2.), mat2(vec2(1., 0.), vec2(0., 1.)));
	assert_eq!(loc.pos, vec2(1., 2.));
	assert_abs_diff_eq!(loc.rot * vec2(1., 0.), vec2(1. / 3., 0.), epsilon = ε);
	assert_abs_diff_eq!(loc.rot * vec2(0., 1.), vec2(0., 1. / 4.), epsilon = ε);

	let loc = put_object(&m, vec2(1., 2.), mat2(vec2(0., 1.), vec2(-1., 0.)));
	assert_eq!(loc.pos, vec2(1., 2.));
	assert_abs_diff_eq!(loc.rot * vec2(1., 0.), vec2(0., 1. / 4.), epsilon = ε);
	assert_abs_diff_eq!(loc.rot * vec2(0., 1.), vec2(-1. / 3., 0.), epsilon = ε);

	let c = 0.5 * std::f32::consts::SQRT_2;
	let loc = put_object(&m, vec2(1., 2.), mat2(vec2(c, c), vec2(-c, c)));
	assert_eq!(loc.pos, vec2(1., 2.));
	assert_abs_diff_eq!(loc.rot * vec2(1., 0.), vec2(1. / 5., 1. / 5.), epsilon = ε);
	assert_abs_diff_eq!(
		loc.rot * vec2(0., 1.),
		vec2(-4. / 15., 3. / 20.),
		epsilon = ε
	);
}

fn draw_cross(gc: &mut Vec<Draw>, pos: Vec2, r: f32) {
	gc.move_to(pos.x - r, pos.y - r);
	gc.line_to(pos.x + r, pos.y + r);
	gc.move_to(pos.x - r, pos.y + r);
	gc.line_to(pos.x + r, pos.y - r);
}

fn draw_ray_2(gc: &mut Vec<Draw>, space: &Space, camera: Location, dir: Vec2) {
	let pos = vec2(0., 0.);
	let (hits, path) = space.trace_dbg(camera, Ray { pos, dir });
	let hits2 = space.trace(camera, Ray { pos, dir });
	for (a, b) in hits.into_iter().zip(hits2.into_iter()) {
		assert_eq!(a.id, b.id);
		assert_eq!(a.pos, b.pos);
		assert_eq!(a.rel, b.rel);
	}

	gc.new_path();
	gc.move_to(pos.x, pos.y);
	for pt in &path.points[1..] {
		gc.line_to(pt.x, pt.y);
	}
	let end_pos = *path
		.points
		.last()
		.expect("the starting point is always in the path");
	let dir_pos = end_pos + 1000.0 * path.end_dir;
	gc.line_to(dir_pos.x, dir_pos.y);
	gc.stroke();
}

fn draw_fan_2(gc: &mut Vec<Draw>, space: &Space, camera: Location, spread: f32) {
	for y in itertools_num::linspace(-spread, spread, 101) {
		draw_ray_2(gc, space, camera, vec2(1., y));
	}
}

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);
	}
}

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