#include <iostream>
#include <fstream>
#include <map>
#include <math.h>
#include "planrec.hxx"
#include "a_h_1_rp.hxx"

#include <ff_to_aptk.hxx>
#include <strips_prob.hxx>
#include <fluent.hxx>
#include <action.hxx>
#include <cond_eff.hxx>
#include <strips_state.hxx>
#include <fwd_search_prob.hxx>
#include <mutex_set.hxx>

#include <h_1.hxx>
#include <rp_heuristic.hxx>

#include <aptk/open_list.hxx>
#include <aptk/string_conversions.hxx>
#include <aptk/at_bfs.hxx>
#include <aptk/at_gbfs.hxx>

#include <fstream>

#include <boost/program_options.hpp>

namespace po = boost::program_options;

using	aptk::STRIPS_Problem;

using	aptk::agnostic::Fwd_Search_Problem;
using	aptk::Action;

using aptk::agnostic::H1_Heuristic;
using	aptk::agnostic::H_Max_Evaluation_Function;

using	aptk::agnostic::H_Add_Evaluation_Function;
using	aptk::agnostic::Relaxed_Plan_Heuristic;

using aptk::agnostic::recog::RP_H1_Heuristic;

using aptk::search::Open_List;
using	aptk::search::bfs::AT_BFS_SQ_SH;
using aptk::search::bfs::AT_GBFS_SQ_SH;
using	aptk::search::Node_Comparer;

typedef	aptk::search::bfs::Node< aptk::State >	Search_Node;
typedef	Node_Comparer< Search_Node >	Tie_Breaking_Algorithm;
typedef	Open_List< Tie_Breaking_Algorithm, Search_Node >	BFS_Open_List;

typedef	H1_Heuristic<Fwd_Search_Problem, H_Max_Evaluation_Function>	H_Max_Fwd;
typedef	AT_BFS_SQ_SH< Fwd_Search_Problem, H_Max_Fwd, BFS_Open_List >	BFS_H_Max_Fwd;

typedef	H1_Heuristic<Fwd_Search_Problem, H_Add_Evaluation_Function>	H_Add_Fwd;
typedef	Relaxed_Plan_Heuristic< Fwd_Search_Problem, H_Add_Fwd >		H_Add_Rp_Fwd;

typedef AT_BFS_SQ_SH<Fwd_Search_Problem,H_Add_Rp_Fwd, BFS_Open_List> BFS_H_RP_Fwd;
typedef AT_GBFS_SQ_SH<Fwd_Search_Problem, H_Add_Rp_Fwd, BFS_Open_List> GBFS_H_RP_Fwd;

typedef RP_H1_Heuristic<Fwd_Search_Problem, H_Max_Evaluation_Function> H_A_MAX_Fwd;

// Greedy Best First Seach and Novelty, Nir's Code
// ----------------------------------------------------------------------

#include "gbfs_f_local/src/new_node_comparer.hxx"
#include "gbfs_f_local/src/dynamic_novelty_partition.hxx"
#include "gbfs_f_local/src/novelty_partition.hxx"
#include "gbfs_f_local/src/novelty_partition_2.hxx"
#include "gbfs_f_local/src/recursive_novelty_partition.hxx"
#include "gbfs_f_local/src/at_gbfs_2h.hxx"

#include "gbfs_f_local/src/at_gbfs_3h.hxx"
#include "gbfs_f_local/src/at_gbfs_alternate.hxx"
#include "gbfs_f_local/src/at_gbfs_2h_exp.hxx"
#include "gbfs_f_local/src/2h_exp_node.hxx"

#include <novelty.hxx>
#include "gbfs_f_local/src/iw.hxx"

using 	aptk::agnostic::Dynamic_Novelty_Partition;
using 	aptk::agnostic::Novelty_Partition;
using 	aptk::agnostic::Recursive_Novelty_Partition;
using 	aptk::agnostic::Novelty_Partition_2;

using	aptk::search::Node_Comparer_2H;
using	aptk::search::Node_Comparer_2H_ehc;
using	aptk::search::Node_Comparer_Second;
using	aptk::search::Node_Comparer_Third;
using	aptk::search::Node_Comparer_3H_FN;
using	aptk::search::gbfs_2h::AT_GBFS_2H;
using	aptk::search::gbfs_3h::AT_GBFS_3H;
using	aptk::search::gbfs_alternate::AT_GBFS_Alternate;
using	aptk::search::gbfs_2h_exp::AT_GBFS_2H_EXP;

using aptk::agnostic::Novelty;
using	aptk::search::brfs::IW;

typedef		aptk::search::gbfs_2h::Node< Fwd_Search_Problem, aptk::State >	Search_Node_gbfs;
typedef		Node_Comparer_2H< Search_Node_gbfs >					Tie_Breaking_Algorithm_gbfs;
typedef		Open_List< Tie_Breaking_Algorithm_gbfs, Search_Node_gbfs >		   BFS_Open_List_gbfs;

typedef   Dynamic_Novelty_Partition<Fwd_Search_Problem, Search_Node_gbfs> H_Dyn_Novel_Fwd;
typedef   Novelty_Partition<Fwd_Search_Problem, Search_Node_gbfs> H_Novel_Partition_Fwd;

typedef   AT_GBFS_2H< Fwd_Search_Problem, H_Dyn_Novel_Fwd, H_Add_Rp_Fwd, BFS_Open_List_gbfs > BWFS_Dynamic_Novel_HRP_Fwd;

// ----------------------------------------------------------------------

void print_plan_to_cout(std::vector< int >& plan, Fwd_Search_Problem& search_prob){
	for ( unsigned k = 0; k < plan.size(); k++ ) {
		const aptk::Action& a = *(search_prob.task().actions()[ plan[k] ]);
		std::cout  << a.signature() << ": ";
    for (auto f: a.prec_vec()){
      std::cout << search_prob.task().fluents()[f]->signature() << ", ";
    }

    std::cout << std::endl;
	}
}

// Do search for plan recognition
void do_pr_search(  Fwd_Search_Problem& search_prob,  std::vector< aptk::Action_Idx >& plan, float& cost,  bool anytime, bool delayed, bool greedy, bool optimal  ) {

  cost = infty;
  anytime = false;

  if (optimal){
    BFS_H_Max_Fwd engine( search_prob );
    engine.set_greedy( greedy );
    engine.set_delay_eval( delayed );
    engine.start();
    engine.find_solution( cost, plan );
    //unsigned h_init = 0;
    //BWFS_Dynamic_Novel_HRP_Fwd engine (search_prob);
    //engine.set_arity(0,h_init);
    //engine.start();
    //engine.find_solution(cost, plan);

  }else {
    GBFS_H_RP_Fwd engine( search_prob );
    engine.start();
    engine.find_solution( cost, plan );
  }
}

void do_heuristic(Fwd_Search_Problem& search_prob, float& cost, std::vector<aptk::Action_Idx>& relaxed_plan, std::vector<unsigned>& observations){

  H_Add_Rp_Fwd h_ff(search_prob);
  //H_A_MAX_Fwd h_amax(search_prob, observations);

  Search_Node* root = new Search_Node( search_prob.init(), 0.0f,-1, NULL );
  std::vector<aptk::Action_Idx> pref_ops;
  h_ff.eval(*(root->state()), cost, pref_ops, relaxed_plan);
  //h_amax.eval(*(root->state()), cost);
}

// Translate the problem to a plan recognition problem
void translate_problem(const aptk::STRIPS_Problem& prob, aptk::STRIPS_Problem& trans_prob, aptk::Fluent_Vec& w_observation, aptk::Fluent_Vec& wo_observation, std::map<std::string, unsigned int>& idx_to_trans_idx, std::vector<unsigned>& observations, std::vector<unsigned> init_state){

	const std::vector<const aptk::Fluent*> f = prob.fluents();
	const std::vector<const aptk::Action*> as = prob.actions();
	aptk::Fluent_Vec init = init_state;
	aptk::Fluent_Vec goal = prob.goal();
	unsigned index = 0;

	// Add each fluent to transformed problem
	for (unsigned int i = 0; i < prob.num_fluents(); i++){
		trans_prob.add_fluent(trans_prob, (*f[i]).signature());
		idx_to_trans_idx.insert(std::pair<std::string, unsigned int>(std::to_string((*f[i]).index()),index++));

	}

	// Add fluents for each observed action
 	for (unsigned int i = 0; i < observations.size(); i++){
		std::string name = "pa" + std::to_string(observations[i]);
		std::string name_neg = "npa" + std::to_string(observations[i]);

		trans_prob.add_fluent(trans_prob, name);
		idx_to_trans_idx.insert(std::pair<std::string, unsigned int>(name ,index++));

		trans_prob.add_fluent(trans_prob, name_neg);
		idx_to_trans_idx.insert(std::pair<std::string, unsigned int>(name_neg ,index++));

		w_observation.push_back(idx_to_trans_idx[name]);
		wo_observation.push_back(idx_to_trans_idx[name_neg]);
		init.push_back(idx_to_trans_idx[name_neg]);

	}

	// Add actions to translated problem
	for (const aptk::Action * a : as){

		auto address = std::find(observations.begin(), observations.end(), a->index());
		int index = std::distance(observations.begin(), address);

		// The actions corrosponding to the first observation adds the corrosponding fluent unconditionally
		if ( address == observations.begin()){
			aptk::Fluent_Vec add = a->add_vec();
			aptk::Fluent_Vec del = a->del_vec();

			add.push_back(idx_to_trans_idx["pa" + std::to_string(observations[0])]);
			del.push_back(idx_to_trans_idx["npa" + std::to_string(observations[0])]);

			trans_prob.add_action(trans_prob, a->signature(), a->prec_vec(), add, del, a->ceff_vec(), a->cost());

		// Any actions corrosponding to any other observation add the corrospoding fluent only if the fluent corrosponding
		// to the observation immediately prior to this one has already been added.
		 } else if (address != observations.end() ){
			aptk::Fluent_Vec ce_add;
			aptk::Fluent_Vec ce_del;
			aptk::Fluent_Vec ce_prec;
			aptk::Conditional_Effect_Vec ceffs( a->ceff_vec() );
			aptk::Conditional_Effect * ceff = new aptk::Conditional_Effect(trans_prob);

			int prev_obs = idx_to_trans_idx["pa" + std::to_string(observations[index - 1])];
			int curr_obs = idx_to_trans_idx["pa" + std::to_string(a->index())];
			int neg_curr_obs = idx_to_trans_idx["npa" + std::to_string(a->index())];

			ce_prec.push_back(prev_obs);
			ce_add.push_back(curr_obs);
			ce_del.push_back(neg_curr_obs);

			ceff->define(ce_prec, ce_add, ce_del);

			ceffs.push_back(ceff);

			trans_prob.add_action(trans_prob, a->signature(), a->prec_vec(), a->add_vec(), a->del_vec(), ceffs, a->cost());

		// Unobserved action is added as is.
		}else
			trans_prob.add_action(trans_prob, a->signature(), a->prec_vec(), a->add_vec(), a->del_vec(), a->ceff_vec(), a->cost());
	}

	// Set names and initial and goal states the same as
	trans_prob.set_domain_name(prob.domain_name());
	trans_prob.set_problem_name(prob.problem_name());
	trans_prob.set_init(trans_prob, init);
	trans_prob.set_goal(trans_prob, goal);

	trans_prob.make_action_tables();
  unsigned int w_back = w_observation.back();
  unsigned int wo_back = wo_observation.back();
	w_observation.clear();
	wo_observation.clear();
	w_observation.push_back((w_back));
	wo_observation.push_back((wo_back));
  //std::cout << "w_observation: " << trans_prob.fluents()[w_observation[0]]->signature() << "\n";
  //std::cout << "wo_observation: " << trans_prob.fluents()[wo_observation[0]]->signature() << "\n";
}

// Main function for plan recognition
std::vector<double> plan_recognition(const aptk::STRIPS_Problem& prob, std::vector<aptk::Fluent_Vec>& goal_set, std::vector<double> priors, std::vector<unsigned>& observations, aptk::Fluent_Vec& init, float beta, bool optimal) {

	aptk::STRIPS_Problem trans_prob;
	aptk::Fluent_Vec w_observation;
	aptk::Fluent_Vec wo_observation;
	std::map<std::string, unsigned int> idx_to_trans_idx;
  std::vector<float> goal_cost_change;
  float sum (0);

	// Translate problem
	translate_problem(prob, trans_prob, w_observation, wo_observation, idx_to_trans_idx, observations, init);

  // For each goal candidate perform two searches:
  // One where observed actions must be used
  // One where observed actions must not be used
	for (unsigned i = 0; i < goal_set.size(); i++){

    //std::cout << "\nGoal No.: " << i << " / " << goal_set.size() << "\n";
		// Set goal to candidate goal and that observations have occured
		aptk::Fluent_Vec g = goal_set[i];
		aptk::Fluent_Vec pos_goal(g);
		pos_goal.insert(pos_goal.end(), w_observation.begin(), w_observation.end());
		trans_prob.set_goal(trans_prob, pos_goal);

		Fwd_Search_Problem pos_search_prob(&trans_prob);
		std::vector< aptk::Action_Idx > plan;
		float pos_cost(0);

    // Perform search with observations
    do_pr_search(pos_search_prob, plan, pos_cost, false, false, false, optimal);

    //std::cout << "Pos Plan: \n";
    //print_plan_to_cout(plan, pos_search_prob);
    //std::cout << std::endl;
    //print_vector(plan);
    plan.clear();

		// Set goal to candidate goal and that observations must not occur
		aptk::Fluent_Vec neg_goal(g);
		neg_goal.insert(neg_goal.end(), wo_observation.begin(), wo_observation.end());
    /*
     *if (observations == obs){
     *  std::cout << "Goals of PICKUP S and STACK S A:" << std::endl;
     *  for (unsigned int i : neg_goal){
     *    std::cout << (*trans_prob.fluents()[i]).signature() << std::endl;
     *  }
     *}
     */
		trans_prob.set_goal(trans_prob, neg_goal);

		Fwd_Search_Problem neg_search_prob(&trans_prob);
		float neg_cost(0);

    // Perform search without observations
    do_pr_search(neg_search_prob, plan, neg_cost, false, false, false, optimal);

    /*
     *std::cout << "Neg Plan: \n";
     *print_plan_to_cout(plan, neg_search_prob);
     *print_vector(plan);
     */
    plan.clear();

    // Calculate liklihood
		const float pOGin = -1.0f*beta*(neg_cost - pos_cost);
		const float pOGinexp = exp(pOGin);
		const float pOG = 1.0f/(1.0f + pOGinexp);
		const float pOGpG = pOG*priors[i];

    //std::cout << i << ":\tprior: " << priors[i] << std::endl;
    //std::cout <<  "\tneg - pos: " << neg_cost << " - " << pos_cost << std::endl;
    //std::cout <<  " \tpogin: " << pOGin << std::endl;
    //std::cout <<  " \tpoginexp: " << pOGinexp << std::endl;
    //std::cout <<  " \tpog: " << pOG << std::endl;
    //std::cout <<  " \tpogpg: " << pOGpG << std::endl;

		sum += pOGpG;

		goal_cost_change.push_back(pOGpG);
	}
  if (sum != 0){
    auto normv = normalize_vector(goal_cost_change);
    return normv;
  }else {
    std::vector<double> new_priors;
    for (unsigned i = 0; i < priors.size(); i++)
      new_priors.push_back(priors[i]);

    return new_priors;
  }
}

void clone_problem(const aptk::STRIPS_Problem& prob, aptk::STRIPS_Problem& trans_prob, std::vector<unsigned>& observations, float cost, std::vector<unsigned> init_state){

	const std::vector<const aptk::Fluent*> f = prob.fluents();
	const std::vector<const aptk::Action*> as = prob.actions();
	aptk::Fluent_Vec init = init_state;
	aptk::Fluent_Vec goal = prob.goal();

	// Add each fluent to transformed problem
	for (unsigned int i = 0; i < prob.num_fluents(); i++){
		trans_prob.add_fluent(trans_prob, (*f[i]).signature());
	}

	// Add actions to translated problem
	for (const aptk::Action * a : as){

		auto address = std::find(observations.begin(), observations.end(), a->index());

		// The actions corrosponding to the observations are given cost "cost", all other actions are the same.
    float a_cost = a->cost();
    //if (address != observations.end() )
      //a_cost = cost;
    //else
      //a_cost = a->cost();

    trans_prob.add_action(trans_prob, a->signature(), a->prec_vec(), a->add_vec(), a->del_vec(), a->ceff_vec(),a_cost);

	}

	// Set names and initial and goal states the same as
	trans_prob.set_domain_name(prob.domain_name());
	trans_prob.set_problem_name(prob.problem_name());
	trans_prob.set_init(trans_prob, init);
	trans_prob.set_goal(trans_prob, goal);

	trans_prob.make_action_tables();

}

void translate_problem_approx_pos(const aptk::STRIPS_Problem& prob, aptk::STRIPS_Problem& trans_prob, std::vector<unsigned>& observations, std::vector<unsigned> init_state){

	const std::vector<const aptk::Fluent*> f = prob.fluents();
	const std::vector<const aptk::Action*> as = prob.actions();
	aptk::Fluent_Vec init = init_state;
	aptk::Fluent_Vec goal = prob.goal();

	// Add each fluent to transformed problem
	for (unsigned int i = 0; i < prob.num_fluents(); i++){
		trans_prob.add_fluent(trans_prob, (*f[i]).signature());

	}

	// Add actions to translated problem
	for (const aptk::Action * a : as){

		auto address = std::find(observations.begin(), observations.end(), a->index());

		// The effects of the actions corrosponding to the observations are made true in the initial state,
		// the number of observations is then added to the length of the relaxed plan.
    if (address != observations.end() ){
      for (unsigned int i = 0; i < a->add_vec().size(); i++ ){
        init.push_back(a->add_vec()[i]);
      }
    }
    trans_prob.add_action(trans_prob, a->signature(), a->prec_vec(), a->add_vec(), a->del_vec(), a->ceff_vec(), a->cost());

	}

	// Set names and initial and goal states the same as
	trans_prob.set_domain_name(prob.domain_name());
	trans_prob.set_problem_name(prob.problem_name());
	trans_prob.set_init(trans_prob, init);
	trans_prob.set_goal(trans_prob, goal);

	trans_prob.make_action_tables();

}

// Purpose: Approximate a planning problem which requires that a certain action sequence may not be
// satisfied. This is acheived by taking the root value of a delete relaxed heuristic. In order for
// the solution to not satisfy the action sequence the heuristic is run once for each action in the
// sequence. The action being considered is then disallowed. The minimum value over these runs is
// taken as the approximate cost of optimal path for the given planning problem without satisfying
// the actions sequence. This should always be lesser or equal to the  approximation_wo_delete_relax_break_sequence
// method.
// Arguments:
//  trans_prob_neg: This is the planning problem to be solved.
//  observations: This is the sequence of actions which must not be satisfied.
// Return:
//  The approximate cost of an optimal path which solves the problem.
float approximation_wo_delete_relax_disallow_action(aptk::STRIPS_Problem& trans_prob_neg, std::vector<unsigned> observations){
  float min(infty);

  // For each action in the observation sequence set its cost to infinity in turn
  for (unsigned o: observations){
    aptk::Action * a = trans_prob_neg.actions()[o];
    float temp_a_cost = a->cost();
    a->set_cost(infty);

		// The heuristic at the root node is taken the approximate cost of an optimal plan
		// which does not use the considered action
		Fwd_Search_Problem neg_search_prob(&trans_prob_neg);
		float neg_cost(0);
    std::vector<aptk::Action_Idx> plan;
    do_heuristic(neg_search_prob, neg_cost, plan, observations);
    plan.clear();

    if (min > neg_cost){
      min = neg_cost;
    }

    a->set_cost(temp_a_cost);
  }

  return min;
}

// Purpose: Approximate a planning problem which requires that a certain action sequence may not be
// satisfied. This is acheived by taking the root value of a delete relaxed heuristic. In order for
// the solution to not satisfy the action sequence the heuristic is run once for each action in the
// sequence. The action being considered is then disallowed until another action in the sequence is
// used. The initial state is then advanced through the action being considered and its cost is added
// to every following run. The minimum value over these runs is taken as an approximate optimal path
// for the given planning problem without satisfying the actions sequence.
// Arguments:
//  trans_prob_neg: This is the planning problem to be solved.
//  observations: This is the sequence of actions which must not be satisfied.
// Return:
//  The approximate cost of an optimal path which solves the problem.
float approximation_wo_delete_relax_break_sequence(aptk::STRIPS_Problem& trans_prob_neg, std::vector<unsigned> observations){

  float min(infty);

  // The "NEXT" fluent is used to determine if the current action may be taken.
  unsigned next_index = trans_prob_neg.add_fluent(trans_prob_neg, "NEXT");

  // The true inial state is stored to reconstruct the problem
  auto true_init  (trans_prob_neg.init());

  // At present the code only allows for uniform cost actions
  // So the amount of observatons seen so far is added to the cost of the plan.
  int obs_count = 0;

  // The "NEXT" fluent is added to the add lists of all actions in the observation sequence.
  for (auto o: observations){
    auto a = trans_prob_neg.actions()[o];
    a->add_vec().push_back(next_index);
  }

  // For each observation a heurisitic value is taken and the minimum is returned.
  for (unsigned o: observations){
    aptk::Action * a = trans_prob_neg.actions()[o];

    // The observation currently considered has the "NEXT" fluent added to its preconditions
    // This means that it may be used only after another observation.
    a->prec_vec().push_back(next_index);

		// The heuristic value of the problem is found
		Fwd_Search_Problem neg_search_prob(&trans_prob_neg);
		float neg_cost(0);
    std::vector<aptk::Action_Idx> plan;
    do_heuristic(neg_search_prob, neg_cost, plan, observations);
    plan.clear();

    // The action cost of the previous observations is added to the cost
    neg_cost += obs_count;

    if (min > neg_cost){
      min = neg_cost;
    }

    // The "NEXT" fluent is removed as a precondition and an effect
    a->prec_vec().pop_back();
    a->add_vec().pop_back();

    // The initial state is progressed through the considered action
    aptk::State root( trans_prob_neg);
    root.set(trans_prob_neg.init() );

    aptk::State* next_root = root.progress_through(*a);
    STRIPS_Problem::set_init(trans_prob_neg, next_root->fluent_vec());

    // The cost is incremented to reflect the changing initial state
    // TODO: This will need to be changed to:
    obs_count += a->cost();
    //obs_count++;
  }

  // The initial state of the problem is reset
  STRIPS_Problem::set_init(trans_prob_neg,true_init);

  return min;
}

std::vector<double> plan_recognition_approx(const aptk::STRIPS_Problem& prob, std::vector<aptk::Fluent_Vec>& goal_set, std::vector<double> priors, std::vector<unsigned>& observations, aptk::Fluent_Vec& init, float beta = 1, unsigned approx_type = 1) {

	aptk::STRIPS_Problem trans_prob_pos;
	aptk::STRIPS_Problem trans_prob_neg;
  std::vector<float> goal_cost_change;
  float sum (0);
  clone_problem(prob, trans_prob_neg, observations, infty, init);
  translate_problem_approx_pos(prob, trans_prob_pos, observations, init);

  // For each goal candidate perform two searches:
  // One where observed actions must be used
  // One where observed actions must not be used
	for (unsigned i = 0; i < goal_set.size(); i++){

    //std::cout << "\nGoal No.: " << i << " / " << goal_set.size() << "\n";

		// Set goal to candidate goal and that observations have occured
		aptk::Fluent_Vec g = goal_set[i];
		aptk::Fluent_Vec pos_goal(g);
		trans_prob_pos.set_goal(trans_prob_pos, pos_goal);

		Fwd_Search_Problem pos_search_prob(&trans_prob_pos);
		float pos_cost(0);
    std::vector<aptk::Action_Idx> plan;

    // Perform search with observations
    do_heuristic(pos_search_prob, pos_cost, plan, observations);
    pos_cost += observations.size();

    //std::cout << "Relaxed Pos Plan: \n";
    //print_plan_to_cout(plan, pos_search_prob);
    plan.clear();

		// Set goal to candidate goal and that observations must not occur
		aptk::Fluent_Vec neg_goal(g);
		trans_prob_neg.set_goal(trans_prob_neg, neg_goal);
    float neg_cost(0);

		if (approx_type == 1){
      Fwd_Search_Problem neg_search_prob(&trans_prob_neg);

      // Perform search without observations
      do_heuristic(neg_search_prob, neg_cost, plan, observations);

      //std::cout << "Relaxed Neg Plan: \n";
      //print_plan_to_cout(plan, pos_search_prob);
      plan.clear();
    } else if (approx_type == 2) {
      neg_cost = approximation_wo_delete_relax_disallow_action(trans_prob_neg, observations);
    } else {
      neg_cost = approximation_wo_delete_relax_break_sequence(trans_prob_neg, observations);
    }

    // Calculate liklihood
		const float pOGin = -1.0f*beta*(neg_cost - pos_cost);
		const float pOGinexp = exp(pOGin);
		const float pOG = 1.0f/(1.0f + pOGinexp);
		const float pOGpG = pOG*priors[i];

    //std::cout << i << ":\tPrior: " << priors[i] << std::endl;
    //std::cout <<  "\tPos - Neg: " << pos_cost << " - " << neg_cost << std::endl;
    //std::cout <<  " \tpOGin: " << pOGin << std::endl;
    //std::cout <<  " \tpOGinexp: " << pOGinexp << std::endl;
    //std::cout <<  " \tpOG: " << pOG << std::endl;
    //std::cout <<  " \tpOGpG: " << pOGpG << std::endl;

		sum += pOGpG;

		goal_cost_change.push_back(pOGpG);
	}

	return normalize_vector(goal_cost_change);
}

std::vector<double> simulate_plan_for_recog(std::vector<aptk::Action_Idx>& plan, aptk::STRIPS_Problem& prob, std::vector<aptk::Fluent_Vec>& goal_set, std::vector<double>& priors, float beta){

  std::vector<unsigned int> observations;
  std::vector<double> probability;

  int index = 0;
  std::set<unsigned> true_goal(prob.goal().begin(), prob.goal().end());
  for (unsigned int i = 0; i < goal_set.size(); i++){
    std::set<unsigned> cand_goal(goal_set[i].begin(),goal_set[i].end());
    if (true_goal == cand_goal)
      index = i;
  }

  for (unsigned int i = 0; i < plan.size(); i++){
    observations.push_back(plan[i]);
    double h_val;
    /*
     *std::cout << "Observations: " << std::endl;
     *for (unsigned int i = 0; i < observations.size(); i++){
     *  std::cout << observations[i] << ": " << prob.actions()[observations[i]]->signature() << std::endl;
     *}
     */
    std::vector<double> belief_probs_at_i = plan_recognition(prob, goal_set, priors, observations, prob.init(), beta, true);

    // Count the number of goals with greater posterior than the goal
    float rank = 1;
    for (unsigned int i = 0; i < belief_probs_at_i.size(); i++){
      if (index != i)
        if (belief_probs_at_i[i] >= belief_probs_at_i[index])
          rank++;
    }

    //h_val = (1 - belief_probs_at_i[index]);
    h_val = rank + (1 - belief_probs_at_i[index]);

    probability.push_back(h_val);
  }

  return probability;
}


#include <boost/python/numeric.hpp>

namespace scitbx { namespace af { namespace boost_python {

#define SCITBX_LOC(pyname, ElementType) \
  boost::python::object \
  flex_##pyname##_as_numpy_array( \
    ref<ElementType, flex_grid<> > const& O, \
    bool optional=false); \
 \
  versa<ElementType, flex_grid<> >* \
  flex_##pyname##_from_numpy_array( \
    boost::python::numeric::array const& arr_obj);

  SCITBX_LOC(bool, bool)
  SCITBX_LOC(int, int)
  SCITBX_LOC(long, long)
  SCITBX_LOC(float, float)
  SCITBX_LOC(double, double)
#define SCITBX_LOC2 std::complex<double>
  SCITBX_LOC(complex_double, SCITBX_LOC2)
#undef SCITBX_LOC2
  SCITBX_LOC(size_t, std::size_t)

#undef SCITBX_LOC

}}} // namespace scitbx::af::boost_python


/*
 [auto_generated]
 boost/numeric/odeint/stepper/dense_output_runge_kutta.hpp

 [begin_description]
 Implementation of the Dense-output stepper for all steppers. Note, that this class does
 not computes the result but serves as an interface.
 [end_description]

 Copyright 2009-2011 Karsten Ahnert
 Copyright 2009-2011 Mario Mulansky

 Distributed under the Boost Software License, Version 1.0.
 (See accompanying file LICENSE_1_0.txt or
 copy at http://www.boost.org/LICENSE_1_0.txt)
 */

#ifndef BOOST_NUMERIC_ODEINT_STEPPER_DENSE_OUTPUT_RUNGE_KUTTA_HPP_INCLUDED
#define BOOST_NUMERIC_ODEINT_STEPPER_DENSE_OUTPUT_RUNGE_KUTTA_HPP_INCLUDED

#include <utility>
#include <stdexcept>

#include <boost/numeric/odeint/util/bind.hpp>

#include <boost/numeric/odeint/util/copy.hpp>

#include <boost/numeric/odeint/util/state_wrapper.hpp>
#include <boost/numeric/odeint/util/is_resizeable.hpp>
#include <boost/numeric/odeint/util/resizer.hpp>

#include <boost/numeric/odeint/stepper/controlled_step_result.hpp>
#include <boost/numeric/odeint/stepper/stepper_categories.hpp>

namespace boost {
namespace numeric {
namespace odeint {

template< class Stepper , class StepperCategory = typename Stepper::stepper_category >
class dense_output_runge_kutta;

/**
 * \brief The class representing dense-output Runge-Kutta steppers.
 * \note In this stepper, the initialize method has to be called before using
 * the do_step method.
 *
 * The dense-output functionality allows to interpolate the solution between
 * subsequent integration points using intermediate results obtained during the
 * computation. This version works based on a normal stepper without step-size
 * control.
 *
 *
 * \tparam Stepper The stepper type of the underlying algorithm.
 */
template< class Stepper >
class dense_output_runge_kutta< Stepper , stepper_tag >
{

public:

    /*
     * We do not need all typedefs.
     */
    typedef Stepper stepper_type;
    typedef typename stepper_type::state_type state_type;
    typedef typename stepper_type::wrapped_state_type wrapped_state_type;
    typedef typename stepper_type::value_type value_type;
    typedef typename stepper_type::deriv_type deriv_type;
    typedef typename stepper_type::wrapped_deriv_type wrapped_deriv_type;
    typedef typename stepper_type::time_type time_type;
    typedef typename stepper_type::algebra_type algebra_type;
    typedef typename stepper_type::operations_type operations_type;
    typedef typename stepper_type::resizer_type resizer_type;
    typedef dense_output_stepper_tag stepper_category;
    typedef dense_output_runge_kutta< Stepper > dense_output_stepper_type;

    /**
     * \brief Constructs the dense_output_runge_kutta class. An instance of the
     * underlying stepper can be provided.
     * \param stepper An instance of the underlying stepper.
     */
    dense_output_runge_kutta( const stepper_type &stepper = stepper_type() )
    : m_stepper( stepper ) , m_resizer() ,
      m_x1() , m_x2() , m_current_state_x1( true ) ,
      m_t() , m_t_old() , m_dt()
    { }

    /**
     * \brief Initializes the stepper. Has to be called before do_step can be
     * used to set the initial conditions and the step size.
     * \param x0 The initial state of the ODE which should be solved.
     * \param t0 The initial time, at which the step should be performed.
     * \param dt0 The step size.
     */
    template< class StateType >
    void initialize( const StateType &x0 , time_type t0 , time_type dt0 )
    {
        m_resizer.adjust_size( x0 , detail::bind( &dense_output_stepper_type::template resize_impl< StateType > , detail::ref( *this ) , detail::_1 ) );
        boost::numeric::odeint::copy( x0 , get_current_state() );
        m_t = t0;
        m_dt = dt0;
    }

    /**
     * \brief Does one time step.
     * \note initialize has to be called before using this method to set the
     * initial conditions x,t and the stepsize.
     * \param system The system function to solve, hence the r.h.s. of the ordinary differential equation. It must fulfill the
     *               Simple System concept.
     * \return Pair with start and end time of the integration step.
     */
    template< class System >
    std::pair< time_type , time_type > do_step( System system )
    {
        m_stepper.do_step( system , get_current_state() , m_t , get_old_state() , m_dt );
        m_t_old = m_t;
        m_t += m_dt;
        toggle_current_state();
        return std::make_pair( m_t_old , m_dt );
    }

    /*
     * The next two overloads are needed to solve the forwarding problem
     */

    /**
     * \brief Calculates the solution at an intermediate point.
     * \param t The time at which the solution should be calculated, has to be
     * in the current time interval.
     * \param x The output variable where the result is written into.
     */
    template< class StateOut >
    void calc_state( time_type t , StateOut &x ) const
    {
        m_stepper.calc_state( x , t , get_old_state() , m_t_old , get_current_state() , m_t );
    }

    /**
     * \brief Calculates the solution at an intermediate point. Solves the forwarding problem
     * \param t The time at which the solution should be calculated, has to be
     * in the current time interval.
     * \param x The output variable where the result is written into, can be a boost range.
     */
    template< class StateOut >
    void calc_state( time_type t , const StateOut &x ) const
    {
        m_stepper.calc_state( x , t , get_old_state() , m_t_old , get_current_state() , m_t );
    }

    /**
     * \brief Adjust the size of all temporaries in the stepper manually.
     * \param x A state from which the size of the temporaries to be resized is deduced.
     */
    template< class StateType >
    void adjust_size( const StateType &x )
    {
        resize_impl( x );
        m_stepper.stepper().resize( x );
    }

    /**
     * \brief Returns the current state of the solution.
     * \return The current state of the solution x(t).
     */
    const state_type& current_state( void ) const
    {
        return get_current_state();
    }

    /**
     * \brief Returns the current time of the solution.
     * \return The current time of the solution t.
     */
    time_type current_time( void ) const
    {
        return m_t;
    }

    /**
     * \brief Returns the last state of the solution.
     * \return The last state of the solution x(t-dt).
     */
    const state_type& previous_state( void ) const
    {
        return get_old_state();
    }

    /**
     * \brief Returns the last time of the solution.
     * \return The last time of the solution t-dt.
     */
    time_type previous_time( void ) const
    {
        return m_t_old;
    }

private:

    state_type& get_current_state( void )
    {
        return m_current_state_x1 ? m_x1.m_v : m_x2.m_v ;
    }

    const state_type& get_current_state( void ) const
    {
        return m_current_state_x1 ? m_x1.m_v : m_x2.m_v ;
    }

    state_type& get_old_state( void )
    {
        return m_current_state_x1 ? m_x2.m_v : m_x1.m_v ;
    }

    const state_type& get_old_state( void ) const
    {
        return m_current_state_x1 ? m_x2.m_v : m_x1.m_v ;
    }

    void toggle_current_state( void )
    {
        m_current_state_x1 = ! m_current_state_x1;
    }

    template< class StateIn >
    bool resize_impl( const StateIn &x )
    {
        bool resized = false;
        resized |= adjust_size_by_resizeability( m_x1 , x , typename is_resizeable<state_type>::type() );
        resized |= adjust_size_by_resizeability( m_x2 , x , typename is_resizeable<state_type>::type() );
        return resized;
    }

    stepper_type m_stepper;
    resizer_type m_resizer;
    wrapped_state_type m_x1 , m_x2;
    bool m_current_state_x1;    // if true, the current state is m_x1
    time_type m_t , m_t_old , m_dt;

};

/**
 * \brief The class representing dense-output Runge-Kutta steppers with FSAL property.
 *
 * The interface is the same as for dense_output_runge_kutta< Stepper , stepper_tag >.
 * This class provides dense output functionality based on methods with step size controlled
 *
 *
 * \tparam Stepper The stepper type of the underlying algorithm.
 */
template< class Stepper >
class dense_output_runge_kutta< Stepper , explicit_controlled_stepper_fsal_tag >
{
public:

    /*
     * We do not need all typedefs.
     */
    typedef Stepper controlled_stepper_type;

    typedef typename controlled_stepper_type::stepper_type stepper_type;
    typedef typename stepper_type::state_type state_type;
    typedef typename stepper_type::wrapped_state_type wrapped_state_type;
    typedef typename stepper_type::value_type value_type;
    typedef typename stepper_type::deriv_type deriv_type;
    typedef typename stepper_type::wrapped_deriv_type wrapped_deriv_type;
    typedef typename stepper_type::time_type time_type;
    typedef typename stepper_type::algebra_type algebra_type;
    typedef typename stepper_type::operations_type operations_type;
    typedef typename stepper_type::resizer_type resizer_type;
    typedef dense_output_stepper_tag stepper_category;
    typedef dense_output_runge_kutta< Stepper > dense_output_stepper_type;

    dense_output_runge_kutta( const controlled_stepper_type &stepper = controlled_stepper_type() )
    : m_stepper( stepper ) , m_resizer() ,
      m_current_state_x1( true ) ,
      m_x1() , m_x2() , m_dxdt1() , m_dxdt2() ,
      m_t() , m_t_old() , m_dt() ,
      m_is_deriv_initialized( false )
    { }

    template< class StateType >
    void initialize( const StateType &x0 , time_type t0 , time_type dt0 )
    {
        m_resizer.adjust_size( x0 , detail::bind( &dense_output_stepper_type::template resize< StateType > , detail::ref( *this ) , detail::_1 ) );
        boost::numeric::odeint::copy( x0 , get_current_state() );
        m_t = t0;
        m_dt = dt0;
        m_is_deriv_initialized = false;
    }

    template< class System >
    std::pair< time_type , time_type > do_step( System system )
    {
        const size_t max_count = 1000;

        if( !m_is_deriv_initialized )
        {
            typename odeint::unwrap_reference< System >::type &sys = system;
            sys( get_current_state() , get_current_deriv() , m_t );
            m_is_deriv_initialized = true;
        }

        controlled_step_result res = fail;
        m_t_old = m_t;
        size_t count = 0;
        do
        {
            res = m_stepper.try_step( system , get_current_state() , get_current_deriv() , m_t ,
                                      get_old_state() , get_old_deriv() , m_dt );
            if( count++ == max_count )
                throw std::overflow_error( "dense_output_controlled_explicit : too much iterations!");
        }
        while( res == fail );
        toggle_current_state();
        return std::make_pair( m_t_old , m_t );
    }

    /*
     * The two overloads are needed in order to solve the forwarding problem.
     */
    template< class StateOut >
    void calc_state( time_type t , StateOut &x ) const
    {
        m_stepper.stepper().calc_state( t , x , get_old_state() , get_old_deriv() , m_t_old ,
                                        get_current_state() , get_current_deriv() , m_t );
    }

    template< class StateOut >
    void calc_state( time_type t , const StateOut &x ) const
    {
        m_stepper.stepper().calc_state( t , x , get_old_state() , get_old_deriv() , m_t_old ,
                                        get_current_state() , get_current_deriv() , m_t );
    }

    template< class StateIn >
    bool resize( const StateIn &x )
    {
        bool resized = false;
        resized |= adjust_size_by_resizeability( m_x1 , x , typename is_resizeable<state_type>::type() );
        resized |= adjust_size_by_resizeability( m_x2 , x , typename is_resizeable<state_type>::type() );
        resized |= adjust_size_by_resizeability( m_dxdt1 , x , typename is_resizeable<deriv_type>::type() );
        resized |= adjust_size_by_resizeability( m_dxdt2 , x , typename is_resizeable<deriv_type>::type() );
        return resized;
    }

    template< class StateType >
    void adjust_size( const StateType &x )
    {
        resize( x );
        m_stepper.stepper().resize( x );
    }

    const state_type& current_state( void ) const
    {
        return get_current_state();
    }

    time_type current_time( void ) const
    {
        return m_t;
    }

    const state_type& previous_state( void ) const
    {
        return get_old_state();
    }

    time_type previous_time( void ) const
    {
        return m_t_old;
    }

    time_type current_time_step( void ) const
    {
        return m_dt;
    }

private:

    state_type& get_current_state( void )
    {
        return m_current_state_x1 ? m_x1.m_v : m_x2.m_v ;
    }

    const state_type& get_current_state( void ) const
    {
        return m_current_state_x1 ? m_x1.m_v : m_x2.m_v ;
    }

    state_type& get_old_state( void )
    {
        return m_current_state_x1 ? m_x2.m_v : m_x1.m_v ;
    }

    const state_type& get_old_state( void ) const
    {
        return m_current_state_x1 ? m_x2.m_v : m_x1.m_v ;
    }

    deriv_type& get_current_deriv( void )
    {
        return m_current_state_x1 ? m_dxdt1.m_v : m_dxdt2.m_v ;
    }

    const deriv_type& get_current_deriv( void ) const
    {
        return m_current_state_x1 ? m_dxdt1.m_v : m_dxdt2.m_v ;
    }

    deriv_type& get_old_deriv( void )
    {
        return m_current_state_x1 ? m_dxdt2.m_v : m_dxdt1.m_v ;
    }

    const deriv_type& get_old_deriv( void ) const
    {
        return m_current_state_x1 ? m_dxdt2.m_v : m_dxdt1.m_v ;
    }

    void toggle_current_state( void )
    {
        m_current_state_x1 = ! m_current_state_x1;
    }

    controlled_stepper_type m_stepper;
    resizer_type m_resizer;
    bool m_current_state_x1;
    wrapped_state_type m_x1 , m_x2;
    wrapped_deriv_type m_dxdt1 , m_dxdt2;
    time_type m_t , m_t_old , m_dt;
    bool m_is_deriv_initialized;

};

} // namespace odeint
} // namespace numeric
} // namespace boost

#endif // BOOST_NUMERIC_ODEINT_STEPPER_DENSE_OUTPUT_RUNGE_KUTTA_HPP_INCLUDED


#include <iostream>
#include <armadillo>
#include <vector>
#include <cstdlib>
#include <fstream>
#include "libsnnlfisrm/snn.h"

/**
 * Calculates initial delays
*/
arma::Col<double> get_initial_delays();

/**
 * Calculates initial weights for lateral synapses
*/
arma::Col<double> get_initial_weights();

int main(int argc, const char **argv) {

    // Test individual LIF/SRM neuron
    int snn_id = 2;
    int n_inputs = 2;
    int n_lateral = 1;
    std::vector<double> init_d({2, 4});
    std::vector<double> init_w({4});
    double tau_m = 1;
    double u_rest = 3;
    double init_v = 20;
    unsigned char t_rest = 2;
    double kappa_naugh = 3;
    double round_zero = 0.1;
    double u_max = 10;

    SpikeResponseModelNeuron neuron(snn_id, n_inputs, init_d, tau_m,
    u_rest, init_v, t_rest, kappa_naugh, round_zero, u_max);

    printf("delay vector: [");
    for(unsigned int i = 0; i < neuron.d_j.size(); i++)
    {
        printf(" %f ", neuron.d_j.at(i));
    }
    printf("]\n");

    // test t_pulse is working and making kappas...
    std::vector<double> u;
    printf("number of kappas: %d\n", (int)neuron.k_filter_list.at(K_LIST_INPUT_SYNAPSES).size());
    neuron.dendrite = std::vector<DelayedSpike>({DelayedSpike(0, true), DelayedSpike(0, 0)});
    printf("membrane potential: %f\n", neuron.membrane_potential());
    u.push_back(neuron.membrane_potential());
    neuron.t_pulse();
    printf("membrane potential: %f\n", neuron.membrane_potential());
    u.push_back(neuron.membrane_potential());
    neuron.t_pulse();
    printf("membrane potential: %f\n", neuron.membrane_potential());
    u.push_back(neuron.membrane_potential());
    neuron.t_pulse();
    printf("membrane potential: %f\n", neuron.membrane_potential());
    u.push_back(neuron.membrane_potential());
    neuron.t_pulse();
    printf("membrane potential: %f\n", neuron.membrane_potential());
    u.push_back(neuron.membrane_potential());
    neuron.dendrite = std::vector<DelayedSpike>({DelayedSpike(0, false), DelayedSpike(0, 0)});
    neuron.t_pulse();
    u.push_back(neuron.membrane_potential());
    printf("membrane potential: %f\n", neuron.membrane_potential());
    u.push_back(neuron.membrane_potential());
    neuron.t_pulse();
    printf("membrane potential: %f\n", neuron.membrane_potential());
    u.push_back(neuron.membrane_potential());
    neuron.t_pulse();
    printf("membrane potential: %f\n", neuron.membrane_potential());
    u.push_back(neuron.membrane_potential());
    neuron.t_pulse();
    printf("membrane potential: %f\n", neuron.membrane_potential());
    u.push_back(neuron.membrane_potential());
    neuron.dendrite = std::vector<DelayedSpike>({DelayedSpike(0, true), DelayedSpike(0, 0)});
    printf("membrane potential: %f\n", neuron.membrane_potential());
    u.push_back(neuron.membrane_potential());
    neuron.t_pulse();
    printf("membrane potential: %f\n", neuron.membrane_potential());
    u.push_back(neuron.membrane_potential());
    neuron.t_pulse();
    printf("membrane potential: %f\n", neuron.membrane_potential());
    u.push_back(neuron.membrane_potential());
    neuron.t_pulse();
    printf("membrane potential: %f\n", neuron.membrane_potential());
    u.push_back(neuron.membrane_potential());
    neuron.t_pulse();
    printf("membrane potential: %f\n", neuron.membrane_potential());
    u.push_back(neuron.membrane_potential());
    neuron.t_pulse();
    printf("membrane potential: %f\n", neuron.membrane_potential());
    u.push_back(neuron.membrane_potential());
    neuron.t_pulse();
    printf("membrane potential: %f\n", neuron.membrane_potential());
    u.push_back(neuron.membrane_potential());
    neuron.t_pulse();
    printf("membrane potential: %f\n", neuron.membrane_potential());
    u.push_back(neuron.membrane_potential());
    neuron.t_pulse();
    printf("membrane potential: %f\n", neuron.membrane_potential());
    u.push_back(neuron.membrane_potential());
    neuron.t_pulse();
    printf("membrane potential: %f\n", neuron.membrane_potential());
    u.push_back(neuron.membrane_potential());
    neuron.t_pulse();
    printf("number of kappas: %d\n", (int)neuron.k_filter_list.at(K_LIST_INPUT_SYNAPSES).size());

    printf("u plot: [ %f", u.at(0));
    for(unsigned int v = 1; v < u.size(); v++)
        printf(", %f", u.at(v));
    printf("]\n");

    // Testing FSTN
    printf("Testing Fisrt spike time neuron\n");
    double alpha = 1;
    FirstSpikeTimeNeuron fstn(0, alpha);
    fstn.dendrite = (double)2;
    fstn.t_pulse();
    fstn.dendrite = (double)2;
    fstn.t_pulse();
    fstn.dendrite = (double)2;
    fstn.t_pulse();
    fstn.dendrite = (double)2;
    fstn.t_pulse();
    fstn.dendrite = (double)2;
    fstn.t_pulse();
    fstn.dendrite = (double)2;
    fstn.t_pulse();
    fstn.dendrite = (double)2;
    fstn.t_pulse();
    fstn.dendrite = (double)2;
    fstn.t_pulse();
    fstn.dendrite = (double)5;
    fstn.t_pulse();
    fstn.dendrite = (double)5;
    fstn.t_pulse();
    fstn.dendrite = (double)5;
    fstn.t_pulse();
    fstn.dendrite = (double)5;
    fstn.t_pulse();
    fstn.dendrite = (double)5;
    fstn.t_pulse();
    fstn.dendrite = (double)5;
    fstn.t_pulse();
    fstn.dendrite = (double)5;
    fstn.t_pulse();
    fstn.dendrite = (double)5;
    fstn.t_pulse();
    fstn.dendrite = (double)5;
    fstn.t_pulse();
    fstn.dendrite = (double)5;
    fstn.t_pulse();
    fstn.dendrite = (double)5;
    fstn.t_pulse();
    fstn.dendrite = (double)5;
    fstn.t_pulse();

    // Testing neural network
    // parameters:
    unsigned int i_layer_size = 2;
    unsigned int h_layer_size = 16;
    std::vector<std::vector<double>> d_init({
        {4, 3, 2, 3, 4, 3, 3, 2, 1, 4, 3, 3, 3, 2, 5, 2},
        {1, 4, 2, 4, 6, 2, 1, 3, 5, 3, 3, 3, 3, 2, 5, 2}
    });
    std::vector<std::vector<double>> w_init({
        {1, 3, 2, -1, 3, 1, -2, 1, -2, 1, 3, 3, 3, 2, 5, 2},
        {3, 4, 2, 3, 4, 12, 1, 1, 2, 1, 3, 3, 3, 2, 5, 2},
        {2, 3, 2, 1, 2, 1, 2, 1, 1, 2, 3, 3, 3, 2, 5, 2},
        {-1, -2, -1, -2, -3, -4, -9, -1, 2, -2, 3, 3, 3, 2, 5, 2},
        {3, 4, 2, 1, 3, 6, 3, 1, 2, 3, 3, 3, 3, 2, 5, 2},
        {1, 3, 1, 3, 5, 3, 1, 0, 3, 1, 3, 3, 3, 2, 5, 2},
        {-2, -3, -1, -3, -5, 3, 2, 4, 5, 4, 3, 3, 3, 2, 5, 2},
        {1, 6, 2, 2, 4, 6, 3, 2, 4, 5, 3, 3, 3, 2, 5, 2},
        {-2, -5, -1, 4, -1, -3, 2, -7, -1, 3, 3, 3, 3, 2, 5, 2},
        {1, 5, 7, 2, 4, 6, 4, 1, 3, 5, 3, 3, 3, 2, 5, 2},
        {1, 5, 7, 2, 4, 6, 4, 1, 3, 5, 3, 3, 3, 2, 5, 2},
        {1, 5, 7, 2, 4, 6, 4, 1, 3, 5, 3, 3, 3, 2, 5, 2},
        {1, 5, 7, 2, 4, 6, 4, 1, 3, 5, 3, 3, 3, 2, 5, 2},
        {1, 5, 7, 2, 4, 6, 4, 1, 3, 5, 3, 3, 3, 2, 5, 2},
        {1, 5, 7, 2, 4, 6, 4, 1, 3, 5, 3, 3, 3, 2, 5, 2},
        {1, 5, 7, 2, 4, 6, 4, 1, 3, 5, 3, 3, 3, 2, 5, 2},
    });
    tau_m = 4;
    u_rest = 0;
    init_v = 20;
    double t_reset = 3;
    double k_nought = 10;
    round_zero = 0.1;
    alpha = 2;

    BaseSNN snn(i_layer_size, d_init, tau_m, u_rest, init_v, t_reset, k_nought,
    round_zero, alpha, u_max);

    snn.re_process({0, 0});
    for(int p = 0; p < 20; p++)
    {
        snn.process();
        for(unsigned int i = 0; i < h_layer_size; i++)
        {
            double u_i = snn.hidden_layer.at(i).membrane_potential();
            printf("U_%u = %f\n", i, u_i);
            if(snn.hidden_layer.at(i).axon.signal)
            {
                // this neuron is the "winner"
                printf("Spike %u winner (Reached threshold)\n", i);
                goto next_test;
            }
        }
    }
    next_test:
    // Test euclidean_distance_matrix function
    double distance_unit = 2;

    printf("Testing matrices.\n");
    arma::Mat<double> e_mat = euclidean_distance_matrix (&snn, distance_unit);
    printf("Matrix Content\n\n");
    for(unsigned int i = 0; i < snn.h_layer_size; i++)
    {
        for(unsigned int j = 0; j < snn.h_layer_size; j++)
        {
            printf("%f\t", e_mat(i, j));
        }
        printf("\n");
    }

    // Test euclidean initial weight function
    double sigma_1 = .5;
    double sigma_2 = 2;
    std::vector<arma::Col<double>> w_mat = initial_weight_euclidean(e_mat, sigma_1, sigma_2);
    printf("\nPrinting weight matrix's contents:\n\n");
    printf("Weight Matrix Content\n\n");
    for(unsigned int i = 0; i < snn.h_layer_size; i++)
    {
        for(unsigned int j = 0; j < snn.h_layer_size; j++)
        {
            printf("%f\t", w_mat.at(i).at(j));
        }
        printf("\n");
    }

    // Testing euclidean_distance_matrix without snn
    std::vector<std::vector<double>> point_list({
        {0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3},
        {0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3}
    });

    printf("Testing matrices.\n");
    arma::Mat<double> e_mat_init = euclidean_distance_matrix(&point_list, distance_unit);
    printf("Matrix Content\n\n");
    for(unsigned int i = 0; i < e_mat_init.n_rows; i++)
    {
        for(unsigned int j = 0; j < e_mat_init.n_cols; j++)
        {
            printf("%f\t", e_mat_init(i, j));
        }
        printf("\n");
    }

    // Testing Random delay generator
    printf("\nTesting Random Delay generator\n");
    std::vector<std::vector<double>> initial_delays;

    unsigned int n_x = 10;
    unsigned int n_y = 10;
    unsigned int delays_per_row = 10;
    unsigned int delays_per_column = 10;

    initial_delays = initial_delay_vector_2d_map(n_x, n_y,
    delays_per_row, delays_per_column);

    printf("Delay Matrix Content\n\n");
    for(unsigned int i = 0; i < initial_delays.size(); i++)
    {
        for(unsigned int j = 0; j < initial_delays.at(i).size(); j++)
        {
            printf("%f\t", initial_delays.at(i).at(j));
        }
        printf("\n");
    }

    // Test point map creator
    printf("\nTesting point map creator.\n");
    unsigned int x_size = 4;
    unsigned int y_size = 7;
    std::vector<std::vector<double>> map_ = euclidean_point_map(x_size, y_size);
    printf("Point Map Matrix Content\n\n");

    for(unsigned int j = 0; j < x_size*y_size; j++)
    {
        printf("(%f, %f)\t", map_.at(0).at(j), map_.at(1).at(j));
    }

    // Testing training algorthm
    unsigned int n_data = 2;
    tau_m = 1.5;
    u_rest = 2;
    init_v = 4.5;
    t_reset = 3;
    k_nought = 2.5;
    round_zero = 0.05;
    alpha = 4;
    // note that n_x * n_y = h_layer_size
    n_x = 10;
    n_y = 10;
    double delay_distance = 1;
    distance_unit = 1;
    sigma_1 = 0.7;
    sigma_2 = 1.6;

    double sigma_neighbor = 1;

    double eta_d = .5;
    unsigned int t_max = 25;
    unsigned int t_delta = 3;
    double ltd_max = -0.45;
    SNN model(n_data, tau_m, u_rest, init_v,
    t_reset, k_nought, round_zero, alpha, n_x, n_y, delay_distance,
    distance_unit, sigma_neighbor, eta_d, t_max, u_max);

    std::vector<std::vector<double>> data = {
        {0, 0, 0, 1, 2, 2, 3, 3, 3, 3, 3, 2, 1, 0, 0, 0, 0, 0, 0, 0, 0},
        {0, 1, 2, 2, 2, 3, 3, 4, 5, 6, 7, 7, 7, 7, 6, 5, 4, 3, 2, 1, 0}
    };
    std::vector<std::vector<double>> data_2 = {
        {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
        {0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 1, 1, 1, 1},
    };

    // save delays into file:
    std::ofstream delay_file;
    delay_file.open("snn_data_delays.txt");

    // print out delays
    delay_file << "Delays before training:" << std::endl;
    for(unsigned int i = 0; i < model.snn->d_ji.size(); i++)
    {
        delay_file << "[ ";
        for(unsigned int j = 0; j < model.snn->d_ji.at(i).size(); j++)
        {
            delay_file << model.snn->d_ji.at(i).at(j) << ", ";
        }
        delay_file << "]" << std::endl;
    }
    // print out weights

    for(unsigned int p = 0; p < 1; p++)
        model.train(data);

    delay_file << std::endl << "Delays after training:"<< std::endl;
     for(unsigned int i = 0; i < model.snn->d_ji.size(); i++)
    {
        delay_file << "[ ";
        for(unsigned int j = 0; j < model.snn->d_ji.at(i).size(); j++)
        {
            delay_file << model.snn->d_ji.at(i).at(j) << ", ";
        }
        delay_file << "]" << std::endl;
    }
    delay_file.close();

    // Armadillo version printout
    arma::arma_version ver;
    printf("\nArmadillo Version: %s\n", ver.as_string().c_str());
    return 0;
}


/* Copyright 2019 The Spin-Scenario Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "grape.h"
#include <kernel/utilities/ssl_plot.h>
using namespace ssl::utility;
#include <boost/algorithm/string.hpp>
#include <boost/lexical_cast.hpp>
#include <boost/format.hpp>
#include <chrono>
namespace ssl {
namespace oc {
int g_oc_iteration_save = 0;
void set_oc_iteration_save(const sol::table &t) {
	  for (auto &kv : t) {
    int val = kv.second.as<int>();
    g_oc_iteration_save = val;
    break;
  }
}
typedef struct {
  size_t i; // step index.
  size_t j; // channel index.
  size_t chs; // channel number.
  double amp;
} amp_constraint_data;

grape::grape(spin_system &sys)
    : rf_(nullptr) {
  sys_ = &sys;
  superop_.rf_ctrl = sys.rf_hamiltonian();
  superop_.L0 = sys.free_hamiltonian();
  superop_.R = sys.relaxation();
  superop_.L0s = sys.free_hamiltonians();
  int bb_size = superop_.L0s.size();
  //std::cout << bb_size << "\n";
  if (bb_size) {
    // superop_.profile = vec::Ones(bb_size);
    superop_.nominal_offset = sys.nominal_broadband();
    superop_.grad_bb = std::vector<std::vector<double>>(bb_size);
  } else {  // no offset case.
    superop_.L0s.push_back(superop_.L0);
    superop_.grad_bb = std::vector<std::vector<double>>(1);
  }
  opt_model_ = _rho2rho;
  axis_ = uxuy_;
}

grape::~grape() {
}

double grape::maxf() const {
  return max_val_;
}

seq_block &grape::optimize(const sol::table &t) {
  assign_pulse(t);
  assign_x();
  assign_nlopt(t);
  assign_state(t);
  assign_aux_var();

  if (opt_model_ == _rho2rho)
      optimizer_->set_max_objective(objfunc_broadband, this);
  if (opt_model_ == _propagator)
    optimizer_->set_max_objective(objfunc_propagator, this);

  print();
  double max_f;
  nlopt::result result = optimizer_->optimize(x_, max_f);
  if (result == nlopt::SUCCESS)
    ssl_color_text("info", "pulse optimization succeed.\n");
  if (result == nlopt::FAILURE)
    ssl_color_text("err", "pulse optimization failed.\n");
  if (result == nlopt::MAXEVAL_REACHED)
    ssl_color_text("warn", "pulse optimization terminated due to maximum iterations limitation.\n");

  max_val_ = max_f;
  rf_->update_raw_data(axis_, x_.data());

  grape::h5write();
  return *rf_;
}

double grape::objfunc_broadband(const std::vector<double> &x, std::vector<double> &grad,
                                void *func) {
  return ((grape *)func)->objfunc_broadband(x, grad);
}
double grape::objfunc_propagator(const std::vector<double> &x, std::vector<double> &grad,
                                 void *func) {
  return ((grape *)func)->objfunc_propagator(x, grad);
}

void grape::assign_state(const sol::table &t) {
  if (is_retrievable("targ_op", t)) {
    opt_model_ = _propagator;
    targ_op_ = sys_->smart_op(retrieve_table_str("targ_op", t));
    targ_op_ = ssl::spinsys::propagator(targ_op_, _pi);
    int n = targ_op_.rows();
    init_op_ = sp_cx_mat(n, n);
    init_op_.setIdentity();

    // in case of unified targ for diff freq offsets.
    if (superop_.L0s.size())
      targ_op_list_ = std::vector<sp_cx_mat>(superop_.L0s.size(), targ_op_);
    return;
  }

  opt_model_ = _rho2rho;
  init_state_ = sys_->smart_state(retrieve_table_str("init_state", t));
  init_state_ = norm_state(init_state_);

  sol::object obj = retrieve_table("targ_state", t);
  if (obj.get_type() == sol::type::string) {
    targ_state_ = sys_->smart_state(retrieve_table_str("targ_state", t));
    targ_state_ = norm_state(targ_state_);

    // in case of unified targ for diff freq offsets.
    if (superop_.L0s.size())
      targ_list_ = std::vector<sp_cx_vec>(superop_.L0s.size(), targ_state_);
  }

  if (obj.get_type() == sol::type::table) {
    targ_list_.clear();
    sol::table targ_table = obj.as<sol::table>();
    for (size_t i = 0; i < targ_table.size(); i++) {
      sol::object val = targ_table[i + 1];
      sp_cx_vec targ = val.as<sp_cx_vec>();
      targ=norm_state(targ);
      targ_list_.push_back(targ);
    }
  }
}

void grape::assign_nlopt(const sol::table &t) {
  // TO BE REMOVE.
  epsilon_ = 0;      // default val.
  epsilon_num_ = 1;  // default val.
  if (is_retrievable("epsilon", t)) {
    epsilon_ = retrieve_table_double("epsilon", t);
  }

  if (is_retrievable("epsilon_num", t)) {
    epsilon_num_ = retrieve_table_int("epsilon_num", t);
  }
  nlopt::algorithm alg = nlopt::LD_LBFGS;

  if (is_retrievable("algorithm", t)) {
    std::string oc = retrieve_table_str("algorithm", t);
    boost::to_upper(oc);
    if (oc == "MNA") alg = nlopt::LD_MMA;
  }
  optimizer_ = new nlopt::opt(alg, x_dim_);

  if (is_retrievable("max_eval", t)) {
    int val = retrieve_table_int("max_eval", t);
    if (val > 0)
      optimizer_->set_maxeval(val);
    else
      throw std::runtime_error("max_eval must > 0.");
  }

  if (is_retrievable("xtol_rel", t)) {
    double val = retrieve_table_double("xtol_rel", t);
    optimizer_->set_xtol_rel(val);
  } else
    optimizer_->set_xtol_rel(1e-8);

  assign_constraint(t);
}
void grape::print() const {
  if (optimizer_->get_algorithm() == nlopt::LD_MMA)
    ssl_color_text("oc", "nlopt algorithm used: LD_MMA.\n");
  if (optimizer_->get_algorithm() == nlopt::LD_LBFGS)
    ssl_color_text("oc", "nlopt algorithm used: LD_LBFGS.\n");

  ssl_color_text("oc", "x dim: " + boost::lexical_cast<std::string>(x_dim_) +  ".\n");
  ssl_color_text("oc", "nlopt xtol_rel: " + boost::lexical_cast<std::string>(optimizer_->get_xtol_rel()) +  ".\n");

}
void grape::assign_x() {
  x_.clear();
  rf_->convert2(_ux_uy);  // must do it.

  if (axis_ == uxuy_)
    x_ = rf_->clone_raw_data();
  else if (axis_ == ux_)
    x_ = rf_->clone_raw_data_ux();
  else if (axis_ == uy_)
    x_ = rf_->clone_raw_data_uy();

  x_dim_ = x_.size();
}

void grape::assign_pulse(const sol::table &t) {
  double width = retrieve_table_double("width", t); // unit in ms.
  size_t nsteps = (size_t) (retrieve_table_int("step", t));
  std::string pattern = "rand_spline";
  if (is_retrievable("init_pattern", t))
    pattern = retrieve_table_str("init_pattern", t);

  double max_init_amp = 1;
  if (is_retrievable("max_init_amp", t))
    max_init_amp = retrieve_table_double("max_init_amp", t);

  double dt = width * 1e-3 / double(nsteps); // into s.

  std::string str_chs;
  if (is_retrievable("limit_channel", t))
      str_chs = retrieve_table_str("limit_channel", t);
  else
  str_chs = boost::algorithm::join(superop_.rf_ctrl.chs, " ");

  std::string code = "user_rf = shapedRF{name = 'opt-rf', width = " +
                boost::lexical_cast<std::string>(width) +
                ",  step = " + boost::lexical_cast<std::string>(nsteps) +
                ",  max_amp = " + boost::lexical_cast<std::string>(max_init_amp) +
                ",  channel = '" + str_chs + "'," + "pattern = '" + pattern +
                "'}";

  g_lua->script(code);

   std::string s = str(boost::format("pulse width - [%.3f] ms, steps - [%d], step width - [%.3f] us.\n") % width % nsteps
                     % (dt * 1e6));
  ssl_color_text("info", s);

  sol::object val = (*g_lua)["user_rf"];
  seq_block &sb = val.as<seq_block &>(); // note the original type should be 'seq_block'.
  shaped_rf &user_rf = (shaped_rf &) sb; // transfrom into 'shaped_rf'.

  rf_ = &user_rf;

  if (is_retrievable("limit_axis", t)) {
    std::string s = retrieve_table_str("limit_axis", t);
    if (s == "x")
      axis_ = ux_;
    else if (s == "y")
      axis_ = uy_;
  }

  rf_->convert2(_ux_uy);  // must do it.

  if(axis_==ux_) {
   std::vector<double>zeros = std::vector<double>(rf_->get_dims() / 2, 0);
    rf_->update_raw_data_uy(zeros.data());
  }
  if (axis_ == uy_) {
    std::vector<double> zeros = std::vector<double>(rf_->get_dims() / 2, 0);
    rf_->update_raw_data_ux(zeros.data());
  }
}

void grape::h5write(std::string file_name) const {
  if (file_name.empty()) {
    std::string time_s = sys_time();
    file_name = "oc_" + time_s + ".h5";
  }
  H5File file(file_name, H5F_ACC_TRUNC);

  rf_->switch_rf_mode("amp/phase");
  rf_->h5write(file, "opt");

  ssl::utility::h5write(file, nullptr, "obj", stl2vec(obj_val_));
  file.close();
}

void grape::assign_aux_var() {
  // traj_ = state_traj(rf_->get_steps());
  int nbb = superop_.L0s.size();
  if (nbb >= 1) {
    omp_set_num_threads(omp_core_num);
    for (size_t i = 0; i < superop_.grad_bb.size(); i++)
      superop_.grad_bb[i] = std::vector<double>(rf_->get_dims(), 0);

    if (opt_model_ == _rho2rho) {
      traj_omp_ = std::vector<state_traj>(omp_core_num);
      for (int i = 0; i < omp_core_num; i++) {
        traj_omp_[i] = state_traj(rf_->get_steps());
      }
    }
    if (opt_model_ == _propagator) {
      op_traj_omp_ = std::vector<op_traj>(omp_core_num);
      for (int i = 0; i < omp_core_num; i++) {
        op_traj_omp_[i] = op_traj(rf_->get_steps());
      }
    }
  }
}

void grape::assign_constraint(const sol::table &t) {
  if (is_retrievable("max_amp", t)) {
    double val = retrieve_table_double("max_amp", t);  // unit in hz.
    val *= 2 * _pi;

    size_t dim = rf_->get_dims();
    size_t nsteps = rf_->get_steps();
    size_t nchannels = rf_->get_channels();
    std::vector<double> up_bound(dim);
    std::vector<double> low_bound(dim);

    // only for ux/uy mode.
    if (rf_->mode() == _ux_uy)
      for (size_t i = 0; i < nsteps; i++) {
        for (size_t j = 0; j < nchannels; j++) {
          // ux.
          up_bound[2 * nchannels * i + 2 * j] = val;
          low_bound[2 * nchannels * i + 2 * j] = -val;

          // uy.
          up_bound[2 * nchannels * i + 2 * j + 1] = val;
          low_bound[2 * nchannels * i + 2 * j + 1] = -val;
        }
      }

    optimizer_->set_upper_bounds(up_bound);
    optimizer_->set_lower_bounds(low_bound);
  }
}

double grape::objfunc_broadband(const std::vector<double> &x, std::vector<double> &grad) {
  auto start = std::chrono::system_clock::now();
  rf_->update_raw_data(axis_, x.data());
  int N = superop_.L0s.size();
  vec phi = vec::Zero(N);

  size_t nsteps = rf_->get_steps();
  size_t nchannels = rf_->get_channels();
  double dt = rf_->get_dt() * 1e-6;  // into s.
  std::vector<std::string> chs = rf_->get_channels_str();

  // double alpha = alpha0;
  vec rf_scaling =
      vec::LinSpaced(epsilon_num_, -epsilon_, epsilon_);  // e.g. -10% - 10%

#pragma omp parallel for
  for (int p = 0; p < N; p++) {
    int id = omp_get_thread_num();
    traj_omp_[id].forward[0] = init_state_;
    traj_omp_[id].backward[nsteps] = targ_list_[p];
    sp_cx_mat L;
    sp_cx_mat L0 = superop_.L0s[p] + ci * superop_.R;

    int dim = rf_->get_dims();
    if (axis_ != uxuy_) dim /= 2;
    superop_.grad_bb[p] = std::vector<double>(dim, 0);
    // start rf inhom.

    for (int q = 0; q < rf_scaling.size();
         q++) {  // for each rf scaling factor, do

      double kx = 1 + rf_scaling[q], ky = 1 + rf_scaling[q];

      sp_cx_vec rho = traj_omp_[id].forward[0];
      for (size_t i = 0; i < nsteps; i++) {
        L = L0;
        for (size_t j = 0; j < nchannels; j++)
          L += update_rf_ham(x.data(), i, j, chs[j], nchannels, kx, ky);
        traj_omp_[id].forward[i + 1] = ssl::spinsys::step(rho, L, dt);
        rho = traj_omp_[id].forward[i + 1];
      }
      rho = traj_omp_[id].backward[nsteps];
      for (int i = nsteps - 1; i >= 0; i--) {
        L = L0.adjoint();
        for (size_t j = 0; j < nchannels; j++)
          L += update_rf_ham(x.data(), i, j, chs[j], nchannels, kx, ky);
        traj_omp_[id].backward[i] = ssl::spinsys::step(rho, L, -dt);
        rho = traj_omp_[id].backward[i];
      }
      sp_cx_mat Gx, Gy, tmp;
      int k = 0;
      for (size_t i = 0; i < nsteps; i++) {
        L = L0;
        for (size_t j = 0; j < nchannels; j++)
          L += update_rf_ham(x.data(), i, j, chs[j], nchannels, kx, ky);
        tmp = traj_omp_[id].backward[i + 1].adjoint() *
              ssl::spinsys::propagator(L, dt);
        for (size_t j = 0; j < nchannels; j++) {
          if (axis_ == uxuy_) {
            Gx = propagator_derivative(L, superop_.rf_ctrl.Lx[j], dt);
            Gy = propagator_derivative(L, superop_.rf_ctrl.Ly[j], dt);

            superop_.grad_bb[p][k] +=
                traced(tmp * Gx * traj_omp_[id].forward[i]).real();
            superop_.grad_bb[p][k + 1] +=
                traced(tmp * Gy * traj_omp_[id].forward[i]).real();
            k += 2;
          } else if (axis_ == ux_) {
            Gx = propagator_derivative(L, superop_.rf_ctrl.Lx[j], dt);

            superop_.grad_bb[p][k] +=
                traced(tmp * Gx * traj_omp_[id].forward[i]).real();
            k += 1;
          } else if (axis_ == uy_) {
            Gy = propagator_derivative(L, superop_.rf_ctrl.Ly[j], dt);

            superop_.grad_bb[p][k] +=
                traced(tmp * Gy * traj_omp_[id].forward[i]).real();
            k += 1;
          }
        }
      }
      phi[p] += transfer_fidelity(traj_omp_[id].forward[nsteps], targ_list_[p]);
    }

    /// transfer_fidelity(targ_list_[p], targ_list_[p]);
  }  // end parallel for.

  double val = phi.sum() / (double)N / (double)rf_scaling.size();

  int k = 0;
  for (size_t i = 0; i < nsteps; i++) {
    for (size_t j = 0; j < nchannels; j++) {
      double gx = 0;
      double gy = 0;
      if (axis_ == uxuy_) {
        for (int p = 0; p < N; p++) {
          gx += superop_.grad_bb[p][k];
          gy += superop_.grad_bb[p][k + 1];
        }
        grad[k] = gx / (double)N / (double)rf_scaling.size();
        grad[k + 1] = gy / (double)N / (double)rf_scaling.size();

        // rf power reduction.
        /*if (alpha != 0) {
          double ux, uy;
          ux = x[2 * nchannels * i + 2 * j];
          uy = x[2 * nchannels * i + 2 * j + 1];
          grad[k] -= 2.0 * alpha * ux * dt;
          grad[k + 1] -= 2.0 * alpha * uy * dt;
        }*/
        k += 2;
      } else if (axis_ == ux_) {
        for (int p = 0; p < N; p++) {
          gx += superop_.grad_bb[p][k];
        }
        grad[k] = gx / (double)N / (double)rf_scaling.size();

        k += 1;
      } else if (axis_ == uy_) {
        for (int p = 0; p < N; p++) {
          gy += superop_.grad_bb[p][k];
        }
        grad[k] = gy / (double)N / (double)rf_scaling.size();

        k += 1;
      }
    }
  }
  auto end = std::chrono::system_clock::now();
  auto duration = std::chrono::duration_cast<std::chrono::microseconds>(end - start);
  if(double(duration.count()) * std::chrono::microseconds::period::num / std::chrono::microseconds::period::den>2){
    g_oc_iteration_save = 1;
  }
  iteration_shape(x, optimizer_->get_numevals());
  //double PHI1 = val;
  //double PHI2 = 0;
  // double PHI2 = alpha * rf_->rf_power();
  // std::cout << boost::format("==> %04d  [%.8f] [%.8f]\n") %
  // (++opt_.iter_count) % PHI1 % PHI2;
  std::cout << boost::format("==> %04d  [%.8f]\n") % (optimizer_->get_numevals()) % val;
 obj_val_.push_back(val);
  //std::cout << "Use Time:" << double(duration.count()) * microseconds::period::num / microseconds::period::den << " s.\n";
  return val;
}
void grape::iteration_shape(const std::vector<double> &x, int iter) {
  if (!g_oc_iteration_save) return;
  rf_->update_raw_data(axis_, x_.data());
  std::string file = "iteration_" + std::to_string(iter) + "_rf_shape.RF";
  std::ofstream ofstr(file.c_str());
  ofstr << "# " << sys_time()<<"\n";
   rf_->write(ofstr);
  ofstr.close();
}
double grape::objfunc_propagator(const std::vector<double> &x,
                                 std::vector<double> &grad) {
  auto start = std::chrono::system_clock::now();
  rf_->update_raw_data(x.data());
  int N = superop_.L0s.size();
  vec phi = vec::Zero(N);

  size_t nsteps = rf_->get_steps();
  size_t nchannels = rf_->get_channels();
  double dt = rf_->get_dt() * 1e-6;  // into s.
  std::vector<std::string> chs = rf_->get_channels_str();

  // double alpha = alpha0;
  vec rf_scaling =
      vec::LinSpaced(epsilon_num_, -epsilon_, epsilon_);  // e.g. -10% - 10%

#pragma omp parallel for
  for (int p = 0; p < N; p++) {
    int id = omp_get_thread_num();
    op_traj_omp_[id].forward[0] = init_op_;
    op_traj_omp_[id].backward[nsteps] = targ_op_list_[p];
    sp_cx_mat L;
    sp_cx_mat L0 = superop_.L0s[p] + ci * superop_.R;

     int dim = rf_->get_dims();
    if (axis_ != uxuy_) dim /= 2;
    superop_.grad_bb[p] = std::vector<double>(dim, 0);
    // start rf inhom.

    for (int q = 0; q < rf_scaling.size();
         q++) {  // for each rf scaling factor, do

      double kx = 1 + rf_scaling[q], ky = 1 + rf_scaling[q];

      sp_cx_mat rho = op_traj_omp_[id].forward[0];
      for (size_t i = 0; i < nsteps; i++) {
        L = L0;
        for (size_t j = 0; j < nchannels; j++)
          L += update_rf_ham(x.data(), i, j, chs[j], nchannels, kx, ky);
        op_traj_omp_[id].forward[i + 1] = ssl::spinsys::propagator(L, dt) * rho;
        rho = op_traj_omp_[id].forward[i + 1];
      }
      rho = op_traj_omp_[id].backward[nsteps];
      for (int i = nsteps - 1; i >= 0; i--) {
        L = L0;
        for (size_t j = 0; j < nchannels; j++)
          L += update_rf_ham(x.data(), i, j, chs[j], nchannels, kx, ky);
        op_traj_omp_[id].backward[i] =
            ssl::spinsys::propagator(L, dt).adjoint() * rho;
        rho = op_traj_omp_[id].backward[i];
      }
      sp_cx_mat Gx, Gy, tmp;
      int k = 0;
      for (size_t i = 0; i < nsteps; i++) {
        L = L0;
        for (size_t j = 0; j < nchannels; j++)
          L += update_rf_ham(x.data(), i, j, chs[j], nchannels, kx, ky);
        tmp = op_traj_omp_[id].backward[i + 1].adjoint() *
              ssl::spinsys::propagator(L, dt);
        for (size_t j = 0; j < nchannels; j++) {
          if (axis_ == uxuy_) {
            Gx = propagator_derivative(L, superop_.rf_ctrl.Lx[j], dt);
            Gy = propagator_derivative(L, superop_.rf_ctrl.Ly[j], dt);

            superop_.grad_bb[p][k] +=
                traced(tmp * Gx * op_traj_omp_[id].forward[i]).real();
            superop_.grad_bb[p][k + 1] +=
                traced(tmp * Gy * op_traj_omp_[id].forward[i]).real();
            k += 2;
          } else if (axis_ == ux_) {
            Gx = propagator_derivative(L, superop_.rf_ctrl.Lx[j], dt);

            superop_.grad_bb[p][k] +=
                traced(tmp * Gx * op_traj_omp_[id].forward[i]).real();
            k += 1;
          } else if (axis_ == uy_) {
            Gy = propagator_derivative(L, superop_.rf_ctrl.Ly[j], dt);
            superop_.grad_bb[p][k] +=
                traced(tmp * Gy * op_traj_omp_[id].forward[i]).real();
            k += 2;
          }
        }
      }
      phi[p] += transfer_fidelity(op_traj_omp_[id].forward[nsteps],
                                  targ_op_list_[p]) /
                transfer_fidelity(targ_op_list_[p], targ_op_list_[p]);
    }

    /// transfer_fidelity(targ_list_[p], targ_list_[p]);
  }  // end parallel for.

  double val = phi.sum() / (double)N / (double)rf_scaling.size();

  int k = 0;
  for (size_t i = 0; i < nsteps; i++) {
    for (size_t j = 0; j < nchannels; j++) {
      double gx = 0;
      double gy = 0;
     if (axis_ == uxuy_) {
        for (int p = 0; p < N; p++) {
          gx += superop_.grad_bb[p][k];
          gy += superop_.grad_bb[p][k + 1];
        }
        grad[k] = gx / (double)N / (double)rf_scaling.size();
        grad[k + 1] = gy / (double)N / (double)rf_scaling.size();

        // rf power reduction.
        /*if (alpha != 0) {
          double ux, uy;
          ux = x[2 * nchannels * i + 2 * j];
          uy = x[2 * nchannels * i + 2 * j + 1];
          grad[k] -= 2.0 * alpha * ux * dt;
          grad[k + 1] -= 2.0 * alpha * uy * dt;
        }*/
        k += 2;
      } else if (axis_ == ux_) {
        for (int p = 0; p < N; p++) {
          gx += superop_.grad_bb[p][k];
        }
        grad[k] = gx / (double)N / (double)rf_scaling.size();

        k += 1;
      } else if (axis_ == uy_) {
        for (int p = 0; p < N; p++) {
          gy += superop_.grad_bb[p][k];
        }
        grad[k] = gy / (double)N / (double)rf_scaling.size();

        k += 1;
      }
    }
  }

  //double PHI1 = val;
  //double PHI2 = 0;
  // double PHI2 = alpha * rf_->rf_power();
  // std::cout << boost::format("==> %04d  [%.8f] [%.8f]\n") %
  // (++opt_.iter_count) % PHI1 % PHI2;

  auto end = std::chrono::system_clock::now();
  auto duration = std::chrono::duration_cast<std::chrono::microseconds>(end - start);
  if(double(duration.count()) * std::chrono::microseconds::period::num / std::chrono::microseconds::period::den>2){
    g_oc_iteration_save = 1;
  }
  iteration_shape(x, optimizer_->get_numevals());

  std::cout << boost::format("==> %04d  [%.8f]\n") % (optimizer_->get_numevals()) % val;
  obj_val_.push_back(val);
  return (val);
}
// void grape::cartesian2polar(double amp, double phase, double gx, double gy,
//                            double &g_amp, double &g_phase) {
//  double ux = amp * cos(phase);
//  double uy = amp * sin(phase);
//  g_amp = gx * ux + gy * uy;
//  g_amp /= amp;
//  g_phase = -gx * uy + gy * ux;
//}

double grape::amplitude_constraint(unsigned n, const double *x, double *grad,
                                   void *data) {
  amp_constraint_data *d = reinterpret_cast<amp_constraint_data *>(data);

  double ux = x[2 * d->chs * d->i + 2 * d->j];
  double uy = x[2 * d->chs * d->i + 2 * d->j + 1];

  return ux * ux + uy * uy - d->amp * d->amp;
}

sp_cx_mat grape::update_rf_ham(const double *x, size_t step, size_t channel, std::string ch_str,
                               size_t nchannels, double kx, double ky) {
  int ch = superop_.rf_ctrl.channel_index(ch_str);
  if (axis_ == uxuy_) {
    double ux = x[2 * nchannels * step + 2 * channel];
    double uy = x[2 * nchannels * step + 2 * channel + 1];
    // Rf inhom
    ux *= kx;
    uy *= ky;
    return ux * superop_.rf_ctrl.Lx[ch] + uy * superop_.rf_ctrl.Ly[ch];
  } else if (axis_ == ux_) {
    double ux = x[nchannels * step + channel];
    // Rf inhom
    ux *= kx;
    return ux * superop_.rf_ctrl.Lx[ch];
  } else if (axis_ == uy_) {
    double uy = x[nchannels * step + channel];
    // Rf inhom
    uy *= ky;
    return uy * superop_.rf_ctrl.Ly[ch];
  }
}
sp_cx_mat grape::propagator_derivative(const sp_cx_mat &L, const sp_cx_mat &Lrf, double dt) {
  sp_cx_mat commu1, commu2, commu3;
  commu1 = commutator(L, Lrf);
  commu2 = commutator(L, commu1);
  commu3 = commutator(L, commu2);
  return -ci * dt * Lrf + 1 / 2.0 * dt * dt * commu1 + 1 / 6.0 * dt * dt * dt * commu2
      - 1 / 24.0 * dt * dt * dt * dt * commu3;
}

void grape::projection(const sol::table &t) {
  epsilon_ = 0;
  epsilon_num_ = 1;  // default val.
  if (is_retrievable("epsilon", t)) {
    epsilon_ = retrieve_table_double("epsilon", t);
  }
  if (is_retrievable("epsilon_num", t)) {
    epsilon_num_ = retrieve_table_int("epsilon_num", t);
  }
  sol::object val = retrieve_table("init_state", t);
  sp_cx_vec init_state = sys_->smart_state(val.as<std::string>());
  init_state = norm_state(init_state);
  //init_state = levante_ernst(init_state);

  val = retrieve_table("rf", t);
  seq_block &sb = val.as<seq_block &>();
 shaped_rf &user_rf = (shaped_rf &) sb; // transfrom into 'shaped_rf'.
  user_rf.convert2(_ux_uy);

  size_t nsteps = user_rf.get_steps();
  size_t nchannels = user_rf.get_channels();
  double dt = user_rf.width_in_ms() * 1e-3 / double(nsteps); // into s.

  std::map<std::string, sp_cx_vec> obsrv_state_map;
  std::vector<std::string> expr;
  std::vector<sp_cx_vec> obsrv_state;
  if (is_retrievable("observ_states", t)) {
    val = retrieve_table("observ_states", t);
    sol::table expr_table = val.as<sol::table>();
    if (expr_table.size() == 0) {
      obsrv_state_map = sys_->cartesian_basis_states();
    } else {
      for (size_t i = 0; i < expr_table.size(); i++) {
        sol::object val = expr_table[i + 1];
        std::string exp = val.as<std::string>();
        sp_cx_vec rho = sys_->smart_state(exp);
        obsrv_state_map.insert(std::pair<std::string, sp_cx_vec>(exp, rho));
      }
    }

  } else if (is_retrievable("ignore_states", t)) {
    val = retrieve_table("ignore_states", t);
    sol::table expr_table = val.as<sol::table>();
    if (expr_table.size() == 0) {
      obsrv_state_map = sys_->cartesian_basis_states();
    } else {
      obsrv_state_map = sys_->cartesian_basis_states();

      for (size_t i = 0; i < expr_table.size(); i++) {
        sol::object val = expr_table[i + 1];
        std::string exp = val.as<std::string>();
        std::map<std::string, sp_cx_vec>::iterator key = obsrv_state_map.find(exp);
        if (key != obsrv_state_map.end()) obsrv_state_map.erase(key);
      }
    }
  } else {
    obsrv_state_map = sys_->cartesian_basis_states();
  }

  std::map<std::string, sp_cx_vec>::iterator iter;
  for (iter = obsrv_state_map.begin(); iter != obsrv_state_map.end(); iter++) {
    obsrv_state.push_back(norm_state(iter->second));
    // obsrv_state.push_back(levante_ernst(iter->second));
    expr.push_back(iter->first);
  }
  std::string dim1 = "observed states\n";
  std::string dim2 = "";
  std::string dim3 = "rf scaling\n";
  std::string s;
  for (size_t p = 0; p < expr.size(); p++)
    s += std::to_string(p + 1) + "-" + expr[p] + "\n";
  dim1 += s;

  std::string str_opt = "step";
  if (is_retrievable("option", t)) str_opt = retrieve_table_str("option", t);

  if (str_opt != "step" && str_opt != "broadband") {
    std::string s = "unknown projection option ** " + str_opt +
               " ** using 'step' or 'broadband' instead.";
    throw std::runtime_error(s.c_str());
  }

  std::vector<double> x = user_rf.clone_raw_data();
  vec rf_scaling =
      vec::LinSpaced(epsilon_num_, -epsilon_, epsilon_);  // e.g. -10% - 10%

  size_t n = rf_scaling.size();
  if (n == 1)
    dim3 += "no rf inhomogeneity";
  else
    dim3 += boost::lexical_cast<std::string>(rf_scaling[0] * 100) + ":" +
            boost::lexical_cast<std::string>(rf_scaling[n - 1] * 100) + " " +
            boost::lexical_cast<std::string>(n) + " %\n";

  cube comp_dist;
  // for BROADBAND case: states, freq offsets, rf scalings
  // for STEP case: states, steps, rf scalings
  std::vector<std::string> chs = rf_->get_channels_str();

  if (str_opt == "step") {
    dim2 = "pulse steps\n";
    dim2 += "interval: " + boost::lexical_cast<std::string>(dt) + " s\n";
    comp_dist = cube(obsrv_state.size(), nsteps,
                     rf_scaling.size());  // states, steps, rf scalings
    sp_cx_mat L;
    sp_cx_mat L0 = superop_.L0 + ci * superop_.R;
    sp_cx_vec *forward = new sp_cx_vec[nsteps + 1];
    forward[0] = init_state;

    for (int q = 0; q < rf_scaling.size();
         q++) {  // for each rf scaling factor, do
      double kx = 1 + rf_scaling[q], ky = 1 + rf_scaling[q];
      sp_cx_vec rho = forward[0];

      for (size_t i = 0; i < nsteps; i++) {
        L = L0;
        for (size_t j = 0; j < nchannels; j++)
          L += update_rf_ham(x.data(), i, j, chs[j], nchannels, kx, ky);
        forward[i + 1] = ssl::spinsys::step(rho, L, dt);
        rho = forward[i + 1];

        for (size_t k = 0; k < obsrv_state.size(); k++) {
          sp_cx_vec compo = obsrv_state[k];
          comp_dist(k, i, q) = transfer_fidelity(rho, compo);
        }
      }
    }
  } else if (str_opt == "broadband") {
    dim2 = "freq offsets\n";
    int n = superop_.nominal_offset.size();
    dim2 += boost::lexical_cast<std::string>(superop_.nominal_offset[0]) + ":" +
            boost::lexical_cast<std::string>(superop_.nominal_offset[n - 1]) + " " +
            boost::lexical_cast<std::string>(n) + " Hz\n";

    comp_dist = cube(obsrv_state.size(), superop_.L0s.size(),
                     rf_scaling.size());  // states, freq offsets, rf scalings

    omp_set_num_threads(omp_core_num);

    std::vector<sp_cx_vec *> forward_trajs_parfor =
        std::vector<sp_cx_vec *>(omp_core_num);
    for (int i = 0; i < omp_core_num; i++) {
      sp_cx_vec *forward = new sp_cx_vec[nsteps + 1];
      forward[0] = init_state;
      forward_trajs_parfor[i] = forward;
    }

#pragma omp parallel for
    for (int p = 0; p < superop_.L0s.size(); p++) {
      int id = omp_get_thread_num();

      for (int q = 0; q < rf_scaling.size();
           q++) {  // for each rf scaling factor, do
        double kx = 1 + rf_scaling[q], ky = 1 + rf_scaling[q];
        sp_cx_mat L;
        sp_cx_mat L0 = superop_.L0s[p] + ci * superop_.R;
        sp_cx_vec rho = forward_trajs_parfor[id][0];

        for (size_t i = 0; i < nsteps; i++) {
          L = L0;
          for (size_t j = 0; j < nchannels; j++)
            L += update_rf_ham(x.data(), i, j, chs[j], nchannels, kx, ky);
          forward_trajs_parfor[id][i + 1] = ssl::spinsys::step(rho, L, dt);
          rho = forward_trajs_parfor[id][i + 1];
        }

        for (size_t k = 0; k < obsrv_state.size(); k++) {
          sp_cx_vec compo = obsrv_state[k];
          // transfer_wave(p, k) = transfer_fidelity(rho, compo);  // /
          // transfer_fidelity(compo, compo);
          comp_dist(k, p, q) = transfer_fidelity(rho, compo);
        }
      }
    }
  }

  std::string time_s = sys_time();
  H5File file("proj_" + time_s + ".h5", H5F_ACC_TRUNC);
  ssl::utility::h5write(file, nullptr, "projection", comp_dist);
  ssl::utility::h5write(file, nullptr, "dim1", dim1);
  ssl::utility::h5write(file, nullptr, "dim2", dim2);
  ssl::utility::h5write(file, nullptr, "dim3", dim3);
  file.close();

  if (rf_scaling.size() == 1) {  // no rf inhomogeneity.
    Eigen::Tensor<double, 2> sub =
        comp_dist.chip(0, 2);  // rf scaling at dim 2.
    Eigen::Map<mat> m(sub.data(), sub.dimension(0), sub.dimension(1));
    mat grid = m.matrix();

    sol::table lines;

    lines = g_lua->create_table();
    for (int i = 0; i < grid.rows(); i++) {  // each state.
      vec line = grid.row(i).transpose();
      lines.add(line);
    }

    std::string fig_spec;
    vec xval;
    if (str_opt == "step") {
      xval = vec::LinSpaced(grid.cols(), 0, user_rf.width_in_ms()); // transfer trajectories of basis operators
	  std::string s = retrieve_table_str("init_state", t);
      fig_spec = "title<initial state: " + s + "> ";
      fig_spec +=
          "xlabel<pulse " //initial state I_{1y}
          "duration "
          "/ ms> ylabel<magnetization>";
      fig_spec +=
          "xrange<0:" + boost::lexical_cast<std::string>(user_rf.width_in_ms()) +
          "> ";
    } else if (str_opt == "broadband") {
      int n = superop_.nominal_offset.size();
      xval = vec::LinSpaced(grid.cols(), superop_.nominal_offset[0],
                            superop_.nominal_offset[n - 1]);
      xval *= 1e-3;
      fig_spec = "xlabel<frequency offset / kHz> ylabel<magnetization>"; // title<scan 1>
      fig_spec += "xrange<" + boost::lexical_cast<std::string>(xval[0]) + ":" +
                  boost::lexical_cast<std::string>(xval[xval.size() - 1]) + "> ";
    }
    if (expr.size() > 5)
        fig_spec += " gnuplot<set key outside>";
      //fig_spec += " gnuplot<set label 'Scan 2' at graph 0.5,0.5 center font 'Arial,26'\n set ytics 0.2\n set key horizontal above>";
	  //fig_spec += " gnuplot<set label 'Scan 2' at graph 0.5,0.5 center font 'Arial,26'\n set ytics 0.2\n unset key>";
	  //fig_spec += " gnuplot<set ytics 0.2\n unset key>";

    fig_spec += " lw<7>";
    fig_spec += " color<YiZhang16,16>";
    std::string lege;
    for (size_t i = 0; i < expr.size(); i++) lege += expr[i] + ";";
    fig_spec += " legend<" + lege + ">";
    plot(fig_spec, line_series(xval, lines));
    return;
  }

  // rf scaling case.

  Eigen::Tensor<double, 2> sub =
      comp_dist.chip(0, 0);  // observed states at dim 0, by default, plot the
                             // map for first state.
  std::string state_name = expr[0];
  Eigen::Map<mat> m(sub.data(), sub.dimension(0), sub.dimension(1));
  mat grid = m.matrix();

  vec2 xrange, yrange;
  yrange[0] = superop_.nominal_offset[0];
  yrange[1] = superop_.nominal_offset[superop_.nominal_offset.size() - 1];
  yrange *= 1e-3;  // into kHz.
  xrange[0] = rf_scaling[0];
  xrange[1] = rf_scaling[rf_scaling.size() - 1];
  xrange *= 100;  // into %

  utility::map gnu(grid, "style<3d>");
  gnu.xrange = xrange;
  gnu.yrange = yrange;
  std::string fig_spec =
      "'ylabel<freq offset / kHz> xlabel<rf inhomogeneity / %> color<Spectral> "
      "gnuplot<set xlabel offset -1,-1; set ylabel offset -2,-2; set palette "
      "negative; set zrange [0.8:1]; set cbrange [0.8:1]> ";

  (*g_lua)["_comp_dist"] = gnu;
  g_lua->script("plot(" + fig_spec + " title<magnetization map [" + state_name +
                "]>', _comp_dist)");
}

} /* namespace oc */
} /* namespace ssl */


/* Software License Agreement (BSD License)
*
* Copyright (c) 2014, Ross Linscott (rossklin@gmail.com)
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are
* met:
*
*     Redistributions of source code must retain the above copyright
*     notice, this list of conditions and the following disclaimer.
*
*     Redistributions in binary form must reproduce the above copyright
*     notice, this list of conditions and the following disclaimer in
*     the documentation and/or other materials provided with the
*     distribution.
*
*     The names of its contributors may not be used to endorse or promote products
*     derived from this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
* "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
* LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
* A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
* HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
* SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
* LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
* DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
* THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

#if defined(NDEBUG)
#undef NDEBUG
#endif

#include <cstdlib>
#include <cmath>
#include <cstring>

#include <boost/random/mersenne_twister.hpp>
#include <boost/random/normal_distribution.hpp>
#include <boost/random/variate_generator.hpp>
#include <boost/date_time/posix_time/posix_time.hpp>

#ifdef STANDALONE
#include <RInside.h>
#endif

#include "../inst/include/SimpleSDESampler.h"

lpoly_evaluator::lpoly_evaluator(NumericMatrix cm, NumericMatrix trm) : terms(trm), coef_matrix(as_ublas_matrix(cm)){}

void lpoly_evaluator::set_h(double v){
  h = v;
}

NumericMatrix lpoly_evaluator::build(NumericMatrix data){
  double value;
  NumericMatrix result(data.nrow(), terms.nrow());
  int i,j,k;

  for (i = 0; i < data.nrow(); i++){
    for (j = 0; j < terms.nrow(); j++){
      value = 1;
      for (k = 0; k < terms.ncol(); k++){
	value *= pow(data(i, k), terms(j,k));
      }
      result(i,j) = value;
    }

  }
  return result;
}

// compute evaluation function
void lpoly_evaluator::operator()(uvector &query, uvector &out, double t) {
  if (query.size() != terms.ncol()){
    Rcpp::Rcout << "lpoly_evaluator::(): dimension mismatch: " << query.size() << " != " << terms.ncol() << endl;
    exit(-1);
  }

  NumericMatrix state(1, query.size(), query.begin());
  NumericMatrix state_model = build(state);
  uvector state_vector = as_ublas_vector(NumericVector(state_model.begin(), state_model.end()));

  uvector dpart = prod(coef_matrix, state_vector);
  out = dpart + (1 / sqrt(h)) * noise(t);
}

uvector lpoly_evaluator::noise(double t){
  int idx = t / h + 0.5;
  int dim = noise_precache.size2();
  uvector x(dim);

  if (idx >= noise_precache.size1()){
    Rcpp::Rcout << "Invalid time index: " << idx << endl;
    exit(1);
  }

  memcpy(&x[0], &noise_precache(idx,0), dim * sizeof(double));
  return x;
}

void lpoly_evaluator::set_noise(umatrix m){
  noise_precache = m;
}

void lpoly_evaluator::build_noise(int n, int d, double sigma){
  int i;
  int size = n * d;
  double *p;

  boost::mt19937 gener;
  boost::normal_distribution<> normal(0,sigma*sigma);
  boost::variate_generator<boost::mt19937&,boost::normal_distribution<> > rng(gener, normal);

  boost::posix_time::time_duration diff_time = boost::posix_time::microsec_clock::local_time() - boost::posix_time::second_clock::local_time();
  unsigned int the_seed = diff_time.total_microseconds();

  rng.engine().seed(the_seed);
  rng.distribution().reset();

  noise_precache.resize(n, d);
  p = &noise_precache(0,0);

  for (i = 0; i < size; i++){
    p[i] = rng();
  }
}

uvector lpoly_evaluator::evalfun(uvector query){
  NumericMatrix state(1, query.size(), query.begin());
  NumericMatrix state_model = build(state);
  uvector state_vector = as_ublas_vector(NumericVector(state_model.begin(), state_model.end()));
  return prod(coef_matrix, state_vector);
}

lpoly_jacobian::lpoly_jacobian(NumericMatrix cm, NumericMatrix trm) : terms(trm), coef_matrix(as_ublas_matrix(cm)){}

// compute jacobian
void lpoly_jacobian::operator()(uvector &q, umatrix &out, double t){
  int i,j,k;
  umatrix term_derivs = umatrix(terms.nrow(), terms.ncol());
  double x;

  for (i = 0; i < terms.nrow(); i++){
    for (j = 0; j < terms.ncol(); j++){
      // compute derivative of term i with respect to variable j
      term_derivs(i,j) = 1;
      for (k = 0; k < terms.ncol(); k++){
	// do something about corny case: q = 0, terms(i,k) = 0, j = k
	// which causes NAN but should actually be zero
	term_derivs(i,j) *= pow(terms(i, k), j == k) * pow(q(k), terms(i,k) - (j == k));
      }
    }
  }

  out = prod(coef_matrix, term_derivs);
}


#define _LIBCPP_ENABLE_THREAD_SAFETY_ANNOTATIONS
#include "keypoint_distribution.h"

#include <boost/histogram/ostream.hpp>
#include <fstream>
#include <gsl/gsl>
#include <opencv2/core/base.hpp>
#include <opencv2/core/mat.hpp>
#include <opencv2/core/persistence.hpp>
#include <opencv2/core/types.hpp>
#include <optional>
#include <sens_loc/analysis/distance.h>
#include <sens_loc/analysis/keypoints.h>
#include <sens_loc/io/histogram.h>
#include <sens_loc/util/thread_analysis.h>
#include <util/batch_visitor.h>
#include <util/common_structures.h>
#include <util/statistic_visitor.h>

using namespace std;
using namespace gsl;

namespace {

struct keypoint_stat_data {
    keypoint_stat_data() = default;

    void insert_points(gsl::span<cv::KeyPoint> points) noexcept {
        lock_guard l{_keypoint_mutex};
        _global_keypoints.insert(end(_global_keypoints), begin(points),
                                 end(points));
    }
    void insert_distances(gsl::span<float> distances) noexcept {
        lock_guard l{_distance_mutex};
        _global_minimal_distances.insert(_global_minimal_distances.end(),
                                         begin(distances), end(distances));
    }

    pair<vector<cv::KeyPoint>, vector<float>> extract() noexcept {
        lock_guard l1{_keypoint_mutex};
        lock_guard l2{_distance_mutex};

        pair p{move(_global_keypoints), move(_global_minimal_distances)};
        _global_keypoints         = vector<cv::KeyPoint>();
        _global_minimal_distances = vector<float>();

        return p;
    }

  private:
    mutex                                  _keypoint_mutex;
    vector<cv::KeyPoint> _global_keypoints GUARDED_BY(_keypoint_mutex);

    mutex                                   _distance_mutex;
    vector<float> _global_minimal_distances GUARDED_BY(_distance_mutex);
};

/// Calculate the 2-dimensional distribution of the keypoints for a dataset.
class keypoint_distribution {
  public:
    keypoint_distribution(keypoint_stat_data& d) noexcept
        : accumulated_data{d} {}

    void operator()(int /*idx*/,
                    optional<vector<cv::KeyPoint>> keypoints,
                    optional<cv::Mat> /*descriptors*/) noexcept {  // NOLINT
        if (keypoints->empty())
            return;

        accumulated_data.insert_points(*keypoints);

        // Calculate the minimal distance of each keypoint to all others
        // and insert that first into a local vector with that information
        // and finally into the global vector with that information.
        // This is the pixel-distance with euclidean norm.
        // NOTE: This is an inefficient implementation if O(n^2) complexity.
        {
            const size_t n_points = keypoints->size();
            if (n_points < 2)
                return;

            vector<cv::KeyPoint>& kp_ref = *keypoints;
            vector<float>         local_minima;
            vector<float>         local_distances;
            local_distances.reserve(n_points);

            // The upper triangle of the distance matrix needs to be calculated.
            // The last row must be ignored and the diagonal element will be 0.
            // Example:
            //    ++p1++p2++p3++
            // p1 |  0   2   4 |
            // p2 |      0   7 |
            // p3 |          0 |
            // The result are 'n_points / 2' number of minimal distances.
            // Because there will
            for (size_t i = 1; i < n_points - 1; ++i) {
                // loop-calculation
                for (size_t k = i + 1; k < n_points; ++k) {
                    const float dx = kp_ref[i].pt.x - kp_ref[k].pt.x;
                    const float dy = kp_ref[i].pt.y - kp_ref[k].pt.y;
                    const float d  = sqrt(dx * dx + dy * dy);
                    Ensures(d >= 0.0F);
                    local_distances.emplace_back(d);
                }

                Ensures(!local_distances.empty());

                // Store the minimal element for statistical evaluation.
                local_minima.emplace_back(
                    *min_element(begin(local_distances), end(local_distances)));

                // Ensure that local_distances is only allocated once.
                local_distances.clear();
            }
            accumulated_data.insert_distances(local_minima);
        }
    }

    size_t postprocess(unsigned int            image_width,
                       unsigned int            image_height,
                       const optional<string>& stat_file,
                       const optional<string>& response_histo,
                       const optional<string>& size_histo,
                       const optional<string>& kp_distance_histo,
                       const optional<string>& kp_distribution_histo) {
        auto [keypoints, distances] = accumulated_data.extract();

        if (keypoints.empty() || distances.empty())
            return 0UL;

        sens_loc::analysis::keypoints kp{image_width, image_height};

        const auto location_bins = 200U;
        kp.configure_distribution(location_bins);
        kp.configure_distribution("normalized width", "normalized height");

        const auto response_bins = 50U;
        kp.configure_response(response_bins, "detector response");

        const auto size_bins = 50U;
        kp.configure_size(size_bins, "keypoint size");
        kp.analyze(keypoints);

        sort(begin(distances), end(distances));
        const auto                   dist_bins = 50UL;
        sens_loc::analysis::distance distance_stat{distances, dist_bins,
                                                   "minimal keypoint distance"};

        if (stat_file) {
            cv::FileStorage kp_statistic{*stat_file,
                                         cv::FileStorage::WRITE |
                                             cv::FileStorage::FORMAT_YAML};
            kp_statistic.writeComment(
                "The following values contain the results of the statistical "
                "analysis for the keypoint distribution and detector results.");
            write(kp_statistic, "characteristics", kp);
            write(kp_statistic, "distance", distance_stat.get_statistic());
            kp_statistic.release();
        } else {
            cout << "==== Response\n"
                 << "count:  " << kp.response().count << "\n"
                 << "min:    " << kp.response().min << "\n"
                 << "max:    " << kp.response().max << "\n"
                 << "median: " << kp.response().median << "\n"
                 << "mean:   " << kp.response().mean << "\n"
                 << "var:    " << kp.response().variance << "\n"
                 << "stddev: " << kp.response().stddev << "\n"
                 << "==== Size\n"
                 << "count:  " << kp.size().count << "\n"
                 << "min:    " << kp.size().min << "\n"
                 << "max:    " << kp.size().max << "\n"
                 << "median: " << kp.size().median << "\n"
                 << "mean:   " << kp.size().mean << "\n"
                 << "var:    " << kp.size().variance << "\n"
                 << "stddev: " << kp.size().stddev << "\n"
                 << "=== Distance\n"
                 << "count:  " << distance_stat.count() << "\n"
                 << "min:    " << distance_stat.min() << "\n"
                 << "max:    " << distance_stat.max() << "\n"
                 << "median: " << distance_stat.median() << "\n"
                 << "mean:   " << distance_stat.mean() << "\n"
                 << "stddev: " << distance_stat.stddev() << "\n";
        }
        if (response_histo) {
            ofstream gnuplot_data{*response_histo};
            gnuplot_data << sens_loc::io::to_gnuplot(kp.response_histo())
                         << endl;
        } else {
            cout << kp.response_histo() << "\n";
        }

        if (size_histo) {
            ofstream gnuplot_data{*size_histo};
            gnuplot_data << sens_loc::io::to_gnuplot(kp.size_histo()) << endl;
        } else {
            cout << kp.size_histo() << "\n";
        }

        if (kp_distance_histo) {
            ofstream gnuplot_data{*kp_distance_histo};
            gnuplot_data << sens_loc::io::to_gnuplot(distance_stat.histogram())
                         << endl;
        } else {
            cout << distance_stat.histogram() << "\n";
        }

        if (kp_distribution_histo) {
            ofstream gnuplot_data{*kp_distribution_histo};
            gnuplot_data << sens_loc::io::to_gnuplot(kp.distribution()) << endl;
        }
        return keypoints.size();
    }

  private:
    keypoint_stat_data& accumulated_data;
};

}  // namespace

namespace sens_loc::apps {
int analyze_keypoint_distribution(
    util::processing_input  in,
    unsigned int            image_width,
    unsigned int            image_height,
    const optional<string>& stat_file,
    const optional<string>& response_histo,
    const optional<string>& size_histo,
    const optional<string>& kp_distance_histo,
    const optional<string>& kp_distribution_histo) {
    using visitor =
        statistic_visitor<keypoint_distribution, required_data::keypoints>;

    keypoint_stat_data d;

    auto f =
        parallel_visitation(in.start, in.end, visitor{in.input_pattern, d});

    size_t n_elements =
        f.postprocess(image_width, image_height, stat_file, response_histo,
                      size_histo, kp_distance_histo, kp_distribution_histo);

    return n_elements > 0UL ? 0 : 1;
}
}  // namespace sens_loc::apps