#include "control.h"
#include <armadillo>
#include <cassert>
#include <cmath>


namespace ct
{
	TransferFn::TransferFn()
		: num(1)
		, den(1)
	{
		num(0) = 1;
		den(0) = 1;
	}

	TransferFn::TransferFn(math::PolyCx num, math::PolyCx den)
		: num(num)
		, den(den)
	{}

	complex TransferFn::dc_gain() const
	{
		arma::cx_vec z(1, arma::fill::zeros);
		arma::cx_vec n = arma::polyval(num.coeffs, z);
		arma::cx_vec d = arma::polyval(den.coeffs, z);
		return n(0) / d(0);
	}

	bool TransferFn::is_proper() const
	{
		return den.degree() >= num.degree();
	}

	bool TransferFn::is_strictly_proper() const
	{
		return den.degree() >= num.degree();
	}

	TransferFn operator + (const TransferFn& g, const TransferFn& h)
	{
		TransferFn tf(g.num * h.den + h.num * g.den, g.den * g.den);
		tf.canonicalize();
		return tf;
	}

	TransferFn operator - (const TransferFn& g, const TransferFn& h)
	{
		TransferFn tf = g + (-1. * h);
		tf.canonicalize();
		return tf;
	}

	TransferFn operator * (const TransferFn& g, const TransferFn& h)
	{
		TransferFn tf(g.num * h.num, g.den * h.den);
		tf.canonicalize();
		return tf;
	}

	TransferFn operator / (const TransferFn& g, const TransferFn& h)
	{
		TransferFn hinv = TransferFn(h.den, h.num);
		hinv.canonicalize();
		TransferFn tf = g * hinv;
		tf.canonicalize();
		return tf;
	}

	TransferFn operator * (const complex k, const TransferFn& h)
	{
		TransferFn tf(k * h.num, h.den);
		tf.canonicalize();
		return tf;
	}

	TransferFn operator / (const TransferFn& h, const complex k)
	{
		TransferFn tf((1. / k) * h.num, h.den);
		tf.canonicalize();
		return tf;
	}

	TransferFn cancel_zp(const TransferFn& tf, double tol)
	{
		math::PolyCx num(tf.num.coeffs.n_elem);
		math::PolyCx den(tf.den.coeffs.n_elem);
		// TODO: return new tf without poles & zeros that are closer to each other than tol
	}

	TransferFn feedback(const TransferFn& tf, complex k)
	{
		const TransferFn unit(math::PolyCx({1.}), math::PolyCx({1.}));
		const TransferFn ol = k * tf;
		const TransferFn bo = unit - ol;
		const TransferFn fb = unit / bo;
		TransferFn tf_cl = ol * fb;
		// const TransferFn tf_cl = (k * tf) / (unit - k * tf);
		tf_cl.canonicalize();
		return tf_cl;
	}

	void TransferFn::canonicalize()
	{
		complex an = den(0);
		if (std::abs(an) > 0)
		{
			num.coeffs /= an;
			den.coeffs /= an;
		}
	}

	TransferFn& TransferFn::operator = (const TransferFn& other)
	{
		if (this == &other)
			return *this;

		num = math::PolyCx(other.num);
		den = math::PolyCx(other.den);
		return *this;
	}

	LocusSeries::LocusSeries(double start, double end, size_t nsamples)
		: n_samples(nsamples)
		, start(start)
		, end(end)
		, in(arma::logspace(log10(start), log10(end), n_samples))
		, out(1, n_samples, arma::fill::zeros)
	{}

	void rlocus(const TransferFn& tf, LocusSeries& ls)
	{
		using namespace arma;

		CT_ASSERT(tf.is_proper());
		CT_ASSERT(tf.den.degree() > 0);

		// prepare output
		ls.out = cx_mat(tf.den.degree(), ls.n_samples, fill::zeros);

		// compute roots
		for (int i = 0; i < ls.n_samples; i++)
			// FIXME: sort the roots
			ls.out.col(i) = (tf.den + ls.in(i) * tf.num).roots();
	}

	void nyquist(const TransferFn& tf, LocusSeries& ls)
	{
		using namespace arma;
		complex j(0., 1.);
		
		// FIXME missing the countour at infinity and the small contours for poles
		// on the imaginary axis

		// contour from 0 to i\infty
		cx_mat num = polyval(tf.num.coeffs, j * conv_to<cx_vec>::from(ls.in));
		cx_mat den = polyval(tf.den.coeffs, j * conv_to<cx_vec>::from(ls.in));

		if (ls.out.n_rows != ls.n_samples)
			ls.out = cx_mat(1, ls.n_samples);

		for (int i = 0; i < ls.n_samples; i++)
			ls.out(i) = num(i) / den(i);
	}

	SSModel::SSModel(size_t n_in, size_t n_out, size_t n_states)
		: n_in(n_in)
		, n_out(n_out)
		, n_states(n_states)
		, A(n_states, n_states)
		, B(n_states, n_in)
		, C(n_out, n_states)
		, D(n_out, n_in)
	{}

	SSModel ctrb_form(const TransferFn& tf)
	{
		// TODO: change to proper by implementing D
		CT_ASSERT(tf.is_strictly_proper());

		int ord = tf.den.degree();
		SSModel ss(1, 1, ord);

		ss.B(ord - 1) = 1;
		for (int i = 0; i < ord; i++)
		{
			ss.A(ord - 1, i) = tf.den(ord - i);
			if (i < tf.num.coeffs.n_elem)
				ss.C(i) = tf.num(i);
		}
	
		return ss;
	}

	TimeSeries::TimeSeries(double start, double end, size_t n_samples)
		: n_samples(n_samples)
		, start(start)
		, end(end)
		, dt((start - end) / (n_samples - 1))
		, time(arma::linspace(start, end, n_samples))
			// FIXME: do not hardcode SISO
		, in(1, n_samples, arma::fill::zeros)
		, out(1, n_samples, arma::fill::zeros)
	{}

	void response(const SSModel& ss, TimeSeries& ts)
	{
		using namespace arma;

		CT_ASSERT(ts.in.n_rows == ss.n_in);
		CT_ASSERT(ts.in.n_cols == ts.n_samples);

		ts.out = cx_mat(ss.n_out, ts.n_samples, fill::zeros);
		ts.state = cx_mat(ss.n_states, ts.n_samples, fill::zeros);
		// FIXME: non-homogeneous initial condition

		// For the current application we want this to be faster rather than
		// accurate. Hence we use a simulation with ZOH input since it is cheap and
		// probably good enough.
		// FIXME: do not invert matrix like that, but I am too lazy to implement
		// regularization
		const int n_steps = 3;
		const cx_mat Ad = expmat(ss.A * ts.dt / n_steps);
		const cx_mat Bd = (Ad - eye(size(Ad))) * inv(ss.A) * ss.B;

		cx_vec x(ss.n_states), xn;
		for (int k = 0; k < ts.n_samples - 1; k++)
		{
			x = ts.state.col(k);
			for (int l = 0; l < n_steps; l++)
				xn = Ad * x + Bd * ts.in.col(k);
			ts.state.col(k + 1) = xn;
		}

		ts.out = ss.C * ts.state + ss.D * ts.in;
	}


	void step(const SSModel& ss, TimeSeries& ts)
	{
		ts.in = arma::cx_mat(ss.n_in, ts.n_samples);
		ts.in.fill(1.);
		response(ss, ts);
	}
}

// vim: ts=2 sw=2 noet: