// Copyright (C) 2005 Zbynek Winkler, zw at robotika cz
// Licensed to the public under the terms of the GNU GPL 2

#include <vector>
#include <algorithm>
//#include <iostream>
#include <assert.h>
#include <fcntl.h>
#include "frame.h"
//#include "frame-io.h"
#include "rnd.h"

#define ASSERT assert

static inline double deg(const double & in_angle) { return in_angle * M_PI / 180.0;}
static inline bool isValid(const double & a) { return !isnan(a) && finite(a); }
int getColor(const point & in_p)
{
	double TILE = 0.05;
	int white = ((int(in_p.m_x/TILE) ^ int(in_p.m_y/TILE)) & 1);
	return white;
}

class sample : public frame
{
	public:
		double m_weight;
		sample() : frame(), m_weight(1)  {}
		sample(const point & in_point, const double & in_heading) : frame(in_point, in_heading), m_weight(1) {}
		sample(const double & in_x, const double & in_y, const double & in_heading) : frame(in_x, in_y, in_heading), m_weight(1) {}

		void set(const sample & s, const double & xNoise, const double & yNoise, const double & aNoise)
		{
			m_point = s.m_point + point(xNoise, yNoise);
			m_heading = s.m_heading + aNoise;
		}
};

const double D = 1;
const double ENCERR = 0.8;
const double GAIN = 1.5;
const double LOSS = 1.0/GAIN;

class MCL : public std::vector<sample>
{
	rnd m_rand;
	double D;
	public:
		MCL(int numsamp) : std::vector<sample>(numsamp), D(1) {} 
		void init(const sample & in_s, const double & dNoise, const double & aNoise)
		{
			for (iterator i = begin(); i != end(); i++)
				i->set(in_s, dNoise*m_rand(), dNoise*m_rand(), aNoise*m_rand());
		}
		MCL& predict(const double & in_left, const double & in_right)
		{
			for (iterator i = begin(); i != end(); i++)
			{
				double r = in_right + in_right*ENCERR*m_rand();
				double l = in_left  + in_left *ENCERR*m_rand();

				i->advanceBy( (l + r)/2 );
				i->turnBy( (r-l)/D );
			}
			return *this;
		}
		MCL& correct(int a, int b, int c)
		{
			point pa(0.0, 0.05);
			point pb(0.0, 0.0);
			point pc(0.0, -0.05);
			for (iterator i = begin(); i != end(); i++)
			{
				int color = getColor(i->transform(pa));
				if (color == a) i->m_weight *= GAIN;
				else if (a != 2) i->m_weight *= LOSS;
				
				color = getColor(i->transform(pb));
				if (color == b) i->m_weight *= GAIN;
				else if (b != 2) i->m_weight *= LOSS;

				color = getColor(i->transform(pc));
				if (color == c) i->m_weight *= GAIN;
				else if (c != 2) i->m_weight *= LOSS;
			}
			return *this;
		}
		MCL& normalize()
		{
			double gain = 0, w = 0;
			for (iterator s = begin(); s != end(); s++)
				w += s->m_weight;
			gain = size() / w;
			ASSERT( isValid(gain) );
			for (iterator i = begin(); i != end(); i++)
				i->m_weight *= gain;
			return *this;
		}
		struct cmp
		{
			double m_pivot;
			cmp(const double & a) : m_pivot(a) {}	
			bool operator () (const sample & s1) const { return s1.m_weight < m_pivot; }
		};
		void partSort(size_t & b_begin, size_t & b_end)
		{
			double w_min = 0.25;
			double w_max = 2.0;
			double inc = (w_max - w_min) / 6;
			
			using namespace std;
			iterator i3 = partition(begin(), end(), cmp(w_min+3*inc));
			iterator i1 = partition(begin(), i3, cmp(w_min+inc));
			b_begin = partition(begin(), i1, cmp(w_min)) - begin();
			partition(i1, i3, cmp(w_min+2*inc));
			
			iterator i5 = partition(i3, end(), cmp(w_min+5*inc));
			b_end = partition(i5, end(), cmp(w_max)) - begin();
			partition(i3, i5, cmp(w_min+4*inc));
		}
		MCL& resample()
		{
			size_t i_min, i_max;
			partSort(i_min, i_max);
				
			size_t num = 0;
			size_t i;
			MCL & a_mcl = *this;
			for (i = i_max; i < size(); i++)
			{
				int count = int(a_mcl[i].m_weight);
				a_mcl[i].m_weight /= count;
				for (size_t j = num; j < num+count-1; j++)
					a_mcl[j] = a_mcl[i];
				num += count-1;
			}
			for (i = num; i < i_min; i++)
			{
				i_max--;
				a_mcl[i_max].m_weight /= 2;
				a_mcl[i] = a_mcl[i_max];
			}
			return *this;
		}
};

int main()
{
	using namespace std;
	
	ASSERT( eq(sin(deg(90)), 1) );
	ASSERT( eq(cos(deg(90)), 0) );
	ASSERT( eq(cos(deg(180)), -1) );
	ASSERT( eq(sin(deg(180)), 0) );

	frame robot(5,5,0);
	robot.advanceBy(1);
	ASSERT( robot.eq(frame(6,5,0)) );

	ASSERT( frame(3,3,0).turnBy(deg(90)).advanceBy(1).eq(frame(3,4,deg(90))) );
	ASSERT( frame(-2,0,deg(43)).advanceBy(1).eq(frame(-2,0,deg(43-90)).slideBy(1).turnBy(deg(90))) );

	sample center(0.05,0.07,deg(20));
	double dNoise = 0.01;
	double aNoise = deg(10);
	MCL mcl(300);
	mcl.init(center, dNoise, aNoise);

	int fd = ::open("mcl.log", O_CREAT|O_WRONLY, 0644);
	::write(fd, &(*mcl.begin()), mcl.size()*sizeof(sample));
	//cout << mcl.begin()->m_point.m_x << "\t" << mcl.begin()->m_point.m_y << "\t" << mcl.begin()->m_heading << "\t" << mcl.begin()->m_weight << endl;
	robot = center;
	for (int i = 0; i < 300; i++)
	{
		robot.advanceBy(0.002);
		mcl.predict(0.002, 0.002);
		mcl.correct(getColor(robot.transform(point(0.0, 0.05))), getColor(robot.transform(point())), getColor(robot.transform(point(0.0, -0.05))));
		mcl.normalize();
		mcl.resample();
		::write(fd, &(*mcl.begin()), mcl.size()*sizeof(sample));
	}
	::close(fd);

	return 0;
}