use crate::hamiltonian::*;

use std::ops::{Add, Sub, Mul, Div, AddAssign};

#[derive(Clone, Copy, Debug)]
pub struct Augmented<T> {
	omega: f64,
	pub a: State<T>,
	pub b: State<T>,
}

impl<T: Copy> Augmented<T> {
	pub fn from(state: State<T>, omega: f64) -> Self {
		Self {
			omega,
			a: state,
			b: state,
		}
	}
}

impl<T> Augmented<T>
where T: Copy + Add<Output = T> + Mul<f64, Output = T>
{
	pub fn average(&self) -> State<T> {
		(self.a + self.b) * 0.5
	}
}

impl<T> Augmented<T>
where
	T: Copy + Add<Output = T> + Sub<Output = T> + Mul<f64, Output = T> + Div<f64, Output = T> + AddAssign,
	State<T>: Hamiltonian
{
	pub fn evolve(&mut self, h: f64) {
		let cbrt2 = 2.0_f64.powf(1.0 / 3.0);
		let z0 = - cbrt2 / (2.0 - cbrt2);
		let z1 = 1.0 / (2.0 - cbrt2);

		self.strang(z1 * h);
		self.strang(z0 * h);
		self.strang(z1 * h);
	}

	fn strang(&mut self, h: f64) {
		self.evolve_a(h / 2.0);
		self.evolve_b(h / 2.0);
		self.evolve_c(h);
		self.evolve_b(h / 2.0);
		self.evolve_a(h / 2.0);
	}

	fn evolve_a(&mut self, h: f64) {
		let dots = State { q: self.a.q, p: self.b.p }.dots();
		self.a.p += dots.p * h;
		self.b.q += dots.q * h;
	}

	fn evolve_b(&mut self, h: f64) {
		let dots = State { q: self.b.q, p: self.a.p }.dots();
		self.a.q += dots.q * h;
		self.b.p += dots.p * h;
	}

	fn evolve_c(&mut self, h: f64) {
		let alpha = 2.0 * self.omega * h;
		let c = alpha.cos();
		let s = alpha.sin();

		let diff = self.a - self.b;
		let r = State {
			q: diff.q * c + diff.p * s,
			p: diff.p * c - diff.q * s,
		};

		let sum = self.a + self.b;
		self.a = (sum + r) / 2.0;
		self.b = (sum - r) / 2.0;
	}
}

impl<T: Add<Output = T>> Add for Augmented<T> {
	type Output = Self;

	fn add(self, rhs: Self) -> Self {
		Self {
			omega: self.omega,
			a: self.a + rhs.a,
			b: self.b + rhs.b,
		}
	}
}

impl<T: Sub<Output = T>> Sub for Augmented<T> {
	type Output = Self;

	fn sub(self, rhs: Self) -> Self {
		Self {
			omega: self.omega,
			a: self.a - rhs.a,
			b: self.b - rhs.b,
		}
	}
}

impl<T: Mul<f64, Output = T>> Mul<f64> for Augmented<T> {
	type Output = Self;

	fn mul(self, rhs: f64) -> Self {
		Self {
			omega: self.omega,
			a: self.a * rhs,
			b: self.b * rhs,
		}
	}
}

impl<T: Div<f64, Output = T>> Div<f64> for Augmented<T> {
	type Output = Self;

	fn div(self, rhs: f64) -> Self {
		Self {
			omega: self.omega,
			a: self.a / rhs,
			b: self.b / rhs,
		}
	}
}