#include "program.hpp"
#include <cassert>
#include <iostream>
#include <armadillo>
#include <stdio.h>
#include <stdlib.h>
#include "GPM.h"

using namespace Ipopt;
using namespace std;
using namespace arma;

/* Constructor. */
Program::Program()
{}

Program::~Program()
{}

////////////////////////////*赋值函数(自编)*/////////////////////////////////
Index var_num, constraint_num, nnz_jac, nnz_h;
vec Fmin, Fmax, Zmin, Zmax, z0, z_ans0, zl_temp0, zu_temp0, lambda_temp0;
uvec first_col, N_Col;
auxdata1 auxdata; 
Grad_struct_info0 Grad_struct_info;
Jacb_struct_info0 Jacb_struct_info;
sp_mat Jacb_struct0, I_Z_struct0;
field<mat> D_phasei;
mat D_phase, W_phase, deta_t_in_col;
int Col_total, length_state, length_control, length_integral, length_parameter, length_path, length_event, interval_num, length_Z0;
char *deri_supplier;
double eps0;
double* Objective_ans0;
int length_F0;
int(*Dynamics_ptr)(mat*, mat*, auxdata1, mat*, mat*, mat*);
int(*Objective_ptr)(input0*, double*, mat*);

void get_var_num(int varnum0) {
	var_num = varnum0;
	z_ans0 = zeros<vec>(var_num);
	zl_temp0 = zeros<vec>(var_num);
	zu_temp0 = zeros<vec>(var_num);
}
void get_constraint_num(int constraint_num0) {
	constraint_num = constraint_num0;
	lambda_temp0 = zeros<vec>(constraint_num);
}
void get_nnz_jac(int nnz_jac0) {
	nnz_jac = nnz_jac0;
}
void get_nnz_h(int nnz_h0) {
	nnz_h = nnz_h0;
}
void get_x0(vec x00) {
	z0 = x00;
}
void get_bounds(vec Fmin0, vec Fmax0, vec Zmin0, vec Zmax0) {
	Fmin = Fmin0;
	Fmax = Fmax0;
	Zmin = Zmin0;
	Zmax = Zmax0;
}
void get_data(auxdata1 auxdata00, mat W_phase00, mat deta_t_in_col00, int Col_total00, int length_state00, int length_control00,
	Grad_struct_info0 Grad_struct_info00, int length_integral00, double eps00, mat D_phase00, uvec first_col00, int length_path00,
	int length_event00, sp_mat Jacb_struct00, field<mat> D_phasei0, uvec N_Col00, sp_mat I_Z_struct00, int interval_num00,
	int length_Z00, int length_parameter00, Jacb_struct_info0 Jacb_struct_info00, char *deri_supplier00, int length_F000,
	int(*Dynamics_ptr00)(mat* StateVar, mat* ControlVar, auxdata1 auxdata, mat* dState, mat* integrand, mat* path),
	int(*Objective_ptr00)(input0* input, double* output, mat* event)) {
	auxdata = auxdata00;
	W_phase = W_phase00;
	deta_t_in_col = deta_t_in_col00;
	Col_total = Col_total00;
	length_state = length_state00;
	length_control = length_control00;
	Grad_struct_info = Grad_struct_info00;
	length_integral = length_integral00;
	eps0 = eps00;
	D_phase = D_phase00;
	first_col = first_col00;
	length_path = length_path00;
	length_event = length_event00;
	Jacb_struct0 = Jacb_struct00;
	D_phasei = D_phasei0;
	N_Col = N_Col00;
	I_Z_struct0 = I_Z_struct00;
	interval_num = interval_num00;
	length_Z0 = length_Z00;
	length_parameter = length_parameter00;
	Jacb_struct_info = Jacb_struct_info00;
	deri_supplier = deri_supplier00;
	length_F0 = length_F000;
	Dynamics_ptr = Dynamics_ptr00;
	Objective_ptr = Objective_ptr00;
}

//初始化静态成员变量
vec Program::x = zeros<vec>(var_num);
vec Program::z_L = zeros<vec>(var_num);
vec Program::z_U = zeros<vec>(var_num);
vec Program::lambda = zeros<vec>(constraint_num);


///////////////////////////*重载函数部分*///////////////////////////////////////////
bool Program::get_nlp_info(Index& n, Index& m, Index& nnz_jac_g,
                         Index& nnz_h_lag, IndexStyleEnum& index_style)
{
  // MyNLP.hpp 中描述的问题有2个变量，x1，&x2。
  n = var_num;

  // 一个等式约束
  m = constraint_num;

  // 雅可比矩阵的非零值数目
  nnz_jac_g = nnz_jac;

  // 海森矩阵的非零值数目
  nnz_h_lag = nnz_h;

  // 我们对行/列条目使用标准的fortran索引样式。
  index_style = TNLP::C_STYLE;
  return true;
}

bool Program::get_bounds_info(Index n, Number* x_l, Number* x_u,
                            Index m, Number* g_l, Number* g_u)
{
  // 在这里，我们在 get_nlp_info 中给 IPOPT 的 n 和 m 被传回给我们。
  // 如果需要的话，我们可以通过assert来确保它们是我们认为的那样。
  assert(n == var_num);
  assert(m == constraint_num);

  // 变量的上下界
  vector_ARM2NLP(Zmin, x_l, var_num);
  vector_ARM2NLP(Zmax, x_u, var_num);

  // 约束的上下界
  vector_ARM2NLP(Fmin, g_l, constraint_num);
  vector_ARM2NLP(Fmax, g_u, constraint_num);
  return true;
}

bool Program::get_starting_point(Index n, bool init_x, Number* x,
                               bool init_z, Number* z_L, Number* z_U,
                               Index m, bool init_lambda,
                               Number* lambda)
{
  
  // 这里，我们假设我们只有x的起始值。
  // 如果你对自己的NLP进行编码，如果你愿意的话，你可以为其他数值提供起始值。
  assert(init_x == true);
  assert(init_z == false);
  assert(init_lambda == false);

  // 自变量的起点值
  vector_ARM2NLP(z0, x, n);
  // 协态变量的起点值
  // 约束乘子的起点值
  return true;
}

bool Program::eval_f(Index n, const Number* x, bool new_x, Number& obj_value)
{
	// 返回目标函数的值
	vec Z1 = zeros(n, 1);;
	vector_NLP2ARM(x, &Z1, n);
	double J = 0;
	Objective_ipopt(Z1, deta_t_in_col, Col_total, auxdata, W_phase, length_state, length_control, length_parameter, Dynamics_ptr, Objective_ptr, &J);
	// cout << J << endl;////////////
	obj_value = J;
	return true;
}

bool Program::eval_grad_f(Index n, const Number* x, bool new_x, Number* grad_f)
{
	// 返回目标函数grad_{x} f(x)的梯度。
	assert(n == var_num);
	sp_mat Grad;
	vec Z1 = zeros(n, 1);;
	vector_NLP2ARM(x, &Z1, n);
	Grad_pro(Z1, Grad_struct_info, deta_t_in_col, Col_total, auxdata, W_phase, length_state, length_control, length_integral, 
			length_parameter, eps0, deri_supplier, Dynamics_ptr, Objective_ptr, &Grad);
	// Grad.print();////////////
	vec Grad_vec = mat(Grad).col(0);
	vector_ARM2NLP(Grad_vec, grad_f, n);
	return true;
}

bool Program::eval_g(Index n, const Number* x, bool new_x, Index m, Number* g)
{
	// 返回约束条件的值：g(x)
	vec Z1 = zeros(var_num, 1);;
	vector_NLP2ARM(x, &Z1, var_num);
	vec F1;
	F_pro(Z1, D_phase, first_col, deta_t_in_col, length_path, Col_total, auxdata, W_phase, var_num, 
		length_event, length_state, length_control, length_integral, length_parameter, Dynamics_ptr, Objective_ptr, &F1);
	//F1.print("F1");/////////////////
	vector_ARM2NLP(F1, g, constraint_num);
	return true;
}

bool Program::eval_jac_g(Index n, const Number* x, bool new_x,
                       Index m, Index nele_jac, Index* iRow, Index *jCol,
                       Number* values)
{
	if (values == NULL) {
		// 返回约束条件的雅各布结构。
		sp_ARM2NLP(Jacb_struct0, iRow, jCol, nnz_jac);
	  }
	else {		
		// 返回约束条件的jacobian值。
		// 赋值
		vec Z1 = zeros(n, 1);;
		vector_NLP2ARM(x, &Z1, n);
		sp_mat Jacb;
		Jacb_pro(Z1, D_phasei, N_Col, deta_t_in_col, I_Z_struct0, Jacb_struct_info, length_F0, length_path, Col_total, auxdata, W_phase, length_event, 
				length_state, length_control, length_integral, length_parameter, interval_num, eps0, deri_supplier, length_Z0, Dynamics_ptr, Objective_ptr, &Jacb);
		// Jacb.print("Jacb");////////////////////////
		values_set(Jacb, Jacb_struct0, values, nnz_jac);
	}
  return true;
}

bool Program::eval_h(Index n, const Number* x, bool new_x,
                   Number obj_factor, Index m, const Number* lambda,
                   bool new_lambda, Index nele_hess, Index* iRow,
                   Index* jCol, Number* values)
{
  if (values == NULL) {
    // 返回该结构。这是一个对称矩阵，只填写左下角的三角形。
    iRow[0] = 0;
    jCol[0] = 0;

  }
  else {
    // 返还值
    values = 0;
  }

  return true;
}

void Program::finalize_solution(SolverReturn status,
                              Index n, const Number* x, const Number* z_L, const Number* z_U,
                              Index m, const Number* g, const Number* lambda,
                              Number obj_value,
			      const IpoptData* ip_data,
			      IpoptCalculatedQuantities* ip_cq)
{	
	Status_Display(status);
	vector_NLP2ARM(x, &z_ans0, var_num);
	vector_NLP2ARM(z_L, &zl_temp0, var_num);
	vector_NLP2ARM(z_U, &zu_temp0, var_num);
	vector_NLP2ARM(lambda, &lambda_temp0, constraint_num);
	Program::x = z_ans0;
	Program::z_L = zl_temp0;
	Program::z_U = zu_temp0;
	Program::lambda = lambda_temp0;
	// 这里是我们将解决方案存储到变量，或写入文件等，以便我们可以使用解决方案。
	// 由于解决方案是显示在控制台的，所以我们目前在这里什么也不做。
}