#pragma once

#include "armadillo/include/armadillo"

#include <utility>
#include <complex>
#include <cassert>
#include <cmath>
#include <cstddef>

#ifndef CONTROL_H
#define CONTROL_H

#define CT_ASSERT(x) (assert(x))
#define CT_SMALL_TOL 1e-6
#define CT_SMALL(x) (std::abs(x) < CT_SMALL_TOL)
#define CT_ASSERT_SMALL(x) (assert(CT_SMALL))

namespace ct
{
	typedef std::complex<double> complex;
	typedef std::size_t size_t;

	namespace math
	{
		template<typename T>
		struct Poly
		{
			arma::Col<T> coeffs;

			Poly(size_t n_elem) : coeffs(n_elem) {}
			Poly(arma::Col<T> c) : coeffs(c) {}

			template<typename U>
			Poly(std::initializer_list<U> l) : coeffs(l.size())
			{
				int i = 0;
				for (const U& u : l)
					coeffs(i++) = T(u);
			}

			int degree() const { return coeffs.n_elem - 1; }
			void add_root(T r) { (*this) = (*this) * Poly({1., -r}); }
			arma::Col<T> roots() const { return arma::roots(coeffs); }

			template<typename IndexT>
			T& operator () (IndexT idx) { return coeffs(idx); };

			template<typename IndexT>
			T operator () (IndexT idx) const { return coeffs(idx); };
		};
		
		template<typename T>
		Poly<T> operator * (const Poly<T>& p, const Poly<T>& q)
		{
			arma::Col<T> coeffs = arma::conv(p.coeffs, q.coeffs);
			CT_ASSERT(!coeffs.has_nan());
			return Poly<T>(coeffs);
		}

		template<typename T, typename R = T>
		Poly<T> operator * (R scalar, const Poly<T>& p)
		{
			Poly<T> q(p.coeffs * scalar);
			CT_ASSERT(!q.coeffs.has_nan());
			return q;
		}

		template<typename T, typename R = T>
		Poly<T> operator / (R scalar, const Poly<T>& p)
		{
			Poly<T> q(p.coeffs / scalar);
			CT_ASSERT(!q.coeffs.has_nan());
			return q;
		}

		template<typename T>
		Poly<T> operator + (const Poly<T>& p, const Poly<T>& q)
		{
			const Poly<T>& big = (p.degree() > q.degree()) ? p : q;
			const Poly<T>& small = (p.degree() > q.degree()) ? q : p;

			Poly<T> s(small);
			int rel = big.degree() - small.degree();
			s.coeffs.insert_rows(0, rel);

			CT_ASSERT(!s.coeffs.has_nan());
			return Poly<T>(big.coeffs + s.coeffs);
		}

		template<typename T>
		Poly<T> operator - (const Poly<T>& p, const Poly<T>& q)
		{
			return p + (-1. * q);
		}

		typedef Poly<complex> PolyCx;
	}

	struct TransferFn
	{
		math::PolyCx num;
		math::PolyCx den;

		TransferFn(void);
		TransferFn(math::PolyCx num, math::PolyCx den);
		TransferFn(const TransferFn& other)
			: num(other.num)
			, den(other.den) {}

		inline void add_pole(complex p) { den.add_root(p); }
		inline void add_zero(complex z) { num.add_root(z); };

		complex dc_gain() const;

		bool is_proper() const;
		bool is_strictly_proper() const;

		void canonicalize();

		TransferFn& operator = (const TransferFn& other);
	};

	TransferFn operator + (const TransferFn& g, const TransferFn& h);
	TransferFn operator - (const TransferFn& g, const TransferFn& h);
	TransferFn operator * (const TransferFn& g, const TransferFn& h);
	TransferFn operator / (const TransferFn& g, const TransferFn& h);
	TransferFn operator * (const complex k, const TransferFn& h);
	TransferFn operator / (const TransferFn& h, const complex k);

	// inline TransferFn operator / (const TransferFn&

	TransferFn cancel_zp(const TransferFn& tf, double tol = CT_SMALL_TOL);
	TransferFn feedback(const TransferFn& tf, complex k = -1);

	struct LocusSeries
	{
		size_t n_samples;
		double start, end;
		arma::vec in;
		arma::cx_mat out;

		LocusSeries() = delete;
		LocusSeries(double start, double end, size_t n_samples);
	};

	void rlocus(const TransferFn& tf, LocusSeries& ls);
	void nyquist(const TransferFn& tf, LocusSeries& ls);

	struct SSModel
	{
		size_t n_in, n_out, n_states; 
		arma::cx_mat A, B, C, D;

		SSModel() = delete;
		SSModel(size_t n_in, size_t n_out, size_t n_states);
	};

	SSModel ctrb_form(const TransferFn& tf);
	// SSModel obsv_form(const TransferFn& tf);
	// SSModel eigm_form(const TransferFn& tf);

	struct TimeSeries
	{
		size_t n_samples;
		double start, end, dt;
		arma::vec time;
		arma::cx_mat in, out;
		arma::cx_mat state;

		TimeSeries() = delete;
		TimeSeries(double start, double end, size_t n_samples);
	};

	void response(const SSModel& ss, TimeSeries& ts);
	void step(const SSModel& ss, TimeSeries& ts);
}

#endif // CONTROL_H
// vim:ts=2 sw=2 noet: