//------------------------------------------//
//  This file is a modified version of      //
//  basics.cpp, which was distributed as    //
//  part of MuJoCo,  Written by Emo Todorov //
//  Copyright (C) 2017 Roboti LLC           //
//  Modifications by Atabak Dehban          //
//------------------------------------------//

#include <iostream>

#include "mujoco.h"
#include "cstdio"
#include "cstdlib"
#include "cstring"
#include "glfw3.h"

// for sleep timers
#include <chrono>
#include <thread>

// Eigen, used by drake
#include <Eigen/Core>
#include <unsupported/Eigen/MatrixFunctions>
#include "drake/systems/controllers/linear_quadratic_regulator.h"

// MuJoCo data structures
mjModel* m = NULL;                  // MuJoCo model
mjData* d = NULL;                   // MuJoCo data
mjvCamera cam;                      // abstract camera
mjvOption opt;                      // visualization options
mjvScene scn;                       // abstract scene
mjrContext con;                     // custom GPU context

// mouse interaction
bool button_left = false;
bool button_middle = false;
bool button_right =  false;
double lastx = 0;
double lasty = 0;

// holders of one step history of time and position to calculate dertivatives
mjtNum position_history = 0;
mjtNum previous_time = 0;

// controller related variables
float_t ctrl_update_freq = 1000;
mjtNum last_update = 0.0;
mjtNum ctrl;

// keyboard callback
void keyboard(GLFWwindow* window, int key, int scancode, int act, int mods)
{
    // backspace: reset simulation
    if( act==GLFW_PRESS && key==GLFW_KEY_BACKSPACE )
    {
        mj_resetData(m, d);
        mj_forward(m, d);
    }
}

// mouse button callback
void mouse_button(GLFWwindow* window, int button, int act, int mods)
{
    // update button state
    button_left =   (glfwGetMouseButton(window, GLFW_MOUSE_BUTTON_LEFT)==GLFW_PRESS);
    button_middle = (glfwGetMouseButton(window, GLFW_MOUSE_BUTTON_MIDDLE)==GLFW_PRESS);
    button_right =  (glfwGetMouseButton(window, GLFW_MOUSE_BUTTON_RIGHT)==GLFW_PRESS);

    // update mouse position
    glfwGetCursorPos(window, &lastx, &lasty);
}

// mouse move callback
void mouse_move(GLFWwindow* window, double xpos, double ypos)
{
    // no buttons down: nothing to do
    if( !button_left && !button_middle && !button_right )
        return;

    // compute mouse displacement, save
    double dx = xpos - lastx;
    double dy = ypos - lasty;
    lastx = xpos;
    lasty = ypos;

    // get current window size
    int width, height;
    glfwGetWindowSize(window, &width, &height);

    // get shift key state
    bool mod_shift = (glfwGetKey(window, GLFW_KEY_LEFT_SHIFT)==GLFW_PRESS ||
                      glfwGetKey(window, GLFW_KEY_RIGHT_SHIFT)==GLFW_PRESS);

    // determine action based on mouse button
    mjtMouse action;
    if( button_right )
        action = mod_shift ? mjMOUSE_MOVE_H : mjMOUSE_MOVE_V;
    else if( button_left )
        action = mod_shift ? mjMOUSE_ROTATE_H : mjMOUSE_ROTATE_V;
    else
        action = mjMOUSE_ZOOM;

    // move camera
    mjv_moveCamera(m, action, dx/height, dy/height, &scn, &cam);
}

// scroll callback
void scroll(GLFWwindow* window, double xoffset, double yoffset)
{
    // emulate vertical mouse motion = 5% of window height
    mjv_moveCamera(m, mjMOUSE_ZOOM, 0, -0.05*yoffset, &scn, &cam);
}

// control loop callback
void mycontroller(const mjModel* m, mjData* d)
{
    // printouts for debugging purposes
    mjtObj geom = mjOBJ_GEOM;
    mjtObj body = mjOBJ_BODY;
    mjtObj dof = mjOBJ_JOINT;
    int ball_id = mj_name2id(m, body, "ballbody");
    int rod_id = mj_name2id(m, geom, "rod");
    int pivot_id = mj_name2id(m, dof, "pivot");
    std::cout << "ball body id: " << ball_id << std::endl;
    std::cout << "rod geom id: " << rod_id << std::endl;
    std::cout << "pivot joint id: " << pivot_id << std::endl;
    std::cout << "ball mass: " << m->body_mass[ball_id] << std::endl;
    std::cout << "length of the rod: " << m->geom_size[rod_id*3+1] << std::endl;
    std::cout << "friction coefficient: " << m->dof_damping[pivot_id] << std::endl;
    std::cout << "gravity: " << m->opt.gravity[2] << std::endl;
    std::cout << "Sensor output: " << d->sensordata[0] << std::endl;

    // System dynamics matrices: https://github.com/RobotLocomotion/drake/blob/master/systems/controllers/zmp_planner.h#L306
    Eigen::Matrix<mjtNum , 2, 2> A_;
    Eigen::Matrix<mjtNum , 2, 1> B_;
    Eigen::Matrix<mjtNum , 2, 2> Q_ = Eigen::Matrix<mjtNum , 2, 2>::Identity();
    Eigen::Matrix<mjtNum , 1, 1> R_ = Eigen::Matrix<mjtNum , 1, 1>::Identity();
    Eigen::Matrix<mjtNum , 2, 1> N = Eigen::Matrix<mjtNum , 2, 1>::Zero();

    // Eqs comming from https://youtu.be/XA6B1IxALhk?t=30m3s
    A_(0,0) = 0.0;
    A_(0,1) = 1.0;
    A_(1,0) = - m->opt.gravity[2]*mju_cos(d->sensordata[0])/(m->geom_size[rod_id*3+1]*2);
    A_(1,1) = - m->dof_damping[pivot_id]/(m->body_mass[ball_id]*mju_pow(m->geom_size[rod_id*3+1]*2,2.0));
    B_(0,0) = 0.0;
    B_(1,0) = 1.0;
    Q_ *= 10.0;

    drake::systems::controllers::LinearQuadraticRegulatorResult lqr_result =
            drake::systems::controllers::LinearQuadraticRegulator(A_, B_, Q_, R_, N);
    Eigen::Matrix<mjtNum, 2, 2> S1_ = lqr_result.S;
    Eigen::Matrix<mjtNum, 1, 2> K_= -lqr_result.K;

    // controller with true values, but it is cheating.
//    ctrl = 3.5*(-d->qvel[0]-10.0*d->qpos[0]);

    // controller with sensor readings
    if (previous_time == 0)
    {
        previous_time = d->time;
        return;
    }
    if (d->time - last_update > 1.0/ctrl_update_freq)
    {
        mjtNum vel = (d->sensordata[0] - position_history)/(d->time-previous_time);
        ctrl = K_(0,0)*d->sensordata[0] + K_(0,1)*vel;
        last_update = d->time;
        position_history = d->sensordata[0];
        previous_time = d->time;
    }
    d->ctrl[0] = ctrl;

    std::cout << "torque effort: " << ctrl << std::endl;
}

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

    // activate software
    mj_activate("../../../mjkey.txt");

    // load and compile model
    char error[1000] = "Could not load binary model";

    // check command-line arguments
    if( argc<2 )
        m = mj_loadXML("../../../models/invertedPendulum.xml", 0, error, 1000);

    else
        if( strlen(argv[1])>4 && !strcmp(argv[1]+strlen(argv[1])-4, ".mjb") )
            m = mj_loadModel(argv[1], 0);
        else
            m = mj_loadXML(argv[1], 0, error, 1000);
    if( !m )
        mju_error_s("Load model error: %s", error);

    // make data
    d = mj_makeData(m);

    // init GLFW
    if( !glfwInit() )
        mju_error("Could not initialize GLFW");

    // create window, make OpenGL context current, request v-sync
    GLFWwindow* window = glfwCreateWindow(1200, 900, "Demo", NULL, NULL);
    glfwMakeContextCurrent(window);
    glfwSwapInterval(1);

    // initialize visualization data structures
    mjv_defaultCamera(&cam);
    mjv_defaultOption(&opt);
    mjv_defaultScene(&scn);
    mjr_defaultContext(&con);
    mjv_makeScene(m, &scn, 2000);                // space for 2000 objects
    mjr_makeContext(m, &con, mjFONTSCALE_150);   // model-specific context

    // install GLFW mouse and keyboard callbacks
    glfwSetKeyCallback(window, keyboard);
    glfwSetCursorPosCallback(window, mouse_move);
    glfwSetMouseButtonCallback(window, mouse_button);
    glfwSetScrollCallback(window, scroll);

    // install control callback
    mjcb_control = mycontroller;

    // initial position
    d->qpos[0] = 1.57;

    // run main loop, target real-time simulation and 60 fps rendering
    mjtNum timezero = d->time;
    double_t update_rate = 0.01;

    // making sure the first time step updates the ctrl previous_time
    last_update = timezero-1.0/ctrl_update_freq;

    // use the first while condition if you want to simulate for a period.
//    while( !glfwWindowShouldClose(window) and d->time-timezero < 1.5)
    while( !glfwWindowShouldClose(window))
    {
        // advance interactive simulation for 1/60 sec
        //  Assuming MuJoCo can simulate faster than real-time, which it usually can,
        //  this loop will finish on time for the next frame to be rendered at 60 fps.
        //  Otherwise add a cpu timer and exit this loop when it is time to render.
        mjtNum simstart = d->time;
        while( d->time - simstart < 1.0/60.0 )
            mj_step(m, d);

        // 15 ms is a little smaller than 60 Hz.
        std::this_thread::sleep_for(std::chrono::milliseconds(15));
       // get framebuffer viewport
        mjrRect viewport = {0, 0, 0, 0};
        glfwGetFramebufferSize(window, &viewport.width, &viewport.height);

          // update scene and render
        mjv_updateScene(m, d, &opt, NULL, &cam, mjCAT_ALL, &scn);
        mjr_render(viewport, &scn, &con);

        // swap OpenGL buffers (blocking call due to v-sync)
        glfwSwapBuffers(window);

        // process pending GUI events, call GLFW callbacks
        glfwPollEvents();

    }

    // free visualization storage
    mjv_freeScene(&scn);
    mjr_freeContext(&con);

    // free MuJoCo model and data, deactivate
    mj_deleteData(d);
    mj_deleteModel(m);
    mj_deactivate();

    // terminate GLFW (crashes with Linux NVidia drivers)
#if defined(__APPLE__) || defined(_WIN32)
    glfwTerminate();
#endif
    return 1;
}


#define NDEBUG
#define _ITERATOR_DEBUG_LEVEL 0

#include <array>
#include <boost/lexical_cast.hpp>
#ifdef DOMAGICK
#include <Magick++.h>
#endif
#include <boost/fusion/adapted/array.hpp>
#include <boost/fusion/adapted/std_array.hpp>
#include <boost/fusion/adapted/std_tuple.hpp>
#include <boost/fusion/adapted/std_pair.hpp>
#include <boost/spirit/home/x3.hpp>
#include <boost/format.hpp>
#include <boost/spirit/include/support_istream_iterator.hpp>

#include <iostream>
#include <fstream>
#include <vector>
#include <map>
#include <string>

using namespace boost::spirit;
using namespace std;

using box = std::array<int, 4>;

box intersect(const box& a, const box& b)
{
	box res = { max(a[0], b[0]), max(a[1], b[1]), min(a[2], b[2]), min(a[3], b[3]) };
	res[0] = min(res[0], res[2]);
	res[1] = min(res[1], res[3]);

	return res;
}

box boxunion(const box& a, const box& b)
{
	box res = { min(a[0], b[0]), min(a[1], b[1]), max(a[2], b[2]), max(a[3], b[3]) };

	return res;
}

int area(const box& a)
{
	return (a[2] - a[0]) * (a[3] - a[1]);
}

// According to https://wandbox.org/permlink/poeusnilIOwmiEBo
namespace boost {
	namespace spirit {
		namespace x3 {
			namespace traits {
				// It can't be specialized because X3 implements is_container as a template aliases,
				// thus we need QUITE TERRIBLE DIRTY hack for fixed length container.
				//template <> struct is_container<Vertex const> : mpl::false_ { };
				//template <> struct is_container<Vertex> : mpl::false_ { };
				namespace detail {
					template <> struct has_type_value_type<box> : mpl::false_ { };
				}
			}
		}
	}
}

template<class T> auto constexpr bracketize(T what)
{
	return '[' > what > ']';
}

// Vector of four ints
auto boxrule = bracketize(x3::int_ > ',' > x3::int_ > ',' > x3::int_ > ',' > x3::int_);
// Comma-separated list of such vectors
auto innerboxlist = (boxrule % ',');
auto fullboxlist = x3::lit("array") > '(' > bracketize(innerboxlist ) > ',' > "dtype" > '=' > (x3::lit("int32") | x3::lit("int64")) > ')';
// The whole first section of the data file
auto boxes = bracketize(*(fullboxlist));

// Array of lists of doubles
auto scorelist = x3::lit("array") > '(' > bracketize(x3::double_ % ',') > ')';
auto scores = bracketize(*(scorelist));

template<class RuleType, class AttrType> void parseToEndWithError(istream& file, const RuleType& rule, AttrType& target)
{
	auto parseriter = boost::spirit::istream_iterator(file);
	boost::spirit::istream_iterator end;

	bool res = phrase_parse(parseriter, end, rule, x3::space - x3::eol, target);

	if (!res)
	{
		std::string val;
		file >> val;
		throw logic_error("Parsing failed. " + (std::string) __func__ + " " + val);
	}
}

struct wordmapper
{
private:
	map<string, int> mapping;
	map<int, string> inversemapping;
public:
	int getMapping(const string& word)
	{
		auto i = mapping.find(word);
		if (i != mapping.end())
		{
			return i->second;
		}

		int index = mapping.size();
		mapping[word] = index;
		inversemapping[index] = word;

		return index;
	}

	string getWord(int mapping)
	{
		return inversemapping[mapping];
	}

	int size()
	{
		return mapping.size();
	}
} wordmap;

// Coming in C++ 17
template <class T>
constexpr std::add_const_t<T>& as_const_cheat(const T& t) noexcept
{
	return t;
}

auto word_ = x3::lexeme[+(x3::char_ - x3::space)];
auto linefile = *word_ % x3::eol;
constexpr int statelimit = 4000;

using stateVec = array<double, statelimit>;

struct hmmtype
{
	vector<vector<box>> ourBoxes;
	vector<vector<double>> scoreVals;
	vector<bool> lineEnd;

	stateVec fwbw[2][statelimit] = { 0 };

	void prepareLineEnd(const vector<vector<int>>& introws)
	{
		lineEnd.reserve(scoreVals.size());
		for (const auto& r : introws)
		{
			for (int i : r)
			{
				lineEnd.push_back(false);
			}
			*(lineEnd.end() - 1) = true;
		}
	}

	void emit(stateVec& state, int pos)
	{
		for (int i = 0; i < scoreVals[pos].size(); i++)
		{
			state[i] *= max(1e-50, scoreVals[pos][i]);
		}
	}

	static double transProb(const box& a, const box& b, bool linebreak)
	{
		bool ok = false;
		int height = max(a[3] - a[1], b[3] - b[1]);
		if (linebreak)
		{
			if (b[3] > a[1]) ok = true;
		}
		else
		{
			if (b[0] > a[0] && b[1] > a[1] - height && b[1] < a[1] + height) ok = true;
		}

		return (ok ? 1 : 1e-2) * (1 - 0.99 * area(intersect(a, b)) / max(min(area(a), area(b)) * 1., 1e-9)) * (linebreak ? 1. : 0.01 + 0.99 * (area(a) + area(b) - area(intersect(a,b))) / max(area(boxunion(a, b)), 1));
	}

	template<int dir> void transition(const stateVec& fState, stateVec& tState, int fromPos, int toPos)
	{
		static_assert(dir == -1 || dir == 1);

		for (int i = 0; i < scoreVals[toPos].size(); i++)
		{
			tState[i] = 0;
		}

		for (int i = 0; i < scoreVals[fromPos].size(); i++)
		{
			for (int j = 0; j < scoreVals[toPos].size(); j++)
			{
				box* fromBox = &ourBoxes[fromPos][i];
				box* toBox = &ourBoxes[toPos][j];
				int fIndex = fromPos;

				if (dir == -1)
				{
					fIndex = toPos;
					swap(fromBox, toBox);
				}

				tState[j] += fState[i] * transProb(*fromBox, *toBox, lineEnd[fIndex]);
			}
		}
	}

	void normalize(stateVec& state, int pos)
	{
		double sum = 0;
		for (int i = 0; i < scoreVals[pos].size(); i++)
		{
			sum += state[i];
		}
		if (isnan(sum)) cerr << "Normalization error at " << pos << "\n";
		if (sum == 0) cerr << "Zero at " << pos << "\n";
		//cout << "Sum at " << pos << ":" << sum << "\n";

		if (sum < 1e-10)
		{
			for (int i = 0; i < scoreVals[pos].size(); i++)
			{
				state[i] *= 1e150;
			}
		}
	}

	void computeFB()
	{
		// FW
		for (int i = 0; i < scoreVals[0].size(); i++)
		{
			fwbw[0][0][i] = 1;
		}

		for (int i = 0; i < scoreVals.size(); i++)
		{
			emit(fwbw[0][i], i);
			if (i != scoreVals.size() - 1)
			{
				transition<1>(fwbw[0][i], fwbw[0][i + 1], i, i + 1);
			}

			normalize(fwbw[0][i + 1], i);
		}

		// BW
		for (int i = 0; i < scoreVals[scoreVals.size() - 1].size(); i++)
		{
			fwbw[1][scoreVals.size() - 1][i] = 1;
		}

		for (int i = scoreVals.size() - 1; i != 0; i--)
		{
			stateVec copy = fwbw[1][i];
			emit(copy, i);
			transition<-1>(copy, fwbw[1][i - 1], i, i - 1);

			normalize(fwbw[1][i - 1], i - 1);
		}
	}

	void fakeFB()
	{
		// FW
		for (int j = 0; j < scoreVals.size(); j++)
		{
			for (int i = 0; i < scoreVals[j].size(); i++)
			{
				fwbw[0][j][i] = 1;
			}

			emit(fwbw[0][j], j);
		}

		// BW
		for (int j = 0; j < scoreVals.size(); j++)
		{
			for (int i = 0; i < scoreVals[j].size(); i++)
			{
				fwbw[1][j][i] = 1;
			}
		}
	}

	stateVec getProbs(int pos) const
	{
		stateVec toret;
		double sum = 0;
		for (int i = 0; i < scoreVals[pos].size(); i++)
		{
			toret[i] = fwbw[0][pos][i] * fwbw[1][pos][i];
			sum += toret[i];
		}

		sum = 1 / sum;
		for (int i = 0; i < scoreVals[pos].size(); i++)
		{
			toret[i] *= sum;
		}

		return toret;
	}

	void softmax(double factor)
	{
		for (auto& scoreList : scoreVals)
		{
			double maxVal = 0;
			for (double& score : scoreList)
			{
				score *= factor;
				maxVal = max(score, maxVal);
			}

			double sum = 0;
			for (double score : scoreList)
			{
				sum += exp(score * score);
				//sum += exp(score - maxVal);
			}

			for (double& score : scoreList)
			{
				score = exp(score * score) / sum;
				//score = exp(score - maxVal) / sum;
			}
		}
	}
} hmm;

template<bool doimg> void writeWithSeparator(const string& separator
#ifdef DOMAGICK
	, Magick::Image img, string path
#endif
)
{
#ifdef DOMAGICK
	std::list<Magick::Drawable> drawList;
	drawList.push_back(Magick::DrawableFillOpacity(0));
	drawList.push_back(Magick::DrawableStrokeWidth(5));
	drawList.push_back(Magick::DrawableStrokeColor("white"));
	array<string, 3> colors = { "red", "green", "blue" };

	if constexpr (doimg)
	{
		img.strokeWidth(0.5);
		//drawList.push_back(Magick::DrawableStrokeOpacity(0.25));
	}
#endif

	for (int i = 0; i < hmm.scoreVals.size(); i++)
	{
		stateVec states = hmm.getProbs(i);
		int maxindex = 0;
		for (int j = 0; j < hmm.scoreVals[i].size(); j++)
		{
			if (states[j] > states[maxindex]) maxindex = j;
		}
		cout << maxindex << ":" << states[maxindex];
		/*for (int j = 0; j < 4; j++)
		{
		cout << ":" << hmm.ourBoxes[i][maxindex][j];
		}*/
		//img.fillColor(Magick::Color::Color(0, 0, 0, 65535 - states[maxindex] * 65535));
		const box& b = hmm.ourBoxes[i][maxindex];

#ifdef DOMAGICK
		if constexpr (doimg)
		{
			drawList.push_back(Magick::DrawableRectangle(b[0], b[1], b[2], b[3]));
			img.draw(drawList);
			drawList.pop_back();

			img.strokeColor(colors[i % 3]);
			img.draw(Magick::DrawableText(b[0] + 20, b[1] + 20, boost::lexical_cast<string>(i) + ":" + boost::str(boost::format("%.2f") % states[maxindex])));
		}
#endif
		cout << " ";
		if (hmm.lineEnd[i]) cout << separator << "\n";
	}
#ifdef DOMAGICK
	if constexpr (doimg)
	{
		img.write(path);
	}
#endif
}

void writeDiffs(const vector<int>& words)
{
	vector<stateVec> wordScoresNew;
	vector<stateVec> wordScoresOld;
	wordScoresNew.resize(wordmap.size());
	wordScoresOld.resize(wordmap.size());
	vector<int> lens;
	lens.resize(wordmap.size());
	fill(lens.begin(), lens.end(), 0);

	for (int i = 0; i < hmm.scoreVals.size(); i++)
	{
		if (!lens[words[i]])
		{
			lens[words[i]] = hmm.scoreVals[i].size();
			for (int j = 0; j < hmm.scoreVals[i].size(); j++)
			{
				wordScoresNew[words[i]][j] = 0;
				wordScoresOld[words[i]][j] = 0;
			}
		}

		stateVec states = hmm.getProbs(i);
		for (int j = 0; j < hmm.scoreVals[i].size(); j++)
		{
			wordScoresNew[words[i]][j] += states[j];
			wordScoresOld[words[i]][j] += hmm.scoreVals[i][j];
		}
	}

	cout << "{";
	bool first = true;
	for (int i = 0; i < wordScoresNew.size(); i++)
	{
		if (!lens[i]) continue;
		if (!first)
		{
			cout << ",\n";
		}

		first = false;
		cout << "\"" << wordmap.getWord(i) << "\" : ";
		cout << "[";
		for (int j = 0; j < lens[i]; j++)
		{
			// Richer debug:
			//cout << wordScoresNew[i][j] << ":" << wordScoresOld[i][j] << ":" << log(sqrt(wordScoresNew[i][j] / wordScoresOld[i][j])) << (j != lens[i] - 1 ? ", " : "]");
			cout << log(sqrt(wordScoresNew[i][j] / wordScoresOld[i][j])) << (j != lens[i] - 1 ? ", " : "]");
		}
	}
	cout << "}";
}

int main(int argc, char** argv)
{
#ifdef DOMAGICK
	Magick::InitializeMagick(0);
#endif
	if (argc < 3)
	{
		return -1;
	}
	ifstream file(argv[1]);

	parseToEndWithError(file, boxes > scores, as_const_cheat(std::forward_as_tuple(hmm.ourBoxes, hmm.scoreVals)));
	cerr << "Read " << hmm.ourBoxes.size() << " box lists and " << hmm.scoreVals.size() << " score lists." << "\n";

	for (int i = 0; i < hmm.ourBoxes.size(); i++)
	{
		if (hmm.ourBoxes[i].size() > statelimit)
		{
			hmm.ourBoxes[i].resize(statelimit);
			hmm.scoreVals[i].resize(statelimit);
			fprintf(stderr, "Capping state count at word %d to limit %d.\n", i, statelimit);
		}
	}

	// Do soft-max with amplification 1
	hmm.softmax(0.5 * sqrt(108));

	file = ifstream(argv[2]);
	vector<vector<int>> introws;
	vector<int> words;
	vector<vector<string>> rows;
	file >> noskipws;
	parseToEndWithError(file, linefile, rows);

	transform(rows.begin(), rows.end(), back_inserter(introws),
		[](vector<string>& row)
	{
		vector<int> introw;
		transform(row.begin(), row.end(), back_inserter(introw),
			[](const string& word)
		{
			return wordmap.getMapping(word);
		});

		return introw;
	});
	for (auto row : introws)
	{
		for (int w : row)
		{
			words.push_back(w);
		}
	}

	if (words.size() != hmm.ourBoxes.size())
	{
		cerr << "Number of words in transcription is " << words.size() << ", but number of box lists is " << hmm.ourBoxes.size() << ". This does not make sense. Terminating.";

		return -1;
	}

	hmm.prepareLineEnd(introws);

#ifdef DOMAGICK
	Magick::Image origImg(argv[3]);
#endif

	hmm.computeFB();
	writeWithSeparator<false>("//");// , origImg, string(argv[3]) + ".N.png");

	// Sanity check, basically compute probs WITHOUT the HMM
	//hmm.fakeFB();
	//writeWithSeparator("//", origImg, string(argv[3]) + ".naive.png");

	writeDiffs(words);
}


//
//=======================================================================
// Copyright (c) 2004 Kristopher Beevers
//
// Distributed under the Boost Software License, Version 1.0. (See
// accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)
//=======================================================================
//

#ifndef BOOST_GRAPH_ASTAR_SEARCH_HPP
#define BOOST_GRAPH_ASTAR_SEARCH_HPP

#include <functional>
#include <vector>
#include <boost/limits.hpp>
#include <boost/graph/named_function_params.hpp>
#include <boost/graph/relax.hpp>
#include <boost/graph/exception.hpp>
#include <boost/graph/breadth_first_search.hpp>
#include <boost/graph/detail/d_ary_heap.hpp>
#include <boost/property_map/property_map.hpp>
#include <boost/property_map/vector_property_map.hpp>

namespace boost {

  template <class Heuristic, class Graph>
  struct AStarHeuristicConcept {
    void constraints()
    {
      function_requires< CopyConstructibleConcept<Heuristic> >();
      h(u);
    }
    Heuristic h;
    typename graph_traits<Graph>::vertex_descriptor u;
  };

  template <class Graph, class CostType>
  class astar_heuristic : public std::unary_function<
    typename graph_traits<Graph>::vertex_descriptor, CostType>
  {
  public:
    typedef typename graph_traits<Graph>::vertex_descriptor Vertex;
    astar_heuristic() {}
    CostType operator()(Vertex u) { return static_cast<CostType>(0); }
  };

  template <class Visitor, class Graph>
  struct AStarVisitorConcept {
    void constraints()
    {
      function_requires< CopyConstructibleConcept<Visitor> >();
      vis.initialize_vertex(u, g);
      vis.discover_vertex(u, g);
      vis.examine_vertex(u, g);
      vis.examine_edge(e, g);
      vis.edge_relaxed(e, g);
      vis.edge_not_relaxed(e, g);
      vis.black_target(e, g);
      vis.finish_vertex(u, g);
    }
    Visitor vis;
    Graph g;
    typename graph_traits<Graph>::vertex_descriptor u;
    typename graph_traits<Graph>::edge_descriptor e;
  };

  template <class Visitors = null_visitor>
  class astar_visitor : public bfs_visitor<Visitors> {
  public:
    astar_visitor() {}
    astar_visitor(Visitors vis)
      : bfs_visitor<Visitors>(vis) {}

    template <class Edge, class Graph>
    void edge_relaxed(Edge e, Graph& g) {
      invoke_visitors(this->m_vis, e, g, on_edge_relaxed());
    }
    template <class Edge, class Graph>
    void edge_not_relaxed(Edge e, Graph& g) {
      invoke_visitors(this->m_vis, e, g, on_edge_not_relaxed());
    }
  private:
    template <class Edge, class Graph>
    void tree_edge(Edge e, Graph& g) {}
    template <class Edge, class Graph>
    void non_tree_edge(Edge e, Graph& g) {}
  };
  template <class Visitors>
  astar_visitor<Visitors>
  make_astar_visitor(Visitors vis) {
    return astar_visitor<Visitors>(vis);
  }
  typedef astar_visitor<> default_astar_visitor;

  namespace detail {

    template <class AStarHeuristic, class UniformCostVisitor,
              class UpdatableQueue, class PredecessorMap,
              class CostMap, class DistanceMap, class WeightMap,
              class ColorMap, class BinaryFunction,
              class BinaryPredicate>
    struct astar_bfs_visitor
    {

      typedef typename property_traits<CostMap>::value_type C;
      typedef typename property_traits<ColorMap>::value_type ColorValue;
      typedef color_traits<ColorValue> Color;
      typedef typename property_traits<DistanceMap>::value_type distance_type;

      astar_bfs_visitor(AStarHeuristic h, UniformCostVisitor vis,
                        UpdatableQueue& Q, PredecessorMap p,
                        CostMap c, DistanceMap d, WeightMap w,
                        ColorMap col, BinaryFunction combine,
                        BinaryPredicate compare, C zero)
        : m_h(h), m_vis(vis), m_Q(Q), m_predecessor(p), m_cost(c),
          m_distance(d), m_weight(w), m_color(col),
          m_combine(combine), m_compare(compare), m_zero(zero) {}

      template <class Vertex, class Graph>
      void initialize_vertex(Vertex u, Graph& g) {
        m_vis.initialize_vertex(u, g);
      }
      template <class Vertex, class Graph>
      void discover_vertex(Vertex u, Graph& g) {
        m_vis.discover_vertex(u, g);
      }
      template <class Vertex, class Graph>
      void examine_vertex(Vertex u, Graph& g) {
        m_vis.examine_vertex(u, g);
      }
      template <class Vertex, class Graph>
      void finish_vertex(Vertex u, Graph& g) {
        m_vis.finish_vertex(u, g);
      }
      template <class Edge, class Graph>
      void examine_edge(Edge e, Graph& g) {
        if (m_compare(get(m_weight, e), m_zero))
          throw negative_edge();
        m_vis.examine_edge(e, g);
      }
      template <class Edge, class Graph>
      void non_tree_edge(Edge, Graph&) {}

      template <class Edge, class Graph>
      void tree_edge(Edge e, Graph& g) {
        m_decreased = relax(e, g, m_weight, m_predecessor, m_distance,
                            m_combine, m_compare);

        if(m_decreased) {
          m_vis.edge_relaxed(e, g);
          put(m_cost, target(e, g),
              m_combine(get(m_distance, target(e, g)),
                        m_h(target(e, g))));
        } else
          m_vis.edge_not_relaxed(e, g);
      }

      template <class Edge, class Graph>
      void gray_target(Edge e, Graph& g) {
        m_decreased = relax(e, g, m_weight, m_predecessor, m_distance,
                            m_combine, m_compare);

        if(m_decreased) {
          put(m_cost, target(e, g),
              m_combine(get(m_distance, target(e, g)),
                        m_h(target(e, g))));
          m_Q.update(target(e, g));
          m_vis.edge_relaxed(e, g);
        } else
          m_vis.edge_not_relaxed(e, g);
      }

      template <class Edge, class Graph>
      void black_target(Edge e, Graph& g) {
        m_decreased = relax(e, g, m_weight, m_predecessor, m_distance,
                            m_combine, m_compare);

        if(m_decreased) {
          m_vis.edge_relaxed(e, g);
          put(m_cost, target(e, g),
              m_combine(get(m_distance, target(e, g)),
                        m_h(target(e, g))));
          m_Q.push(target(e, g));
          put(m_color, target(e, g), Color::gray());
          m_vis.black_target(e, g);
        } else
          m_vis.edge_not_relaxed(e, g);
      }

      AStarHeuristic m_h;
      UniformCostVisitor m_vis;
      UpdatableQueue& m_Q;
      PredecessorMap m_predecessor;
      CostMap m_cost;
      DistanceMap m_distance;
      WeightMap m_weight;
      ColorMap m_color;
      BinaryFunction m_combine;
      BinaryPredicate m_compare;
      bool m_decreased;
      C m_zero;

    };

  } // namespace detail

  template <typename VertexListGraph, typename AStarHeuristic,
            typename AStarVisitor, typename PredecessorMap,
            typename CostMap, typename DistanceMap,
            typename WeightMap, typename ColorMap,
            typename VertexIndexMap,
            typename CompareFunction, typename CombineFunction,
            typename CostInf, typename CostZero>
  inline void
  astar_search_no_init
    (VertexListGraph &g,
     typename graph_traits<VertexListGraph>::vertex_descriptor s,
     AStarHeuristic h, AStarVisitor vis,
     PredecessorMap predecessor, CostMap cost,
     DistanceMap distance, WeightMap weight,
     ColorMap color, VertexIndexMap /*index_map*/,
     CompareFunction compare, CombineFunction combine,
     CostInf /*inf*/, CostZero zero)
  {
    typedef typename graph_traits<VertexListGraph>::vertex_descriptor
      Vertex;
    typedef boost::vector_property_map<std::size_t> IndexInHeapMap;
    IndexInHeapMap index_in_heap;
    typedef d_ary_heap_indirect<Vertex, 4, IndexInHeapMap, CostMap, CompareFunction>
      MutableQueue;
    MutableQueue Q(cost, index_in_heap, compare);

    detail::astar_bfs_visitor<AStarHeuristic, AStarVisitor,
        MutableQueue, PredecessorMap, CostMap, DistanceMap,
        WeightMap, ColorMap, CombineFunction, CompareFunction>
      bfs_vis(h, vis, Q, predecessor, cost, distance, weight,
              color, combine, compare, zero);

    breadth_first_visit(g, s, Q, bfs_vis, color);
  }

