use flo_draw::*;
use flo_canvas::*;
use glm::*;
use num_traits::identities::Zero;

pub fn main() {
	let space = Coil {
		coil_scale: 2.0,
		coil_r: 300.0,
		coil_w: 50.0,
		coil_m: 10.0,
	};
	with_2d_graphics(move || {
		let canvas = create_drawing_window("Refraction");
		canvas.draw(|gc| {
			gc.canvas_height(1000.0);

			gc.new_path();
			gc.circle(0.0, 0.0, space.coil_r + space.coil_w + space.coil_m);
			gc.circle(0.0, 0.0, space.coil_r + space.coil_w);
			gc.circle(0.0, 0.0, space.coil_r - space.coil_w);
			gc.circle(0.0, 0.0, space.coil_r - space.coil_w - space.coil_m);
			gc.winding_rule(WindingRule::EvenOdd);
			gc.fill_color(Color::Rgba(0.8, 0.8, 0.8, 1.0));
			gc.fill();

			gc.line_width(1.0);
			gc.stroke_color(Color::Rgba(1.0, 0.5, 0.0, 1.0));
			for y in itertools_num::linspace(-1.0, 1.0, 101) {
				let base = vec2(-500.0, 0.0);
				let dir = vec2(1.0, y);
				let path = trace_iter(&space, base, dir, 1.0);
				gc.new_path();
				gc.move_to(base.x, base.y);
				for pt in path.take(10000) {
					gc.line_to(pt.x, pt.y);
					if any(greaterThan(abs(pt), Vec2::from_s(1000.0))) {
						break
					}
				}
				gc.stroke();
			}
		});
	});
}

struct Coil {
	coil_m: f32,
	coil_scale: f32,
	coil_r: f32,
	coil_w: f32,
}

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

		let s = smoothbox (r, vec2(self.coil_r - self.coil_w, self.coil_r + self.coil_w), self.coil_m);
		let t = mix(1.0, self.coil_r / r / self.coil_scale, s);

		Decomp2{
			ortho: mat2(
				dir.x, - dir.y,
				dir.y, dir.x,
			),
			diag : vec2(1.0, t),
		}
	}
}

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

type Tens2 = [Mat2; 2];

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

	fn metric(&self, pos: Vec2) -> Mat2 {
		let h = self.halfmetric(pos);
		transpose(&h.ortho) * diagonal(h.diag * h.diag) * h.ortho
	}

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

	fn globalize(&self, at: Vec2, v: Vec2) -> Vec2 {
		let h = self.halfmetric(at);
		transpose(&h.ortho) * diagonal( Vec2::from_s(1.0) / h.diag) * h.ortho * v
	}
}

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

fn trace(space: &impl Metric, base: Vec2, dir: Vec2, distance: f32, dt: f32) -> Vec<Vec2> {
	let steps = floor(distance / dt) as usize;
	let mut result = Vec::with_capacity(steps);
	let mut p = base;
	let mut v = normalize(dir);
	for _ in 0..steps {
		let a: Vec2 = -convolute(krist(space, p), v);
		v = v + a * dt;
		p = p + v * dt;
		result.push(p);
	}
	result
}

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

#[test]
fn t_iter() {
	let space = Coil {
		coil_scale: 2.0,
		coil_r: 300.0,
		coil_w: 50.0,
		coil_m: 10.0,
	};
	let base = vec2(-500.0, 0.0);
	let dir = vec2(1.0, 0.3);
	let dt = 1.0;
	let steps = 1000;
	let a = trace(&space, base, dir, dt * (steps as f32), dt);
	let b: Vec<Vec2> = trace_iter(&space, base, dir, dt).take(steps).collect();
	assert_eq!(a, b);
}

fn krist(space: &impl Metric, pos: Vec2) -> Tens2 {
	// Γ^i_k_l = .5 * g^i^m * (g_m_k,l + g_m_l,k - g_k_l,m)
	let g = inverse(&space.metric(pos)); // с верхними индексами
	let d = space.dmetric(pos);
	let mut ret: Tens2 = [Mat2::zero(); 2];
	// ret[i][l][k] = sum((m) => .5f * g[m][i] * (d[k][l][m] + d[l][k][m] - d[m][k][l]))
	for i in 0..2 {
		for l in 0..2 {
			for k in 0..2 {
				let mut v = 0.0;
				for m in 0..2 {
					v += g[m][i] * (d[l][k][m] + d[k][m][l] - d[m][k][l]);
				}
				ret[i][l][k] = 0.5 * v;
			}
		}
	}
	ret
}

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

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

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

fn diagonal(v: Vec2) -> Mat2 {
	mat2(v.x, 0.0, 0.0, v.y)
}

fn sqr(x: f32) -> f32 {
	return x * x;
}

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

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 = min(slope1, slope2);
	smoothstep(glm::clamp(lin, 0.0, 1.0))
}