#pragma once

#include <armadillo>
#include <complex>
#include <cassert>
#include <cmath>

#ifndef CONTROL_H
#define CONTROL_H

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

namespace ct
{
	using complex = std::complex<double>;

	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) {}

			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);
			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);
			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(big.degree() + 1);

			for (int i = 0; i <= small.degree(); i++)
				s(i) = big(i) + small(i);

			for (int i = small.degree() + 1; i <= big.degree(); i++)
				s(i) = big(i);

			return s;
		}
	}

	struct TransferFn
	{
		complex gain;
		math::Poly<complex> num;
		math::Poly<complex> den;

		TransferFn(complex gain);

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

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

	TransferFn feedback(const TransferFn& tf, double 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);


	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: