#include "TuLiPControl.h"

/* not implemented yet */
void abstract(Controller* controller)
{
	return;
}

/* compute the control input u(k) and store it in controller.u */
void compute_control(Controller* controller)
{

	Polytope* currentRegion = cartesian_prod(regions[controller->dRegion],input_bound);
	Polytope* projGoalRegion = regions[controller->goal];
	/* todo: check if the computation is successful */
	get_input_helper(n, p, controller->nSteps,
					 A, B,
					 currentRegion->A, currentRegion->b, currentRegion->k,
					 projGoalRegion->A, projGoalRegion->b, projGoalRegion->k,
					 controller->hatx, projGoalRegion->center, controller->u);
	free_polytope(currentRegion);
}

Controller* instantiate_controller()
{
	Controller *c = malloc(sizeof(Controller));
	int i;
	c->u = (pfloat*)malloc(p*sizeof(pfloat));
	c->y = (pfloat*)malloc(m*sizeof(pfloat));
	c->dInput = (idxint*)malloc(nInputVariable*sizeof(idxint));
	for (i = 0;i<nInputVariable;i++)
	{
		*(c->dInput+i) = 0;
	}
	c->hatx = (pfloat*)malloc(n*sizeof(pfloat));
	return c;
}

void init_controller(Controller* controller)
{
	int i;
	init_input_bound();
	init_region();
	printf("successful init input and ppp polytopes\n");
	fflush(stdout);
	for(i=0;i<n;i++)
	{
		*(controller->hatx+i) = x0[i];
	}
	controller->dRegion = dRegion0;
	controller->nSteps = totalSteps;
	controller->fsm=init_fsm();
	printf("successful init fsm\n");
	fflush(stdout);
	controller->goal = fsm_transition(controller->fsm,controller->dInput);
	compute_control(controller);
	controller->nSteps--;
}

void free_controller(Controller* controller)
{
	free(controller->u);
	free(controller->y);
	free(controller->hatx);
	free_fsm(controller->fsm);
	free(controller);
	free_region();
	free_input_bound();
}

void transition_function(Controller* controller)
{
	estimate_function(controller);
	/* arrives at the goal */
	if (controller->nSteps<=0)
	{
		/* todo:check if the goal is really reached */
		controller->dRegion = controller->goal;
		/* reset nSteps */
		controller->nSteps = totalSteps;
		controller->goal = fsm_transition(controller->fsm, controller->dInput);
	}
	compute_control(controller);
	controller->nSteps--;
	/* display_controller(controller); */
}


void output_function(Controller* controller, pfloat u[])
{
	int i;
	for(i=0;i<p;i++)
	{
		u[i] = *(controller->u+i);
	}
}

void input_function(Controller* controller,const pfloat y[], const idxint dInput[])
{
	int i;
	for(i=0;i<m;i++)
	{
		*(controller->y+i) = y[i];
	}
	for(i=0;i<nInputVariable;i++)
	{
		*(controller->dInput+i) = dInput[i];
	}
}

/* assuming perfect full state observation
   possibly consider noisy observations in the future */
void estimate_function(Controller* controller)
{
	int i;
	for(i=0;i<n;i++)
	{
		*(controller->hatx+i) = *(controller->y+i);
	}
}


void display_controller(Controller* controller)
{
	int i;
	printf("State Estimation: ");
	for(i=0;i<n;i++)
	{
		printf("%5.2f, ",*(controller->hatx+i));
	}
	printf("\nControl Inputs: ");
	for(i=0;i<p;i++)
	{
		printf("%5.2f, ",*(controller->u+i));
	}
	printf("\nSensor Measurements: ");
	for(i=0;i<m;i++)
	{
		printf("%5.2f, ",*(controller->y+i));
	}
	printf("\nState of the Mealy machine: %d",controller->fsm->currentState);
	printf("\nDiscrete Region: %d",controller->dRegion);
	printf("\nNext goal: %d",controller->goal);
	printf("\nSteps left to reach goal: %d\n",controller->nSteps);
	fflush(stdout);
	return;
}

idxint get_input_helper(const idxint n, const idxint p, const idxint N,
						const pfloat* A, const pfloat* B,
						const pfloat* A1, const pfloat* b1, const idxint l1,
						const pfloat* A2, const pfloat* b2, const idxint l2,
						const pfloat* x0, const pfloat* xc, pfloat* u0)
{
	idxint i, j, k, l, sparsepr;

	idxint exitflag = ECOS_FATAL;
	pwork* mywork;
	/*the state is [x0,u0,x1,u1,...,uN-1,xN,a,b,t], where a,b,t are auxiliary variables.
	 * a,b,t satifies the following equalities and inequality:
	 * a = (t-1)/2, b = (t+1)/2
	 * norm([u0,...,uN-1,xN,a],2)<= b
	 * these equalities and inequality implies that u0^2+...+(uN-1)^2 + xN^2<=t*/
	idxint ecos_n = n*(N+1) + p*N + 3;
	idxint ecos_m = l1 * N + l2 + p*N + n + 2;
	idxint ecos_p = n*(N+1) + 2;
	idxint ecos_l = l1 * N + l2;
	idxint ecos_ncones = 1;
	idxint ecos_q[1];

	idxint nnzG;
	pfloat *ecos_Gpr;
	idxint *ecos_Gjc;
	idxint *ecos_Gir;

	idxint nnzA;
	idxint nnzAblock;
	pfloat *ecos_Apr;
	idxint *ecos_Ajc;
	idxint *ecos_Air;

	pfloat *ecos_c;
	pfloat *ecos_h;
	pfloat *ecos_b;

	ecos_q[0] = p*N + n + 2;

	/*generating matrix G of ECOS*/
	nnzG = N*p + n + N*(n+p)*l1 + n*l2 + 2;
	ecos_Gpr = (pfloat*)malloc(sizeof(pfloat)*nnzG);
	if (ecos_Gpr == NULL)
		return ECOS_FATAL;
	ecos_Gjc = (idxint*)malloc(sizeof(idxint)*(ecos_n+1));
	if (ecos_Gjc == NULL)
		return ECOS_FATAL;
	ecos_Gir = (idxint*)malloc(sizeof(idxint)*nnzG);
	if (ecos_Gir == NULL)
		return ECOS_FATAL;
	sparsepr = 0;

	for(k=0; k<N; k++)
	{
		for(j = 0; j < n+p; j++)
		{
			*(ecos_Gjc+j+k*(n+p)) = sparsepr;
			for (i=0; i<l1;i++)
			{
				*(ecos_Gpr+sparsepr) = -1*(*(A1+j*l1+i));
				*(ecos_Gir+sparsepr) = i+k*l1;
				sparsepr++;
			}
			if (j >= n)
			{
				*(ecos_Gpr+sparsepr) = -1;
				*(ecos_Gir+sparsepr) = j-n+k*p+N*l1+l2+2;
				sparsepr++;
			}
		}
	}

	for(j=N*(n+p); j<N*(n+p)+n;j++)
	{
		*(ecos_Gjc+j) = sparsepr;
		for(i=0;i<l2;i++)
		{
			*(ecos_Gpr+sparsepr) = -1*(*(A2+(j-N*(n+p))*l2+i));
			*(ecos_Gir+sparsepr) = i+N*l1;
			sparsepr++;
		}
		*(ecos_Gpr+sparsepr) = -1;
		*(ecos_Gir+sparsepr) = j-N*(n+p)+N*l1+l2+2+N*p;
		sparsepr++;
	}

	*(ecos_Gjc+ecos_n-3) = nnzG-2;
	*(ecos_Gpr+nnzG-2) = -1;
	*(ecos_Gir+nnzG-2) = N*l1+l2+1;

	*(ecos_Gjc+ecos_n-2) = nnzG-1;
	*(ecos_Gpr+nnzG-1) = -1;
	*(ecos_Gir+nnzG-1) = N*l1+l2;

	*(ecos_Gjc+ecos_n-1) = nnzG;
	*(ecos_Gjc+ecos_n) = nnzG;

	/*generating matrix A of ECOS*/
	nnzA = (N+1)*n + N*n*n + N*n*p + 4;
	/*The number of none zero entries in the block [I,0;-A,-B]*/
	nnzAblock = n*n + n*p + n;
	ecos_Apr = (pfloat*)malloc(sizeof(pfloat)*nnzA);
	if (ecos_Apr == NULL)
		return ECOS_FATAL;
	ecos_Ajc = (idxint*)malloc(sizeof(idxint)*(ecos_n+1));
	if (ecos_Ajc == NULL)
		return ECOS_FATAL;
	ecos_Air = (idxint*)malloc(sizeof(pfloat)*nnzA);
	if (ecos_Air == NULL)
		return ECOS_FATAL;

	sparsepr = 0;

	for(j = 0; j < n; j++)
	{
		*(ecos_Ajc+j)=sparsepr;
		*(ecos_Apr + sparsepr) = 1;
		*(ecos_Air + sparsepr) = j;
		sparsepr++;
		for (i = 0;i < n; i++)
		{
			*(ecos_Apr+sparsepr) = -1*(*(A+j*n+i));
			*(ecos_Air+sparsepr) = n+i;
			sparsepr++;
		}
	}

	for (j = n; j < n+p; j++)
	{
		*(ecos_Ajc+j)=sparsepr;
		for (i = 0;i < n; i++)
		{
			*(ecos_Apr + sparsepr) = -1*(*(B + (j-n)*n + i));
			*(ecos_Air + sparsepr) = n+i;
			sparsepr++;
		}
	}

	*(ecos_Ajc+n+p) = sparsepr;
	for (k = 1; k < N; k++)
	{
		for (l = 1; l <= n+p; l++)
		{
			*(ecos_Ajc+k*(n+p)+l) = *(ecos_Ajc + l) + *(ecos_Ajc + k*(n+p));
		}
		for (l = 0; l < nnzAblock; l++)
		{
			*(ecos_Apr+sparsepr) = *(ecos_Apr+sparsepr%nnzAblock);
			*(ecos_Air+sparsepr) = *(ecos_Air+sparsepr%nnzAblock) + k*n;
			sparsepr++;
		}
	}

	for(j=(n+p)*N;j<(n+p)*N+n;j++)
	{
		*(ecos_Ajc+j) = sparsepr;
		*(ecos_Apr+sparsepr) = 1;
		*(ecos_Air+sparsepr) = N*n+j-(n+p)*N;
		sparsepr++;
	}

	*(ecos_Ajc+ecos_n-3) = nnzA-4;
	*(ecos_Apr+nnzA-4) = 1;
	*(ecos_Air+nnzA-4) = (N+1)*n;

	*(ecos_Ajc+ecos_n-2) = nnzA-3;
	*(ecos_Apr+nnzA-3) = 1;
	*(ecos_Air+nnzA-3) = (N+1)*n+1;

	*(ecos_Ajc+ecos_n-1) = nnzA-2;
	*(ecos_Apr+nnzA-2) = -0.5;
	*(ecos_Air+nnzA-2) = (N+1)*n;
	*(ecos_Apr+nnzA-1) = -0.5;
	*(ecos_Air+nnzA-1) = (N+1)*n+1;

	*(ecos_Ajc+ecos_n) = nnzA;

	/*The objective function is equivalent to -2xc'*xN + t + xc'*xc
	 * (the constant xc'*xc is omitted )*/
	ecos_c = (pfloat*)malloc(sizeof(pfloat)*ecos_n);
	if (ecos_c == NULL)
		return ECOS_FATAL;
	for(i = 0; i < ecos_n; i++)
	{
		if (i == ecos_n - 1)
		{
			*(ecos_c + i) = 1;
		}
		else if (i >= (n+p)*N && i< (n+p)*N+n)
		{
			*(ecos_c + i) = -2*(*(xc + i - (n+p)*N));
		}
		else
		{
			*(ecos_c+i) = 0;
		}
	}

	/* h is of the form [-b1,...,-b1,-b2,0,...,0] */
	ecos_h = (pfloat*)malloc(sizeof(pfloat)*ecos_m);
	if (ecos_h == NULL)
		return ECOS_FATAL;
	for (i = 0; i < ecos_m; i++)
	{
		if (i < l1*N)
		{
			*(ecos_h + i) = -1 * (*(b1 + i%l1));
		}
		else if (i < l1*N + l2)
		{
			*(ecos_h + i) = -1 * (*(b2 + i - l1*N));
		}
		else
		{
			*(ecos_h + i) = 0;
		}
	}

	/*b is of the form [x0,0,...,0,-0.5,0.5]*/
	ecos_b = (pfloat*)malloc(sizeof(pfloat)*ecos_p);
	if (ecos_b == NULL)
		return ECOS_FATAL;
	for (i = 0; i < ecos_p; i++)
	{
		if (i < n)
		{
			*(ecos_b + i) = *(x0 + i);
		}
		else
		{
			*(ecos_b + i) = 0;
		}
	}
	*(ecos_b + ecos_p - 2) = -0.5;
	*(ecos_b + ecos_p - 1) = 0.5;

	mywork = ECOS_setup(ecos_n, ecos_m, ecos_p, ecos_l, ecos_ncones, ecos_q, ecos_Gpr, ecos_Gjc, ecos_Gir, ecos_Apr, ecos_Ajc, ecos_Air, ecos_c, ecos_h, ecos_b);
	if( mywork != NULL ){
		mywork->stgs->verbose = 0;
		/* solve */
		exitflag = ECOS_solve(mywork);

		for (i = 0; i < p; i++)
		{
			*(u0+i) = mywork->x[i+n];
		}
		ECOS_cleanup(mywork, 0);
	}

	free(ecos_Gpr);
	free(ecos_Gjc);
	free(ecos_Gir);
	free(ecos_Apr);
	free(ecos_Ajc);
	free(ecos_Air);
	free(ecos_c);
	free(ecos_h);
	free(ecos_b);

	return exitflag;
}