  // Non-named parameter interface
  template <typename VertexListGraph, typename AStarHeuristic,
            typename AStarVisitor, typename PredecessorMap,
            typename CostMap, typename DistanceMap,
            typename WeightMap, typename VertexIndexMap,
            typename ColorMap,
            typename CompareFunction, typename CombineFunction,
            typename CostInf, typename CostZero>
  inline void
  astar_search
    (VertexListGraph &g,
     typename graph_traits<VertexListGraph>::vertex_descriptor s,
     AStarHeuristic h, AStarVisitor vis,
     PredecessorMap predecessor, CostMap cost,
     DistanceMap distance, WeightMap weight,
     VertexIndexMap index_map, ColorMap color,
     CompareFunction compare, CombineFunction combine,
     CostInf inf, CostZero zero)
  {

    typedef typename property_traits<ColorMap>::value_type ColorValue;
    typedef color_traits<ColorValue> Color;
    typename graph_traits<VertexListGraph>::vertex_iterator ui, ui_end;
    for (tie(ui, ui_end) = vertices(g); ui != ui_end; ++ui) {
      put(color, *ui, Color::white());
      put(distance, *ui, inf);
      put(cost, *ui, inf);
      put(predecessor, *ui, *ui);
      vis.initialize_vertex(*ui, g);
    }
    put(distance, s, zero);
    put(cost, s, h(s));

    astar_search_no_init
      (g, s, h, vis, predecessor, cost, distance, weight,
       color, index_map, compare, combine, inf, zero);

  }

  namespace detail {
    template <class VertexListGraph, class AStarHeuristic,
              class CostMap, class DistanceMap, class WeightMap,
              class IndexMap, class ColorMap, class Params>
    inline void
    astar_dispatch2
      (VertexListGraph& g,
       typename graph_traits<VertexListGraph>::vertex_descriptor s,
       AStarHeuristic h, CostMap cost, DistanceMap distance,
       WeightMap weight, IndexMap index_map, ColorMap color,
       const Params& params)
    {
      dummy_property_map p_map;
      typedef typename property_traits<CostMap>::value_type C;
      astar_search
        (g, s, h,
         choose_param(get_param(params, graph_visitor),
                      make_astar_visitor(null_visitor())),
         choose_param(get_param(params, vertex_predecessor), p_map),
         cost, distance, weight, index_map, color,
         choose_param(get_param(params, distance_compare_t()),
                      std::less<C>()),
         choose_param(get_param(params, distance_combine_t()),
                      closed_plus<C>()),
         choose_param(get_param(params, distance_inf_t()),
                      std::numeric_limits<C>::max BOOST_PREVENT_MACRO_SUBSTITUTION ()),
         choose_param(get_param(params, distance_zero_t()),
                      C()));
    }

    template <class VertexListGraph, class AStarHeuristic,
              class CostMap, class DistanceMap, class WeightMap,
              class IndexMap, class ColorMap, class Params>
    inline void
    astar_dispatch1
      (VertexListGraph& g,
       typename graph_traits<VertexListGraph>::vertex_descriptor s,
       AStarHeuristic h, CostMap cost, DistanceMap distance,
       WeightMap weight, IndexMap index_map, ColorMap color,
       const Params& params)
    {
      typedef typename property_traits<WeightMap>::value_type D;
      std::vector<D> distance_map;
      std::vector<D> cost_map;
      std::vector<default_color_type> color_map;

      detail::astar_dispatch2
        (g, s, h,
         choose_param(cost, vector_property_map<D, IndexMap>(index_map)),
         choose_param(distance, vector_property_map<D, IndexMap>(index_map)),
         weight, index_map,
         choose_param(color, vector_property_map<default_color_type, IndexMap>(index_map)),
         params);
    }
  } // namespace detail

  // Named parameter interface
  template <typename VertexListGraph,
            typename AStarHeuristic,
            typename P, typename T, typename R>
  void
  astar_search
    (VertexListGraph &g,
     typename graph_traits<VertexListGraph>::vertex_descriptor s,
     AStarHeuristic h, const bgl_named_params<P, T, R>& params)
  {

    detail::astar_dispatch1
      (g, s, h,
       get_param(params, vertex_rank),
       get_param(params, vertex_distance),
       choose_const_pmap(get_param(params, edge_weight), g, edge_weight),
       choose_const_pmap(get_param(params, vertex_index), g, vertex_index),
       get_param(params, vertex_color),
       params);

  }

} // namespace boost

#endif // BOOST_GRAPH_ASTAR_SEARCH_HPP


#include <utility/Configuration_Chain.hpp>
#include <utility/Configurations.hpp>
#include <data/Spin_System.hpp>
#include <engine/Vectormath.hpp>
#include <engine/Manifoldmath.hpp>

#include <Eigen/Dense>

#include <random>
#include <iostream>
#include <string>
#include <vector>

namespace Utility
{
	namespace Configuration_Chain
	{
		void Add_Noise_Temperature(std::shared_ptr<Data::Spin_System_Chain> c, int idx_1, int idx_2, scalar temperature)
		{
			for (int img = idx_1 + 1; img <= idx_2 - 1; ++img)
			{
				Configurations::Add_Noise_Temperature(*c->images[img], temperature, img);
			}
		}

		void Homogeneous_Rotation(std::shared_ptr<Data::Spin_System_Chain> c, int idx_1, int idx_2)
		{
			int nos = c->images[0]->nos;
			int noi = idx_2 - idx_1 + 1;

			scalar angle, rot_angle;
			Vector3 axis, rot_axis, a, b, temp;

			for (int i = 0; i < nos; ++i)
			{
				a = (*c->images[idx_1]->spins)[i];
				b = (*c->images[idx_2]->spins)[i];

				rot_angle = Engine::Manifoldmath::dist_greatcircle(a, b);
				rot_axis = a.cross(b);

				// If they are not strictly parallel we can rotate
				if (rot_axis.norm() > 1e-8)
				{
					rot_axis.normalize();

					for (int img = idx_1+1; img < idx_2; ++img)
					{
						angle = (img-idx_1)*rot_angle/noi;
						Engine::Vectormath::rotate(a, rot_axis, angle, temp);

						(*c->images[img]->spins)[i] = temp;
					}
				}
				// Otherwise we simply leave the spin untouched
				else
				{
					for (int img = idx_1+1; img < idx_2; ++img)
					{
						(*c->images[img]->spins)[i] = a;
					}
				}
			}

		}

		void Homogeneous_Rotation(std::shared_ptr<Data::Spin_System_Chain> c, vectorfield A, vectorfield B)
		{
			(*c->images[0]->spins) = A;
			(*c->images[c->noi - 1]->spins) = B;

			int nos = c->images[0]->nos;

			scalar angle, rot_angle;
			Vector3 axis, rot_axis, a, b, temp;

			for (int i = 0; i < c->images[0]->nos; ++i)
			{
				a = A[i];
				b = B[i];

				rot_angle = Engine::Manifoldmath::dist_greatcircle(a, b);
				rot_axis = a.cross(b);

				// If they are not strictly parallel we can rotate
				if (rot_axis.norm() > 1e-8)
				{
					rot_axis.normalize();

					for (int img = 1; img < c->noi - 1; ++img)
					{
						angle = (img)*rot_angle / (c->noi - 1);
						Engine::Vectormath::rotate(a, rot_axis, angle, temp);

						(*c->images[img]->spins)[i] = temp;
					}
				}
				// Otherwise we simply leave the spin untouched
				else
				{
					for (int img = 1; img < c->noi - 1; ++img)
					{
						(*c->images[img]->spins)[i] = a;
					}
				}
			}

			/*
				do i=1,NOS
				  Start(i,:)  = Start(i,:)/length(Start(i,:))
				  Finish(i,:) = Finish(i,:)/length(Finish(i,:))

				  IMAGES(idx_start,i,:)  = Start(i,:)
				  IMAGES(idx_finish,i,:) = Finish(i,:)

				  r = max(-1.0,min(1.0, dot_product(Start(i,:), Finish(i,:)) )) !! this prevents NaNs from ocurring
				  rot_angle(i)  = acos(r)
				  rot_axis(i,:) = cross_product(Start(i,:), Finish(i,:))
				  if (abs(length(rot_axis(i,:))) > 1e-18) rot_axis(i,:) = rot_axis(i,:)/length(rot_axis(i,:))
				enddo
			*/

			/*
				do i=idx_start+1,idx_finish-1
				  !! loop over spins in image
				  do j=1,NOS
					 angle = (i-1)*rot_angle(j)/(idx_finish-idx_start)
					 IMAGES(i,j,:) = rotate_spin(Start(j,:),rot_axis(j,:),angle)

				  enddo
				enddo
			*/
		}
	}//end namespace Configuration_Chain
}//end namespace Utility


