#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 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;
		const TransferFn tf_cl = ol * fb;
		// const TransferFn tf_cl = (k * tf) / (unit - k * tf);
		return tf_cl;
	}

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

	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))
		// FIXME: do not hardcode one root locus
		, 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++)
			ls.out.col(i) = (tf.den + ls.in(i) * tf.num).roots();
		
		// sort the roots
	}

	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 is cheap and probably
		// good enough.
		// FIXME: do not invert matrix like that
		const cx_mat Ad = expmat(ss.A * ts.dt);
		const cx_mat Bd = (Ad - eye(size(Ad))) * pinv(ss.A) * ss.B;

		for (int k = 0; k < ts.n_samples - 1; k++)
			ts.state.col(k + 1) = Ad * ts.state.col(k) + Bd * ts.in.col(k);

		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: