/*
 Copyright (C) 2020 Quaternion Risk Management Ltd
 All rights reserved.
*/

#include <boost/make_shared.hpp>
#include <ql/math/distributions/chisquaredistribution.hpp>
#include <ql/math/distributions/normaldistribution.hpp>
#include <ql/processes/eulerdiscretization.hpp>
#include <qle/models/crcirpp.hpp>
#include <qle/processes/crcirppstateprocess.hpp>

namespace QuantExt {

CrCirppStateProcess::CrCirppStateProcess(CrCirpp* const model,
                                         CrCirppStateProcess::Discretization disc)
    : StochasticProcess(boost::shared_ptr<StochasticProcess::discretization>(new EulerDiscretization)), model_(model),
      discretization_(disc) {}

Size CrCirppStateProcess::size() const { return 2; }

Disposable<Array> CrCirppStateProcess::initialValues() const {
    Array res(size(), 0.0);
    res[0] = model_->parametrization()->y0(0); // y0
    res[1] = 1.0; // S(0,0) = 1
    return res;
}

Disposable<Array> CrCirppStateProcess::drift(Time t, const Array& x) const {
    QL_FAIL("not implemented");
}

Disposable<Matrix> CrCirppStateProcess::diffusion(Time t, const Array& x) const {
    QL_FAIL("not implemented");
}

Disposable<Array> CrCirppStateProcess::evolve(Time t0, const Array& x0, Time dt, const Array& dw) const {
    Array retVal(size());
    Real kappa, theta, sigma, y0;
    kappa = model_->parametrization()->kappa(t0);
    theta = model_->parametrization()->theta(t0);
    sigma = model_->parametrization()->sigma(t0);
    y0 = model_->parametrization()->y0(t0);

    const Real sdt = std::sqrt(dt);
    switch (discretization_) {

    case BrigoAlfonsi: {
        // see D. Brigo and F. Mercurio. Interest Rate Models: Theory and Practice, 2nd
        // Edition. Springer, 2006.
        // Ensures non-negative values for \sigma^2 \leq 2*\kappa*\theta
        Real temp = (1 - kappa / 2.0 * dt);
        Real temp2 = temp * std::sqrt(x0[0]) + sigma * sdt * dw[0] / (2.0 * temp);
        retVal[0] = temp2 * temp2 + (kappa * theta - sigma * sigma / 4) * dt;
        break;
    }
    default:
        QL_FAIL("unknown discretization schema");
    }

    // second element is S(0,i)
    Real SM_ti = model_->defaultCurve()->survivalProbability(t0 + dt);
    Real SM_ti_pre = model_->defaultCurve()->survivalProbability(t0);
    Real Pcir_ti = model_->zeroBond(0, t0 + dt, y0);
    Real Pcir_ti_pre = model_->zeroBond(0, t0, y0);
    retVal[1] = x0[1] * SM_ti / SM_ti_pre * Pcir_ti_pre / Pcir_ti * exp(-x0[0] * dt);

    return retVal;
}

} // namespace QuantExt


/*
 * Software License Agreement (BSD License)
 *
 * Copyright (c) 2013, The University of Texas at Dallas
 * All rights reserved.
 *
 * 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.
 *     * Neither the name of The University of Texas at Dallas nor the
 *       names of its contributors may 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.
 *
 */

/*
 * Filename: connectivity.cpp
 *
 * Description: This file contains the ROS node that calculates the control variables
 * for a collection of iRobot Creates.
 *
 * Log
 * ----
 * 2013-09-15 File created by Hazen Eckert
 *
 */

// ROS includes
#include "ros/ros.h"
#include "ros/assert.h"
#include "dynamic_reconfigure/server.h"
#include "create_driver/vicon_driver.h"

// Library includes
#include <string>
#include <vector>
#include <map>
#include <set>
#include <algorithm>
#include <Eigen/Core>
#include <Eigen/LU>
#include <Eigen/Eigenvalues>

// Local Package includes
#include "connectivity_controller/EigenConnectivityMsgs.h"
#include "connectivity_controller/ConnectivityVariablesConfig.h"

using namespace std;
using namespace Eigen;

#define RHO1 0.7
#define RHO2 2.3

// Collision avoidance variables
double R = 0.7;
double r = 0.33;

double psi( double x, double rho1, double rho2 )
{
	if ( x <= rho1 ) {
		return 1;
	} else if ( rho1 < x && x < rho2 ) {
		return ( exp(-1/(rho2-x)) / ( exp(-1/(rho2-x)) + exp(1/(rho1-x)) ) );
	} else if ( rho2 <= x ) {
		return 0;
	}
	return 0;
}

double psiPrime( double d )
{
	if ( d <= RHO1 ) {
		return 0;
	} else if ( RHO1 < d && d < RHO2 ) {
		return (-0.5)*( pow((d-RHO1),(-2))+pow((d-RHO2),(-2)) ) / (1+  cosh(  (1/(-d+RHO1))+(1/(-d+RHO2)) )   );
	} else if ( RHO2 <= d ) {
		return 0;
	}
	return 0;
}

void reconfigureCallback(connectivity_controller::ConnectivityVariablesConfig &config, uint32_t level) {
	R = config.r_max;
	r = config.r_min;

}

int main(int argc, char **argv)
{
	// ROS Initalization
	ros::init(argc, argv, "connectivity");

	ros::NodeHandle n;
	ros::NodeHandle private_n("~");

	dynamic_reconfigure::Server<connectivity_controller::ConnectivityVariablesConfig> server;
  	dynamic_reconfigure::Server<connectivity_controller::ConnectivityVariablesConfig>::CallbackType f;
  	f = boost::bind(&reconfigureCallback, _1, _2);
 	server.setCallback(f);

	// ROS Parameters

	// List of robot names
	vector<string> robot_names;
	XmlRpc::XmlRpcValue robot_list;
	private_n.getParam("robot_list", robot_list);
	ROS_ASSERT(robot_list.getType() == XmlRpc::XmlRpcValue::TypeArray);

	for (int i = 0; i < robot_list.size(); i++)
	{
		ROS_ASSERT(robot_list[i].getType() == XmlRpc::XmlRpcValue::TypeString);
		robot_names.push_back(static_cast<string>(robot_list[i]));
	}

	// List of formation x positions
	vector<double> formation_x;
	XmlRpc::XmlRpcValue x_list;
	private_n.getParam("formation_x_list", x_list);
	ROS_ASSERT(x_list.getType() == XmlRpc::XmlRpcValue::TypeArray);

	for (int i = 0; i < x_list.size(); i++)
	{
		ROS_ASSERT(x_list[i].getType() == XmlRpc::XmlRpcValue::TypeDouble);
		formation_x.push_back(static_cast<double>(x_list[i]));
	}

	// List of formation y positions
	vector<double> formation_y;
	XmlRpc::XmlRpcValue y_list;
	private_n.getParam("formation_y_list", y_list);
	ROS_ASSERT(y_list.getType() == XmlRpc::XmlRpcValue::TypeArray);

	for (int i = 0; i < y_list.size(); i++)
	{
		ROS_ASSERT(y_list[i].getType() == XmlRpc::XmlRpcValue::TypeDouble);
		formation_y.push_back(static_cast<double>(y_list[i]));
	}

	// Number of robots
	const int num_robots = robot_names.size();

	// ROS Subscribers
	map<string, create_driver::ViconStream> vicon;
	vector<string>::iterator name_it;
	vector<ros::Subscriber> sub;

	for (name_it = robot_names.begin(); name_it != robot_names.end(); name_it++)
	{
		vicon.insert(pair<string, create_driver::ViconStream>(*name_it, create_driver::ViconStream()));
		sub.push_back(n.subscribe(*name_it + "/tf", 10, &create_driver::ViconStream::callback, &vicon[*name_it]));
	}

	// ROS Publishers
	ros::Publisher pub = n.advertise<connectivity_controller::EigenConnectivityMsgs>("/connectivity", 10);

	// ROS loop
	ros::Rate loop_rate(250); // 250 Hz

	double rho2[] = { 2.3, 2.2, 1.9, 1.5, 1.7, 2.0 };

	while (ros::ok())
	{
		// Retrieve Vicon Data
		ros::spinOnce();

		// Calculations
		MatrixXd adjacency_matrix = MatrixXd::Zero(num_robots, num_robots);
		MatrixXd P = MatrixXd::Zero(num_robots, num_robots - 1);
		MatrixXd L = MatrixXd::Zero(num_robots, num_robots);
		MatrixXd S = MatrixXd::Zero(num_robots, num_robots);

		// Adjacency Matrix
		for (int i = 0; i < num_robots; i++)
		{
			for (int j = 0; j < num_robots; j++)
			{
				if ( i == j )
				{
					adjacency_matrix(i,j) = 0.0;
				} else {
					adjacency_matrix(i,j) = psi(sqrt(
					pow(vicon[robot_names[j]].x() - vicon[robot_names[i]].x(), 2) +
					pow(vicon[robot_names[j]].y() - vicon[robot_names[i]].y(), 2)
					), RHO1, rho2[i]);
				}
			}
		}
		ROS_WARN_STREAM("Adj matrix:"  << adjacency_matrix );

		// L Matrix
		for (int i = 0; i < num_robots; i++)
		{
			for (int j = 0; j < num_robots; j++)
			{
				if ( i == j ) {
					double rowSum = 0;
					for (int k = 0; k < num_robots; k++) {
						rowSum += adjacency_matrix(i,k);
					}
					L(i,j) = rowSum - adjacency_matrix(i,j);
				} else {
					L(i,j) = -adjacency_matrix(i,j);
				}
			}
		}

		ROS_WARN_STREAM("L matrix:"  << L );

		// S Matrix
		S = MatrixXd::Identity(num_robots, num_robots) - L/((double)num_robots); // S = I - L/n

		ROS_WARN_STREAM("S matrix:"  << S );

		// Get eigenvalues and eigen vectors
		EigenSolver<MatrixXd> es(S.transpose());
		VectorXcd eigenvalues = es.eigenvalues();
		MatrixXcd eigenvectors = es.eigenvectors();

		ROS_WARN_STREAM("Eigenvalues"  << eigenvalues );
		ROS_WARN_STREAM("Eigenvectors:"  << eigenvectors );

		std::complex<double> max_eigenvalue(0.0,0.0);
		VectorXcd gamma = VectorXcd::Zero(num_robots);

		for (int i = 0; i < num_robots; i++)
		{
			if (abs(eigenvalues(i)) > abs(max_eigenvalue))
			{
				max_eigenvalue = eigenvalues(i);
				gamma = eigenvectors.col(i);
			}
		}

		//ROS_WARN_STREAM("Max Eigenvalue:"  << max_eigenvalue );
		ROS_WARN_STREAM("Gamma:"  << gamma );

		MatrixXcd inv_gamma1 = MatrixXcd::Zero(num_robots, num_robots);
		MatrixXcd inv_gamma2 = MatrixXcd::Zero(num_robots, num_robots);
		MatrixXcd inv_gamma1m = MatrixXcd::Zero(num_robots, num_robots);;
		MatrixXcd inv_gamma2m = MatrixXcd::Zero(num_robots, num_robots);;

		double control_input_CC_x[num_robots];
		double control_input_CC_y[num_robots];

		for (int i = 0; i < num_robots; i++)
		{
			VectorXcd gammai = VectorXcd::Zero(num_robots);
			for (int m = 0; m < num_robots; m++)
			{
				gammai(m) = abs(gamma(i));
			}

			//ROS_WARN_STREAM("gammai:"  << gammai );

			inv_gamma1.row(i) = gamma.transpose()/gamma(i);

			//ROS_WARN_STREAM("inv_gamma1.row(" << i << "):"  << inv_gamma1.row(i) );

			VectorXcd temp_inv_gamma2 = VectorXcd::Zero(num_robots);
			for (int n = 0; n < num_robots; n++)
			{
				temp_inv_gamma2(n) = 1/gamma(n).real();
			}
			inv_gamma2.row(i) = gamma(i)*temp_inv_gamma2.transpose();

			//ROS_WARN_STREAM("inv_gamma2.row(" << i << "):"  << inv_gamma2.row(i) );

			double upper_bound_g = 1000.0;
			double lower_bound_g = 1.1;
			for (int k = 0; k < num_robots; k++)
			{
				double gain_param = inv_gamma1(i,k).real();
				if (gain_param > lower_bound_g && gain_param < upper_bound_g)
				{
					inv_gamma1m(i,k) = pow(pow(gain_param,2)-pow(lower_bound_g,2), 2)/pow(pow(gain_param,2) - pow(upper_bound_g,2),2);
				}

				gain_param = inv_gamma2(i,k).real();
				if (gain_param > lower_bound_g && gain_param < upper_bound_g)
				{
					inv_gamma2m(i,k) = pow(pow(gain_param,2)-pow(lower_bound_g,2), 2)/pow(pow(gain_param,2) - pow(upper_bound_g,2),2);
				}
			}

			//ROS_WARN_STREAM("inv_gamma1m.row(" << i << "):"  << inv_gamma1m.row(i) );
			//ROS_WARN_STREAM("inv_gamma2m.row(" << i << "):"  << inv_gamma2m.row(i) );

			set<int> neighbors_out;
			set<int> neighbors_in;
			for (int j = 0; j < num_robots; j++)
			{
				if (adjacency_matrix(i,j) > 0)
					neighbors_out.insert(j);
				if (adjacency_matrix(j,i) > 0)
					neighbors_in.insert(j);
			}
			vector<int> neighbors(num_robots*2);
			vector<int>::iterator it = set_union(neighbors_out.begin(), neighbors_out.end(),
													neighbors_in.begin(), neighbors_in.end(),
														neighbors.begin());
			neighbors.resize(it - neighbors.begin());

			int Ni = neighbors.size();
			////////////////////////////////////////////
			std::stringstream ss1;
			for (std::set<int>::iterator it = neighbors_out.begin(); it != neighbors_out.end(); ++it)
			{
				ss1 << ' ' << *it;
			}
			//ROS_WARN_STREAM("neighbors_out:"  << ss1.str() );

			std::stringstream ss2;
			for (std::set<int>::iterator it = neighbors_in.begin(); it != neighbors_in.end(); ++it)
			{
				ss2 << ' ' << *it;
			}
			//ROS_WARN_STREAM("neighbors_in:"  << ss2.str() );

			std::stringstream ss3;
			for (std::vector<int>::iterator it = neighbors.begin(); it != neighbors.end(); ++it)
			{
				ss3 << ' ' << *it;
			}
			//ROS_WARN_STREAM("neighbors:"  << ss3.str() );
			/////////////////////////////////////
			double sumx = 0.0;
			double sumy = 0.0;

			for (int k = 0; k < Ni; k++)
			{
				sumx = sumx + (-1)*( abs(inv_gamma1m(i,neighbors[k]) ) + abs(inv_gamma2m(i,neighbors[k]) )  ) * (vicon[robot_names[i]].x() - vicon[robot_names[neighbors[k]]].x());
				sumy = sumy + (-1)*( abs(inv_gamma1m(i,neighbors[k]) ) + abs(inv_gamma2m(i,neighbors[k]) )  ) * (vicon[robot_names[i]].y() - vicon[robot_names[neighbors[k]]].y());
			}

			control_input_CC_x[i] = sumx;
			control_input_CC_y[i] = sumy;

		}

		// P Matrix
		if (num_robots == 6)
		{
			P << 	1/sqrt(2), 		-1/sqrt(6.0),		-1/2.0/sqrt(3), 	-1/2.0/sqrt(5), 	-1/sqrt(30),
					0.0,       		sqrt(2.0/3), 		-1/2.0/sqrt(3), 	-1/2.0/sqrt(5), 	-1/sqrt(30),
					0.0, 		 	0.0, 				sqrt(3.0)/2,    	-1/(2*sqrt(5)),		-1/sqrt(30),
					0.0, 			0.0, 				0.0, 				2/sqrt(5), 			-1/sqrt(30),
					0.0, 			0.0, 				0.0,				0.0, 				sqrt(5.0/6),
					-1/sqrt(2), 	-1/sqrt(6),			-1/2.0/sqrt(3), 	-1/2.0/sqrt(5), 	-1/sqrt(30);
		} else if (num_robots == 4) {
			P << 	-0.5,        		-0.5,              	-0.5,
					 0.833333333333333, -0.166666666666667, -0.166666666666667,
					-0.166666666666667,  0.833333333333333, -0.166666666666667,
					-0.166666666666667, -0.166666666666667,  0.833333333333333;
		}
		// M Matrix
		MatrixXd M = P.transpose() * L * P;
		//std::cerr << "adjacency_matrix"  << adjacency_matrix  ;
		// dLdx and dLdy for each robot
		MatrixXd dLdx[num_robots];
		MatrixXd dLdy[num_robots];

		for (int i = 0; i < num_robots; i++)
		{
			dLdx[i] = MatrixXd::Zero(num_robots, num_robots);
			dLdy[i] = MatrixXd::Zero(num_robots, num_robots);
		}

		for(int k = 0; k < num_robots; k++)
		{
			for(int i = 0; i < num_robots; i++)
			{
				for (int j = i + 1; j < num_robots; j++)
				{
					double D = sqrt(
					pow(vicon[robot_names[i]].x() - vicon[robot_names[j]].x(), 2) +
					pow(vicon[robot_names[i]].y() - vicon[robot_names[j]].y(), 2)
					);

					double dAdD = psiPrime(D);

					if( k == i ) {
						double dDdx = (vicon[robot_names[i]].x()-vicon[robot_names[j]].x()) / D;
						double dDdy = (vicon[robot_names[i]].y()-vicon[robot_names[j]].y()) / D;

						dLdx[k](i,j) = dAdD*dDdx;
						dLdy[k](i,j) = dAdD*dDdy;

					}

					if( k == j ) {
						double dDdx = -(vicon[robot_names[i]].x()-vicon[robot_names[j]].x()) /D;
						double dDdy = -(vicon[robot_names[i]].y()-vicon[robot_names[j]].y()) / D;

						dLdx[k](i,j) = dAdD*dDdx;
						dLdy[k](i,j) = dAdD*dDdy;

					}

					dLdx[k](j,i) = dLdx[k](i,j);
					dLdy[k](j,i) = dLdy[k](i,j);

				}
			}

			for(int i = 0; i < num_robots; i++)
			{
				double sumx = 0;
				double sumy = 0;
				for(int j = 0; j<num_robots; j++)
				{
					sumx += dLdx[k](i,j);
					sumy += dLdy[k](i,j);
				}
				dLdx[k](i,i) = -sumx;
	            dLdy[k](i,i) = -sumy;
			}
		}

		// dMdx and dMdy for each robot
		MatrixXd dMdx[num_robots];
		MatrixXd dMdy[num_robots];

		for (int i = 0; i < num_robots; i++)
		{
			dMdx[i] = P.transpose() * dLdx[i] * P;
			dMdy[i] = P.transpose() * dLdy[i] * P;
			//std:cerr << dMdx[i]  ;
		}

		// Collision Avoidance for each robot starts here:
		MatrixXd CAMatrix_x = MatrixXd::Zero(num_robots,num_robots);
		MatrixXd CAMatrix_y = MatrixXd::Zero(num_robots,num_robots);

		for(int i = 0; i < num_robots; i++)
		{
			for (int j = i + 1; j < num_robots; j++)
			{
				double D = sqrt(
				pow(vicon[robot_names[i]].x() - vicon[robot_names[j]].x(), 2) +
				pow(vicon[robot_names[i]].y() - vicon[robot_names[j]].y(), 2)
				);
				if( D < R && D > r){
					CAMatrix_x(i,j) = -4*(pow(R,2)-pow(r,2))
										*(pow(D,2)-pow(R,2))
										/pow((pow(D,2)-pow(r,2)),3)
										*(vicon[robot_names[i]].x()
										- vicon[robot_names[j]].x());

		        	CAMatrix_y(i,j) = -4*(pow(R,2)-pow(r,2))
		        						*(pow(D,2)-pow(R,2))
		        						/pow((pow(D,2)-pow(r,2)),3)
		        						*(vicon[robot_names[i]].y()
		        						- vicon[robot_names[j]].y());
            	} else {
            		CAMatrix_x(i,j) = 0;
            		CAMatrix_y(i,j) = 0;
            	}

            	CAMatrix_x(j,i) = -CAMatrix_x(i,j);
	       		CAMatrix_y(j,i) = -CAMatrix_y(i,j);
			}
		}

		double control_input_CA_x[num_robots];
		double control_input_CA_y[num_robots];

		for(int i = 0; i < num_robots; i++)
		{
			double sumx = 0;
			double sumy = 0;
			for (int j = 0; j < num_robots; j++)
			{
				sumx = sumx + CAMatrix_x(i,j);
				sumy = sumy + CAMatrix_y(i,j);
			}
			control_input_CA_x[i] = sumx;
			control_input_CA_y[i] = sumy;
		}

		double control_input_FC_x[num_robots];
		double control_input_FC_y[num_robots];

		for (int i = 0; i < num_robots; i++)
		{
			control_input_FC_x[i] = formation_x[i] - vicon[robot_names[i]].x();
			control_input_FC_y[i] = formation_y[i] - vicon[robot_names[i]].y();
		}

		// Trace of Mx and My for each robot
		double trace_Mx[num_robots];
		double trace_My[num_robots];

		for (int i = 0; i < num_robots; i++)
		{
			//std::cerr << (M.inverse()*dMdx[i])  ;
			trace_Mx[i] = (M.inverse()*dMdx[i]).trace();
			trace_My[i] = (M.inverse()*dMdy[i]).trace();
		}

		double determinant_M = M.determinant();

		// Send msg

		connectivity_controller::EigenConnectivityMsgs msg;

		msg.trMx.assign(trace_Mx, trace_Mx + num_robots);
		msg.trMy.assign(trace_My, trace_My + num_robots);

		msg.u_cc_x.assign(control_input_CC_x, control_input_CC_x + num_robots);
		msg.u_cc_y.assign(control_input_CC_y, control_input_CC_y + num_robots);

		msg.u_ca_x.assign(control_input_CA_x, control_input_CA_x + num_robots);
		msg.u_ca_y.assign(control_input_CA_y, control_input_CA_y + num_robots);

		msg.u_fc_x.assign(control_input_FC_x, control_input_FC_x + num_robots);
		msg.u_fc_y.assign(control_input_FC_y, control_input_FC_y + num_robots);
		msg.detM = determinant_M;

		pub.publish(msg);

		// Sleep
		loop_rate.sleep();
	}

	return 0;
}


#define PY_SSIZE_T_CLEAN
#include <Python.h>
#include <vector>
#include <boost/timer/timer.hpp>
#include <GeomLib.hpp>
#include <MPILib/include/MPINetworkCode.hpp>
#include <MPILib/include/RateAlgorithmCode.hpp>
#include <MPILib/include/SimulationRunParameter.hpp>
#include <MPILib/include/report/handler/InactiveReportHandler.hpp>
#include <MPILib/include/WilsonCowanAlgorithm.hpp>
#include <MPILib/include/RateFunctorCode.hpp>
#include <MPILib/include/utilities/ProgressBar.hpp>
#include <MPILib/include/BasicDefinitions.hpp>
#include <MPILib/include/WilsonCowanParameter.hpp>
#include <MPILib/include/WilsonCowanAlgorithm.hpp>
#include <MPILib/include/MiindTvbModelAbstract.hpp>

MPILib::Rate External_RateFunction(MPILib::Time t){
	return 1.0;
}
/* This class is designed to be used in conjunction with
 * Miind_WilsonCowan.py only. Any updates to this shared library
 * should be communicated to the TVB team and a new .so file should be provided
 * to be included in the tvb-library source tree (tvb-library/tvb/simulator/models/).
 */
class MiindWilsonCowan : public MPILib::MiindTvbModelAbstract<double, MPILib::utilities::CircularDistribution>{
private:
	std::vector<double> E_Initials = std::vector<double> (); // initial excitatory values
	std::vector<double> I_Initials = std::vector<double> (); // initial inhibitory values

public:

	MiindWilsonCowan(int num_nodes, double simulation_length, double dt) :
		MiindTvbModelAbstract(num_nodes, simulation_length) {
			this->_time_step = dt;
		}

	// In general, we won't need to set initial values but we implement it here
	// so that we can match TVB's Wilson Cowan example.
	void setInitialValues(std::vector<double> E_vals, std::vector<double> I_vals) {
		for(int i=0; i<_num_nodes; i++) {
			E_Initials.push_back(E_vals[i]);
			I_Initials.push_back(I_vals[i]);
		}
	}

	void init(std::vector<double> params)
	{
		// (TVB : tau_e) population time constant
		Time   E_tau       = params[4];
		// (TVB : c_e) maximum rate reached by the sigmoid function
		Rate   E_max_rate  = params[6];
		// (TVB : a_e) noise term modulates the input (a term in Wilson Cowan)
		double E_noise     = params[10];
		// (TVB : b_e) bias term shifts the input (theta term in Wilson Cowan)
		double E_bias      = params[12];
		// (TVB : r_e) smoothing term over output which models refractory dynamics
		double E_smoothing = params[8];
		// (TVB : tau_i) population time constant
		Time   I_tau       = params[5];
		// (TVB : c_i) maximum rate reached by the sigmoid function
		Rate   I_max_rate  = params[7];
		// (TVB : a_i) noise term modulates the input
		double I_noise     = params[11];
		// (TVB : b_i) bias term shifts the input
		double I_bias      = params[13];
		// (TVB : r_e) smoothing term over output which models refractory dynamics
		double I_smoothing = params[9];
		// (TVB : c_ee) Weight of E->E connection
		double E_E_Weight  = params[0];
		// (TVB : -c_ie) Weight of E->I connection
		double I_E_Weight  = -params[1];
		 // (TVB : c_ei) Weight of I->E connection
		double E_I_Weight  = params[2];
		// (TVB : -c_ii) Weight of I->I connection
		double I_I_Weight  = -params[3];
		// (TVB : P) Some additional drive to excitatory pop
		double P_E_Weight  = params[14];
		// (TVB : Q) Some additional drive to inhibitory pop
		double Q_I_Weight  = params[15];

		try {
			std::vector<NodeId> E_ids = std::vector<NodeId>();
			std::vector<NodeId> I_ids = std::vector<NodeId>();

			for(int i=0; i<_num_nodes; i++) {
				MPILib::WilsonCowanParameter E_param = MPILib::WilsonCowanParameter(
								E_tau, E_max_rate, E_noise, E_bias, E_Initials[i], E_smoothing);
				MPILib::WilsonCowanAlgorithm E_alg(E_param);

				MPILib::WilsonCowanParameter I_param = MPILib::WilsonCowanParameter(
								I_tau, I_max_rate, I_noise, I_bias, I_Initials[i], I_smoothing);
				MPILib::WilsonCowanAlgorithm I_alg(I_param);

				MPILib::Rate RateFunction_P(MPILib::Time);
				MPILib::RateFunctor<double> rate_functor_p(External_RateFunction);

				MPILib::Rate RateFunction_Q(MPILib::Time);
				MPILib::RateFunctor<double> rate_functor_q(External_RateFunction);

				MPILib::NodeId id_E = network.addNode(E_alg, MPILib::EXCITATORY_DIRECT);
				MPILib::NodeId id_I = network.addNode(I_alg, MPILib::INHIBITORY_DIRECT);
				MPILib::NodeId id_P = network.addNode(rate_functor_p, MPILib::NEUTRAL);
				MPILib::NodeId id_Q = network.addNode(rate_functor_q, MPILib::NEUTRAL);

				E_ids.push_back(id_E);
				I_ids.push_back(id_I);

				network.makeFirstInputOfSecond(id_E,id_E,E_E_Weight);
				network.makeFirstInputOfSecond(id_I,id_E,I_E_Weight);
				network.makeFirstInputOfSecond(id_E,id_I,E_I_Weight);
				network.makeFirstInputOfSecond(id_I,id_I,I_I_Weight);
				network.makeFirstInputOfSecond(id_P,id_E,P_E_Weight);
				network.makeFirstInputOfSecond(id_Q,id_I,Q_I_Weight);

			}

			// Set each node to have an external successor (for coupling input from TVB)
			for(auto& id : E_ids) {
				network.setNodeExternalSuccessor(id);
				network.setNodeExternalPrecursor(id, 1);
			}

			for(auto& id : I_ids) {
				network.setNodeExternalSuccessor(id);
				network.setNodeExternalPrecursor(id, 1);
			}

			std::string sim_name = "miind_wc";
			MPILib::report::handler::InactiveReportHandler handler =
											MPILib::report::handler::InactiveReportHandler();

			SimulationRunParameter par_run( handler,(_simulation_length/_time_step)+1,0,
											_simulation_length,_time_step,_time_step,sim_name,_time_step);

			network.configureSimulation(par_run);

		} catch(std::exception& exc){
			std::cout << exc.what() << std::endl;
		}
	}
};

static MiindWilsonCowan *model;

static PyObject *miind_init(PyObject *self, PyObject *args)
{
    int nodes;
    double sim_time;
    double time_step;

    if (!PyArg_ParseTuple(args, "idd", &nodes, &sim_time, &time_step))
	return NULL;

    model = new MiindWilsonCowan(nodes, sim_time, time_step);

    Py_INCREF(Py_None);
    return Py_None;
}

static PyObject *miind_initParams(PyObject *self, PyObject *args)
{
    PyObject *float_list;
    int pr_length;

    if (!PyArg_ParseTuple(args, "O", &float_list))
        return NULL;
    pr_length = PyObject_Length(float_list);
    if (pr_length < 0)
        return NULL;

    std::vector<double> params(pr_length);

    for (int index = 0; index < pr_length; index++) {
        PyObject *item;
        item = PyList_GetItem(float_list, index);
        if (!PyFloat_Check(item))
            params[index] = 0.0;
        params[index] = PyFloat_AsDouble(item);
    }

    model->init(params);

    Py_INCREF(Py_None);
    return Py_None;
}

static PyObject *miind_setInitialValues(PyObject *self, PyObject *args)
{
    PyObject *float_listE;
    PyObject *float_listI;
    int pr_length;

    if (!PyArg_ParseTuple(args, "OO", &float_listE, &float_listI))
        return NULL;
    pr_length = PyObject_Length(float_listE);
    if (pr_length < 0)
        return NULL;

    std::vector<double> insE(pr_length);
    std::vector<double> insI(pr_length);

    for (int index = 0; index < pr_length; index++) {
        PyObject *item;
        item = PyList_GetItem(float_listE, index);
        if (!PyFloat_Check(item))
            insE[index] = 0.0;
        insE[index] = PyFloat_AsDouble(item);
    }

    for (int index = 0; index < pr_length; index++) {
        PyObject *item;
        item = PyList_GetItem(float_listI, index);
        if (!PyFloat_Check(item))
            insI[index] = 0.0;
        insI[index] = PyFloat_AsDouble(item);
    }

    model->setInitialValues(insE,insI);

    Py_INCREF(Py_None);
    return Py_None;
}

static PyObject *miind_getTimeStep(PyObject *self, PyObject *args)
{
    return Py_BuildValue("d", model->getTimeStep());
}

static PyObject *miind_getSimulationLength(PyObject *self, PyObject *args)
{
    return Py_BuildValue("d", model->getSimulationLength());
}

static PyObject *miind_startSimulation(PyObject *self, PyObject *args)
{
    model->startSimulation();
    Py_INCREF(Py_None);
    return Py_None;
}

static PyObject *miind_evolveSingleStep(PyObject *self, PyObject *args)
{
    PyObject *float_list;
    int pr_length;

    if (!PyArg_ParseTuple(args, "O", &float_list))
        return NULL;
    pr_length = PyObject_Length(float_list);
    if (pr_length < 0)
        return NULL;

    std::vector<double> activities(pr_length);

    for (int index = 0; index < pr_length; index++) {
        PyObject *item;
        item = PyList_GetItem(float_list, index);
        if (!PyFloat_Check(item))
            activities[index] = 0.0;
        activities[index] = PyFloat_AsDouble(item);
    }

    std::vector<double> out_activities = model->evolveSingleStep(activities);

    PyObject* tuple = PyTuple_New(pr_length);

    for (int index = 0; index < pr_length; index++) {
        PyTuple_SetItem(tuple, index, Py_BuildValue("d", out_activities[index]));
    }

    return tuple;
}

static PyObject *miind_endSimulation(PyObject *self, PyObject *args)
{
    model->endSimulation();
    Py_INCREF(Py_None);
    return Py_None;
}

static PyMethodDef MiindModelMethods[] = {
    {"init",  miind_init, METH_VARARGS, "Init Miind model."},
    {"initParams",  miind_initParams, METH_VARARGS, "Init Miind model parameters."},
    {"getTimeStep",  miind_getTimeStep, METH_VARARGS, "Get time step."},
    {"getSimulationLength",  miind_getSimulationLength, METH_VARARGS, "Get sim time."},
    {"startSimulation",  miind_startSimulation, METH_VARARGS, "Start simulation."},
    {"evolveSingleStep",  miind_evolveSingleStep, METH_VARARGS, "Evolve one time step."},
    {"endSimulation",  miind_endSimulation, METH_VARARGS, "Clean up."},
    {"setInitialValues", miind_setInitialValues, METH_VARARGS, "Set initial values."},
    {NULL, NULL, 0, NULL}        /* Sentinel */
};

static struct PyModuleDef miindmodule = {
    PyModuleDef_HEAD_INIT,
    "libmiindwc",   /* name of module */
    NULL, /* module documentation, may be NULL */
    -1,       /* size of per-interpreter state of the module,
                 or -1 if the module keeps state in global variables. */
    MiindModelMethods
};

PyMODINIT_FUNC
PyInit_libmiindwc(void)
{
    return PyModule_Create(&miindmodule);
}


/* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
 * Filename    : sphdec.cpp                                                                            *
 * Project     : Planewalker - Schnorr-Euchner sphere decoder simulation for space-time lattice codes  *
 * Authors     : Pasi Pyrrö, Oliver Gnilke                                                             *
 * Version     : 1.0                                                                                *
 * Copyright   : Aalto University ~ School of Science ~ Department of Mathematics and Systems Analysis *
 * Date        : 9.1.2017                                                                              *
 * Language    : C++ (2011 or newer standard)                                                          *
 * Description : The core algorithm implementations for the simulation + wrapper function for them     *
 * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */

#define ARMA_NO_DEBUG /* disable Armadillo bound checks for addiotional speed */

#include <iostream>
#include <armadillo> /* linear algebra library */
#include <vector>
#include <complex>
#include <string>

#include "sphdec.hpp"
#include "misc.hpp"
#include "algorithms.hpp"

using namespace std;
using namespace arma;

/* Decision feedback equalization on xt[i]
   i.e. Use the Babai nearest plane algorithm as a starting point in dimension i */
inline void dfe(int i, int q, vec &xt, const vec &y, vec &delta, const vec &ksi, const mat &R, const vector<int> &S) {
    /* calculate the "Babai coeffient" and "round" it to nearest symbol in the symbol set */
    xt[i] = nearest_symbol((y[i]-ksi[i])/R(i,i), S);
    if (xt[i] <= S[0]) { /* only feasible way is positive direction */
        delta[i] = 2;
    } else if (xt[i] >= S[q - 1]) { /* only feasible way is negative direction */
        delta[i] = -2;
    } else { /* there's room to enumerate within signal set */
        delta[i] = 2*sesd_sign(y[i]-ksi[i]-R(i,i)*xt[i]);
    }
}

/* Schnorr-Euchner enumeration step for the sphere decoder i.e. zig-zag around the dfe point.
   Makes sure we loop over the points in dimension i in non descending order.
   This enables early stopping criteria when we find a point in dimension i outside the radius,
   i.e. no points in that sub search tree are inside the sphere because they can only be further away or at the same distance. */
inline void se_enum(int i, vec &xt, vec &delta) {
    /* add delta to xt and pick new delta (difference to next symbol in S we want to try) */
    xt[i] = xt[i] + delta[i];
    delta[i] = -delta[i] - 2*sesd_sign(delta[i]);
}

/* Perform some basic checks to ensure that the sphere decoder works as intended */
inline bool check(double radius, const mat &R, const vec &y) {
    if (radius <= 0){
        log_msg("sphdec: negative squared initial radius given!", "Error");
        log_msg("aborting simulation round...");
        return false;
    }

    if (R.n_rows != R.n_cols) {
        log_msg("sphdec: R is not a square matrix!", "Error");
        log_msg("aborting simulation round...");
        return false;
    }

    if (y.n_elem != R.n_cols) {
        log_msg("sphdec: vector y dimension mismatch!", "Error");
        log_msg("aborting simulation round...");
        return false;
    }
    return true;
}

/* Basic sphere decoder algorithm */
vector<int> sphdec(const vec &y, const mat &R, const vector<int> &S, int &counter, double radius){

    /* Initialize */
    int k = params["no_of_matrices"];
    int q = params["x-PAM"];
    int i = k-1; /* We start from dimension k-1 and iterate all the way down to dimension 0 */

    vector<int> x(k); /* point to decode */
    vec xt(k), ksi(k), delta(k), dist(k);
    xt.zeros(); ksi.zeros(); delta.zeros(); dist.zeros();
    counter = 0; /* counts how many search tree nodes (loop iterations) we went through */

    vec distances(k, fill::zeros);

    double xidist = 0.0;
    bool found = false;

    if (!check(radius, R, y))
        return vector<int>(0);

    /* Initialize xt[k-1] and delta[k-1] */
    dfe(i, q, xt, y, delta, ksi, R, S);

    while (!exit_flag) {

        counter++;
        /* Step 3. */
        xidist = pow(y[i]-ksi[i]-R(i,i)*xt[i], 2);

        distances[i] = dist[i] + xidist;

        /***** Uncomment line below to debug *****/
        // cout << i << ": " << vec2str(xt, k) << ", xidist = " << xidist + dist[i] << ", C = " << radius << endl;

        if (radius < dist[i] + xidist) { // current point xt is outside the sphere
            // Step 4.
            if (i == k-1) {
                break;
            } else {
                // Step 6.
                i++;
                se_enum(i, xt, delta);
            }
        } else { // we are inside the sphere
            if (xt[i] < S[0] || xt[i] > S[q - 1]){ // we are outside the signal set boundaries
                if ((xt[i] < S[0] && (xt[i] + delta[i]) > S[q - 1]) || (xt[i] > S[q - 1] && (xt[i] + delta[i]) < S[0])){
                    // Step 4.
                    if (i == k-1) {
                        break;
                    } else {
                        i++;
                        se_enum(i, xt, delta);
                    }
                } else {
                    se_enum(i, xt, delta);
                }
            } else {
                if (i > 0) {
                    ksi[i-1] = 0;
                    for (int j = i; j < k; j++)
                        ksi[i-1] += R(i-1, j)*xt[j];
                    dist[i-1] = dist[i] + xidist;
                    i--;
                    dfe(i, q, xt, y, delta, ksi, R, S);
                } else { // lattice point is found (Step 5)
                    radius = dist[0] + xidist;
                    found = true;
                    x = conv_to<vector<int>>::from(xt);
                    i++;
                    se_enum(i, xt, delta);
                }
            }
        }
    }
    if (found)
        return x;
    else {
        log_msg("Initial squared radius used: " + to_string(radius), "Alert");
        log_msg("distances considered: " + vec2str(distances, distances.size()), "Alert");
        log_msg("sphdec: point not found!", "Alert");
        return vector<int>(0);
    }
}

