/** CoupledPendulums.cpp
 * Implementation of the coupled inverted pendulums benchmark.
 *
 *  @date: 30.11.2010
 *  @author: Heiko Hamann
*/


#include "CoupledPendulums.h"


CoupledPendulums::CoupledPendulums() {
  CoupledPendulums(2);
}


CoupledPendulums::CoupledPendulums(int myCrabNum) {
  crabNum = myCrabNum;
  phi.resize(myCrabNum);
  omega.resize(myCrabNum);
  v.resize(myCrabNum);
  u.resize(myCrabNum);
  x.resize(myCrabNum);
  fitness = 0.0;

  Kp = 0.005;
  Kl = 0.05;
  g = 9.81;
  length = 0.5;
  Vmotor = 7.0;
  Vbreak = 8.5;
  deltaUV = 0.05;
  
  deltaT = 0.01;
  
  chainMin = 0.05;
  chainMax = 0.35;
  worldWidth = 2.0;

  stepsPerRun = 4000;
  currentStep = 0;

  initialPosition();
}

CoupledPendulums::~CoupledPendulums() {

}


void CoupledPendulums::initialPosition() {

  if(crabNum>0)
    phi[0] = 0.8*PI;
  if(crabNum>1)
    phi[1] = 0.9*PI;
  if(crabNum>2)
    phi[2] = 1.0*PI;
  if(crabNum>3)
    phi[3] = 1.1*PI;
  if(crabNum>4)
    phi[4] = 1.2*PI;
  
  if(crabNum>0)
    x[0] = 0.0+0.25*(-4.0)*(chainMin+chainMax);
  if(crabNum>1)
    x[1] = 0.0+0.25*(-2.0)*(chainMin+chainMax);
  if(crabNum>2)
    x[2] = 0.0+0.25*0.0*(chainMin+chainMax);
  if(crabNum>3)
    x[3] = 0.0+0.25*2.0*(chainMin+chainMax);
  if(crabNum>4)
    x[4] = 0.0+0.25*4.0*(chainMin+chainMax);
  
  for(unsigned int j=0;j<crabNum;j++) {
    omega[j] = 0.0;
    v[j] = 0.0;
    u[j] = 0.0;
  }

  fitness = 0.0;
  currentStep = 0;
}



float CoupledPendulums::motor(float u, float v) {
  if(fabs(u-v)>deltaUV) {
    if(u>v) {
      if(v>=0)
	return Vmotor;
      else
	return Vbreak;
    } else {
      // u<=v
      if(v>=0)
	return -Vbreak;
      else
	return -Vmotor;
    }
  } else {
    // fabs(u-v)<=deltaUV
    if(u>=v) {
      if(v>=0)
	return Vmotor/deltaUV*(u-v);
      else
	return Vbreak/deltaUV*(u-v);
    } else {
      // u<v
      if(v>=0)
	return Vbreak/deltaUV*(u-v);
      else
	return Vmotor/deltaUV*(u-v);
    }
  }
}


float CoupledPendulums::changeAngularVel(float phi, float omega, 
					 float u, float v) {

  if(omega==0.0)
    return 3.0*g/(2.0*length)*sin(phi)-3.0/(2.0*length)*motor(u,v)*cos(phi);
  else
    return 3.0*g/(2.0*length)*sin(phi)-3.0/(2.0*length)*motor(u,v)*cos(phi)
      -Kp*omega*fabs(omega)-Kl*omega/fabs(omega);
  
}



int CoupledPendulums::checkForBorders() {

  for(unsigned int j=0;j<crabNum;j++) {
    // no room to the right:
    if(j<crabNum-1) {
      if( (fabs(x[j]-x[j+1])<chainMin) 
	  && ( (v[j]>0.0) || (v[j]==0.0 && u[j]>0.0) ) ) {
	//printf("%d knocks over %d to the right. (%f)\n", j, j+1, fabs(x[j]-x[j+1]));
	v[j]=0.0;
	return 0;
      }
    }
    if(j>0) {
      if( (fabs(x[j]-x[j-1])>chainMax)  
	  && ( (v[j]>0.0) || (v[j]==0.0 && u[j]>0.0) ) ) {
	v[j]=0.0;
	//printf("%d snaps the chain by going right. (%f)\n", j, fabs(x[j]-x[j-1]));
	return 0;
      }
    }
    // no room to the left:
    if(j<crabNum-1) {
      if( (fabs(x[j]-x[j+1])>chainMax) 
	  && ( (v[j]<0.0) || (v[j]==0.0 && u[j]<0.0) ) ) {
	v[j]=0.0;
	//printf("%d snaps the chain by going left. (%f)\n", j, fabs(x[j]-x[j+1]));
	return 0;
      }
    }
    if(j>0) {
      if( (fabs(x[j]-x[j-1])<chainMin) 
	  && ( (v[j]<0.0) || (v[j]==0.0 && u[j]<0.0) ) ) {
	v[j]=0.0;
	//printf("%d knocks over %d to the left. (%f)\n", j, j-1, fabs(x[j]-x[j-1]));
	return 0;
      }
    }
    
    
    if(omega[j]>5.0*PI) {
      omega[j]=5.0*PI;
      return 0;
    }
    if(omega[j]<-5.0*PI) {
      omega[j]=-5.0*PI;
      return 0;
    }
    if(x[j]<-worldWidth/2.0) {
      x[j]=-worldWidth/2.0;
      v[j]=0.0;
      return 0;
    }
    if(x[j]>worldWidth/2.0) {
      x[j]=worldWidth/2.0;
      v[j]=0.0;
      return 0;
    }
    if(v[j]<-2.0) {
      v[j]=-2.0;
      return 0;
    }
    if(v[j]>2.0) {
      v[j]=2.0;
      return 0;
    }
  }
  return 1;
}


std::vector< std::vector<int> > CoupledPendulums::getSensorValues() {

  std::vector< std::vector<int> > sensorValues;
  sensorValues.resize(crabNum, std::vector<int>(10));
  float distLeft;
  float distRight;

  for(unsigned int i=0;i<crabNum;i++) {
    
    distLeft = 1.0;
    distRight = 1.0;
    
    if(crabNum>1) {
      if(i==0)
	distRight = fabs(x[0]-x[1]);
      else if(i==crabNum-1)
	distLeft = fabs(x[crabNum-1]-x[crabNum-2]);
      else {
	distLeft = fabs(x[i]-x[i-1]);
	distRight = fabs(x[i]-x[i+1]);
      }
    }

    while(phi[i]<0.0)
      phi[i] += 2.0*PI;
    while(phi[i]>2.0*PI)
      phi[i] -= 2.0*PI;
    
    if(phi[i]<0.5*PI) {
      sensorValues[i][0] = 127-(int)(127.0*phi[i]/(0.5*PI));
      sensorValues[i][1] = 0;
      sensorValues[i][2] = 0;
      sensorValues[i][3] = 0;
    } else if(phi[i]<PI) {
      sensorValues[i][0] = 0;
      sensorValues[i][1] = 0;
      sensorValues[i][2] = 127-(int)(127.0*(phi[i]-0.5*PI)/(0.5*PI));
      sensorValues[i][3] = 0;
    } else if(phi[i]<1.5*PI) {
      sensorValues[i][0] = 0;
      sensorValues[i][1] = (int)(-127.0*(PI-phi[i])/(0.5*PI));
      sensorValues[i][2] = 0;
      sensorValues[i][3] = 0;
    } else {
      sensorValues[i][0] = 0;
      sensorValues[i][1] = 0;
      sensorValues[i][2] = 0;
      sensorValues[i][3] = (int)(-127.0*(1.5*PI-phi[i])/(0.5*PI));
    }
    
    if(x[i]<0.0) {
      if((worldWidth/4.0)+x[i] < distLeft)
	sensorValues[i][4] = (int)(-1.0*127.0*x[i]/(worldWidth/2.0));
      else
	sensorValues[i][4] = (int)(127.0*distLeft/(worldWidth/2.0));
      if(distRight < (worldWidth/4.0)) {
	sensorValues[i][5] = (int)(127.0*distRight/(worldWidth/2.0));
      } else {
	sensorValues[i][5] = 0;
      }
    } else {
      if((worldWidth/4.0)-x[i] < distRight) {
	sensorValues[i][5] = (int)(127.0*x[i]/(worldWidth/2.0));
      } else {
	sensorValues[i][5] = (int)(127.0*distRight/(worldWidth/2.0));
      }
      if(distLeft < (worldWidth/4.0))
	sensorValues[i][4] = (int)(127.0*distLeft/(worldWidth/2.0));
      else
	sensorValues[i][4] = 0;
    }
    if(v[i]<0.0) {
      sensorValues[i][6] = (int)(-1.0*127.0*v[i]/2.0);
      sensorValues[i][7] = 0;
  } else {
      sensorValues[i][6] = 0;
      sensorValues[i][7] = (int)(127.0*v[i]/2.0);
    }
    if(omega[i]<0.0) {
      sensorValues[i][8] = (int)(-1.0*127.0*omega[i]/(5.0*PI));
      sensorValues[i][9] = 0;
    } else {
      sensorValues[i][8] = 0;
      sensorValues[i][9] = (int)(127.0*omega[i]/(5.0*PI));
    }
    
    // noise:
    /*
      for(int i=0;i<10;i++) {
      sensorValues[i] += rand()%6 - 3;
      if(sensorValues[i]<0)
      sensorValues[i] = 0;
      }
    */
  }

  return sensorValues;
}



