#pragma once

#include <psemek/math/vector.hpp>

#include <iostream>
#include <compare>

namespace psemek::math
{

	template <typename T, std::size_t N>
	struct dual
	{
		T scalar = T{0};
		vector<T, N> delta = vector<T, N>::zero();

		dual() = default;

		dual(T scalar)
			: scalar{std::move(scalar)}
		{}

		dual(T scalar, vector<T, N> delta)
			: scalar{std::move(scalar)}
			, delta{std::move(delta)}
		{}

		dual & operator += (dual const & d);
		dual & operator -= (dual const & d);
		dual & operator *= (dual const & d);
		dual & operator /= (dual const & d);
	};

	template <typename T, std::size_t N>
	dual<T, N> operator + (dual<T, N> const & d)
	{
		return d;
	}

	template <typename T, std::size_t N>
	dual<T, N> operator - (dual<T, N> const & d)
	{
		return {-d.scalar, -d.delta};
	}

	template <typename T, std::size_t N>
	dual<T, N> operator + (dual<T, N> const & d1, dual<T, N> const & d2)
	{
		return {d1.scalar + d2.scalar, d1.delta + d2.delta};
	}

	template <typename T, std::size_t N>
	dual<T, N> operator - (dual<T, N> const & d1, dual<T, N> const & d2)
	{
		return {d1.scalar - d2.scalar, d1.delta - d2.delta};
	}

	template <typename T, std::size_t N>
	dual<T, N> operator * (dual<T, N> const & d1, dual<T, N> const & d2)
	{
		return {d1.scalar * d2.scalar, d1.scalar * d2.delta + d2.scalar * d1.delta};
	}

	template <typename T, std::size_t N>
	dual<T, N> inverse(dual<T, N> const & d)
	{
		return {T{1} / d.scalar, - d.delta / d.scalar / d.scalar};
	}

	template <typename T, std::size_t N>
	dual<T, N> operator / (dual<T, N> const & d1, dual<T, N> const & d2)
	{
		return d1 * inverse(d2);
	}

	template <typename T, std::size_t N>
	dual<T, N> & dual<T, N>::operator += (dual<T, N> const & d)
	{
		*this = *this + d;
		return *this;
	}

	template <typename T, std::size_t N>
	dual<T, N> & dual<T, N>::operator -= (dual const & d)
	{
		*this = *this - d;
		return *this;
	}

	template <typename T, std::size_t N>
	dual<T, N> & dual<T, N>::operator *= (dual const & d)
	{
		*this = *this * d;
		return *this;
	}

	template <typename T, std::size_t N>
	dual<T, N> & dual<T, N>::operator /= (dual const & d)
	{
		*this = *this / d;
		return *this;
	}

	template <typename T, std::size_t N>
	dual<T, N> operator + (dual<T, N> const & d, T const & v)
	{
		return {d.scalar + v, d.delta};
	}

	template <typename T, std::size_t N>
	dual<T, N> operator + (T const & v, dual<T, N> const & d)
	{
		return {d.scalar + v, d.delta};
	}

	template <typename T, std::size_t N>
	dual<T, N> operator - (dual<T, N> const & d, T const & v)
	{
		return {d.scalar - v, d.delta};
	}

	template <typename T, std::size_t N>
	dual<T, N> operator - (T const & v, dual<T, N> const & d)
	{
		return {v - d.scalar, -d.delta};
	}

	template <typename T, std::size_t N>
	dual<T, N> operator * (dual<T, N> const & d, T const & v)
	{
		return {d.scalar * v, d.delta * v};
	}

	template <typename T, std::size_t N>
	dual<T, N> operator * (T const & v, dual<T, N> const & d)
	{
		return {d.scalar * v, d.delta * v};
	}

	template <typename T, std::size_t N>
	dual<T, N> operator / (dual<T, N> const & d, T const & v)
	{
		return {d.scalar / v, d.delta / v};
	}

	template <typename T, std::size_t N>
	dual<T, N> operator / (T const & v, dual<T, N> const & d)
	{
		return v * inverse(d);
	}

	// Exists mostly to make max/min/clamp/etc work
	template <typename T, std::size_t N>
	std::partial_ordering operator <=> (dual<T, N> const & d1, dual<T, N> const & d2)
	{
		return d1.scalar <=> d2.scalar;
	}

	template <typename T, std::size_t N>
	dual<T, N> abs(dual<T, N> const & d)
	{
		return d.scalar > T{0} ? d : -d;
	}

	template <typename T, std::size_t N>
	dual<T, N> exp(dual<T, N> const & d)
	{
		auto e = std::exp(d.scalar);
		return {e, e * d.delta};
	}

	template <typename T, std::size_t N>
	dual<T, N> log(dual<T, N> const & d)
	{
		return {std::log(d.scalar), d.delta / d.scalar};
	}

	template <typename T, std::size_t N>
	dual<T, N> sqrt(dual<T, N> const & d)
	{
		auto s = std::sqrt(d.scalar);
		return {s, d.delta / (T{2} * s)};
	}

	template <typename T, std::size_t N>
	dual<T, N> sin(dual<T, N> const & d)
	{
		return {std::sin(d.scalar), std::cos(d.scalar) * d.delta};
	}

	template <typename T, std::size_t N>
	dual<T, N> cos(dual<T, N> const & d)
	{
		return {std::cos(d.scalar), - std::sin(d.scalar) * d.delta};
	}

	template <typename T, std::size_t N>
	dual<T, N> pow(dual<T, N> const & d1, dual<T, N> const & d2)
	{
		return exp(log(d1) * d2);
	}

	template <typename T, std::size_t N>
	std::ostream & operator << (std::ostream & os, dual<T, N> const & d)
	{
		os << d.scalar;
		for (std::size_t i = 0; i < N; ++i)
			os << " + " << d.delta[i] << " e_" << i;
		return os;
	}

}