/* Sphere decoder algorithm that considers the spherical shaping of the codebook */
vector<int> sphdec_spherical_shaping(const vec &y, const mat &HR, const mat &R, const vector<int> &S,
                                     int &counter, double P, double radius){

    /* Initialize */
    int k = params["no_of_matrices"];
    int q = params["x-PAM"];
    int i = k-1; /* We start from dimension k-1 and iterate all the way down to dimension 0 */

    vector<int> x(k); /* point to decode */
    vec xt(k), ksi(k), delta(k), dist(k), curr(k), ener(k);
    xt.zeros(); ksi.zeros(); delta.zeros(); dist.zeros(); curr.zeros(); ener.zeros();
    counter = 0; /* counts how many search tree nodes (loop iterations) we went through */

    vec distances(k, fill::zeros);

    double xidist = 0.0, xiener = 0.0;
    bool found = false;

    if (!check(radius, HR, y))
        return vector<int>(0);

    /* Initialize xt[k-1] and delta[k-1] */
    dfe(i, q, xt, y, delta, ksi, HR, S);

    while (!exit_flag) {

        counter++;
        // Step 3.
        xidist = pow(y[i]-ksi[i]-HR(i,i)*xt[i], 2);

        distances[i] = dist[i] + xidist;

        /***** Uncomment line below to debug *****/
        // cout << i << ": " << vec2str(xt, k) << ", xidist = " << xidist + dist[i] << ", C = " << radius << endl;

        if (radius < dist[i] + xidist) { /* current point xt is outside the sphere */
            // Step 4.
            if (i == k-1) {
                break;
            } else {
                // Step 6.
                i++;
                se_enum(i, xt, delta);
            }
        } else { /* we are inside the sphere */

            xiener = pow(xt[i]*R(i,i) + curr[i], 2);

            // we are outside the signal set boundaries
            if (xt[i] < S[0] || xt[i] > S[q - 1] || ener[i] + xiener > P + 10e-6) {
                if ((xt[i] < S[0] && (xt[i] + delta[i]) > S[q - 1]) || (xt[i] > S[q - 1] && (xt[i] + delta[i]) < S[0])){
                    // Step 4.
                    if (i == k-1) {
                        break;
                    } else {
                        i++;
                        se_enum(i, xt, delta);
                    }
                } else
                    se_enum(i, xt, delta);
            } else {
                if (i > 0) {
                    ksi[i-1] = 0;
                    curr[i-1] = 0;
                    for (int j = i; j < k; j++){
                        ksi[i-1] += HR(i-1, j)*xt[j];
                        curr[i-1] += xt[j]*R(i-1,j);
                    }
                    dist[i-1] = dist[i] + xidist;
                    ener[i-1] = ener[i] + xiener;
                    i--;
                    dfe(i, q, xt, y, delta, ksi, HR, S);
                } else { // lattice point is found (Step 5)
                    radius = dist[0] + xidist;
                    found = true;
                    x = conv_to<vector<int>>::from(xt);
                    i++;
                    se_enum(i, xt, delta);
                }
            }
        }
    }
    if (found)
        return x;
    else {
        log_msg("Initial squared radius used: " + to_string(radius), "Alert");
        // mat tmp = y-HR*xt;
        // cx_mat distasd = cx_mat(tmp, mat(tmp.n_rows, tmp.n_cols, fill::zeros));
        // log_msg("Current squared distance: " + to_string(frob_norm_squared(distasd)), "Alert");
        // log_msg("xidist: " + to_string(xidist), "Alert");
        log_msg("distances considered: " + vec2str(distances, distances.size()), "Alert");
        log_msg("sphdec: point not found!", "Alert");
        return vector<int>(0);
    }
}

/* Wrapper function for the sphere decoder to handle complex to real matrix conversion,
   QR-decomposition and other mappings */
vector<int> sphdec_wrapper(const vector<cx_mat> &bases, const mat Rorig, const cx_mat &H,
                           const cx_mat &X, const cx_mat &N, const vector<int> &symbset, int &visited_nodes, double radius) {

    /* read simulation parameters */
    int n = params["no_of_receiver_antennas"];
    int t = params["time_slots"];
    int k = params["no_of_matrices"];
    double P = dparams["spherical_shaping_max_power"];

    vector<int> x(k);                          /* decoded point */
    cx_mat Y(n, t);                            /* received code block */
    mat B(2*t*n, k), Q, R;                     /* real matrices for QR-decompostion */
    vec y(2*t*n);                              /* raw input vector for the sphere decoder (unmapped) */
    vec y2;                                    /* input vector for the sphere decoder (mapped to same space a R) */

    Y = H*X + N;                               /* Calculate simulated code block that we would receive */
    y = to_real_vector(Y);                     /* convert Y to real vector */

    for (int i = 0; i < k; i++)
        B.col(i) = to_real_vector(H*bases[i]); /* B = (HX1 HX2 ... HXk) = generator matrix of faded lattice */

    qr_econ(Q, R, B);                          /* QR-decomposition of B (omits zero rows in R) */
    process_qr(Q, R);                          /* Make sure R has positive diagonal elements */

    y2 = Q.st()*y;                             /* Map y to same basis as R */

    /* decide which sphere decoder algorithm to use */
    if (P <= 0)
        x = sphdec(y2, R, symbset, visited_nodes, radius);
    else
        x = sphdec_spherical_shaping(y2, R, Rorig, symbset, visited_nodes, P, radius);

    return x;
}


#if HAVE_CONFIG_H
# include <config.h>
#endif

#include <boost/bind.hpp>
using boost::cref;

#include "Cosmology.h"
#include "PowerSpectrum.h"
#include "Quadrature.h"
#include "CDE.h"
#include "BSPTN.h"
#include "SPT.h"
#include "Spline.h"
#include "SpecialFunctions.h"
#include "LinearPS.h"

#include <cerrno>
#include <cstdio>
#include <cstdlib>
#include <sys/stat.h>
#include <math.h>       /* pow */

/////////////////CLUSTERING DARK ENERGY (AND MG RESUMMED SPECTRUM) LIBRARY ///////////////////
// Load Linear PS
CDE::CDE(const Cosmology& C, const PowerSpectrum& P_l, double epsrel_)
: C(C), P_l(P_l)
{
    epsrel = epsrel;
}

/* Tree level cross bispectrum for  isosceles or equilateral configurations */
// vars specifies omega_m, cs^2, w , p3 and scale factor
double CDE::Bcde(int a, double vars[], double k, double p, double x) const {
IOW iow;
iow.initn_cde_btree(vars[4], k, p, x, vars[0], vars[1], vars[2],vars[3]);
double d =  sqrt(pow2(k) + pow2(p) + 2*p*k*x);
double omega_m = vars[0]/pow3(vars[4]);
double A = -3.*(1.+vars[2]);
double omegaf = pow(vars[4],A);
double omega_q= (1.-vars[0])*omegaf;
double cs2 = vars[1];
double norm = 2./pow4(dnorm_spt);
double terms[3];
if(d < 1e-6)
    return 0;
else
    switch(a) {
      // theta_m density_m density_m
        case 1:
						return    norm*(P_l(k)*P_l(p)*G1_nk*F1p_nk[0]*F2A_nk[0] +  P_l(d)*P_l(k)*G1_nk*F2D_nk[0]*F1kmp_nk[0]
                          + P_l(p)*P_l(d)*G2_nk[0]*F1p_nk[0]*F1kmp_nk[0]);
            break;
      // theta_q density_m density_m
        case 2:
            return    norm*(P_l(k)*P_l(p)*G1q_nk*F1p_nk[0]*F2A_nk[0] +  P_l(d)*P_l(k)*G1q_nk*F2D_nk[0]*F1kmp_nk[0]
                      + P_l(p)*P_l(d)*G2q_nk[0]*F1p_nk[0]*F1kmp_nk[0]);
            break;
      // theta_m \delta_q \delta_q
        case 3:
            return    norm*(P_l(k)*P_l(p)*G1_nk*F1pq_nk[0]*F2Aq_nk[0] +  P_l(d)*P_l(k)*G1_nk*F2Dq_nk[0]*F1kmpq_nk[0]
                      + P_l(p)*P_l(d)*G2_nk[0]*F1pq_nk[0]*F1kmpq_nk[0]);
            break;
       // theta_q \delta_q \delta_q
        case 4:
            return    norm*(P_l(k)*P_l(p)*G1q_nk*F1p_nk[0]*F2A_nk[0] +  P_l(d)*P_l(k)*G1q_nk*F2D_nk[0]*F1kmp_nk[0]
                      + P_l(p)*P_l(d)*G2q_nk[0]*F1p_nk[0]*F1kmp_nk[0]);
           break;
      // theta_m \delta_m \delta_q
        case 5:
            return    norm*(P_l(k)*P_l(p)*G1_nk*F1p_nk[0]*F2Aq_nk[0] +  P_l(d)*P_l(k)*G1_nk*F2D_nk[0]*F1kmpq_nk[0]
                      + P_l(p)*P_l(d)*G2_nk[0]*F1p_nk[0]*F1kmpq_nk[0]);
            break;
     // theta_q \delta_m \delta_q
        case 6:
            return    norm*(P_l(k)*P_l(p)*G1q_nk*F1p_nk[0]*F2Aq_nk[0] +  P_l(d)*P_l(k)*G1q_nk*F2D_nk[0]*F1kmpq_nk[0]
                      + P_l(p)*P_l(d)*G2q_nk[0]*F1p_nk[0]*F1kmpq_nk[0]);
           break;
     // theta_m \delta_t \delta_t [\delta_t = \omega_m*(\delta_m + \omega_q/\omega_m*(1+cs^2)*\delta_q) ]
        case 7:
        // theta_m density_m density_m
  				 terms[0]= norm*(P_l(k)*P_l(p)*G1_nk*F1p_nk[0]*F2A_nk[0] +  P_l(d)*P_l(k)*G1_nk*F2D_nk[0]*F1kmp_nk[0]
                         + P_l(p)*P_l(d)*G2_nk[0]*F1p_nk[0]*F1kmp_nk[0]);

        // theta_m \delta_q \delta_q
           terms[1]= norm*(P_l(k)*P_l(p)*G1_nk*F1pq_nk[0]*F2Aq_nk[0] +  P_l(d)*P_l(k)*G1_nk*F2Dq_nk[0]*F1kmpq_nk[0]
                     + P_l(p)*P_l(d)*G2_nk[0]*F1pq_nk[0]*F1kmpq_nk[0]);

        // theta_m \delta_m \delta_q
           terms[2]= norm*(P_l(k)*P_l(p)*G1_nk*F1p_nk[0]*F2Aq_nk[0] +  P_l(d)*P_l(k)*G1_nk*F2D_nk[0]*F1kmpq_nk[0]
                     + P_l(p)*P_l(d)*G2_nk[0]*F1p_nk[0]*F1kmpq_nk[0]);

          return pow2(omega_m)*(terms[0] + 2.*omega_q/omega_m*(1.+cs2)*terms[2] + pow2(omega_q/omega_m*(1.+cs2))*terms[1]);
           break;

           // theta_q \delta_t \delta_t [\delta_t = \omega_m*(\delta_m + \omega_q/\omega_m*(1+cs^2)*\delta_q) ]
       case 8:
          // theta_q density_m density_m
            terms[0]= norm*(P_l(k)*P_l(p)*G1q_nk*F1p_nk[0]*F2A_nk[0] +  P_l(d)*P_l(k)*G1q_nk*F2D_nk[0]*F1kmp_nk[0]
                      + P_l(p)*P_l(d)*G2q_nk[0]*F1p_nk[0]*F1kmp_nk[0]);

          // theta_q \delta_m \delta_q
            terms[1]=  norm*(P_l(k)*P_l(p)*G1q_nk*F1p_nk[0]*F2Aq_nk[0] +  P_l(d)*P_l(k)*G1q_nk*F2D_nk[0]*F1kmpq_nk[0]
                      + P_l(p)*P_l(d)*G2q_nk[0]*F1p_nk[0]*F1kmpq_nk[0]);

          // theta_q \delta_q \delta_q
            terms[2]=  norm*(P_l(k)*P_l(p)*G1q_nk*F1p_nk[0]*F2A_nk[0] +  P_l(d)*P_l(k)*G1q_nk*F2D_nk[0]*F1kmp_nk[0]
                      + P_l(p)*P_l(d)*G2q_nk[0]*F1p_nk[0]*F1kmp_nk[0]);

            return pow2(omega_m)*(terms[0] + 2.*omega_q/omega_m*(1.+cs2)*terms[2] + pow2(omega_q/omega_m*(1.+cs2))*terms[1]);
            break;
        case 9:
          // d_m d_m d_m
          return    norm*(P_l(k)*P_l(p)*F1_nk*F1p_nk[0]*F2A_nk[0] +  P_l(d)*P_l(k)*F1_nk*F2D_nk[0]*F1kmp_nk[0]
                        + P_l(p)*P_l(d)*F2_nk[0]*F1p_nk[0]*F1kmp_nk[0]);
          break;
      default:
            warning("SPT: invalid indices, a = %d\n", a);
            return 0;
    }
}