float CoupledPendulums::getFitness() {

  return fitness;

}




int CoupledPendulums::performOneStep(std::vector< std::vector<int> > motorControl) {

  float phiD1, phiD2, phiD3;
  float omegaD1, omegaD2, omegaD3;
  float vD1, vD2, vD3;
  float xD1, xD2, xD3;

  for(unsigned int j=0;j<crabNum;j++) {

    float a1 = (float)motorControl[j][0]/127.0;
    float a2 = (float)motorControl[j][1]/127.0;
    
    u[j] = 2.0*(a1-a2);
    
    // runge-kutta:
    phiD1 = omega[j];
    xD1 = v[j];
    omegaD1 = changeAngularVel(phi[j], omega[j], u[j], v[j]);
    vD1 = motor(u[j],v[j]);
    omegaD2 = changeAngularVel(phi[j]+phiD1*deltaT*0.5, 
			       omega[j]+omegaD1*deltaT*0.5, 
			       u[j], 
			       v[j]+vD1*deltaT*0.5);
    phiD2 = omega[j]+omegaD1*deltaT*0.5;
    xD2 = v[j]+vD1*deltaT*0.5;
    vD2 = motor(u[j],v[j]+vD1*deltaT*0.5);
    omegaD3 = changeAngularVel(phi[j]-phiD1*deltaT+phiD2*deltaT*2.0, 
			       omega[j]-omegaD1*deltaT+omegaD2*deltaT*2.0, 
			       u[j], 
			       v[j]+vD1*deltaT+vD2*deltaT*2.0);
    phiD3 = omega[j]-omegaD1*deltaT+omegaD2*deltaT*2.0;
    xD3 = v[j]-vD1*deltaT+vD2*deltaT*2.0;
    vD3 = motor(u[j],v[j]-vD1*deltaT+vD2*deltaT*2.0);
    
    phi[j] += deltaT*(1.0/6.0*phiD1+4.0/6.0*phiD2+1.0/6.0*phiD3);
    x[j] += deltaT*(1.0/6.0*xD1+4.0/6.0*xD2+1.0/6.0*xD3);
    omega[j] += deltaT*(1.0/6.0*omegaD1+4.0/6.0*omegaD2+1.0/6.0*omegaD3);
    v[j] += deltaT*(1.0/6.0*vD1+4.0/6.0*vD2+1.0/6.0*vD3);
    
    
    fitness += fabs(phi[j]-PI)/(PI*crabNum*(float)stepsPerRun);
    
  } // for crab num

  currentStep++;
  return checkForBorders();
}




std::vector< std::vector<float> > CoupledPendulums::getState() {

  std::vector< std::vector<float> > state;
  state.resize(crabNum, std::vector<float>(5));

  for(unsigned int i=0;i<crabNum;i++) {
    state[i][0] = u[i];
    state[i][1] = x[i];
    state[i][2] = phi[i];
    state[i][3] = omega[i];
    state[i][4] = v[i];
  }

  return state;
}



void CoupledPendulums::writeTrajectoryToFile() {

  FILE *f;
  char filename[256];

  f = fopen("results.sim","w");
  fprintf(f, "number_of_cats        = %d\n", crabNum);
  fprintf(f, "pendulum_length       = %f\n", length); 
  fprintf(f, "width_of_cat          = %f\n", chainMin);
  fprintf(f, "length_of_chain       = %f\n", chainMax-chainMin);
  fprintf(f, "width_of_half_world   = %f\n", worldWidth/2.0);
  fprintf(f, "info_text             = AHHS2\n");
  fclose(f);

  for(unsigned int i=0;i<crabNum;i++) {
    sprintf(filename, "out%d.txt", i);
    f = fopen(filename,"a");
    fprintf(f, "%d %e %e %e %e %e %e\n",
	    currentStep, u[i], x[i], phi[i], omega[i], v[i], motor(u[i],v[i]));
    fclose(f);
  }
}