/* Copyright 2015,2016 Tao Xu
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied. See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

#include "Gbm.h"

#include <boost/scoped_array.hpp>
#include <boost/shared_ptr.hpp>
#include <vector>

#include "Concurrency.h"
#include "Config.h"
#include "DataSet.h"
#include "GbmFun.h"
#include "Tree.h"
#include "TreeRegressor.h"
#include <gflags/gflags.h>

namespace boosting {

using namespace std;

Gbm::Gbm(const GbmFun& fun, const DataSet& ds, const Config& cfg)
  : fun_(fun), ds_(ds), cfg_(cfg) {
}

class ParallelEval : public apache::thrift::concurrency::Runnable {
 public:
  ParallelEval(
    CounterMonitor& monitor,
    const int numExamples,
    const int numFeatures,
    const GbmFun& fun,
    const std::unique_ptr<TreeNode<uint16_t>>& weakModel,
    const DataSet& ds,
    const vector<double>& targets,
    boost::scoped_array<double>& F,
    boost::scoped_array<double>& subLoss,
    const int workIdx,
    const int totalWorkers)
    : monitor_(monitor), numExamples_(numExamples),
      numFeatures_(numFeatures), fun_(fun), weakModel_(weakModel),
      ds_(ds), targets_(targets), F_(F),
      subLoss_(subLoss), workIdx_(workIdx),
      totalWorkers_(totalWorkers) {
  }

  void run() {
    //boost::scoped_array<uint16_t> fvec(new uint16_t[numFeatures_]);
    for (int i = 0; i < numExamples_; i++) {
      if (i % totalWorkers_ == workIdx_) {
        //ds_.getFeatureVec(i, fvec);
        //double score = weakModel_->eval(fvec);
        double score = ds_.getPrediction(weakModel_.get(), i);
        F_[i] += score;
        subLoss_[workIdx_] += fun_.getExampleLoss(targets_[i], F_[i]);
      }
    }
    monitor_.decrement();
  }

 private:
  CounterMonitor& monitor_;
  const int numExamples_;
  const int numFeatures_;
  const GbmFun& fun_;
  const std::unique_ptr<TreeNode<uint16_t>>& weakModel_;
  const DataSet& ds_;
  const vector<double> targets_;
  boost::scoped_array<double>& F_;
  boost::scoped_array<double>& subLoss_;
  const int workIdx_;
  const int totalWorkers_;
};

void Gbm::getModel(
  vector<TreeNode<double>*>* model,
  double fimps[]) {

  const int numExamples = ds_.getNumExamples();

  boost::scoped_array<double> F(new double[numExamples]);
  boost::scoped_array<double> y(new double[numExamples]);

  double f0 = fun_.getF0(ds_.targets_);
  for (int i = 0; i < numExamples; i++) {
    F[i] = f0;
  }

  model->push_back(new LeafNode<double>(f0));

  double initLoss = fun_.getInitLoss(ds_.targets_);

  LOG(INFO) << "init avg loss " << initLoss / numExamples;

  for (int it = 0; it < cfg_.getNumTrees(); it++) {

    LOG(INFO) << "------- iteration " << it << " -------";

    fun_.getGradient(ds_.targets_, F, y);
    TreeRegressor regressor(ds_, y, fun_);

    std::unique_ptr<TreeNode<uint16_t>> weakModel(
      regressor.getTree(cfg_.getNumLeaves(), cfg_.getExampleSamplingRate(),
                        cfg_.getFeatureSamplingRate(), fimps));

    weakModel->scale(cfg_.getLearningRate());

    model->push_back(mapTree(weakModel.get()));

    VLOG(1) << toPrettyJson(weakModel->toJson(cfg_));
    double newLoss = 0.0;

    if (FLAGS_num_threads > 1) {
      CounterMonitor monitor(FLAGS_num_threads);
      boost::scoped_array<double> subLoss(new double[FLAGS_num_threads]);
      for (int wid = 0; wid < FLAGS_num_threads; wid++) {
        subLoss[wid] = 0.0;
        Concurrency::threadManager->add(
          boost::shared_ptr<apache::thrift::concurrency::Runnable>(
            new ParallelEval(monitor, numExamples, ds_.numFeatures_,
                             fun_, weakModel,
                             ds_, ds_.targets_, F, subLoss,
                             wid, FLAGS_num_threads)));
      }
      monitor.wait();

      for (int wid = 0; wid < FLAGS_num_threads; wid++) {
        newLoss += subLoss[wid];
      }
    } else {
      //boost::scoped_array<uint16_t> fvec(new uint16_t[ds_.numFeatures_]);
      for (int i = 0; i < numExamples; i++) {
        // ds_.getFeatureVec(i, fvec);
        // double score = weakModel->eval(fvec);
        double score = ds_.getPrediction(weakModel.get(), i);
        F[i] += score;
        newLoss += fun_.getExampleLoss(ds_.targets_[i], F[i]);
      }
    }

    LOG(INFO) << "total avg loss " << newLoss/numExamples
              << " reduction: " << 1.0 - newLoss/initLoss;
  }
}

TreeNode<double>* Gbm::mapTree(const TreeNode<uint16_t>* rt) {
  const PartitionNode<uint16_t>* pnode =
    dynamic_cast<const PartitionNode<uint16_t>*>(rt);
  if (pnode != NULL) {
    int fid = pnode->getFid();
    PartitionNode<double>* newNode = new PartitionNode<double>(
      fid, ds_.features_[fid].transitions[pnode->getFv()]);
    newNode->setVote(pnode->getVote());
    newNode->setLeft(mapTree(pnode->getLeft()));
    newNode->setRight(mapTree(pnode->getRight()));
    return newNode;
  } else {
    const LeafNode<uint16_t>* lfnode =
      dynamic_cast<const LeafNode<uint16_t>*>(rt);
    return new LeafNode<double>(lfnode->getVote());
  }
}

}


//----------------------------------------------------------------------
//	File:			segmentation.cpp
//	Author:			Elena Garces and Adolfo Munoz
//	Last modified:	18/07/2012
//	Description:	funcions that control the clustering
//----------------------------------------------------------------------
// This file is part of Intrinsic Images by Clustering.
//
//    Intrinsic Images by Clustering is free software: you can redistribute it
//    and/or modify it under the terms of the GNU General Public License as
//    published by the Free Software Foundation, either version 3 of the License,
//     or (at your option) any later version.
//
//    Intrinsic Images by Clustering is distributed in the hope that it will
//    be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
//    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//    GNU General Public License for more details.
//
//    You should have received a copy of the GNU General Public License
//    along with Intrinsic Images by Clustering.  If not,
//	  see <http://www.gnu.org/licenses/>.
////----------------------------------------------------------------------

#include "segmentacion.h"
#include "auxiliar.h"

#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/symmetric.hpp>

#include <algorithm>
#include <list>
#include <iostream>
#include <fstream>
#include "CImg.h"
using namespace std;
using namespace cimg_library;

/*-------------------------------------------------------------------------------------------------*/
void Segmentacion::updateLabels()
/*-------------------------------------------------------------------------------------------------*/
{

	int width = this->imageMapLabels.width();
	int height = this->imageMapLabels.height();

	vector<Cluster> newClusters;
	vector<int> founds;

	cimg_forXY(imageMapLabels,x,y)
	{

		int comp = *imageMapLabels.data(x,y,0,0);
		if ( comp != BACKG_CLUSTER )
		{
			vector<int>::iterator i = find(founds.begin(), founds.end(), comp);

			if (i == founds.end())
			{
				/* add */
				founds.push_back(comp);
				*imageMapLabels.data(x,y,0,0) = int(founds.size())-1;
				Cluster nuevo(comp, y*width+x);
				newClusters.push_back(nuevo);
			}
			else
			{
				/* already exists, update*/
				int index = std::distance (founds.begin(), i );
				*imageMapLabels.data(x,y,0,0) = index;
				newClusters[index].addPixel(y*width+x);

			}
		}
	}
	setListOfClusters(newClusters);
}

int Segmentacion::mergeSegment(const CImg<double>& image, int pos, int min_size_clus)
{

	int width = image.width();
	int n_changes=0;

	vector<int> pixels = getListOfClusters()[pos].getPixels();

	int idSetCentral = getListOfClusters()[pos].idCluster;

	for (int p = 0; p < (int) pixels.size(); p++)
	{
		int x = pixels[p] % width;	int y = pixels[p] / width;
		int idMostSimilar=-1;
		double diffMostSimilar=10000000.;

		int dist = 1;
		for (int j = y-dist; j <= y+dist; j++) {
			for (int i = x-dist; i <= x+dist; i++) {

				int set = getCluster(i,j);
				if (set >= 0 && set != idSetCentral && getListOfClusters().at(*imageMapLabels.data(i,j,0,0)).size() >= min_size_clus)
				{
					double squaredDiff = 0.0;
					cimg_forC(image,c) squaredDiff += pow(image(i, j,0,c) -  image(x, y,0,c), 2);
					squaredDiff=sqrt(squaredDiff);

					if (squaredDiff < diffMostSimilar)
					{
						diffMostSimilar = squaredDiff;
						idMostSimilar = set;
					}
				}
			}
		}
		if (idMostSimilar >=0)
		{
			n_changes++;
			setCluster(x, y, idMostSimilar);
		}
	}
		return n_changes;
}

/*-------------------------------------------------------------------------------------------------*/
/*  */

int Segmentacion::detectSpurious(const CImg<double>& image, int max_size_clus, float percent_external)
/*-------------------------------------------------------------------------------------------------*/
{
	int total_changes=0;

	int width = image.width();
	for (int k = 0; k < (int) listOfClusters.size(); k++)
	{
		if (listOfClusters[k].getPixels().size() < max_size_clus)
		{
			vector<int> pixels = listOfClusters[k].getPixels();
			int votesBoundary = 0;
			for (int p = 0; p < (int) pixels.size(); p++)
			{
				int x = pixels[p] % width;
				int y = pixels[p] / width;
				int idSetCentral = getCluster(x,y); /* id del conjunto del pixel central */

				if ( (getCluster(x-1, y) >= 0 && getCluster(x-1, y) != idSetCentral) ||
					(getCluster(x+1, y) >= 0 && getCluster(x+1, y) != idSetCentral) ||
					(getCluster(x, y+1) >= 0 && getCluster(x, y+1) != idSetCentral) ||
					(getCluster(x, y-1) >= 0 && getCluster(x, y-1) != idSetCentral))
					votesBoundary++;
			}

			if (votesBoundary >= percent_external * pixels.size())
			{
				int changes = mergeSegment(image, k, max_size_clus);
				total_changes+=changes;
			}
		}
	}

	return total_changes;
}

void Segmentacion::mergeSmallClusters(const CImg<double>& image, int max_it,
								   int max_size_clus, float percent_external)
{
	int n_changes = 1; 	int it = 1;
	while (it < max_it && n_changes > 0)
	{
		n_changes = detectSpurious(image, max_size_clus, percent_external);
		#ifdef VERBOSE
		cout << "Clustering: --> \t\tIteration  " << it << " changes " << n_changes << endl;
		#endif
		updateLabels();
		it++;
	}
}

template <typename matrizBoost>
void analysis_gradient_clusters(matrizBoost& matrixRatios, const Segmentacion& seg,
			const CImg<double>& input, const MatteImage& matte)
	{
		int num_sets = int(seg.getListOfClusters().size());
		list<edge> edges;
		ublas::symmetric_matrix<int> edges_counter(num_sets,num_sets);

		for (int i=0; i < num_sets; i++)
			for (int j=i; j < num_sets; j++)
			{
				edges_counter(i,j) = 0;
				matrixRatios(i,j) = 0.0;
			}

		for (int i=0; i < num_sets; i++)
			matrixRatios(i,i) = 0.0;

		buildGraph(edges, input, 2, matte);
		for (list<edge>::const_iterator i=edges.begin(); i != edges.end(); i++)
		{
			int aSeg = seg.getCluster((*i).a);
			int bSeg = seg.getCluster((*i).b);

			if (aSeg != bSeg && aSeg >= 0 && bSeg >= 0)
			{
				int xa = (*i).a % input.width();
				int ya = (*i).a / input.width();
				int xb = (*i).b % input.width();
				int yb = (*i).b / input.width();

				double diffgrad =0.0;
				cimg_forC(input,c)
				{
					diffgrad+= fabs(input(xa, ya,0,c) - input(xb, yb,0,c));
				}
				diffgrad/=input.spectrum();

				if (aSeg <= bSeg)
				{
					edges_counter(aSeg,bSeg)++;
					matrixRatios(aSeg,bSeg)+=diffgrad;
				} else
				{
					edges_counter(bSeg,aSeg)++;
					matrixRatios(bSeg,aSeg)+=diffgrad;
				}
			}
		}

		for (int i=0; i < num_sets; i++)
			for (int j=(i+1); j < num_sets; j++)
			{
				if (matrixRatios(i,j)!=0)
				{
					if (edges_counter(i,j) > 15)
						matrixRatios(i,j)/=edges_counter(i,j);
					else matrixRatios(i,j)=0.0;
				}
			}
	}

bool func_ord_e(const edge &a, const edge &b)
	{ return (a.w < b.w);}

int Segmentacion::mergeSmoothBoundaries(CImg<double> *im_input, const MatteImage &matte, vector<tpPair> &paresR)
{

	int width = im_input->width();
	int height = im_input->height();

	for (unsigned int p=0; p < paresR.size(); p++)
	{
		paresR[p].i= imageMapLabels( paresR[p].i % im_input->width(), paresR[p].i / im_input->width(),0,0);
		paresR[p].j= imageMapLabels( paresR[p].j % im_input->width(), paresR[p].j / im_input->width(),0,0);
	}

	float max_grad = percent_grad_max*im_input->max();

	/* MERGE, ANALYZING GRADIENTS */
	int num_sets = int(getListOfClusters().size());
	ublas::symmetric_matrix<float> matrixRatios(num_sets,num_sets);
	analysis_gradient_clusters(matrixRatios, *(this),*im_input,matte);

	vector<edge> edges;

	int n=0;
	for (int i=0; i < num_sets; i++)
		for (int j=(i+1); j < num_sets; j++)
		{
			if (matrixRatios(i,j) != 0)
			{
				double diff=matrixRatios(i,j);
				if (diff  < max_grad)
					edges.push_back(edge(i,j,diff));
			}
		}

	std::sort(edges.begin(), edges.end(), func_ord_e);

	vector<bool> tratados(getListOfClusters().size(),false);
	vector<bool> referenciado(getListOfClusters().size(),false);

	int n_changes=0;
	for (vector<edge>::const_iterator i=edges.begin(); i != edges.end(); i++)
	{
		int peq, grand;
		if (getListOfClusters()[(*i).a].size() < getListOfClusters()[(*i).b].size())
		{
			peq = (*i).a;
			grand = (*i).b;
		} else  {
			peq = (*i).b;
			grand = (*i).a;
		}

		if (tratados[grand] == false)
		{

			n_changes++;
			// remove the smaller cluster
			for (unsigned int p=0; p < getListOfClusters()[peq].size(); p++)
			{
				setCluster(getListOfClusters()[peq].getPixels()[p], grand);
			}
			tratados[peq]=true;
			listOfClusters[grand].addPixel(listOfClusters[peq].getPixels());

			for (unsigned int i=0; i < paresR.size(); i++)
			{
				if (paresR[i].valid && (paresR[i].i==peq || paresR[i].j == peq))
				{
					paresR[i].valid=false;
				}
			}
		}
	}

	vector<tpPair> paresRnew;
	for (unsigned int i=0; i < paresR.size(); i++)
	{
			if (paresR[i].valid)
			{
				tpPair aux;
				aux.i = getListOfClusters()[paresR[i].i].getPixels()[0];
				aux.j = getListOfClusters()[paresR[i].j].getPixels()[0];
				aux.valid = true;
				paresRnew.push_back(aux);
			}

	}
	paresR = paresRnew;
	updateLabels();

	//* --------------------------------------------------*/

	return n_changes;

}

void Segmentacion::updateReflectancePairs(vector<tpPair>& eqReflectance,  CImg<double> *luminance, int MIN_SIZE_CLUS)
{
	vector<tpPair> eqReflectancenew;
	float MIN_VARIATION_L = (float) 0.07*luminance->max();

	if (eqReflectance.size() > 0)
	{
		vector<float> averageLuminance(getListOfClusters().size(), -1);
		int index=0;
		for (unsigned int i=0; i < eqReflectance.size(); i++)
		{
			tpPair aux;
			aux.valid=true;
			aux.i = imageMapLabels(eqReflectance[i].i % width(),eqReflectance[i].i / width(),0,0);
			aux.j = imageMapLabels(eqReflectance[i].j % width(),eqReflectance[i].j / width(),0,0);

			if (aux.i > aux.j) std::swap(aux.i, aux.j);

			if (aux.i != aux.j && getListOfClusters().at(aux.i).size() > MIN_SIZE_CLUS && getListOfClusters().at(aux.j).size() > MIN_SIZE_CLUS)
			{
				averageLuminance[aux.i] = (averageLuminance[aux.i] < 0) ? computeMean(luminance, getListOfClusters().at(aux.i).getPixels()) : averageLuminance[aux.i];
				averageLuminance[aux.j] = (averageLuminance[aux.j] < 0) ? computeMean(luminance, getListOfClusters().at(aux.j).getPixels()) : averageLuminance[aux.j];

				if (averageLuminance[aux.i]  > 0 && averageLuminance[aux.j] > 0)
				{
					double diff = fabs(averageLuminance[aux.i]-averageLuminance[aux.j]);
					if (diff < MIN_VARIATION_L)	eqReflectancenew.push_back(aux);

				}
			}
		}

	}

	eqReflectance=eqReflectancenew;

}

vector<tpPair> Segmentacion::searchPairsRcte(Segmentacion& seg_old)
{

	vector<tpPair> pairsReflectance;
	vector< vector<int> > vec(seg_old.getListOfClusters().size(), vector<int>(0));
	for (unsigned int s = 0; s < getListOfClusters().size(); s++)
	{
		int p = getListOfClusters()[s].getPixels()[0];
		int clus_old=seg_old.getCluster(p);
		if (getListOfClusters()[s].getPixels().size() > 100)
			vec[clus_old].push_back(s);
	}

	for (unsigned int s = 0; s < vec.size(); s++)
	{
		for (unsigned int subS = 0; subS < vec[s].size(); subS++)
		{

			for (unsigned int sS = subS+1; sS < vec[s].size(); sS++)
			{
				int clusa=vec[s][subS];
				int clusb=vec[s][sS];

				tpPair aux;
				aux.i = getListOfClusters()[clusa].getPixels()[0];
				aux.j = getListOfClusters()[clusb].getPixels()[0];

				pairsReflectance.push_back(aux);
			}
		}
	}

	return pairsReflectance;
}

void Segmentacion::scanline(unsigned int coriginal, unsigned int cnew, int x, int y, Segmentacion& sol)
{
	int x2;

	if (getCluster(x,y) != BACKG_CLUSTER)
	{
		for (x2 = x; (getCluster(x2,y)==coriginal);x2++)	sol.setCluster(x2,y,cnew);

		for (x2 = (x - 1); (getCluster(x2,y)==coriginal);x2--) sol.setCluster(x2,y,cnew);

		for (x2 = (x - 1); (sol.getCluster(x2,y)==cnew) && (getCluster(x2,y)==coriginal);x2--)
		{
			if ( (sol.getCluster(x2,y-1)==NON_ASIGN) && (getCluster(x2,y-1)==coriginal) ) scanline(coriginal,cnew,x2,y-1,sol);
			if ( (sol.getCluster(x2,y+1)==NON_ASIGN) && (getCluster(x2,y+1)==coriginal) ) scanline(coriginal,cnew,x2,y+1,sol);
		}
		for (x2 = x; (sol.getCluster(x2,y)==cnew) && (getCluster(x2,y)==coriginal);x2++)
		{
			if ( (sol.getCluster(x2,y-1)==NON_ASIGN) && (getCluster(x2,y-1)==coriginal) ) scanline(coriginal,cnew,x2,y-1,sol);
			if ( (sol.getCluster(x2,y+1)==NON_ASIGN) && (getCluster(x2,y+1)==coriginal) ) scanline(coriginal,cnew,x2,y+1,sol);
		}

		for (x2 = (x - 1); (sol.getCluster(x2,y)==cnew) && (getCluster(x2,y)==coriginal);x2--)
		{
			if ( (sol.getCluster(x2+1,y-1)==NON_ASIGN) && (getCluster(x2+1,y-1)==coriginal) ) scanline(coriginal,cnew,x2+1,y-1,sol);
			if ( (sol.getCluster(x2-1,y-1)==NON_ASIGN) && (getCluster(x2-1,y-1)==coriginal) ) scanline(coriginal,cnew,x2-1,y-1,sol);

			if ( (sol.getCluster(x2+1,y+1)==NON_ASIGN) && (getCluster(x2+1,y+1)==coriginal) ) scanline(coriginal,cnew,x2+1,y+1,sol);
			if ( (sol.getCluster(x2-1,y+1)==NON_ASIGN) && (getCluster(x2-1,y+1)==coriginal) ) scanline(coriginal,cnew,x2-1,y+1,sol);

		}

		for (x2 = x; (sol.getCluster(x2,y)==cnew) && (getCluster(x2,y)==coriginal);x2++)
		{
			if ( (sol.getCluster(x2+1,y-1)==NON_ASIGN) && (getCluster(x2+1,y-1)==coriginal) ) scanline(coriginal,cnew,x2+1,y-1,sol);
			if ( (sol.getCluster(x2-1,y-1)==NON_ASIGN) && (getCluster(x2-1,y-1)==coriginal) ) scanline(coriginal,cnew,x2-1,y-1,sol);

			if ( (sol.getCluster(x2+1,y+1)==NON_ASIGN) && (getCluster(x2+1,y+1)==coriginal) ) scanline(coriginal,cnew,x2+1,y+1,sol);
			if ( (sol.getCluster(x2-1,y+1)==NON_ASIGN) && (getCluster(x2-1,y+1)==coriginal) ) scanline(coriginal,cnew,x2-1,y+1,sol);

		}
	} else sol.setCluster(x,y,BACKG_CLUSTER);
}

Segmentacion Segmentacion::get_separate_noncontiguous_clusters()
{
	Segmentacion sol(width(),height());
	sol.imageMapLabels.fill(NON_ASIGN);
	unsigned int cnew = 0;
	int x=0, y=0;

	while (y<height())
	{
		while ( (sol.getCluster(x,y)!=NON_ASIGN) && (y<height()) )
		{
			x++;
			if (x>=width()) { x = 0; y++;}
		}

		if (y<height())	scanline(getCluster(x,y),cnew,x,y,sol);
		cnew++;
	}
	return sol;
}


// Copyright (C) 2002-2006 Klaas Gadeyne <first dot last at gmail dot com>
//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation; either version 2.1 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with this program; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
//
#include "rng.h"

#include <../config.h>
#ifdef __RNGWRAPPER_BOOST__  // BOOST RANDOM LIBRARY
// THE BOOST RANDOM NUMBER GENERATION LIBRARY

#include <boost/random/mersenne_twister.hpp>
#include <boost/random/variate_generator.hpp>
#include <boost/random/normal_distribution.hpp>
#include <boost/random/uniform_01.hpp>
#include <boost/random/uniform_real.hpp>

static boost::mt19937 Boost_Rng; // Source of randomness

static boost::uniform_real<double> Uniform_Distribution; // Uniform distribution
static boost::variate_generator<boost::mt19937&,boost::uniform_real<double> > roll(Boost_Rng,Uniform_Distribution);

double BFL::rnorm(const double& mu,const double& sigma)
{
  boost::normal_distribution<double> TestDist(mu,sigma);
  boost::variate_generator <boost::mt19937 &,boost::normal_distribution<double> > TestGen(Boost_Rng,TestDist);
  return TestGen();
}

double BFL::runif()
{
  return roll();
}

double BFL::runif(const double &min, const double& max)
{
  boost::uniform_real<double> Uniform_DistributionMinMax(min,max); // Uniform distribution
  boost::variate_generator<boost::mt19937&,boost::uniform_real<double> > roll(Boost_Rng,Uniform_DistributionMinMax);
  return roll();
}

#endif // __RNGWRAPPER_BOOST__


/*Copyright (c) 2021 James Gayvert

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.
*/

/*
 * qchem_interface.cpp
 */

#include "qchem_interface.h"

#include <algorithm>
#include <chrono>
#include <Eigen/Dense>
#include <fstream>
#include <iostream>
#include <sstream>
#include <string>

#include "gto_ordering.h"
#include "opencap_exception.h"
#include "utils.h"

void qchem_parse_fchk_dms(std::string dmat_filename,std::vector<std::vector<Eigen::MatrixXd>> &alpha_opdms,
		std::vector<std::vector<Eigen::MatrixXd>> &beta_opdms, size_t nstates, size_t ntdm, BasisSet &bs,
		bool symmetric_rdm, bool do_spin)
{
	alpha_opdms = std::vector< std::vector<Eigen::MatrixXd>>(nstates, std::vector<Eigen::MatrixXd> (nstates));
	beta_opdms = std::vector< std::vector<Eigen::MatrixXd>>(nstates, std::vector<Eigen::MatrixXd> (nstates));
	int lt_number = bs.Nbasis*(bs.Nbasis+1)/2;
	int lt_tdm = nstates*(nstates+1)/2 - nstates;
	std::ifstream is(dmat_filename);
	std::vector<Eigen::MatrixXd> alpha_tdms;
	std::vector<Eigen::MatrixXd> beta_tdms;
    if (is.good())
    {
    	std::string line, rest;
    	std::getline(is, line);
    	auto start = std::chrono::high_resolution_clock::now();
    	for (size_t i=0;i<nstates;i++)
    	{
    			//alpha first, then beta
    			for(size_t spin=0;spin<=do_spin;spin++)
    			{
					while(line.find("State Density")== std::string::npos)
					{
						std::getline(is,line);
	    	    		if (is.peek()==EOF)
	    	    			opencap_throw("Error: Reached end of file before densities for "+
	    	    					std::to_string(nstates) + " states were found.");
					}
					//last part of line should be number of elements to read
					int num_elements = stoi(split(line,' ').back());
					if(sqrt(num_elements)!=bs.Nbasis && num_elements!=lt_number)
						opencap_throw("Error: dimensions of TDMs do not match specified basis set.");
					size_t lines_to_read = num_elements%5==0 ? (num_elements/5) : num_elements/5+1;
					std::vector<double> matrix_elements;
					for (size_t k=1;k<=lines_to_read;k++)
					{
						std::getline(is,line);
						std::vector<std::string> tokens = split(line,' ');
						for (auto token:tokens)
							matrix_elements.push_back(std::stod(token));
					}
					Eigen::MatrixXd st_opdm(bs.Nbasis,bs.Nbasis);
					if(symmetric_rdm)
						fill_LT<double>(matrix_elements,st_opdm);
					else
						fill_mat<double>(matrix_elements,st_opdm);
					to_opencap_ordering(st_opdm,bs,get_qchem_ids(bs));
					if(spin==0)
					{
						alpha_opdms[i][i]=st_opdm;
						if(!do_spin)
							beta_opdms[i][i]=st_opdm;
					}
					else
						beta_opdms[i][i]=st_opdm;
    			}
		}
        is.seekg (0, ios::beg);
    	for(size_t i=0;i<ntdm;i++)
    	{
			//alpha first, then beta
			for (size_t spin=0;spin<=do_spin;spin++)
			{
				while(line.find("Transition DM")== std::string::npos)
				{
					std::getline(is,line);
					if (is.peek()==EOF)
						opencap_throw("Error: Reached end of file before densities for "+
								std::to_string(nstates) + " states were found.");
				}
				//last part of line should be number of elements to read
				size_t num_elements = stoi(split(line,' ').back());
				if(sqrt(num_elements)!=bs.Nbasis && num_elements!=lt_number)
					opencap_throw("Error: dimensions of TDMs do not match specified basis set.");
				size_t lines_to_read = num_elements%5==0 ? (num_elements/5) : num_elements/5+1;
				std::vector<double> matrix_elements;
				for (size_t k=1;k<=lines_to_read;k++)
				{
					std::getline(is,line);
					std::vector<std::string> tokens = split(line,' ');
					for (auto token:tokens)
						matrix_elements.push_back(std::stod(token));
				}
				Eigen::MatrixXd st_opdm(bs.Nbasis,bs.Nbasis);
				fill_mat<double>(matrix_elements,st_opdm);
				to_opencap_ordering(st_opdm,bs,get_qchem_ids(bs));
				if(spin==0)
				{
					alpha_tdms.push_back(st_opdm);
					if(!do_spin)
						beta_tdms.push_back(st_opdm);
				}
				else
					beta_tdms.push_back(st_opdm);
			}
    	}
    }
    is.close();
    if(ntdm == lt_tdm)
    {
        std::cout << "Warning: TDM M-->N is assumed to be conjugate transpose of "
        << "TDM N-->M where M>N" << std::endl;
    	size_t dm_idx = 0;
    	for(size_t i=0;i<nstates;i++)
    	{
    		for(size_t j=i+1;j<nstates;j++)
    		{
    			alpha_opdms[i][j] = alpha_tdms[dm_idx];
    			alpha_opdms[j][i] = alpha_tdms[dm_idx].adjoint();
    			beta_opdms[i][j] = beta_tdms[dm_idx].adjoint();
    			beta_opdms[j][i] = beta_tdms[dm_idx];
				dm_idx++;
    		}
    	}
    }
    else
    {
    	size_t dm_idx = 0;
    	for(size_t i=0;i<nstates;i++)
    	{
    		for(size_t j=0;j<nstates;j++)
    		{
    			if(i!=j)
    			{
    				alpha_opdms[i][j] = alpha_tdms[dm_idx];
    				beta_opdms[i][j] = beta_tdms[dm_idx];
    				dm_idx++;
    			}
    		}
    	}
    }
}

void qchem_read_dms(std::vector<std::vector<Eigen::MatrixXd>> &alpha_dms,
		std::vector<std::vector<Eigen::MatrixXd>> &beta_dms,
		std::string fchk_filename, BasisSet &bs,size_t num_states)
{
	std::ifstream is(fchk_filename);
	size_t nstates = 0;
	size_t num_tdm = 0;
	bool sep_alpha_beta;
	bool symmetric_rdm = false;
	bool symmetric_tdm = false;
	std::string line, rest;
	if (is.good())
	{
		while(is.peek() != EOF )
		{
			std::getline(is, line);
			if(line.find("State Density")!= std::string::npos)
			{
				if(line.find("Alpha")!=std::string::npos)
				{
					sep_alpha_beta=true;
					int num_elements = stoi(split(line,' ').back());
					if(sqrt(num_elements)!=bs.Nbasis)
					{
						int lt_number = bs.Nbasis*(bs.Nbasis+1)/2;
						if(lt_number == num_elements)
							symmetric_rdm = true;
						else
							opencap_throw("Error: dimensions of DMs do not match specified basis set.");
					}
					break;
				}
				else
				{
					sep_alpha_beta=false;
					int num_elements = stoi(split(line,' ').back());
					if(sqrt(num_elements)!=bs.Nbasis)
					{
						int lt_number = bs.Nbasis*(bs.Nbasis+1)/2;
						if (lt_number == num_elements)
							symmetric_rdm = true;
						else
							opencap_throw("Error: dimensions of DMs do not match specified basis set.");
					}
					break;
				}
			}
		}
		is.seekg (0, ios::beg);
		while(is.peek()!=EOF)
		{
			std::getline(is, line);
			if((line.find("Alpha")!= std::string::npos && line.find("State Density")!= std::string::npos) && sep_alpha_beta)
				nstates++;
			else if(line.find("State Density")!= std::string::npos && !sep_alpha_beta)
				nstates++;
			if((line.find("Alpha")!= std::string::npos && line.find("Transition DM")!= std::string::npos) && sep_alpha_beta)
				num_tdm++;
			else if(line.find("Transition DM")!= std::string::npos && !sep_alpha_beta)
				num_tdm++;
		}
		is.close();
		if(nstates==0)
		{
			opencap_throw("Error: Unable to find any densities in:" +fchk_filename);
		}
		else if(nstates!=num_states)
		{
			opencap_throw("Error: number of states found: " + std::to_string(nstates) + ", does not match "
					"number of states specified in the input: " + std::to_string(num_states)+ ". Exiting...");
		}
		qchem_parse_fchk_dms(fchk_filename,alpha_dms,beta_dms,nstates,num_tdm,bs,symmetric_rdm,sep_alpha_beta);
	}
	else
	    opencap_throw("Error: I couldn't read:" + fchk_filename);
}

Eigen::MatrixXd qchem_read_overlap(std::string dmat_filename, BasisSet bs)
{
	size_t num_bf = bs.Nbasis;
    std::ifstream is(dmat_filename);
	Eigen::MatrixXd smat(num_bf,num_bf);
	smat=Eigen::MatrixXd::Zero(num_bf,num_bf);
    if (is.good())
    {
    	std::string line, rest;
    	while (line.find("Overlap Matrix")== std::string::npos)
    	{
        	std::getline(is, line);
    		if (is.peek()==EOF)
    			opencap_throw("Error: Reached end of file before Overlap Matrix was found.");
    	}
    	size_t num_elements = stoi(split(line,' ').back());
    	int lt_number = bs.Nbasis*(bs.Nbasis+1)/2;
    	if(num_elements!=lt_number)
    		opencap_throw("Error: Dimensions of overlap matrix do not match basis.");
		size_t lines_to_read = num_elements%5==0 ? (num_elements/5) : num_elements/5+1;
		std::vector<double> matrix_elements;
		for (size_t k=1;k<=lines_to_read;k++)
		{
			std::getline(is,line);
			std::vector<std::string> tokens = split(line,' ');
			for (auto token:tokens)
				matrix_elements.push_back(std::stod(token));
		}
		fill_LT<double>(matrix_elements,smat);
		to_opencap_ordering(smat,bs,get_qchem_ids(bs));
    }
    else
    	opencap_throw("Error: I couldn't read:" + dmat_filename);
    return smat;
}

Eigen::MatrixXd read_qchem_tddft_energies(size_t nstates,std::string output_file)
{
	Eigen::MatrixXd ZERO_ORDER_H(nstates,nstates);
	ZERO_ORDER_H=Eigen::MatrixXd::Zero(nstates,nstates);
	std::ifstream is(output_file);
    if (is.good())
    {
    	std::string line, rest;
    	std::getline(is, line);
    	size_t state_idx = 1;
    	while (state_idx<nstates)
    	{
    		if (is.peek()==EOF)
    			opencap_throw("Error: Reached end of file before "+ std::to_string(nstates) + " energies were found. "
    					"Only " + std::to_string(state_idx) + " states were found. Exiting...");
    		if (line.find("Total energy in the final basis set")!= std::string::npos)
    			ZERO_ORDER_H(0,0) = std::stod(split(line,' ')[8]);
			if (line.find("Total energy for state")!= std::string::npos)
			{
					ZERO_ORDER_H(state_idx,state_idx) = std::stod(split(line,' ')[5]);
					state_idx++;
					std::getline(is,line);
			}
    		else
    			std::getline(is,line);
    	}
    }
    else
    	opencap_throw("Error: I couldn't read:" + output_file);
    return ZERO_ORDER_H;
}

Eigen::MatrixXd read_qchem_eom_energies(size_t nstates,std::string output_file)
{
	Eigen::MatrixXd ZERO_ORDER_H(nstates,nstates);
	ZERO_ORDER_H=Eigen::MatrixXd::Zero(nstates,nstates);
	std::ifstream is(output_file);
    if (is.good())
    {
    	std::string line, rest;
    	std::getline(is, line);
    	size_t state_idx = 1;
    	while (state_idx<=nstates)
    	{
    		if (is.peek()==EOF)
    			opencap_throw("Error: Reached end of file before "+ std::to_string(nstates) + " energies were found. "
    					"Only " + std::to_string(state_idx) + " states were found. Exiting...");
			if (line.find("Total energy")!= std::string::npos && line.find("Excitation energy")!= std::string::npos)
			{
					ZERO_ORDER_H(state_idx-1,state_idx-1) = std::stod(split(line,' ')[3]);
					state_idx++;
					std::getline(is,line);
			}
    		else
    			std::getline(is,line);
    	}
    }
    else
    	opencap_throw("Error: I couldn't read:" + output_file);
    return ZERO_ORDER_H;
}

std::vector<Atom> read_geometry_from_fchk(std::string fchk_filename)
{
	std::vector<Atom> atoms;
	std::vector<size_t> atom_nums;
	std::vector<double> coords;
	std::ifstream is(fchk_filename);
    if (is.good())
    {
    	std::string line, rest;
    	std::getline(is, line);
		while(line.find("Atomic numbers")== std::string::npos)
		{
			std::getline(is,line);
    		if (is.peek()==EOF)
    			opencap_throw("Error: Reached end of file before atomic numbers were found.");
		}
		size_t num_elements = stoi(split(line,' ').back());
		size_t lines_to_read = num_elements%6==0 ? (num_elements/6) : num_elements/6+1;
		for (size_t k=1;k<=lines_to_read;k++)
		{
			std::getline(is,line);
			std::vector<std::string> tokens = split(line,' ');
			for (auto token:tokens)
				atom_nums.push_back(std::stoi(token));
		}
		while(line.find("Current cartesian coordinates")== std::string::npos)
		{
			std::getline(is,line);
    		if (is.peek()==EOF)
    			opencap_throw("Error: Reached end of file before Current Cartesian Coordinates were found.");
		}
		num_elements = stoi(split(line,' ').back());
		lines_to_read = num_elements%5==0 ? (num_elements/5) : num_elements/5+1;
		for (size_t k=1;k<=lines_to_read;k++)
		{
			std::getline(is,line);
			std::vector<std::string> tokens = split(line,' ');
			for (auto token:tokens)
				coords.push_back(std::stod(token));
		}
		//ok now lets populate our atoms
		for(size_t i=0;i<atom_nums.size();i++)
			atoms.push_back(Atom(atom_nums[i],coords[i*3],coords[i*3+1],coords[i*3+2]));
    }
    else
    	opencap_throw("Error: I couldn't read:" + fchk_filename);
    return atoms;
}

BasisSet read_basis_from_fchk(std::string fchk_filename, std::vector<Atom> atoms)
{
	auto start = std::chrono::high_resolution_clock::now();
	std::vector<int> shell_types;
	std::vector<int> prims_per_shell;
	std::vector<int> atom_ids;
	std::vector<double> exps;
	std::vector<double> coeffs;
	// needed for SP specification
	std::vector<double> p_coeffs;
	bool SP_basis_function = true;
	BasisSet bs;
	for(auto atm:atoms)
		bs.centers.push_back(atm.coords);
	std::vector<shell_id> ids;
	std::ifstream is(fchk_filename);
    if (is.good())
    {
    	std::string line, rest;
    	std::getline(is, line);
    	//First lets figure out if there are SP functions
    	//shell types
		while(line.find("Shell types")== std::string::npos)
		{
			std::getline(is,line);
    		if (is.peek()==EOF)
    			opencap_throw("Error: Reached end of file before shell types were found.");
		}
		size_t num_elements = stoi(split(line,' ').back());
		size_t lines_to_read = num_elements%6==0 ? (num_elements/6) : num_elements/6+1;
		for (size_t k=1;k<=lines_to_read;k++)
		{
			std::getline(is,line);
			std::vector<std::string> tokens = split(line,' ');
			for (auto token:tokens)
			{
				shell_types.push_back(std::stoi(token));
			}
		}
		//prims per shell
		while(line.find("Number of primitives per shell")== std::string::npos)
		{
			std::getline(is,line);
    		if (is.peek()==EOF)
    			opencap_throw("Error: Reached end of file before Number of primitives per shell were found.");
		}
		num_elements = stoi(split(line,' ').back());
		lines_to_read = num_elements%6==0 ? (num_elements/6) : num_elements/6+1;
		for (size_t k=1;k<=lines_to_read;k++)
		{
			std::getline(is,line);
			std::vector<std::string> tokens = split(line,' ');
			for (auto token:tokens)
				prims_per_shell.push_back(std::stoi(token));
		}
		//atom ids
		while(line.find("Shell to atom map")== std::string::npos)
		{
			std::getline(is,line);
    		if (is.peek()==EOF)
    			opencap_throw("Error: Reached end of file before Shell to atom map was found.");
		}
		num_elements = stoi(split(line,' ').back());
		lines_to_read = num_elements%6==0 ? (num_elements/6) : num_elements/6+1;
		for (size_t k=1;k<=lines_to_read;k++)
		{
			std::getline(is,line);
			std::vector<std::string> tokens = split(line,' ');
			for (auto token:tokens)
				atom_ids.push_back(std::stoi(token));
		}
		//prims
		while(line.find("Primitive exponents")== std::string::npos)
		{
			std::getline(is,line);
    		if (is.peek()==EOF)
    			opencap_throw("Error: Reached end of file before Primitive exponents were found.");
		}
		num_elements = stoi(split(line,' ').back());
		lines_to_read = num_elements%5==0 ? (num_elements/5) : num_elements/5+1;
		for (size_t k=1;k<=lines_to_read;k++)
		{
			std::getline(is,line);
			std::vector<std::string> tokens = split(line,' ');
			for (auto token:tokens)
				exps.push_back(std::stod(token));
		}
		//coeffs
		while(line.find("Contraction coefficients")== std::string::npos)
		{
			std::getline(is,line);
    		if (is.peek()==EOF)
    			opencap_throw("Error: Reached end of file before Contraction coefficients were found.");
		}
		num_elements = stoi(split(line,' ').back());
		lines_to_read = num_elements%5==0 ? (num_elements/5) : num_elements/5+1;
		for (size_t k=1;k<=lines_to_read;k++)
		{
			std::getline(is,line);
			std::vector<std::string> tokens = split(line,' ');
			for (auto token:tokens)
				coeffs.push_back(std::stod(token));
		}
		while(line.find("P(S=P) Contraction coefficients")== std::string::npos)
		{
			std::getline(is,line);
    		if (is.peek()==EOF)
    		{
    			SP_basis_function = false;
    			break;
    		}
		}
		if (SP_basis_function)
		{
			num_elements = stoi(split(line,' ').back());
			lines_to_read = num_elements%5==0 ? (num_elements/5) : num_elements/5+1;
			for (size_t k=1;k<=lines_to_read;k++)
			{
				std::getline(is,line);
				std::vector<std::string> tokens = split(line,' ');
				for (auto token:tokens)
					p_coeffs.push_back(std::stod(token));
			}
		}
    }
    else
    	opencap_throw("Error: I couldn't read:" + fchk_filename);
    size_t prim_idx=0;
    for(size_t i=0;i<shell_types.size();i++)
    {
    	//SP
    	if(shell_types[i]==-1)
    	{
    		if(p_coeffs.size()==0)
    			opencap_throw("Error: missing section P(S=P) Contraction coefficients.");
    		Shell s_shell(0,atoms[atom_ids[i]-1].coords);
    		Shell p_shell(1,atoms[atom_ids[i]-1].coords);
			int num_prims = prims_per_shell[i];
			for(int j=1;j<=num_prims;j++)
			{
				s_shell.add_primitive(exps[prim_idx],coeffs[prim_idx]);
				p_shell.add_primitive(exps[prim_idx],p_coeffs[prim_idx]);
				prim_idx++;
			}
			bs.add_shell(s_shell);
			bs.add_shell(p_shell);
    	}
    	else
    	{
			Shell new_shell(abs(shell_types[i]),atoms[atom_ids[i]-1].coords);
			if(shell_types[i]>0 && new_shell.l>1)
				new_shell.pure=false;
			int num_prims = prims_per_shell[i];
			for(int j=1;j<=num_prims;j++)
			{
				new_shell.add_primitive(exps[prim_idx],coeffs[prim_idx]);
				prim_idx++;
			}
			bs.add_shell(new_shell);
    	}
    }
    bs.normalize();
	auto stop = std::chrono::high_resolution_clock::now();
	auto total_time = std::chrono::duration<double>(stop-start).count();
    return bs;
}


/*

Copyright (c) 2005-2016, University of Oxford.
All rights reserved.

University of Oxford means the Chancellor, Masters and Scholars of the
University of Oxford, having an administrative office at Wellington
Square, Oxford OX1 2JD, UK.

This file is part of Chaste.

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 Oxford 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.

*/

#ifndef AVERAGEDSOURCEPDE_HPP_
#define AVERAGEDSOURCEPDE_HPP_

#include "ChasteSerialization.hpp"
#include <boost/serialization/base_object.hpp>
#include <boost/serialization/vector.hpp>

#include "AbstractCellPopulation.hpp"
#include "TetrahedralMesh.hpp"
#include "AbstractLinearEllipticPde.hpp"

/**
 *  An Elliptic PDE which calculates the source term by adding the number of cells
 *  in the element containing that point and scaling by the element area.
 */
template<unsigned DIM>
class AveragedSourcePde : public AbstractLinearEllipticPde<DIM,DIM>
{
    friend class TestCellBasedPdes;

private:

    /** Needed for serialization.*/
    friend class boost::serialization::access;
    /**
     * Serialize the PDE and its member variables.
     *
     * @param archive the archive
     * @param version the current version of this class
     */
    template<class Archive>
    void serialize(Archive & archive, const unsigned int version)
    {
       archive & boost::serialization::base_object<AbstractLinearEllipticPde<DIM, DIM> >(*this);
       archive & mCoefficient;
       archive & mCellDensityOnCoarseElements;
    }

protected:

    /** The cell population member. */
    AbstractCellPopulation<DIM>& mrCellPopulation;

    /** Coefficient of consumption of nutrient by cells. */
    double mCoefficient;

    /** Vector of averaged cell densities on elements of the coarse mesh. */
    std::vector<double> mCellDensityOnCoarseElements;

public:

    /**
     * Constructor.
     *
     * @param rCellPopulation reference to the cell population
     * @param coefficient the coefficient of consumption of nutrient by cells (defaults to 0.0)
     */
    AveragedSourcePde(AbstractCellPopulation<DIM>& rCellPopulation, double coefficient=0.0);

    /**
     * @return const reference to the cell population (used in archiving).
     */
    const AbstractCellPopulation<DIM>& rGetCellPopulation() const;

    /**
     * @return mCoefficient
     */
    double GetCoefficient() const;

    /**
     * Set up the source terms.
     *
     * @param rCoarseMesh reference to the coarse mesh
     * @param pCellPdeElementMap optional pointer to the map from cells to coarse elements
     */
    void virtual SetupSourceTerms(TetrahedralMesh<DIM,DIM>& rCoarseMesh, std::map<CellPtr, unsigned>* pCellPdeElementMap=NULL);

    /**
     * Overridden ComputeConstantInUSourceTerm() method.
     *
     * @param rX The point in space
     * @param pElement the element
     *
     * @return the constant in u part of the source term, i.e g(x) in
     *  Div(D Grad u)  +  f(x)u + g(x) = 0.
     */
    double ComputeConstantInUSourceTerm(const ChastePoint<DIM>& rX, Element<DIM,DIM>* pElement);

    /**
     * Overridden ComputeLinearInUCoeffInSourceTerm() method.
     *
     * @param rX The point in space
     * @param pElement the element
     *
     * @return the coefficient of u in the linear part of the source term, i.e f(x) in
     *  Div(D Grad u)  +  f(x)u + g(x) = 0.
     */
    double ComputeLinearInUCoeffInSourceTerm(const ChastePoint<DIM>& rX, Element<DIM,DIM>* pElement);

    /**
     * Overridden ComputeDiffusionTerm() method.
     *
     * @param rX The point in space at which the diffusion term is computed
     *
     * @return a matrix.
     */
    c_matrix<double,DIM,DIM> ComputeDiffusionTerm(const ChastePoint<DIM>& rX);

    /**
     * @return the uptake rate.
     *
     * @param elementIndex the element we wish to return the uptake rate for
     */
    double GetUptakeRateForElement(unsigned elementIndex);
};

#include "SerializationExportWrapper.hpp"
EXPORT_TEMPLATE_CLASS_SAME_DIMS(AveragedSourcePde)

namespace boost
{
namespace serialization
{
/**
 * Serialize information required to construct an AveragedSourcePde.
 */
template<class Archive, unsigned DIM>
inline void save_construct_data(
    Archive & ar, const AveragedSourcePde<DIM>* t, const unsigned int file_version)
{
    // Save data required to construct instance
    const AbstractCellPopulation<DIM>* p_cell_population = &(t->rGetCellPopulation());
    ar & p_cell_population;
}

/**
 * De-serialize constructor parameters and initialise an AveragedSourcePde.
 */
template<class Archive, unsigned DIM>
inline void load_construct_data(
    Archive & ar, AveragedSourcePde<DIM>* t, const unsigned int file_version)
{
    // Retrieve data from archive required to construct new instance
    AbstractCellPopulation<DIM>* p_cell_population;
    ar >> p_cell_population;

    // Invoke inplace constructor to initialise instance
    ::new(t)AveragedSourcePde<DIM>(*p_cell_population);
}
}
} // namespace ...

#endif /*AVERAGEDSOURCEPDE_HPP_*/