//////////////////////////////////////////////////////////////////////////
////////////////           PSOPT  Example             ////////////////////
//////////////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////////
//////// Title:                 Rayleigh problem          ////////////////
//////// Last modified:         11 July 2011              ////////////////
//////// Reference:             Betts (2010)              ////////////////
//////// (See PSOPT handbook for full reference)          ////////////////
//////////////////////////////////////////////////////////////////////////
////////     Copyright (c) Victor M. Becerra, 2011        ////////////////
//////////////////////////////////////////////////////////////////////////
//////// 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"

//////////////////////////////////////////////////////////////////////////
///////////////////  Define the end point (Mayer) cost function //////////
//////////////////////////////////////////////////////////////////////////

adouble endpoint_cost(adouble* initial_states, adouble* final_states,
                      adouble* parameters,adouble& t0, adouble& tf,
                      adouble* xad, int iphase, Workspace* workspace)
{
   return 0.0;
}

//////////////////////////////////////////////////////////////////////////
///////////////////  Define the integrand (Lagrange) cost function  //////
//////////////////////////////////////////////////////////////////////////

adouble integrand_cost(adouble* states, adouble* controls,
                       adouble* parameters, adouble& time, adouble* xad,
                       int iphase, Workspace* workspace)
{
  adouble x1 = states[0];
  adouble u = controls[0];

  return  (u*u + x1*x1);
}

//////////////////////////////////////////////////////////////////////////
///////////////////  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 u = controls[0];

   double p = 0.14;

   derivatives[ 0 ] = x2;
   derivatives[ 1 ] = -x1 + x2*(1.4-p*x2*x2) + 4.0*u;

   path[ 0 ] = u + x1/6.0;
}


////////////////////////////////////////////////////////////////////////////
///////////////////  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)
{
   adouble x10 = initial_states[ 0 ];
   adouble x20 = initial_states[ 1 ];

   e[ 0 ] = x10;
   e[ 1 ] = x20;

}


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

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


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


int main(void)
{

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

    Alg  algorithm;
    Sol  solution;
    Prob problem;

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

    problem.name        							= "Rayleigh problem";
    problem.outfilename                	 	= "rayleigh.txt";

////////////////////////////////////////////////////////////////////////////
////////////  Declare 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 					= 1;
    problem.phases(1).nevents   					= 2;
    problem.phases(1).npath     					= 1;
    problem.phases(1).nodes                  << 80;

    psopt_level2_setup(problem, algorithm);

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

    DMatrix x, u, t;
    DMatrix lambda, H, mu;

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


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

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

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

    problem.phases(1).bounds.lower.events(0) 		= -5.0;
    problem.phases(1).bounds.lower.events(1) 		= -5.0;

    problem.phases(1).bounds.upper.events(0) 		= -5.0;
    problem.phases(1).bounds.upper.events(1) 		= -5.0;

    problem.phases(1).bounds.lower.path(0) 			= -100.0;
    problem.phases(1).bounds.upper.path(0) 			= 0.0;

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


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


    problem.integrand_cost 								= &integrand_cost;
    problem.endpoint_cost 									= &endpoint_cost;
    problem.dae 												= &dae;
    problem.events 											= &events;
    problem.linkages											= &linkages;

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

    DMatrix state_guess(2,30);

    state_guess   << -5.0*ones(1,30),
                     -5.0*ones(1, 30);

    problem.phases(1).guess.controls       = zeros(1,30);
    problem.phases(1).guess.states         = state_guess;
    problem.phases(1).guess.time           = linspace(0.0, 4.5, 30);




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


    algorithm.nlp_method                  = "IPOPT";
    algorithm.scaling                     = "automatic";
    algorithm.derivatives                 = "automatic";
    algorithm.collocation_method          = "Hermite-Simpson";
    algorithm.nlp_iter_max                = 1000;
    algorithm.nlp_tolerance               = 1.e-10;

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


    psopt(solution, problem, algorithm);


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

    x      = solution.get_states_in_phase(1);
    u      = solution.get_controls_in_phase(1);
    t      = solution.get_time_in_phase(1);
    lambda = solution.get_dual_costates_in_phase(1);
    mu      = solution.get_dual_path_in_phase(1);
    H      = solution.get_dual_hamiltonian_in_phase(1);

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


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

    plot(t,x,problem.name, "time (s)", "states", "x1 x2" );
    plot(t,u,problem.name ,"time (s)", "control", "u" );
    plot(t,lambda,problem.name ,"time (s)", "costates", "p1 p2");
    plot(t,mu, problem.name,"time (s)", "mu", "mu");

    plot(t,x,problem.name, "time (s)", "states", "x1 x2", "pdf", "rayleigh_states.pdf" );
    plot(t,u,problem.name ,"time (s)", "control", "u", "pdf", "rayleigh_control.pdf" );
    plot(t,lambda,problem.name ,"time (s)", "costates", "l1 l2", "pdf", "rayleigh_costates.pdf");
    plot(t,mu, problem.name,"time (s)", "mu", "mu", "pdf", "rayleigh_mu.pdf");

}

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

