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

	pub(crate) 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
	}
}

pub type Tens2 = [Mat2; 2];

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

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

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

	fn part_derivs_at(&self, pos: Vec2) -> Tens2 {
		part_deriv(|p| self.at(p), pos, 1.0 / 1024.0) // division by such eps is exact which is good for overall precision
	}

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

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

	fn globalize(&self, at: Vec2, v: Vec2) -> Vec2 {
		Mat2::from(self.sqrt_at(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 = -contract2(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_vec_at(base, dir),
		dt,
	}
}

pub 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 = &space.inverse_at(pos); // с верхними индексами
	let d = space.part_derivs_at(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)),
	]
}

/// Сворачивает тензор t с вектором u
pub fn contract(t: Tens2, u: Vec2) -> Mat2 {
	mat2(t[0] * u, t[1] * u).transpose()
}

/// Сворачивает тензор t с вектором v дважды, по второму и третьему индексам.
pub fn contract2(t: Tens2, v: Vec2) -> Vec2 {
	contract(t, v) * 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);
}