// 1-loop resummed spectrum

// Spline linear growth factors and no wiggle spectra
Spline F1_spline, G1_spline, EHUNWcde, RENWcde;
double ANWcde,sigmav;

double KMIN_cde = 1e-4;
double KMAX_cde = 50.;
void CDE::F1_init(int a, double scalef, double omega0, double par1, double par2, double par3) const{
  IOW iow;
  vector<double> kval_table, F1_table, G1_table;
  int n3 = 300;
  for(int i = 0; i<n3; i++){
  double k = KMIN_cde*exp(i*log(KMAX_cde/(KMIN_cde))/(n3-1.));
  iow.initn_lin(a, scalef, k, omega0, par1, par2, par3);
  double ling1=F1_nk;
  double ling2=G1_nk;

  kval_table.push_back(k);
  F1_table.push_back(ling1);
  G1_table.push_back(ling2);

  }
  F1_spline = LinearSpline(kval_table,F1_table);
  G1_spline = LinearSpline(kval_table,G1_table);
}

static double ANWcde_integrand(double q, double k){
return pow2(F1_spline(k))*RENWcde(k)*(1.-j0(k*q))*pow2(q);
}

static double nowig_integrand(const PowerSpectrum& P_l, double k, double q){
double lambda = 0.25*pow(k/0.05,0.04);
return P_l(q)/EHUNWcde(q)*exp(-pow2(log(k/q)/lambda)/2.)/q;
}

////////////////
// FOR RSD ONLY

static double cde_exp(const PowerSpectrum& P_l, double q){
  return  P_l(q)*pow2(F1_spline(q))/(6.*pow2(M_PI));
}

static void sigv_init(const PowerSpectrum& P_l){
    sigmav = Integrate(bind(cde_exp,cref(P_l), _1), KMIN_cde, KMAX_cde, 1e-3);
  }
//////////////

// Initialize the NW power spectrum, A^{nw,0l} and sigma_v

// run this before running other kernel initializations (initn or initn_cde)
// a chooses cde(1) or mg(0)
void CDE::initcde(int a, double scalef, double omega0, double par1, double par2, double par3) const{
NoWigglePS now(C, 0. , EisensteinHu);
IOW iow;

// linear spectrum initialization in range [1e-4,50]
F1_init(a, scalef, omega0, par1, par2, par3);

// sigma_v iinitialization
//sigv_init(cref(P_l));

// EHUNWcde initialization
int n3 = 500;
vector<double> kval_table, ehu_table, now_table;;
for(int i = 0; i<n3; i++ ){
double  k = KMIN_cde*exp(i*log(KMAX_cde/KMIN_cde)/(n3-1.));
double  ehups = now.Evaluate(k);
  kval_table.push_back(k);
  ehu_table.push_back(ehups);
    }

EHUNWcde = LinearSpline(kval_table,ehu_table);

// Renormalized linear spectrum and ANWcde initialization
const double qmin = 10.;
const double qmax = 300.;
double c[2] = {qmin,KMIN_cde};
double d[2] = {qmax,KMAX_cde};

for(int i = 0; i<n3; i++ ){
  double k = kval_table[i];

  double lambda = 0.25*pow(k/0.05,0.04);
  double nowps = EHUNWcde(k)/sqrt(2*M_PI*pow2(lambda))*Integrate<ExpSub>(bind(nowig_integrand,cref(P_l),k,_1),KMIN_cde,KMAX_cde,epsrel);
  now_table.push_back(nowps);
    }

  RENWcde = LinearSpline(kval_table,now_table);
  ANWcde = 1./((pow3(qmax)-pow3(qmin))*pow2(M_PI))*Integrate<2>(bind(ANWcde_integrand,_1,_2),c,d,1e-3,1e-3);

}

// NUMERICAL KERNEL RESUMMED 1-LOOP TERMS
//Integrating over angle - P22 numerical:
static double F2F(int y, int a,  const PowerSpectrum& P_l, double k, double r){
	double temp_ps;
	double myresult=0.;
	switch (a) {
			case 1:
			for( int i = 0; i < n1; i++ )
				{
        double d = 1+ r*r - 2*r*x128[i];
		 		if(d < 1e-5){
			 	temp_ps=0.;
		 		}
		  	else {
		  	temp_ps = RENWcde(k*sqrt(d)) * pow2(F2_nk[i*n2 + y]);
				}
        myresult += w128[i] * temp_ps;
				}
        return 2. * r * r * myresult;
        break;
			case 2:
      myresult = 0.;
			for( int i = 0; i < n1; i++ )
				{
        double d = 1+ r*r - 2*r*x128[i];
				if(d < 1e-5){
				temp_ps=0.;
				}
				else {
				temp_ps = RENWcde(k*sqrt(d)) * G2_nk[i*n2 + y]*F2_nk[i*n2 + y];
				}
        myresult += w128[i] * temp_ps;
				}
        return  2. * r * r * myresult;
        break;
			case 3:
			for( int i = 0; i < n1; i++ )
				{
				double d = 1+ r*r - 2*r*x128[i];
				if(d < 1e-5){
				temp_ps=0.;
				}
				else {
				temp_ps = RENWcde(k*sqrt(d)) * pow2(G2_nk[i*n2 + y]);
				}
        myresult += w128[i] * temp_ps;
				}
        return  2. * r * r * myresult;
        break;
			}
	}

  //Integrating over angle -P13 numerical:
  static double F3F(int y, int a,  const PowerSpectrum& P_l, double k, double r){
    double myresult = 0.;
    // Set the limits of angular integrations
  	switch (a) {
  			case 1:
  			 for( int i = 0; i < n1; i++ )
  				{
  				myresult += w128[i] * F1_nk * F3_nk[i*n2 + y];
  				}
        return   6. * r * r * myresult;
        break;
  			case 2:
  			 for( int i = 0; i < n1; i++ )
  				{
  					myresult += w128[i] * (G1_nk*F3_nk[i*n2 + y] + F1_nk*G3_nk[i*n2 + y]);
  				}
        return   3. * r * r * myresult;
        break;
  			case 3:
  				for( int i = 0; i < n1; i++ )
  				{
  					myresult += w128[i] * G1_nk * G3_nk[i*n2 + y];
  				}
          return    6. * r * r * myresult;
          break;

  		}
    }

/*Selection function for numerical kernel magnitude */

inline int num_selec_mag(double kmin, double kmax, double y){
      double YMIN =QMINp/kmax;
      double YMAX = QMAXp/kmin;
    	int mag_int;
  //  	mag_int = (y-YMIN)*(n2*1.-1)/(YMAX-YMIN); // linear
    //mag_int = sqrt((y-QMINp/kmax)*kmin/QMAXp)*(n2-1); //quadratic
    mag_int = (int)round((n2-1)*log(y/YMIN)/log(YMAX/YMIN)); //exponential
    		return mag_int;
}

/* P^{(22)} */
// a = 1 : delta delta
// a = 2 : delta theta
// a = 3 : theta theta
double CDE::P22nres(double kmin, double kmax,  int a, double k) const {
    	double KMAX = QMAXp/k;
    	double KMIN = QMINp/k;
      int y1 = num_selec_mag(kmin, kmax, KMIN);
      int y2 = num_selec_mag(kmin, kmax, KMAX);
      double y[n2];
      double integrand[n2];
      for (int i = y1; i<=y2; i++){
      y[i] = QMINp/kmax * exp(i*log(QMAXp*kmax/(QMINp*kmin))/(n2*1.-1.)); // exponential sampling
  //    y[i] = i*1./(n2-1.)*(QMAXp/kmin-QMINp/kmax)+QMINp/kmax; // linear sampling
      integrand[i] = RENWcde(k*y[i]) * F2F(i, a, cref(P_l), k, y[i]);
}
double res = 0.;
  for( int i = y1+1; i <= y2; ++ i ){
res += 0.5 * (y[i] - y[i-1])*(integrand[i] + integrand[i-1]);
}
return  k*k*k/(4*M_PI*M_PI)/pow4(dnorm_spt) * res;
}

double CDE::P13nres(double kmin, double kmax,  int a, double k) const {
  double KMAX = QMAXp/k;
  double KMIN = QMINp/k;
  int y1 = num_selec_mag(kmin, kmax, KMIN);
  int y2 = num_selec_mag(kmin, kmax, KMAX);
  double y[n2];
  double integrand[n2];
  for (int i = y1; i<=y2; i++){
  y[i] = QMINp/kmax * exp(i*log(QMAXp*kmax/(QMINp*kmin))/(n2*1.-1.));
//  y[i] = i*1./(n2-1.)*(QMAXp/kmin-QMINp/kmax)+QMINp/kmax; // linear sampling
  integrand[i] = RENWcde(k*y[i]) * F3F(i, a, cref(P_l), k, y[i]);
  }
double res = 0.;
  for( int i = y1+1; i <= y2; ++ i ){
res += 0.5 * (y[i] - y[i-1])*(integrand[i] + integrand[i-1]);
}
  return  k*k*k/(4*M_PI*M_PI)/pow4(dnorm_spt) * RENWcde(k) * res;
}

double CDE::PNWloopn(double kmin, double kmax, double k, int a ) const{
  switch(a) {
      case 1:
          return   pow2(F1_spline(k)/dnorm_spt)*RENWcde(k) + P13nres(kmin, kmax, a, k) + P22nres(kmin, kmax, a, k);
          break;
      case 2:
          return  F1_spline(k)*G1_spline(k)/pow2(dnorm_spt)*RENWcde(k) + P13nres(kmin, kmax, a, k) + P22nres(kmin, kmax, a, k);
          break;
      case 3:
          return  pow2(G1_spline(k)/dnorm_spt)*RENWcde(k) + P13nres(kmin, kmax, a, k) + P22nres(kmin, kmax, a, k);
          break;
      default:
          warning("CDE: invalid indices, a = %d\n", a);
          return 0;
  }
}

/* Resummed 1-loop spectra a la Fonseca */

double CDE::Presumn(double kmin, double kmax, double k, int a) const{
  SPT spt(C,P_l,epsrel);
  switch(a){
  case 1:
   return   PNWloopn(kmin,kmax,k,1) + exp(-0.5*pow2(k)*ANWcde/pow2(dnorm_spt))*(spt.PLOOPn(kmin,kmax,1,k)-PNWloopn(kmin,kmax,k,1) + 0.5*pow2(k)*ANWcde/pow2(dnorm_spt)*pow2(F1_spline(k)/dnorm_spt)*(P_l(k) - RENWcde(k)));
      break;
  case 2:
   return PNWloopn(kmin,kmax,k,2) + exp(-0.5*pow2(k)*ANWcde/pow2(dnorm_spt))*(spt.PLOOPn(kmin,kmax,2,k)-PNWloopn(kmin,kmax,k,2) + 0.5*pow2(k)*ANWcde/pow2(dnorm_spt)*G1_spline(k)*F1_spline(k)/pow2(dnorm_spt)*(P_l(k) - RENWcde(k)));
      break;
  case 3:
      return PNWloopn(kmin,kmax,k,3) + exp(-0.5*pow2(k)*ANWcde/pow2(dnorm_spt))*(spt.PLOOPn(kmin,kmax,3,k)-PNWloopn(kmin,kmax,k,3) + 0.5*pow2(k)*ANWcde/pow2(dnorm_spt)*pow2(G1_spline(k)/dnorm_spt)*(P_l(k) - RENWcde(k)));
      break;
  case 4:
      return pow2(F1_spline(k)/dnorm_spt)*RENWcde(k);
      //IR CHECK
  case 5:
      return P13nres(kmin, kmax, 1, k) + P22nres(kmin, kmax, 1, k) + exp(-0.5*pow2(k)*ANWcde/pow2(dnorm_spt))*(spt.P22n(kmin,kmax,1,k) + spt.P13n(kmin,kmax,1,k) -(P13nres(kmin, kmax, 1, k) + P22nres(kmin, kmax, 1, k)));
  default:
          warning("CDE: invalid indices, a = %d\n", a);
          return 0;
}
}


/*
Copyright (c) 2019 Matthew H. Reilly (kb1vc)
All rights reserved.

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.

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.
*/
#ifndef POINT_HDR_DEF
#define POINT_HDR_DEF
#include <string>
#include <regex>
#include <cmath>
#include <boost/format.hpp>

/**
 * \class GeoProf::Point
 *
 * \brief Point implements navigation operations on a geographic point
 * on the surface of the earth. The Point object includes methods to
 * calculate the distance and bearing from one point to another, or the
 * location (point) at a given distance and bearing from a starting location.
 *
 * \author $Author: kb1vc $
 *
 * \date $Date: 2005/04/14 14:16:20 $
 *
 * Contact: kb1vc@kb1vc.org
 *
 */
namespace GeoProf {
  class Point {
  public:
    /**
     * @brief Create a point object from latitude, and longitude
     *
     * @param _lat the latitude of the point -- negative values are south of the equator
     * @param _lon the longitude of the point -- negative values are west of Grenwich
     *
     * Note that bearings between the poles are unlikely to make sense.
     * In places where every direction is "south" or "north" the math
     * is often beyond sensible limits.  If your application requires
     * pole-to-pole paths, then look elsewhere, or just point south.
     */
    Point(double _lat = 0.0, double _lon = 0.0) {
      lat = _lat;
      lon = _lon;
    }

    /**
     * @brief Create a point object from another point object
     *
     * @param orig the point that we're cloning.
     */
    Point(const Point & orig) {
      lat = orig.lat;
      lon = orig.lon;
    }

    /**
     * @brief Create a point from a Maidenhead Grid specifier
     *
     * @param grid A Maidenhead Grid specifier of the form XXnnxx (capital
     * letter pair, digit pair, letter pair)
     */
    Point(const std::string & grid) {
      grid2Pt(grid);
    }

    /**
     * @brief Return degrees latitude (negative is south of the equator)
     *
     * @return latitude
     */
    double getLatitude() const { return lat; }

    /**
     * @brief Return degrees longitude (negative is west of the Grenwich)
     *
     * @return longitude
     */
    double getLongitude() const { return lon; }

    /**
     * @brief Calculate the bearing (in degrees -- 0 is north) from this point to anothe point
     * along the shorter great circle route
     *
     * @param other the other point
     * @return bearing in degrees.
     */
    double bearingTo(const Point & other) const;

    /**
     * @brief Calculate the great circle distance (in meters) from this point to another point
     *
     * @param other the other point
     * @return great circle distance in meters
     */
    double distanceTo(const Point & other) const;

    /**
     * @brief Calculate the bearing (in degrees -- 0 is north) and distance (in meters)
     * from this point to another point
     * along the shorter great circle route
     *
     * @param other the other point
     * @param bearing (output) direction to the other point along the shorter great circle path
     * @param reverse_bearing direction of travel from other point to this point (in degrees, 0 is north)
     * @param distance (output) distance in meters to the other point along the shorter great circle path
     */
    void bearingDistanceTo(const Point & other, double & bearing, double & reverse_bearing, double & distance) const;

    /**
     * @brief Return the point at which one would arrive after traveling on the
     * great circle path from this point at the specified bearing and for the specified distance.
     *
     * @param bearing direction of travel (in degrees, 0 is north)
     * @param distance of travel in meters
     * @param next (output) the point at which we'll arrive.
     */
    void stepTo(double bearing, double distance, Point & next) const;

    /**
     * @brief convert this point to its Maidenhead Grid specifier
     *
     * @param grid (output) The Maidenhead Grid string XXnnxx
     */
    void pt2Grid(std::string & grid) const;

    /**
     * @brief Set this point to the location of a Maidenhead Grid specifier
     *
     * @param grid A Maidenhead Grid specifier of the form XXnnxx (capital
     * letter pair, digit pair, letter pair)
     */
    void grid2Pt(const std::string & grid);

    /**
     * @brief Set this point from a string describing the location in
     * degrees-minutes-seconds of latitude and logitude. Specification is
     * in lat degrees/minutes/seconds and lon degrees/minutes/seconds as
     * doubles.  North/South/East/West
     *
     * @param lat_d latitude degrees
     * @param lat_m latitude minutes
     * @param lat_s latitude seconds
     * @param ns true if latitude is north of the equator
     * @param lon_d longitude degrees
     * @param lon_m longitude minutes
     * @param lon_s longitude seconds
     * @param ew true if latitude is east of Grenwich
     */
    void dms2Pt(double lat_d, double lat_m, double lat_s, char ns,
		 double lon_d, double lon_m, double lon_s, char ew);

    /**
     * @brief Correct the calculated bearing and distance by iterating
     * around the calculated values to find the minimum error.
     *
     * @param other the other point
     * @param bearing direction to the other point along the shorter great circle path
     * @param distance distance in meters to the other point along the shorter great circle path
     * @param new_bearing (output) corrected bearing
     * @param new_distance (output) corrected range
     *
     */
    void correctBearingDistanceTo(const Point & other,
				  const double bearing,
				  const double distance,
				  double & new_bearing,
				  double & new_distance) const;

    std::string toString() const {
      char ns = ' ';
      if (lat > 0.) {
	ns = 'N';
      }
      if (lat < 0.0) {
	ns = 'S';
      }

      char ew = ' ';
      if (lon > 0.0) {
	ew = 'E';
      }
      if (lon < 0.0) {
	ew = 'W';
      }

      return (boost::format("%g %c %g %c") % lat % ns % lon % ew).str();
    }
  private:
    /// Latitude South is negative, North is positive.
    double lat;

    /// Longitude West is negative, East is positive.
    double lon;

    static std::regex grid_regexp; //  ("[A-R][A-R][0-9][0-9][A-X][A-X]", std::regex_constants::icase);

    /**
     * @brief Validate a string as a Maidenhead Grid specifier
     *
     * @param grid A Maidenhead Grid specifier of the form XXnnxx (capital
     * letter pair, digit pair, letter pair)
     * @return true if this is a properly formatted grid, false otherwise.
     */
    bool checkGrid(const std::string & grid) const;

    /**
     * @brief Helper for translating char positions in a grid
     * locator into double offsets from 0 degrees.
     */
    double gridDiff(char v, char s, double mul) const;

    /**
     * @brief four quadrant arc tangent returning result in the range 0..2pi
     *
     * @param y rise
     * @param x run
     * @return arctan in range 0..2pi
     */
    double atan2Pt(double y, double x) const;

    /**
     * @brief translate an angle into the range 0..2pi
     *
     * @param ang angle in radians
     * @return angle in range 0..2pi
     */
    double inSpan(double ang) const;

    // recursive helper to correctBearingDistanceTo
    bool recCorrectBearingDistanceTo(const Point & other,
				     const double bearing,
				     const double distance,
				     const double b_span,
				     const double d_span,
				     double & new_bearing,
				     double & new_distance
				     ) const;

    // Useful constants
    static const double clarke_66_al;    /*Clarke 1866 ellipsoid*/
    static const double clarke_66_bl;
    static const double rad_per_deg;
  };

}
#endif


// Copyright 2018 The Simons Foundation, Inc. - 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.
//
// by G. Mazzola, May-Aug 2018

#include <Eigen/Dense>
#include <iostream>
#include <vector>
#include "Utils/all_utils.hpp"
#include "Utils/lookup.hpp"

#ifndef NETKET_JASTROW_HPP
#define NETKET_JASTROW_HPP

namespace netket {

/** Jastrow machine class.
 *
 */
template <typename T>
class Jastrow : public AbstractMachine<T> {
  using VectorType = typename AbstractMachine<T>::VectorType;
  using MatrixType = typename AbstractMachine<T>::MatrixType;
  using VectorRefType = typename AbstractMachine<T>::VectorRefType;
  using VectorConstRefType = typename AbstractMachine<T>::VectorConstRefType;
  using VisibleConstType = typename AbstractMachine<T>::VisibleConstType;

  const AbstractHilbert &hilbert_;

  // number of visible units
  int nv_;

  // number of parameters
  int npar_;

  // weights
  MatrixType W_;

  // buffers
  VectorType thetas_;
  VectorType thetasnew_;

 public:
  using StateType = typename AbstractMachine<T>::StateType;
  using LookupType = typename AbstractMachine<T>::LookupType;

  // constructor
  explicit Jastrow(const AbstractHilbert &hilbert)
      : hilbert_(hilbert), nv_(hilbert.Size()) {
    Init();
  }

  void Init() {
    if (nv_ < 2) {
      throw InvalidInputError(
          "Cannot construct Jastrow states with less than two visible units");
    }

    W_.resize(nv_, nv_);
    W_.setZero();

    npar_ = (nv_ * (nv_ - 1)) / 2;

    thetas_.resize(nv_);
    thetasnew_.resize(nv_);

    InfoMessage() << "Jastrow WF Initizialized with nvisible = " << nv_
                  << " and nparams = " << npar_ << std::endl;
  }

  int Nvisible() const override { return nv_; }

  int Npar() const override { return npar_; }

  void InitRandomPars(int seed, double sigma) override {
    VectorType par(npar_);

    netket::RandomGaussian(par, seed, sigma);

    SetParameters(par);
  }

  VectorType GetParameters() override {
    VectorType pars(npar_);

    int k = 0;

    for (int i = 0; i < nv_; i++) {
      for (int j = i + 1; j < nv_; j++) {
        pars(k) = W_(i, j);
        k++;
      }
    }

    return pars;
  }

  void SetParameters(VectorConstRefType pars) override {
    int k = 0;

    for (int i = 0; i < nv_; i++) {
      W_(i, i) = T(0.);
      for (int j = i + 1; j < nv_; j++) {
        W_(i, j) = pars(k);
        W_(j, i) = W_(i, j);  // create the lower triangle
        k++;
      }
    }
  }

  void InitLookup(VisibleConstType v, LookupType &lt) override {
    if (lt.VectorSize() == 0) {
      lt.AddVector(v.size());
    }
    if (lt.V(0).size() != v.size()) {
      lt.V(0).resize(v.size());
    }

    lt.V(0) = (W_.transpose() * v);  // does not matter the transpose W is symm
  }

  // same as for the RBM
  void UpdateLookup(VisibleConstType v, const std::vector<int> &tochange,
                    const std::vector<double> &newconf,
                    LookupType &lt) override {
    if (tochange.size() != 0) {
      for (std::size_t s = 0; s < tochange.size(); s++) {
        const int sf = tochange[s];
        lt.V(0) += W_.row(sf) * (newconf[s] - v(sf));
      }
    }
  }

  T LogVal(VisibleConstType v) override { return 0.5 * v.dot(W_ * v); }

  // Value of the logarithm of the wave-function
  // using pre-computed look-up tables for efficiency
  T LogVal(VisibleConstType v, const LookupType &lt) override {
    return 0.5 * v.dot(lt.V(0));
  }

  // Difference between logarithms of values, when one or more visible variables
  // are being flipped
  VectorType LogValDiff(
      VisibleConstType v, const std::vector<std::vector<int>> &tochange,
      const std::vector<std::vector<double>> &newconf) override {
    const std::size_t nconn = tochange.size();
    VectorType logvaldiffs = VectorType::Zero(nconn);

    thetas_ = (W_.transpose() * v);
    T logtsum = 0.5 * v.dot(thetas_);

    for (std::size_t k = 0; k < nconn; k++) {
      if (tochange[k].size() != 0) {
        thetasnew_ = thetas_;
        Eigen::VectorXd vnew(v);

        for (std::size_t s = 0; s < tochange[k].size(); s++) {
          const int sf = tochange[k][s];

          thetasnew_ += W_.row(sf) * (newconf[k][s] - v(sf));
          vnew(sf) = newconf[k][s];
        }

        logvaldiffs(k) = 0.5 * vnew.dot(thetasnew_) - logtsum;
      }
    }
    return logvaldiffs;
  }

  T LogValDiff(VisibleConstType v, const std::vector<int> &tochange,
               const std::vector<double> &newconf,
               const LookupType &lt) override {
    T logvaldiff = 0.;

    if (tochange.size() != 0) {
      T logtsum = 0.5 * v.dot(lt.V(0));
      thetasnew_ = lt.V(0);
      Eigen::VectorXd vnew(v);

      for (std::size_t s = 0; s < tochange.size(); s++) {
        const int sf = tochange[s];

        thetasnew_ += W_.row(sf) * (newconf[s] - v(sf));
        vnew(sf) = newconf[s];
      }

      logvaldiff = 0.5 * vnew.dot(thetasnew_) - logtsum;
    }

    return logvaldiff;
  }

  VectorType DerLog(VisibleConstType v) override {
    VectorType der(npar_);

    int k = 0;

    for (int i = 0; i < nv_; i++) {
      for (int j = i + 1; j < nv_; j++) {
        der(k) = v(i) * v(j);
        k++;
      }
    }

    return der;
  }

  const AbstractHilbert &GetHilbert() const noexcept override {
    return hilbert_;
  }

  void to_json(json &j) const override {
    j["Name"] = "Jastrow";
    j["Nvisible"] = nv_;
    j["W"] = W_;
  }

  void from_json(const json &pars) override {
    if (pars.at("Name") != "Jastrow") {
      throw InvalidInputError(
          "Error while constructing Jastrow from Json input");
    }

    if (FieldExists(pars, "Nvisible")) {
      nv_ = pars["Nvisible"];
    }
    if (nv_ != hilbert_.Size()) {
      throw InvalidInputError(
          "Number of visible units is incompatible with given "
          "Hilbert space");
    }

    Init();

    if (FieldExists(pars, "W")) {
      W_ = pars["W"];
    }
  }
};
}  // namespace netket

#endif


// Copyright (c) 2020 Marcus Valtonen Örnhag
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

#include <Eigen/Dense>
#include <math.h>  // copysign

#include "radial.hpp"

namespace DronePoseLib {
    Eigen::MatrixXd radialdistort(const Eigen::MatrixXd& x, double kappa) {
        // We expect inhomogenous input data
        assert(x.rows() == 2);

        Eigen::ArrayXd ru2 = x.colwise().squaredNorm();
        Eigen::ArrayXd ru = ru2.sqrt();

        Eigen::ArrayXd rd;
        // Compute distorted radius
        if (kappa == 0) {
            rd = ru;
        } else {
            rd = 0.5 / kappa / ru - copysign(1.0, kappa) * (0.25 / std::pow(kappa, 2) / ru2 - 1.0 / kappa).sqrt();
        }

        // compute distorted coordinates
        Eigen::MatrixXd y(2, x.cols());
        y = (rd / ru).replicate(1, x.rows()).transpose() * x.array();

        // TODO(marcusvaltonen): Avoid divion by zero (centre coordinate - usually synthethic images)
        // y(isnan(y(:))) = 0;

        return y;
    }

    Eigen::MatrixXd radialundistort(const Eigen::MatrixXd& x, double kappa) {
        // We expect inhomogenous input data
        assert(x.rows() == 2);

        Eigen::VectorXd rd2 = x.colwise().squaredNorm();

        // Compute undistorted coordinates
        Eigen::MatrixXd y(3, x.cols());
        y.topRows(2) = x;
        y.bottomRows(1) = (Eigen::VectorXd::Ones(x.cols()) + kappa * rd2).transpose();

        return y.colwise().hnormalized();
    }
}  // namespace DronePoseLib


/**
 * @file 	main.cpp
 * @author 	Fabian Wegscheider
 * @date 	Jul 19, 2017
 */

#include <iostream>
#include <boost/program_options.hpp>
#include <boost/timer/timer.hpp>
#include <omp.h>
#include "GraphParser.h"
#include "PrimeNumbers.h"
#include "SteinerTreeHeuristic.h"

using std::string;
using std::cout;
using std::cerr;
using std::endl;
using namespace boost;
namespace po = boost::program_options;

using weight_type = double;

using CSR_Graph = compressed_sparse_row_graph<directedS, no_property,
		property<edge_weight_t, weight_type>>;

/**
 * The main function which reads in a graph from a .gph file, considers all
 * vertices with prime indices as terminals and calculates a steiner tree
 * using an improved shortest-path-heuristic based on Dijkstra
 * @param numargs number of inputs on command line
 * @param args array of inputs on command line
 * @return whether the program operated successfully
 */
int main(int numargs, char* args[]) {

	timer::cpu_timer overall_timer;

	// default values for unspecified options
	int N_THREADS = 1;
	int N_TERMINALS = 100;

	/* parsing command line options */
	string input;
	bool print_tree_selected;

	try {
		po::options_description desc("Allowed options");
		desc.add_options()
				("help,h", "produce help message")
				("showtree,s", "print tree")
				("threads,t", po::value<int>(&N_THREADS)->default_value(1),
						"set number of threads")
				("terminals,n", po::value<int>(&N_TERMINALS)->default_value(100),
						"set number of starting terminals")
				("input-file", po::value<string>(), "input file");

		po::positional_options_description p;
		p.add("input-file", 1);
		p.add("terminals", -1);
		po::variables_map vm;
		po::store(po::command_line_parser(numargs, args).
				options(desc).positional(p).run(), vm);
		po::notify(vm);

		if (vm.count("help")) {
			cout << desc << "\n";
			exit(EXIT_SUCCESS);
		}

		if (vm.count("input-file")) {
			input = vm["input-file"].as< string >();
		} else {
			cerr << "Usage: ./ex10 <file.gph> (<number of starting terminals>) "
					<< "(other options)" << endl;
			cerr << "run with -h to see information about options" << endl;
			exit(EXIT_FAILURE);
		}

		print_tree_selected = vm.count("showtree");

	} catch (...) {
		cerr << "Usage: ./ex10 <file.gph> (<number of starting terminals>) "
				<< "(other options)" << endl;
		cerr << "run with -h to see information about options" << endl;
		exit(EXIT_FAILURE);
	}
	/* end of parsing command line options */

	/* parsing .gph file */
	int num_vertices;
	int num_edges;

	GraphParser parser(input);
	if (!parser.opened_successfully) {
		cerr << "file could not be read" << endl;
		exit(EXIT_FAILURE);
	}

	// first line is read to get number of vertices and edges
	if (!parser.read_first_line(num_vertices, num_edges)) {
		cerr << "error while reading file, not the right format" << endl;
		exit(EXIT_FAILURE);
	}

	// we use arrays for edges and weights so that we can use a boost graph
	Edge* edges = new Edge[num_edges];
	weight_type* weights = new double[num_edges];

	// rest of the file is read and parsed to a graph
	if (!parser.read_edge_data(edges, weights)) {
		cerr << "error while reading file, not the right format" << endl;
		exit(EXIT_FAILURE);
	}
	/* end of file parsing */

	/* initialization and computing */

	// here the graph is constructed using a CSR graph from boost
	const CSR_Graph g(boost::edges_are_unsorted, &edges[0], &edges[0] + num_edges,
			weights, num_vertices);

	// all primes in {2,...,num_vertices} are considered to be terminals
	vector<int> terminals = PrimeNumbers::find_primes(num_vertices);

	timer::cpu_timer algo_timer;

	const int iterations = std::min((int) terminals.size(), N_TERMINALS);

	// The tree is implicitly stored by remembering the predecessor of each
	// vertex, a -1 means that the vertex is not in the tree
	vector<int> best_tree(num_vertices);

	weight_type min_value = std::numeric_limits<weight_type>::infinity();
	int min_root = -1;

	// heuristic is called for the specified number of terminals and the best
	// solution is kept. It is done in parallel if more than one thread is chosen
	#pragma omp parallel for schedule(dynamic) num_threads(N_THREADS)
	for (int i = 0; i < iterations; ++i) {

		vector<int> tree(num_vertices, -1);
		weight_type objective_value = SteinerTreeHeuristic::compute_steiner_tree(g,
			num_vertices, terminals[i], tree, terminals);

		#pragma omp critical
		if (objective_value < min_value) {

			min_value = objective_value;
			min_root = i;
			best_tree = tree;
		}
	}

	string algo_time = algo_timer.format(3, "%w");
	/* end of computation */

	assert(SteinerTreeHeuristic::test_tree(g, num_vertices, terminals[min_root],
			best_tree, terminals));

	// output is printed, with tree if option was chosen
	cout << "TLEN: " << min_value << endl;
	if (print_tree_selected) {
		string edge_string = SteinerTreeHeuristic::print_tree(g,
				num_vertices, terminals[min_root], best_tree, terminals);
		cout << "TREE: " << edge_string << endl;
	}
	cout << "TIME: " << overall_timer.format(3, "%t") << endl;
	cout << "WALL: " << algo_time << endl;

	delete[] edges;
	delete[] weights;

	exit(EXIT_SUCCESS);
}


#ifndef STAN_MATH_TORSTEN_LINODE_HPP
#define STAN_MATH_TORSTEN_LINODE_HPP

#include <Eigen/Dense>
#include <stan/math/torsten/ev_manager.hpp>
#include <boost/math/tools/promotion.hpp>
#include <stan/math/torsten/ev_solver.hpp>
#include <stan/math/torsten/to_array_2d.hpp>
#include <stan/math/torsten/pmx_linode_model.hpp>
#include <stan/math/torsten/pmx_check.hpp>
#include <stan/math/prim/err/check_square.hpp>
#include <vector>

namespace torsten {

  namespace {
    template<typename T>
    using torsten_matrix_dyn_t = Eigen::Matrix<T,-1,-1>;
  }

/**
 * Computes the predicted amounts in each compartment at each event
 * for a compartment model, described by a linear system of ordinary
 * differential equations. Uses the stan::math::matrix_exp
 * function.
 *
 * @tparam T0 type of scalar for time of events.
 * @tparam T1 type of scalar for amount at each event.
 * @tparam T2 type of scalar for rate at each event.
 * @tparam T3 type of scalar for inter-dose inteveral at each event.
 * @tparam T4 type of scalar for matrix describing linear ODE system.
 * @tparam T5 type of scalars for bio-variability parameters.
 * @tparam T6 type of scalars for tlag parameters
 * @param[in] time times of events
 * @param[in] amt amount at each event
 * @param[in] rate rate at each event
 * @param[in] ii inter-dose interval at each event
 * @param[in] evid event identity:
 *                    (0) observation
 *                    (1) dosing
 *                    (2) other
 *                    (3) reset
 *                    (4) reset AND dosing
 * @param[in] cmt compartment number at each event
 * @param[in] addl additional dosing at each event
 * @param[in] ss steady state approximation at each event (0: no, 1: yes)
 * between time-points
 * @param[in] system square matrix describing the linear system of ODEs
 * @param[in] bio-variability at each event
 * @param[in] lag times at each event
 * @return a matrix with predicted amount in each compartment
 * at each event.
 */
template <typename T0, typename T1, typename T2, typename T3,
          typename T4, typename T5, typename T6>
stan::matrix_return_t<T0, T1, T2, T3, T4, T5, T6>
pmx_solve_linode(const std::vector<T0>& time,
                 const std::vector<T1>& amt,
                 const std::vector<T2>& rate,
                 const std::vector<T3>& ii,
                 const std::vector<int>& evid,
                 const std::vector<int>& cmt,
                 const std::vector<int>& addl,
                 const std::vector<int>& ss,
                 const std::vector< Eigen::Matrix<T4, -1, -1> >& system,
                 const std::vector<std::vector<T5> >& biovar,
                 const std::vector<std::vector<T6> >& tlag) {
  using std::vector;
  using Eigen::Dynamic;
  using Eigen::Matrix;
  using boost::math::tools::promote_args;

  static const char* function("pmx_solve_linode");
  for (size_t i = 0; i < system.size(); i++)
    stan::math::check_square(function, "system matrix", system[i]);
  int nCmt = system[0].cols();

  using ER = NONMENEventsRecord<T0, T1, T2, T3>;
  using EM = EventsManager<ER, NonEventParameters<T0, T4, torsten_matrix_dyn_t, std::tuple<T5, T6> >>;
  const ER events_rec(nCmt, time, amt, rate, ii, evid, cmt, addl, ss);

  Matrix<typename EM::T_scalar, Dynamic, Dynamic> pred =
    Matrix<typename EM::T_scalar, Dynamic, Dynamic>::Zero(events_rec.num_event_times(), EM::nCmt(events_rec));

  using model_type = torsten::PMXLinODEModel<typename EM::T_par>;
  EventSolver<model_type, EM> pr;
  pr.pred(0, events_rec, pred, dsolve::PMXAnalyiticalIntegrator(), system, biovar, tlag, nCmt);
  return pred;
}

/**
 * Overload function to allow user to pass a matrix for
 * system.
 */
template <typename T0, typename T1, typename T2, typename T3,
          typename T4, typename T_biovar, typename T_tlag>
stan::matrix_return_t<T0, T1, T2, T3, T4, T_biovar, T_tlag>
pmx_solve_linode(const std::vector<T0>& time,
                 const std::vector<T1>& amt,
                 const std::vector<T2>& rate,
                 const std::vector<T3>& ii,
                 const std::vector<int>& evid,
                 const std::vector<int>& cmt,
                 const std::vector<int>& addl,
                 const std::vector<int>& ss,
                 const Eigen::Matrix<T4, -1, -1>& system,
                 const std::vector<T_biovar>& biovar,
                 const std::vector<T_tlag>& tlag) {
  std::vector<Eigen::Matrix<T4, -1, -1> > system_{system};
  auto biovar_ = torsten::to_array_2d(biovar);
  auto tlag_ = torsten::to_array_2d(tlag);

  return pmx_solve_linode(time, amt, rate, ii, evid, cmt, addl, ss,
                          system_, biovar_, tlag_);
}

/**
 * Overload function to allow user to pass a matrix for
 * system.
 */
template <typename T0, typename T1, typename T2, typename T3,
          typename T4, typename T_biovar, typename T_tlag,
          typename = require_any_not_std_vector_t<T_biovar, T_tlag> >
stan::matrix_return_t<T0, T1, T2, T3, T4, T_biovar, T_tlag>
pmx_solve_linode(const std::vector<T0>& time,
                 const std::vector<T1>& amt,
                 const std::vector<T2>& rate,
                 const std::vector<T3>& ii,
                 const std::vector<int>& evid,
                 const std::vector<int>& cmt,
                 const std::vector<int>& addl,
                 const std::vector<int>& ss,
                 const std::vector< Eigen::Matrix<T4, -1, -1> >& system,
                 const std::vector<T_biovar>& biovar,
                 const std::vector<T_tlag>& tlag) {
  auto biovar_ = torsten::to_array_2d(biovar);
  auto tlag_ = torsten::to_array_2d(tlag);

  return pmx_solve_linode(time, amt, rate, ii, evid, cmt, addl, ss,
                          system, biovar_, tlag_);
}

  // old version by using transpose
template <typename T0, typename T1, typename T2, typename T3,
          typename T4, typename T5, typename T6>
stan::matrix_return_t<T0, T1, T2, T3, T4, T5, T6>
linOdeModel(const std::vector<T0>& time,
            const std::vector<T1>& amt,
            const std::vector<T2>& rate,
            const std::vector<T3>& ii,
            const std::vector<int>& evid,
            const std::vector<int>& cmt,
            const std::vector<int>& addl,
            const std::vector<int>& ss,
            const std::vector< Eigen::Matrix<T4, Eigen::Dynamic, Eigen::Dynamic> >& system,
            const std::vector<std::vector<T5> >& biovar,
            const std::vector<std::vector<T6> >& tlag) {
  auto x = pmx_solve_linode(time, amt, rate, ii, evid, cmt, addl, ss, system, biovar, tlag);
  return x.transpose();
}

template <typename T0, typename T1, typename T2, typename T3,
          typename T4, typename T_biovar, typename T_tlag>
stan::matrix_return_t<T0, T1, T2, T3, T4, T_biovar, T_tlag>
linOdeModel(const std::vector<T0>& time,
                 const std::vector<T1>& amt,
                 const std::vector<T2>& rate,
                 const std::vector<T3>& ii,
                 const std::vector<int>& evid,
                 const std::vector<int>& cmt,
                 const std::vector<int>& addl,
                 const std::vector<int>& ss,
                 const Eigen::Matrix<T4, -1, -1>& system,
                 const std::vector<T_biovar>& biovar,
                 const std::vector<T_tlag>& tlag) {
  auto x = pmx_solve_linode(time, amt, rate, ii, evid, cmt, addl, ss, system, biovar, tlag);
  return x.transpose();
}

  template <typename T0, typename T1, typename T2, typename T3,
            typename T4, typename T_biovar, typename T_tlag,
            typename = require_any_not_std_vector_t<T_biovar, T_tlag> >
  stan::matrix_return_t<T0, T1, T2, T3, T4, T_biovar, T_tlag>
  linOdeModel(const std::vector<T0>& time,
              const std::vector<T1>& amt,
              const std::vector<T2>& rate,
              const std::vector<T3>& ii,
              const std::vector<int>& evid,
              const std::vector<int>& cmt,
              const std::vector<int>& addl,
              const std::vector<int>& ss,
              const std::vector< Eigen::Matrix<T4, -1, -1> >& system,
              const std::vector<T_biovar>& biovar,
              const std::vector<T_tlag>& tlag) {
    auto x = pmx_solve_linode(time, amt, rate, ii, evid, cmt, addl, ss,
                            system, biovar, tlag);
    return x.transpose();
  }

}
#endif


#ifndef qlex_svi_interpolation_hpp
#define qlex_svi_interpolation_hpp

#include <ql/utilities/null.hpp>
#include <ql/utilities/dataformatters.hpp>
#include <ql/math/interpolation.hpp>
#include <ql/math/optimization/method.hpp>
#include <ql/math/optimization/simplex.hpp>
#include <ql/math/optimization/levenbergmarquardt.hpp>
#include <ql/pricingengines/blackformula.hpp>
#include <ql/math/optimization/projectedcostfunction.hpp>
#include <ql/math/optimization/constraint.hpp>
#include <ql/math/randomnumbers/haltonrsg.hpp>

#include <termstructures/volatility/svi.hpp>

#include <boost/make_shared.hpp>
#include <boost/assign/list_of.hpp>

using namespace QuantLib;

namespace QLExtension {

namespace detail {
	class SVICoeffHolder {
	public:
		SVICoeffHolder(const Time t, const Real &forward, std::vector<Real> params)
			: t_(t), forward_(forward), params_(params),
			weights_(std::vector<Real>()), error_(Null<Real>()),
			maxError_(Null<Real>()), SVIEndCriteria_(EndCriteria::None) {
			QL_REQUIRE(t > 0.0, "expiry time must be positive: " << t
				<< " not allowed");
		}
		virtual ~SVICoeffHolder() {}

		/*! Expiry, Forward */
		Real t_;
		const Real forward_;
		/*! Parameters */
		std::vector<Real> params_;
		std::vector<Real> weights_;
		/*! Interpolation results */
		Real error_, maxError_;
		EndCriteria::Type SVIEndCriteria_;
	};

	template <class I1, class I2>
	class SVIInterpolationImpl : public Interpolation::templateImpl<I1, I2>,
								 public SVICoeffHolder
	{
	public:
		SVIInterpolationImpl(
			const I1 &xBegin, const I1 &xEnd, const I2 &yBegin, Time t,
			const Real &forward, std::vector<Real> params, bool vegaWeighted,
			const boost::shared_ptr<EndCriteria> &endCriteria,
			const boost::shared_ptr<OptimizationMethod> &optMethod,
			const Real errorAccept, const bool useMaxError, const Size maxGuesses)
			: Interpolation::templateImpl<I1, I2>(xBegin, xEnd, yBegin),
			SVICoeffHolder(t, forward, params),
			endCriteria_(endCriteria), optMethod_(optMethod),
			errorAccept_(errorAccept), useMaxError_(useMaxError),
			maxGuesses_(maxGuesses),  vegaWeighted_(vegaWeighted) {
				// if no optimization method or endCriteria is provided, we provide one
				if (!optMethod_)
					optMethod_ = boost::shared_ptr<OptimizationMethod>(
					new LevenbergMarquardt(1e-8, 1e-8, 1e-8));
				// optMethod_ = boost::shared_ptr<OptimizationMethod>(new
				//    Simplex(0.01));
				if (!endCriteria_) {
					endCriteria_ = boost::shared_ptr<EndCriteria>(
						new EndCriteria(60000, 100, 1e-8, 1e-8, 1e-8));
				}
				// equal weighted
				this->weights_ = std::vector<Real>(xEnd - xBegin, 1.0 / (xEnd - xBegin));
				completed_ = false;
			}

		// Optimization
		void update() {
			// we must update weights if it is vegaWeighted
			if (vegaWeighted_) {
				// std::vector<Real>::const_iterator x = this->xBegin_;
				// std::vector<Real>::const_iterator y = this->yBegin_;
				// std::vector<Real>::iterator w = weights_.begin();
				this->weights_.clear();
				Real weightsSum = 0.0;
				for (Size i = 0; i<Size(this->xEnd_ - this->xBegin_); ++i) {
					Real stdDev = std::sqrt((this->yBegin_[i]) * (this->yBegin_[i]) * this->t_);
					this->weights_.push_back(
						blackFormulaStdDevDerivative(this->xBegin_[i], forward_, stdDev));
					weightsSum += this->weights_.back();
				}
				// weight normalization
				std::vector<Real>::iterator w = this->weights_.begin();
				for (; w != this->weights_.end(); ++w)
					*w /= weightsSum;
			}

			try {
				SVIError costFunction(this);
				Constraint constraint = SVIConstraint();

				Array guess(params_.size());
				for (Size i = 0; i < params_.size(); i++)
				{
					guess[i] = this->params_[i];
				}
				Size iterations = 0;
				Real tmpInterpolationError;
				EndCriteria::Type tmpEndCriteria;
				do {
					Problem problem(costFunction, constraint, guess);
					tmpEndCriteria = optMethod_->minimize(problem, *endCriteria_);
					guess = problem.currentValue();
					tmpInterpolationError = useMaxError_ ? interpolationMaxError()
						: interpolationError();

					switch (tmpEndCriteria) {
					case EndCriteria::None:
					case EndCriteria::MaxIterations:
					case EndCriteria::Unknown:
						completed_ = false;
						break;
					default:
						completed_ = true;
					}
				} while (++iterations < maxGuesses_ &&
					tmpInterpolationError > errorAccept_);

				for (Size i = 0; i < guess.size(); ++i)
					this->params_[i] = guess[i];

				this->error_ = interpolationError();
				this->maxError_ = interpolationMaxError();
			}
			catch (const std::exception& e)
			{
				QL_FAIL("error in SVI calibration: " << e.what());
			}
		}

		Real value(Real x) const {
			QL_REQUIRE(x > 0.0, "strike must be positive: " << io::rate(x)
				<< " not allowed");
			return sviVolatility(x, forward_, t_, params_[0], params_[1], params_[2], params_[3], params_[4]);
		}

		Real primitive(Real) const { QL_FAIL("SVI primitive not implemented"); }
		Real derivative(Real) const { QL_FAIL("SVI derivative not implemented"); }
		Real secondDerivative(Real) const {
			QL_FAIL("SVI secondDerivative not implemented");
		}

		// calculate total squared weighted difference (L2 norm)
		Real interpolationSquaredError() const {
			Real error, totalError = 0.0;
			//std::vector<Real>::const_iterator x = this->xBegin_;
			//std::vector<Real>::const_iterator y = this->yBegin_;
			// std::vector<Real>::const_iterator w = this->weights_.begin();
			for (Size i = 0; i<Size(this->xEnd_ - this->xBegin_); ++i) {
				error = (value(this->xBegin_[i]) - this->yBegin_[i]);
				totalError += error * error * (this->weights_[i]);
			}
			return totalError;
		}

		// calculate weighted differences
		Disposable<Array> interpolationErrors(const Array &) const {
			Array results(Size(this->xEnd_ - this->xBegin_));
			//std::vector<Real>::const_iterator x = this->xBegin_;
			//std::vector<Real>::const_iterator y = this->yBegin_;
			//std::vector<Real>::const_iterator w = this->weights_.begin();
			Array::iterator r = results.begin();
			for (Size i = 0; i < Size(this->xEnd_ - this->xBegin_); i++)
			{
				*r = (value(this->xBegin_[i]) - this->yBegin_[i]) * std::sqrt(this->weights_[i]);
			}
			return results;
		}

		Real interpolationError() const {
			Size n = this->xEnd_ - this->xBegin_;
			Real squaredError = interpolationSquaredError();
			return std::sqrt(n * squaredError / (n - 1));
		}

		Real interpolationMaxError() const {
			Real error, maxError = QL_MIN_REAL;
			I1 i = this->xBegin_;
			I2 j = this->yBegin_;
			for (; i != this->xEnd_; ++i, ++j) {
				error = std::fabs(value(*i) - *j);
				maxError = std::max(maxError, error);
			}
			return maxError;
		}
	public:
		boost::shared_ptr<EndCriteria> endCriteria_;
		boost::shared_ptr<OptimizationMethod> optMethod_;
		const Real errorAccept_;
		const bool useMaxError_;
		const Size maxGuesses_;
		bool vegaWeighted_;
		bool completed_;
		NoConstraint constraint_;

	private:
		class SVIError : public CostFunction {
		public:
			SVIError(SVIInterpolationImpl *svi) : svi_(svi) {}

			// calculate total squared weighted difference
			Real value(const Array &x) const {
				for (Size i = 0; i < svi_->params_.size(); i++)
				{
					svi_->params_[i] = x[i];
				}

				return svi_->interpolationSquaredError();
			}

			Disposable<Array> values(const Array &x) const {
				for (Size i = 0; i < svi_->params_.size(); i++)
				{
					svi_->params_[i] = x[i];
				}

				return svi_->interpolationErrors(x);
			}

		private:
			SVIInterpolationImpl *svi_;
		};

		class SVIConstraint : public Constraint {
		private:
			class Impl : public Constraint::Impl {
			public:
				bool test(const Array& params) const {
					if (params[1] < 0)
						return false;
					if (std::fabs(params[2]) >= 1)
						return false;
					if (params[4] <= 0)
						return false;
					//if (params[0] + params[1] * params[4] * std::sqrt(1-params[2]*params[2]) < 0)
					//	return false;

					return true;
				}
			};
		public:
			SVIConstraint()
				: Constraint(boost::shared_ptr<Constraint::Impl>(
				new SVIConstraint::Impl)) {}
		};
	};
}	// end of namespace detail

	//! %SVI smile interpolation between discrete volatility points.
	class SVIInterpolation : public Interpolation {
	public:
		template <class I1, class I2>
		SVIInterpolation(const I1 &xBegin,  // x = strikes
			const I1 &xEnd,
			const I2 &yBegin,  // y = volatilities
			Time t,            // option expiry
			const Real& forward,
			Real a,
			Real b,
			Real rho,
			Real m,
			Real sigma,
			bool vegaWeighted = true,
			const boost::shared_ptr<EndCriteria>& endCriteria
			= boost::shared_ptr<EndCriteria>(),
			const boost::shared_ptr<OptimizationMethod>& optMethod
			= boost::shared_ptr<OptimizationMethod>(),
			const Real errorAccept = 0.0020,
			const bool useMaxError = false,
			const Size maxGuesses = 1) {

			impl_ = boost::shared_ptr<Interpolation::Impl>(
				new detail::SVIInterpolationImpl<I1, I2>(
				xBegin, xEnd, yBegin, t, forward,
				boost::assign::list_of(a)(b)(rho)(m)(sigma),
				vegaWeighted, endCriteria, optMethod, errorAccept, useMaxError,
				maxGuesses));
			coeffs_ = boost::dynamic_pointer_cast<
				detail::SVICoeffHolder >(impl_);
		}
		Real expiry()  const { return coeffs_->t_; }
		Real forward() const { return coeffs_->forward_; }
		Real a()   const { return coeffs_->params_[0]; }
		Real b()    const { return coeffs_->params_[1]; }
		Real rho()      const { return coeffs_->params_[2]; }
		Real m()     const { return coeffs_->params_[3]; }
		Real sigma()     const { return coeffs_->params_[4]; }
		Real rmsError() const { return coeffs_->error_; }
		Real maxError() const { return coeffs_->maxError_; }
		const std::vector<Real>& interpolationWeights() const {
			return coeffs_->weights_;
		}
		EndCriteria::Type endCriteria() { return coeffs_->SVIEndCriteria_; }

	private:
		boost::shared_ptr<detail::SVICoeffHolder> coeffs_;
	};
}	// end of namespace QLExtension

#endif