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


namespace ct
{
	TransferFn::TransferFn(complex gain)
		: gain(gain)
		, num(math::Poly<complex>{1})
		, den(math::Poly<complex>{1})
	{
		num(0) = 1.;
		den(0) = 1.;
	}

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

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

	TransferFn feedback(const TransferFn& tf, double k)
	{
		// TransferFn tf_cl;
	}

	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) * tf.gain;
		}
	
		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: