#pragma once

#include <psemek/math/detail/array_2d.hpp>
#include <psemek/math/vector.hpp>
#include <psemek/math/math.hpp>
#include <psemek/util/span.hpp>

#include <iostream>
#include <iomanip>
#include <algorithm>

namespace psemek::math
{

	template <typename T, std::size_t R, std::size_t C>
	struct matrix
	{
		using scalar_type = T;
		static constexpr std::size_t static_rows = R;
		static constexpr std::size_t static_columns = C;

		using dynamic_dimensions_type = detail::dynamic_dimensions_2d<R, C>;

		typename detail::array_2d<T, R, C>::type coords;

		// Internal
		matrix(dynamic_dimensions_type dimensions);

		matrix();

		template <typename ... Args>
		matrix(Args && ... args) requires (R != dynamic && C != dynamic && sizeof...(Args) == R * C && detail::all_convertible_to<T, Args...>::value);

		explicit matrix(std::size_t rows) requires(R == dynamic && C != dynamic);
		explicit matrix(std::size_t columns) requires(R != dynamic && C == dynamic);
		matrix(std::size_t rows, std::size_t columns) requires(R == dynamic && C == dynamic);

		std::size_t rows() const
		{
			return coords.rows;
		}

		std::size_t columns() const
		{
			return coords.columns;
		}

		auto operator[](std::size_t i)
		{
			return coords[i];
		}

		auto operator[](std::size_t i) const
		{
			return coords[i];
		}

		util::span<T> values()
		{
			return {&coords[0][0], rows() * columns()};
		}

		util::span<T const> values() const
		{
			return {&coords[0][0], rows() * columns()};
		}

		matrix copy() const
		{
			if constexpr (R == dynamic || C == dynamic)
			{
				matrix r;
				r.coords = coords.copy();
				return r;
			}
			else
			{
				return *this;
			}
		}

		matrix & operator *= (T const & s);
		matrix & operator /= (T const & s);

		matrix & operator += (matrix const & v);
		matrix & operator -= (matrix const & v);

		static matrix zero(dynamic_dimensions_type dimensions = {});
		static matrix identity(dynamic_dimensions_type dimensions = {});
		static matrix scalar(T const & s, dynamic_dimensions_type dimensions = {});
		static matrix diagonal(vector<T, std::min(R, C)> const & d);
	};

	template <typename T, std::size_t R, std::size_t C>
	void check_dynamic_size(matrix<T, R, C> const & m1, matrix<T, R, C> const & m2)
	{
		check_dynamic_size(m1.coords.dimensions(), m2.coords.dimensions());
	}

	template <typename T, std::size_t R, std::size_t C>
	matrix<T, R, C>::matrix(dynamic_dimensions_type dimensions)
		: coords(dimensions)
	{}

	template <typename T, std::size_t R, std::size_t C>
	matrix<T, R, C>::matrix()
		: matrix(dynamic_dimensions_type{})
	{}

	template <typename T, std::size_t R, std::size_t C>
	template <typename ... Args>
	matrix<T, R, C>::matrix(Args && ... args) requires (R != dynamic && C != dynamic && sizeof...(Args) == R * C && detail::all_convertible_to<T, Args...>::value)
		: coords({})
	{
		auto out = values().begin();
		((*out++ = args), ...);
	}

	template <typename T, std::size_t R, std::size_t C>
	matrix<T, R, C>::matrix(std::size_t rows) requires(R == dynamic && C != dynamic)
		: coords({.rows = rows})
	{}

	template <typename T, std::size_t R, std::size_t C>
	matrix<T, R, C>::matrix(std::size_t columns) requires(R != dynamic && C == dynamic)
		: coords({.columns = columns})
	{}

	template <typename T, std::size_t R, std::size_t C>
	matrix<T, R, C>::matrix(std::size_t rows, std::size_t columns) requires(R == dynamic && C == dynamic)
		: coords({.rows = rows, .columns = columns})
	{}

	template <typename T, std::size_t R, std::size_t C>
	matrix<T, R, C> matrix<T, R, C>::zero(dynamic_dimensions_type dimensions)
	{
		matrix<T, R, C> m(dimensions);
		for (auto & v : m.values())
			v = T(0);
		return m;
	}

	template <typename T, std::size_t R, std::size_t C>
	matrix<T, R, C> matrix<T, R, C>::identity(dynamic_dimensions_type dimensions)
	{
		return scalar(T(1), dimensions);
	}

	template <typename T, std::size_t R, std::size_t C>
	matrix<T, R, C> matrix<T, R, C>::scalar(T const & s, dynamic_dimensions_type dimensions)
	{
		matrix<T, R, C> m = zero(dimensions);
		for (std::size_t i = 0; i < std::min(m.rows(), m.columns()); ++i)
			m[i][i] = s;
		return m;
	}

