//////////////////////////////////////////////////////////////////////////
//////////////////         notorious.cxx       ///////////////////////////
//////////////////////////////////////////////////////////////////////////
////////////////           PSOPT  Example             ////////////////////
//////////////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////////
//////// Title:  Bock's notorious parameter estimation problem ///////////
//////// Last modified: 08 April 2011                     ////////////////
//////// Reference:     Schittkowskit (2002)     	      ////////////////
//////// (See PSOPT handbook for full reference)           ///////////////
//////////////////////////////////////////////////////////////////////////
////////     Copyright (c) Victor M. Becerra, 2009        ////////////////
//////////////////////////////////////////////////////////////////////////
//////// This is part of the PSOPT software library, which ///////////////
//////// is distributed under the terms of the GNU Lesser ////////////////
//////// General Public License (LGPL)                    ////////////////
//////////////////////////////////////////////////////////////////////////

#include "psopt.h"

using namespace PSOPT;




//////////////////////////////////////////////////////////////////////////
///////////////////  Define the observation function //////////
//////////////////////////////////////////////////////////////////////////



void  observation_function( adouble* observations,
                            adouble* states, adouble* controls,
                            adouble* parameters, adouble& time, int k,
                            adouble* xad, int iphase, Workspace* workspace)
{

      observations[ 0 ] = states[ 0 ];
      observations[ 1 ] = states[ 1 ];
}




//////////////////////////////////////////////////////////////////////////
///////////////////  Define the DAE's ////////////////////////////////////
//////////////////////////////////////////////////////////////////////////


void dae(adouble* derivatives, adouble* path, adouble* states,
         adouble* controls, adouble* parameters, adouble& time,
         adouble* xad, int iphase, Workspace* workspace)
{

   adouble x1 = states[ 0 ];
   adouble x2 = states[ 1 ];


   adouble p  = parameters[ 0 ];
   adouble t  = time;

   double mu = 60.0;

   derivatives[0] =  x2;
   derivatives[1] =  mu*mu*x1 - (mu*mu + p*p)*sin(p*t);


}

////////////////////////////////////////////////////////////////////////////
///////////////////  Define the events function ////////////////////////////
////////////////////////////////////////////////////////////////////////////

void events(adouble* e, adouble* initial_states, adouble* final_states,
            adouble* parameters,adouble& t0, adouble& tf, adouble* xad,
            int iphase, Workspace* workspace)
{

   e[ 0 ] =  initial_states[0];
   e[ 1 ] =  initial_states[1];


}


///////////////////////////////////////////////////////////////////////////
///////////////////  Define the phase linkages function ///////////////////
///////////////////////////////////////////////////////////////////////////

void linkages( adouble* linkages, adouble* xad, Workspace* workspace)
{
  // No linkages as this is a single phase problem
}

using namespace std;

////////////////////////////////////////////////////////////////////////////
///////////////////  Define the main routine ///////////////////////////////
////////////////////////////////////////////////////////////////////////////

int main(void)
{
   int nobs =200;

   // Generate true solution at sampling points and add noise

   double sigma = 0.05;

   MatrixXd y1m, y2m;


//   theta = randu(1,nobs);

   MatrixXd noise1= GaussianRandom(1,nobs);
   
   MatrixXd noise2= GaussianRandom(1,nobs);
   
   MatrixXd theta = (MatrixXd::Random(1,nobs)+ones(1,nobs))/2.0;

   sort(theta);
   
   MatrixXd ss = (pi*theta).array().sin();
   MatrixXd cc = (pi*theta).array().cos();

   y1m =  ss + sigma*noise1;

   y2m =  pi*cc + sigma*noise2;
   


////////////////////////////////////////////////////////////////////////////
///////////////////  Declare key structures ////////////////////////////////
////////////////////////////////////////////////////////////////////////////

    Alg  algorithm;
    Sol  solution;
    Prob problem;

////////////////////////////////////////////////////////////////////////////
///////////////////  Register problem name  ////////////////////////////////
////////////////////////////////////////////////////////////////////////////

    problem.name          		= "Bocks notorious parameter estimation problem";
    problem.outfilename                 = "notorious.txt";

////////////////////////////////////////////////////////////////////////////
////////////  Define problem level constants & do level 1 setup ////////////
////////////////////////////////////////////////////////////////////////////

    problem.nphases   							= 1;
    problem.nlinkages                   	= 0;

    psopt_level1_setup(problem);


/////////////////////////////////////////////////////////////////////////////
/////////   Define phase related information & do level 2 setup /////////////
/////////////////////////////////////////////////////////////////////////////

    problem.phases(1).nstates   					= 2;
    problem.phases(1).ncontrols 					= 0;
    problem.phases(1).nevents   					= 2;
    problem.phases(1).npath     					= 0;
    problem.phases(1).nparameters        		= 1;
    problem.phases(1).nodes    		    		<< 80;
    problem.phases(1).nobserved              = 2;
    problem.phases(1).nsamples               = nobs;


    psopt_level2_setup(problem, algorithm);

////////////////////////////////////////////////////////////////////////////
////////////  Enter estimation information                      ////////////
////////////////////////////////////////////////////////////////////////////


    problem.phases(1).observation_nodes      = theta;
    problem.phases(1).observations           << y1m, 
                                                y2m;



////////////////////////////////////////////////////////////////////////////
///////////////////  Enter problem bounds information //////////////////////
////////////////////////////////////////////////////////////////////////////


    problem.phases(1).bounds.lower.states(0) =  -10.0;
    problem.phases(1).bounds.lower.states(1) =  -100.0;



    problem.phases(1).bounds.upper.states(0) =  10.0;
    problem.phases(1).bounds.upper.states(1) =  100.0;


    problem.phases(1).bounds.lower.parameters(0) =  -10.0;
    problem.phases(1).bounds.upper.parameters(0) =   10.0;

    problem.phases(1).bounds.lower.events(0) =  0.0;
    problem.phases(1).bounds.upper.events(0) =  0.0;

    problem.phases(1).bounds.lower.events(1) =  pi;
    problem.phases(1).bounds.upper.events(1) =  pi;



    problem.phases(1).bounds.lower.StartTime    = 0.0;
    problem.phases(1).bounds.upper.StartTime    = 0.0;

    problem.phases(1).bounds.lower.EndTime      = 1.0;
    problem.phases(1).bounds.upper.EndTime      = 1.0;

////////////////////////////////////////////////////////////////////////////
///////////////////  Register problem functions  ///////////////////////////
////////////////////////////////////////////////////////////////////////////

    problem.dae 							= &dae;
    problem.events 						= &events;
    problem.linkages						= &linkages;
    problem.observation_function 	= & observation_function;

////////////////////////////////////////////////////////////////////////////
///////////////////  Define & register initial guess ///////////////////////
////////////////////////////////////////////////////////////////////////////

    int nnodes = problem.phases(1).nodes(0);


    problem.phases(1).guess.states         = zeros(2,nnodes);
    problem.phases(1).guess.time           = linspace(0.0, 1.0, nnodes);
    problem.phases(1).guess.parameters(0)  = 0.0;


////////////////////////////////////////////////////////////////////////////
///////////////////  Enter algorithm options  //////////////////////////////
////////////////////////////////////////////////////////////////////////////

    algorithm.nlp_method                  = "IPOPT";
    algorithm.scaling                     = "automatic";
    algorithm.derivatives                 = "automatic";
    algorithm.collocation_method          = "trapezoidal";
    algorithm.nlp_iter_max                =  200;
    algorithm.nlp_tolerance               = 1.e-4;
//    algorithm.mesh_refinement             = "automatic";
//    algorithm.ode_tolerance               = 1.e-6;



////////////////////////////////////////////////////////////////////////////
///////////////////  Now call PSOPT to solve the problem   //////////////////
////////////////////////////////////////////////////////////////////////////


    psopt(solution, problem, algorithm);

////////////////////////////////////////////////////////////////////////////
///////////////////  Declare DMatrix objects to store results //////////////
////////////////////////////////////////////////////////////////////////////

    DMatrix states, x1, x2, p, t;


////////////////////////////////////////////////////////////////////////////
///////////  Extract relevant variables from solution structure   //////////
////////////////////////////////////////////////////////////////////////////

    states = solution.get_states_in_phase(1);
    t      = solution.get_time_in_phase(1);
    p      = solution.get_parameters_in_phase(1);
    x1     = states.row(0);
    x2     = states.row(1);


////////////////////////////////////////////////////////////////////////////
///////////  Save solution data to files if desired ////////////////////////
////////////////////////////////////////////////////////////////////////////

    Save(states,"states.dat");
    Save(t,"t.dat");
    Print(p,"Estimated parameter");
    printf( "\nParameter error %e\n", abs(p(0)-pi) );


////////////////////////////////////////////////////////////////////////////
///////////  Plot some results if desired (requires gnuplot) ///////////////
////////////////////////////////////////////////////////////////////////////

    plot(theta,y1m,t,x1,problem.name, "time (s)", "observed x1", "x1m x1");
    plot(theta,y2m,t,x2,problem.name, "time (s)", "observed x2", "x2m x2");



}

////////////////////////////////////////////////////////////////////////////
///////////////////////      END OF FILE     ///////////////////////////////
////////////////////////////////////////////////////////////////////////////