	template <typename T, std::size_t R, std::size_t C>
	matrix<T, R, C> matrix<T, R, C>::diagonal(vector<T, std::min(R, C)> const & d)
	{
		dynamic_dimensions_type dimensions;
		if constexpr (R == dynamic)
		{
			dimensions.rows = d.dimension();
		}
		if constexpr (C == dynamic)
		{
			dimensions.columns = d.dimension();
		}

		matrix<T, R, C> m = zero(dimensions);
		for (std::size_t i = 0; i < std::min(R, C); ++i)
			m[i][i] = d[i];
		return m;
	}

	template <typename T1, typename T, std::size_t R, std::size_t C>
	matrix<T1, R, C> cast(matrix<T, R, C> const & m)
	{
		matrix<T1, R, C> r(m.coords.dimensions());

		auto out = r.values().begin();
		for (auto const & value : m.values())
			*out++ = static_cast<T1>(value);

		return r;
	}

	template <typename T, std::size_t R, std::size_t C>
	std::strong_ordering operator <=> (matrix<T, R, C> const & m1, matrix<T, R, C> const & m2)
	{
		check_dynamic_size(m1, m2);
		auto begin1 = m1.values().begin();
		auto begin2 = m2.values().begin();
		auto end1 = m1.values().end();

		// Intentionally don't use std::lexicographical_compare_three_way to fake strong ordering
		// (actual ordering might be partial & weak due to NaN's and signed zeros
		while (begin1 != end1)
		{
			if (*begin1 < *begin2)
				return std::strong_ordering::less;
			else if (*begin1 > *begin2)
				return std::strong_ordering::greater;
			++begin1;
			++begin2;
		}
		return std::strong_ordering::equal;
	}

	template <typename T, std::size_t R, std::size_t C>
	bool operator == (matrix<T, R, C> const & m1, matrix<T, R, C> const & m2)
	{
		return (m1 <=> m2) == std::strong_ordering::equal;
	}

	template <typename T, std::size_t R, std::size_t C>
	bool operator != (matrix<T, R, C> const & m1, matrix<T, R, C> const & m2)
	{
		return !(m1 == m2);
	}

	template <typename T, std::size_t R, std::size_t C>
	bool operator < (matrix<T, R, C> const & m1, matrix<T, R, C> const & m2)
	{
		return (m1 <=> m2) == std::strong_ordering::less;
	}

	template <typename T, std::size_t R, std::size_t C>
	bool operator > (matrix<T, R, C> const & m1, matrix<T, R, C> const & m2)
	{
		return (m2 < m1);
	}

	template <typename T, std::size_t R, std::size_t C>
	bool operator <= (matrix<T, R, C> const & m1, matrix<T, R, C> const & m2)
	{
		return !(m2 < m1);
	}

	template <typename T, std::size_t R, std::size_t C>
	bool operator >= (matrix<T, R, C> const & m1, matrix<T, R, C> const & m2)
	{
		return !(m1 < m2);
	}

	template <typename T, std::size_t R, std::size_t C>
	matrix<T, R, C> operator * (matrix<T, R, C> const & m, T const & s)
	{
		matrix<T, R, C> r(m.coords.dimensions());
		auto out = r.values().begin();
		for (auto const & value : m.values())
			*out++ = value * s;
		return r;
	}

	template <typename T, std::size_t R, std::size_t C>
	matrix<T, R, C> operator * (T const & s, matrix<T, R, C> const & m)
	{
		matrix<T, R, C> r(m.coords.dimensions());
		auto out = r.values().begin();
		for (auto const & value : m.values())
			*out++ = s * value;
		return r;
	}

	template <typename T, std::size_t R, std::size_t C>
	matrix<T, R, C> operator / (matrix<T, R, C> const & m, T const & s)
	{
		matrix<T, R, C> r(m.coords.dimensions());
		auto out = r.values().begin();
		for (auto const & value : m.values())
			*out++ = value / s;
		return r;
	}

	template <typename T, std::size_t R, std::size_t C>
	matrix<T, R, C> & matrix<T, R, C>::operator *= (T const & s)
	{
		*this = *this * s;
		return *this;
	}

	template <typename T, std::size_t R, std::size_t C>
	matrix<T, R, C> & matrix<T, R, C>::operator /= (T const & s)
	{
		*this = *this / s;
		return *this;
	}

	template <typename T, std::size_t R, std::size_t C>
	matrix<T, R, C> operator - (matrix<T, R, C> const & m)
	{
		matrix<T, R, C> r(m.coords.dimensions());
		auto out = r.values().begin();
		for (auto const & value : m.values())
			*out++ = -value;
		return r;
	}

	template <typename T, std::size_t R, std::size_t C>
	matrix<T, R, C> operator + (matrix<T, R, C> const & m1, matrix<T, R, C> const & m2)
	{
		check_dynamic_size(m1, m2);
		matrix<T, R, C> r(m1.coords.dimensions());
		auto in1 = m1.values().begin();
		auto in2 = m2.values().begin();
		for (auto & value : r.values())
			value = (*in1++) + (*in2++);
		return r;
	}

	template <typename T, std::size_t R, std::size_t C>
	matrix<T, R, C> operator - (matrix<T, R, C> const & m1, matrix<T, R, C> const & m2)
	{
		check_dynamic_size(m1, m2);
		matrix<T, R, C> r(m1.coords.dimensions());
		auto in1 = m1.values().begin();
		auto in2 = m2.values().begin();
		for (auto & value : r.values())
			value = (*in1++) - (*in2++);
		return r;
	}

	template <typename T, std::size_t R, std::size_t C>
	matrix<T, R, C> & matrix<T, R, C>::operator += (matrix<T, R, C> const & m)
	{
		*this = *this + m;
		return *this;
	}

	template <typename T, std::size_t R, std::size_t C>
	matrix<T, R, C> & matrix<T, R, C>::operator -= (matrix<T, R, C> const & m)
	{
		*this = *this - m;
		return *this;
	}

	template <typename T, std::size_t R, std::size_t C>
	vector<T, R> operator * (matrix<T, R, C> const & m, vector<T, C> const & v)
	{
		detail::check_dynamic_size(m.coords.dimensions().columns, v.coords.dimensions());
		auto r = vector<T, R>::zero(m.coords.dimensions().rows);
		for (std::size_t i = 0; i < m.rows(); ++i)
		{
			for (std::size_t j = 0; j < m.columns(); ++j)
				r[i] += m[i][j] * v[j];
		}
		return r;
	}

	template <typename T, std::size_t R, std::size_t C>
	vector<T, C> operator * (vector<T, R> const & v, matrix<T, R, C> const & m)
	{
		detail::check_dynamic_size(v.coords.dimensions(), m.coords.dimensions().rows);
		auto r = vector<T, C>::zero(m.coords.dimensions().columns);
		for (std::size_t i = 0; i < m.rows(); ++i)
		{
			for (std::size_t j = 0; j < m.columns(); ++j)
				r[j] += v[i] * m[i][j];
		}
		return r;
	}

	template <typename T, std::size_t R, std::size_t K, std::size_t C>
	matrix<T, R, C> operator * (matrix<T, R, K> const & m1, matrix<T, K, C> const & m2)
	{
		detail::check_dynamic_size(m1.coords.dimensions().columns, m2.coords.dimensions().rows);
		auto r = matrix<T, R, C>::zero();
		for (std::size_t i = 0; i < m1.rows(); ++i)
		{
			for (std::size_t j = 0; j < m2.columns(); ++j)
			{
				for (std::size_t k = 0; k < m1.columns(); ++k)
					r[i][j] += m1[i][k] * m2[k][j];
			}
		}
		return r;
	}

	template <typename T, std::size_t R, std::size_t C>
	matrix<T, C, R> transpose(matrix<T, R, C> const & m)
	{
		matrix<T, C, R> r({.rows = m.coords.dimensions().columns, .columns = m.coords.dimensions().rows});
		for (std::size_t i = 0; i < m.rows(); ++i)
			for (std::size_t j = 0; j < m.columns(); ++j)
				r[j][i] = m[i][j];
		return r;
	}

	template <typename T, std::size_t N, typename ... Rows>
	auto by_rows(vector<T, N> const & v, Rows const & ... rows)
	{
		(check_dynamic_size(v, rows), ...);

		vector<T, N> m[] = {v, rows...};

		matrix<T, sizeof...(Rows) + 1, N> r({.columns = v.coords.dimensions()});
		for (std::size_t i = 0; i < r.rows(); ++i)
			for (std::size_t j = 0; j < r.columns(); ++j)
				r[i][j] = m[i][j];
		return r;
	}

	template <typename T, std::size_t N, typename ... Columns>
	auto by_columns(vector<T, N> const & v, Columns const & ... columns)
	{
		(check_dynamic_size(v, columns), ...);
		vector<T, N> m[] = {v, columns...};

		matrix<T, N, sizeof...(Columns) + 1> r({.rows = v.coords.dimensions()});
		for (std::size_t i = 0; i < r.rows(); ++i)
			for (std::size_t j = 0; j < r.columns(); ++j)
				r[i][j] = m[j][i];
		return r;
	}

	template <typename T, std::size_t R, std::size_t C>
	vector<T, C> row(matrix<T, R, C> const & m, std::size_t i)
	{
		vector<T, C> r(m.coords.dimensions().columns);
		for (std::size_t j = 0; j < m.columns(); ++j)
			r[j] = m[i][j];
		return r;
	}

	template <typename T, std::size_t R, std::size_t C>
	vector<T, R> column(matrix<T, R, C> const & m, std::size_t j)
	{
		vector<T, R> r(m.coords.dimensions().rows);
		for (std::size_t i = 0; i < m.rows(); ++i)
			r[i] = m[i][j];
		return r;
	}

	template <typename T, std::size_t R, std::size_t C>
	T frobenius_norm_sqr(matrix<T, R, C> const & m)
	{
		T r{0};
		for (auto const & value : m.values())
			r += sqr(value);
		return r;
	}

	template <typename T, std::size_t R, std::size_t C>
	T linf_norm(matrix<T, R, C> const & m)
	{
		using std::abs;
		using std::max;
		T r = abs(m[0][0]);
		for (auto const & value : m.values())
			make_max(r, abs(value));
		return r;
	}

	template <typename T, std::size_t R, std::size_t C>
	T frobenius_norm(matrix<T, R, C> const & m)
	{
		T const max = linf_norm(m);
		if (max == T(0)) return max;

		T sum{};
		for (auto const & value : m.values())
			sum += sqr(value / max);

		using std::sqrt;
		return max * sqrt(sum);
	}

	template <typename T, std::size_t R, std::size_t C>
	T trace(matrix<T, R, C> const & m)
	{
		T r{0};
		for (std::size_t i = 0; i < std::min(m.rows(), m.columns()); ++i)
			r += m[i][i];
		return r;
	}

	template <std::size_t R1, std::size_t C1, std::size_t R2, std::size_t C2, typename T, std::size_t R, std::size_t C>
	requires (R1 != dynamic && C1 != dynamic && R2 != dynamic && C2 != dynamic)
	matrix<T, R2 - R1, C2 - C1> submatrix(matrix<T, R, C> const & m)
	{
		static_assert(R1 < R2);
		static_assert(C1 < C2);
		static_assert(R2 <= R);
		static_assert(C2 <= C);

		matrix<T, R2 - R1, C2 - C1> result;
		for (std::size_t i = 0; i < result.rows(); ++i)
			for (std::size_t j = 0; j < result.columns(); ++j)
				result[i][j] = m[i + R1][j + C1];
		return result;
	}

	template <typename T>
	matrix<T, 3, 3> cross_product_matrix(vector<T, 3> const & v)
	{
		auto r = matrix<T, 3, 3>::zero();
		r[0][1] = -v[2];
		r[0][2] = v[1];
		r[1][0] = v[2];
		r[1][2] = -v[0];
		r[2][0] = -v[1];
		r[2][1] = v[0];
		return r;
	}

	template <typename T, std::size_t R, std::size_t C>
	matrix<T, R, C> outer_product(vector<T, R> const & v1, vector<T, C> const & v2)
	{
		matrix<T, R, C> r({.rows = v1.coords.dimensions(), .columns = v2.coords.dimensions()});
		for (std::size_t i = 0; i < r.rows(); ++i)
			for (std::size_t j = 0; j < r.columns(); ++j)
				r[i][j] = v1[i] * v2[j];
		return r;
	}

	template <typename T, std::size_t R, std::size_t C>
	std::ostream & operator << (std::ostream & os, matrix<T, R, C> const & m)
	{
		for (std::size_t i = 0; i < m.rows(); ++i)
		{
			for (std::size_t j = 0; j < m.columns(); ++j)
				os << m[i][j] << ' ';
			os << '\n';
		}
		return os;
	}

	template <typename T, std::size_t R, std::size_t C>
	struct setw
	{
		matrix<T, R, C> const & m;
		int w;
	};

	template <typename T, std::size_t R, std::size_t C>
	std::ostream & operator << (std::ostream & os, setw<T, R, C> const & w)
	{
		for (std::size_t i = 0; i < w.m.rows(); ++i)
		{
			for (std::size_t j = 0; j < w.m.columns(); ++j)
				os << std::fixed << std::right << std::setw(w.w) << w.m[i][j] << ' ';
			os << '\n';
		}
		return os;
	}

}

namespace std
{

	template <typename T, std::size_t R, std::size_t C, typename Char>
	struct formatter<::psemek::math::matrix<T, R, C>, Char>
		: formatter<T, Char>
	{
		using formatter<T, Char>::parse;

		template <typename FormatContext>
		auto format(::psemek::math::matrix<T, R, C> const & m, FormatContext & ctx) const
		{
			for (std::size_t i = 0; i < m.rows(); ++i)
			{
				for (std::size_t j = 0; j < m.columns(); ++j)
				{
					ctx.advance_to(formatter<T, Char>::format(m[i][j], ctx));
					ctx.advance_to(std::format_to(ctx.out(), " "));
				}
				ctx.advance_to(std::format_to(ctx.out(), "\n"));
			}
			return ctx.out();
		}
	};

}
