#include "process.h"
#include "data.h"
#include "linear.h"
#include "Sample.h"
#include "RIRL.h"
#include "DETree.h"

#include <iostream>
#include <cstdlib>
#include <vector>
#include <algorithm>
#include <random>
#include <chrono>
#include <fstream>
#include <limits>
#include <cfloat>
#include <boost/archive/text_oarchive.hpp>
#include <boost/archive/text_iarchive.hpp>
#include <boost/serialization/vector.hpp>

using namespace std;

Process::Process()
{

}

Process::~Process()
{

}

//each sample is a tuple <x,y>
vector<vector<double> > Process::getContinuousSamples(Data &data,vector<double> previousPoint,vector<double> currentPoint){
    int numberOfSamples = data.getNumberOfSamples();

    double startPointX = previousPoint.at(0);
    double startPointY = previousPoint.at(1);

    double endPointX = currentPoint.at(0);
    double endPointY = currentPoint.at(1);

    double deltaX = abs((startPointX - endPointX)) / ((double) (numberOfSamples));
    double deltaY = abs((startPointY - endPointY)) / ((double) (numberOfSamples));

    double xSample = startPointX;
    double ySample = startPointY;

    vector<vector<double> > samplesList;

    for (int i = 0 ; i < numberOfSamples ; i++){
        vector<double> sample;
        sample.push_back(xSample);
        sample.push_back(ySample);

        samplesList.push_back(sample);

        xSample = xSample + deltaX;

        if(deltaX != 0.0){
            ySample = ((startPointY - endPointY) / (startPointX - endPointX)) * xSample
                    + (startPointY - ((startPointY - endPointY) /
                                      (startPointX - endPointX)) * startPointX);
        }else if(deltaX == 0.0 && deltaY != 0.0){
            ySample = ySample + deltaY;
        }else{
            ySample = ySample;
        }
    }

    return samplesList;
}

double Process::getDistanceToPowerTwo(vector <double> sample){
    return (pow ((sample.at(0) - 10), 2) + pow ((sample.at(1) - 10), 2));
}

vector<double> Process::getObsSampleList(Data &data,vector<double> previousPoint,vector<double> currentPoint){
    //for debugging
    //    if(currentPoint == previousPoint){
    //        cout << "we are turnning!" << endl;
    //    }

    //list of <x,y> between two points
    vector< vector<double> > continuousSampleList = getContinuousSamples(data, previousPoint, currentPoint);

    vector<double> obsList;

    // parameters of the normal distribution
    double sigma = data.getSigma();
    double mean = data.getMean();
    double obs;
    double accNoise = 0.0;
    for (int i = 0 ; i < (int)continuousSampleList.size() ; i++){
        vector<double> sample = continuousSampleList.at(i);

        double r2 = getDistanceToPowerTwo(sample);
        double p = data.getP(); // presure in intensity formula

        unsigned seed = std::chrono::system_clock::now().time_since_epoch().count();
        default_random_engine generator(seed);
        normal_distribution<double> distribution(mean,sigma);
        double noise = distribution(generator);
        //        double noise = 0.0; // for debugging noise free
        //        cout << (p / (4 * M_PI * r2)) << " + " << noise << endl;

        obs = (p / (4 * M_PI * r2)) + noise;
        accNoise = accNoise + noise * noise;
        obsList.push_back(obs);
    }

    data.updateNoise(accNoise);
    return obsList;
}

//changed
// don't make continus trajectory any more just the discrete one. center to center move
void Process::bulidData(Data &data, int lenghtOfTrajectory, bool forTrainingObs){

    vector<double> initialPoint;
    vector<double> firstPoint;

    vector<int> firstState;
    int initialAction;
    int action;
    vector<int> initialState; //at 0 facing east
    initialState.push_back(0);
    initialState.push_back(0);

    if(forTrainingObs == true){
        initialAction = randomPolicy(data);
    }else{
        //call rational policy here
        map<vector<int>,map<int,double>> ExpertPolicy = data.getExpertPolicy();
        vector<int> listOfActions = data.getListOfActions();
        double max = -DBL_MAX;

        for(int j = 0 ; j < (int)(listOfActions.size()) ; j++ ){
            int a = listOfActions.at(j);
            if(max < ExpertPolicy[initialState][a]){
                max = ExpertPolicy[initialState][a];
                initialAction = a;
            }
        }
        //print the T for debugging
        cout << "[< " << initialState.at(0) << " , "
             << initialState.at(1) << " > , " <<
                initialAction << " ]" << endl;
    }
    initialPoint = getCenterPointInState(initialState.at(0));
    firstPoint = initialPoint;

    firstState = initialState;
    action = initialAction;

    int counter = 0;

    bool done = false;

    while(!done){
        vector<double> secondPoint;

        vector<int> secondState = transitionFunction(data, firstState, action);

        if(firstState == secondState){
            secondPoint = firstPoint;
        }else{
            secondPoint = getCenterPointInState(secondState.at(0));
        }

        vector<double> obs = getObsSampleList(data,firstPoint, secondPoint);
        updateFlatObsList(data,firstState,action,obs,forTrainingObs);

        firstState = secondState;
        firstPoint = secondPoint;
        counter++;

        if(counter == lenghtOfTrajectory){
            done = true;
        }else{
            if(forTrainingObs == true){
                action = randomPolicy(data);
            }else{
                //call rational policy here
                map<vector<int>,map<int,double>> ExpertPolicy = data.getExpertPolicy();
                vector<int> listOfActions = data.getListOfActions();
                double max = -DBL_MAX;

                for(int j = 0 ; j < (int)(listOfActions.size()) ; j++ ){
                    int a = listOfActions.at(j);
                    if(max < ExpertPolicy[firstState][a]){
                        max = ExpertPolicy[firstState][a];
                        action = a;
                    }
                }
                //print the T for debugging
                cout << "[< " << firstState.at(0) << " , "
                     << firstState.at(1) << " > , " <<
                        action << " ]" << endl;
            }
        }
    }
}

// copress that to above function
void Process::updateFlatObsList(Data &data, vector<int> state, int action, vector<double> obs, bool forTrainingObs){
    int numberOfSamples = data.getNumberOfSamples();

    vector<double> time_x = data.getTimeChunk();
    Sample obsTuple;
    Maths::Regression::Linear linearRegression(numberOfSamples, time_x, obs); //Intensity_y = Obs
    double slope =linearRegression.getSlope();
    obsTuple.values.push_back(slope);
    double intercept =linearRegression.getIntercept();
    obsTuple.values.push_back(intercept);
    if(forTrainingObs){
        obsTuple.values.push_back(state.at(0)); // adding state.position
        obsTuple.values.push_back(state.at(1)); // adding state.orientation
        obsTuple.values.push_back(action); // adding action
    }

    obsTuple.p = data.getNormalizerFactorForP(); //uniform
    data.updateFlatObsList(obsTuple);
}

//changed
//just consider the center
//rawstate = S_0 to S_n
vector<double> Process::getCenterPointInState(int rawState){
    vector<double> result;

    double x = (rawState + 1) * 20 +10;
    double y = 30;

    result.push_back(x);
    result.push_back(y);

    return result;
}

vector<int> Process::calcStateBoundaries(Data &data, int state){
    vector<int> result;

    int stepSize = data.getStepSize();
    int remainder = state % stepSize;
    int quotient = state / stepSize;

    int stateLowerX = remainder * stepSize;
    int stateUpperX = remainder  * stepSize + stepSize;
    int stateLowerY = quotient * stepSize - stepSize;
    int stateUpperY = quotient * stepSize;

    result.push_back(stateLowerX);
    result.push_back(stateUpperX);
    result.push_back(stateLowerY);
    result.push_back(stateUpperY);

    return result;
}

//return a random action
//done
int Process::randomPolicy(Data &data){
    vector<int> listOfActions = data.getListOfActions();
    int randomAction = listOfActions.at(getRandomIntNumber(listOfActions.size()));
    return randomAction;
}

//I have to chagne it later
double Process::probablityOfNextStateGivenCurrentStateAction(Data &data,vector<int> nextState, vector<int> currentState, int currentAction){
    double stochasticity = data.getStochasticity();
    vector<int> idealNextState = returnNextState(data,currentState,currentAction);
    double pr = 0.0;
    if (nextState == idealNextState){
        pr = 1.0 - stochasticity;
    }/*else if(areAdj(currentState,nextState) &&
             currentState.at(0) != 0 &&
             currentState.at(0) != data.getNumberOfRawStates()-1){ //if in the middle of hallway
        pr = stochasticity / 2.0; //confirm it with ken
    }else if(areAdj(currentState,nextState)){ //if at the end of the hallway
        pr = stochasticity / 1.0; //confirm it with ken
    }*/else if(nextState == currentState){
        pr = stochasticity;
    }

    return pr;
}

//changed
vector<int> Process::returnNextState(Data &data,vector<int> currentState, int currentAction){
    vector<int> nextState;
    if(/*(currentState.at(1) == 0 && currentAction == -1) ||*/ //facing east going backward
            (currentState.at(1) == 1 && currentAction == 1)){ //facing west going forward
        int nextPosition = currentState.at(0) - 1;
        //check if the next position in the range of possible positions
        if(nextPosition < 0){
            nextPosition = 0;
        }else if(nextPosition > data.getNumberOfRawStates() - 1){
            nextPosition = data.getNumberOfRawStates() - 1;
        }
        nextState.push_back(nextPosition);
        nextState.push_back(currentState.at(1));
    }else if(currentAction == 0){ //turning
        int nextOrientation;
        if(currentState.at(1) == 0){
            nextOrientation = 1;
        }else{
            nextOrientation = 0;
        }
        nextState.push_back(currentState.at(0));
        nextState.push_back(nextOrientation);
    }else if((currentState.at(1) == 0 && currentAction == 1) /*|| //facing east going forward
                                                                                (currentState.at(1) == 1 && currentAction == -1)*/){ //facing west going backward
        int nextPosition = currentState.at(0) + 1;
        //check if the next position (rawState) in the range of possible positions
        if(nextPosition < 0){
            nextPosition = 0;
        }else if(nextPosition > data.getNumberOfRawStates() - 1){
            nextPosition = data.getNumberOfRawStates() - 1;
        }
        nextState.push_back(nextPosition);
        nextState.push_back(currentState.at(1));
    }else{
        cout << "missing something!" << endl;
    }

    return nextState;
}

//changed
vector<int> Process::transitionFunction(Data &data,vector<int> currentState, int currentAction){
    vector<int> nextState;
    double stochasticity = data.getStochasticity();
    //    vector<int> listOfActions = data.getListOfActions();
    double stochasticityIndicator = getRandomDoubleNumber(0,1);

    if(stochasticityIndicator < stochasticity ){
        //        int randomAction = listOfActions.at(getRandomIntNumber(listOfActions.size()));
        //        while(randomAction == currentAction){
        //            randomAction = listOfActions.at(getRandomIntNumber(listOfActions.size()));
        //        }
        //        //generate next state here
        //        nextState = returnNextState(data, currentState, randomAction);
        nextState = currentState;
    }else{
        //generate next state here
        nextState = returnNextState(data, currentState, currentAction);
    }

    return nextState;
}

double Process::getRandomDoubleNumber(int lowerBound, int upperBound){
    double min = (double) lowerBound;
    double max = (double) upperBound;
    return (max - min) * ((double)rand() / (double)RAND_MAX) + min;
}

int Process::getRandomIntNumber(int upperBound){

    return  rand() % upperBound;
}

vector<double> Process::rowNormalizer(vector<int> row){
    vector<double> result;
    int sum = 0;
    for (int i = 0 ; i < (int)row.size() ; i++){
        sum = sum + row.at(i);
    }

    if (sum != 0){
        for (int i = 0 ; i < (int)row.size() ; i++){
            result.push_back((double) row.at(i) / (double) sum);
        }
    }else{
        for (int i = 0 ; i < (int)row.size() ; i++){
            result.push_back(0);
        }
    }

    return result;
}

void Process::trainObs(Data &data){
    vector<double> * low  = data.getLow();
    vector<double> * high = data.getHigh();
    DETree observationModel(data.getFlatObsList(), low, high);
    data.setObsModel(observationModel);
}

//changed
void Process::generateTrajectories(Data &data, int lenghtOfT, bool forTrainingObs){
    data.setSampleLength(lenghtOfT);
    //setting up expert's policy
    if(!forTrainingObs){
        map<vector<int>,map<int,double>> qValuesForExpert = qValueSoftMaxSolver(data, 0.001, data.getExpertWeights());
        constructExpertPolicy(data,qValuesForExpert);
    }
    cout << "Generating data!" << endl;
    data.setNormalizerFactorForP(lenghtOfT);
    bulidData(data, lenghtOfT, forTrainingObs);
    cout << "Data generated!" << endl;
    // claculating noise in relative mean square error fashion
    double rmsNoise = data.getNoise() / ((double)lenghtOfT * (double)data.getNumberOfSamples());
    cout << "Average Noise Added to observation = " << rmsNoise << endl;

    if(forTrainingObs){
        saveObsModel(data);
        cout << "Saving obsModel!\n";
        saveObsModel(data);
        cout << "obsModel saved!\n";
    }
}

//changed
vector<double> Process::getFeatures(Data &data, vector<int> state, int action){
    vector<double> features(data.getNumberOfFeatures(), 0.0);
    int rawState = state.at(0);
    int orientation = state.at(1);
    int activatedFeatureForTurningIndex;
    //    if(rawState == 9){
    //        cout << (data.getNumberOfRawStates() - 1) << endl;
    //    }
    if(action == 0){
        if (rawState < (data.getNumberOfRawStates() - 1) / 2 && orientation == 1){ //first half should face west
            activatedFeatureForTurningIndex = rawState;
            features.at(activatedFeatureForTurningIndex) = 1.0;
        }else if(rawState > (data.getNumberOfRawStates() - 1) / 2 && orientation == 0) { //second half should face east
            activatedFeatureForTurningIndex = data.getNumberOfRawStates() - rawState - 1;
            features.at(activatedFeatureForTurningIndex) = 1.0;
        }else if(rawState == (data.getNumberOfRawStates() - 1) / 2){
            activatedFeatureForTurningIndex = (data.getNumberOfRawStates() - 1) / 2;
            features.at(activatedFeatureForTurningIndex) = 1.0;
        }else{
            //            cout << "I missing some Feature!!" << endl;
            activatedFeatureForTurningIndex = features.size()-2;
            features.at(activatedFeatureForTurningIndex) = 1.0;
        }
    }else{
        //construct the feature vector here
        //last feature is for moving forward
        //I may change this part if i want to add moving backward!
        //I have to check if I am bangging the wall it should be no reward
        if(rawState == 0 || rawState == data.getNumberOfRawStates() - 1){
            //by not activating feature ==> no reward by going into the wall
            //            activatedFeatureForTurningIndex = features.size()-2;
            //            features.at(activatedFeatureForTurningIndex) = 1.0;
        }else{
            features.at(features.size()-1) = 1.0;
        }
    }

    return features;
}

//new
map<vector<int>,map<int,double>> Process::qValueSoftMaxSolver(Data &data, double err, vector<double> weights){

    map<vector<int>,double> V; //[State] //[State] ==> vector<int>
    map<vector<int>,map<int,double>> Q; //[State][Action] //[State] ==> vector<int> & [Action] ==> int
    vector<vector<int>> listOfStates = data.getListOfStates();
    vector<int> listOfActions = data.getListOfActions();

    for(int i = 0 ; i < (int)(listOfStates.size()) ; i++ ){ //foreach (State s ; model.S()) {
        vector<int> s = listOfStates.at(i);
        V[s] = 0.0;

        for(int j = 0 ; j < (int)(listOfActions.size()) ; j++ ){//foreach (a; model.A(s)) {
            int a = listOfActions.at(j);
            Q[s][a] = 0;
        }
    }
    //    V = V.rehash;
    //    Q = Q.rehash;
    double delta = 0;
    int iteration = 0; //size_t
    while (true) {
        delta = 0;

        for(int i = 0 ; i < (int)(listOfStates.size()) ; i++ ){ //foreach (State s ; model.S()) {
            vector<int> s = listOfStates.at(i);

            for(int j = 0 ; j < (int)(listOfActions.size()) ; j++ ){//foreach (a; model.A(s)) {
                int a = listOfActions.at(j);
                double r = reward(data,s,a,weights);//add reward function later
                //double[State] T = model.T(s, a);

                double expected_rewards = 0;
                for (int k = 0 ; k < int(listOfStates.size()) ; k++){ //foreach (s_prime, p; T){
                    vector<int> s_prime = listOfStates.at(k);
                    //p ==> probabilityOfTheNextState
                    double p  = probablityOfNextStateGivenCurrentStateAction(data,s_prime, s, a);
                    expected_rewards += p*V[s_prime];
                }

                Q[s][a] = r + data.getGamma() * expected_rewards;
            }
        }

        map<vector<int>,double> v;
        v = V;
        for(int i = 0 ; i < (int)(listOfStates.size()) ; i++ ){ //foreach (State s ; model.S()) {
            vector<int> s = listOfStates.at(i);
            double maxx = -DBL_MAX;
            for(int j = 0 ; j < (int)(listOfActions.size()) ; j++ ){//foreach (a; model.A(s)) {
                int a = listOfActions.at(j);
                maxx = max(maxx, Q[s][a]);
            }

            double e_sum = DBL_MIN;

            for(int j = 0 ; j < (int)(listOfActions.size()) ; j++ ){//foreach (a; model.A(s)) {
                int a = listOfActions.at(j);
                e_sum += exp(Q[s][a] - maxx);
            }

            v[s] = maxx - log(e_sum);
            delta = max(delta, abs(V[s] - v[s]));
        }

        V = v;
        //cout << "Current Iteration: " <<  delta << endl;

        if (delta < err || iteration > data.getMaxIter()) {
            map<vector<int>,map<int,double>> returnval; //[state][action]

            //doing this part to avoid overfelow problem when do e to the power in softmax!!
            for(int i = 0 ; i < (int)(listOfStates.size()) ; i++ ){ //foreach(state, actvalue; Q) {
                vector<int> state = listOfStates.at(i);

                for(int j = 0 ; j < (int)(listOfActions.size()) ; j++ ){//foreach(action, value; actvalue) {
                    int action = listOfActions.at(j);
                    double value = Q[state][action];
                    returnval[state][action] = value - v[state];
                }
            }
            //				writeln("Finished Q-Value");
            return returnval;
        }
        iteration ++;
    }
}

//new
//this function calculate the reward reciving state and action
double Process::reward(Data &data,vector<int> state, int action, vector<double> weights){
    double reward;
    vector<double> features = getFeatures(data,state,action);
    reward = calcInnerProduct(features,weights);
    return reward;
}

//construct the expert policy using the actual weights
//using softmax
void Process::constructExpertPolicy(Data &data, map<vector<int>,map<int,double>> Q_value){
    map<vector<int>,map<int,double>> policy; //[State][Action]
    vector<vector<int>> listOfStates = data.getListOfStates();
    vector<int> listOfActions = data.getListOfActions();

    for(int i = 0 ; i < (int)(listOfStates.size()) ; i++ ){ //foreach (State s ; model.S()) {
        vector<int> s = listOfStates.at(i);
        double normalizer = 0.0; //[State]

        for(int j = 0 ; j < (int)(listOfActions.size()) ; j++ ){//foreach (Action a; model.A(s)) {
            int a = listOfActions.at(j);
            normalizer = normalizer + exp(Q_value[s][a]);
        }

        for(int j = 0 ; j < (int)(listOfActions.size()) ; j++ ){//foreach (Action a; model.A(s)) {
            int a = listOfActions.at(j);
            policy[s][a] = exp(Q_value[s][a]) / normalizer;
        }
    }
    data.setExpertPolicy(policy);
}

double Process::calcInnerProduct(vector<double> v1, vector<double> v2){

    double result = 0.0;
    if(int(v1.size()) != int(v2.size())){
        cout << "Something is wrong in Process::calcInnerProduct" << endl;
        return 0.0;
    }
    for (int i = 0 ; i < int(v1.size()) ; i++){
        result = result + v1.at(i) * v2.at(i);
    }
    return result;
}

void Process::printPolicy(map<vector<int>,map<int,double>> policy){
    //    std::map< std::string, std::map<std::string, std::string> > m;

    for (auto i : policy){
        for (auto j : i.second){
            cout << " < " << i.first.at(0) << " , " << i.first.at(1) << " > , " << j.first << " ==> Pr = " << j.second << endl;
        }
        cout << "--------------------------------------------" << endl;
    }

}

vector<double> Process::multiply(double scalar, vector<double> v){
    vector<double> result;
    for (int i = 0 ; i < int(v.size()) ; i++){
        result.push_back(scalar * v.at(i));
    }
    return result;
}

vector<double> Process::divide(double scalar, vector<double> v){
    vector<double> result;
    for (int i = 0 ; i < int(v.size()) ; i++){
        result.push_back(v.at(i) / scalar);
    }
    return result;
}

vector<double> Process::add(vector<double> v1, vector<double> v2){
    vector<double> result;
    for (int i = 0 ; i < int(v1.size()) ; i++){
        result.push_back(v1.at(i) + v2.at(i));
    }
    return result;
}

double Process::l1norm(vector<double> v){
    double returnval = 0;
    for (int i = 0 ; i < int(v.size()) ; i++){
        returnval += abs(v.at(i));
    }
    return returnval;
}

double Process::l2norm(vector<double> v){
    return sqrt(calcInnerProduct(v,v));
}

void Process::loadObsModel(Data &data){
    vector<Sample> flatObsList;
    string fileAddress = data.getObsModelAddress();
    std::ifstream ifs(fileAddress);
    boost::archive::text_iarchive ia(ifs);
    ia >> flatObsList;
    vector<double> * low  = data.getLow();
    vector<double> * high = data.getHigh();
    DETree observationModel(flatObsList, low, high);
    data.setObsModel(observationModel);
}

void Process::saveObsModel(Data &data){
    vector<Sample> flatObsList = data.getFlatObsList();
    ofstream myfile;
    string fileAddress = data.getObsModelAddress();
    myfile.open (fileAddress.c_str());
    std::ofstream ofs(fileAddress.c_str());
    boost::archive::text_oarchive oa(ofs);
    // write class instance to archive
    oa << flatObsList;
}

//changed
bool  Process::areAdj(vector<int> currentState, vector<int> nextState){
    int rawState1 = currentState.at(0);
    int rawState2 = nextState.at(0);
    if(rawState1 - rawState2 == 0){
        return true;
    }else if(rawState2 - rawState1 == 1 && currentState.at(1) == nextState.at(1)){
        return true;
    }
    return false;
}

vector<double> Process::normalize(vector<double> v){
    double sum = 0;
    vector<double> result;
    for (int i = 0 ; i < int(v.size()) ; i++){
        sum = sum + v.at(i);
    }
    for (int i = 0 ; i < int(v.size()) ; i++){
        result.push_back(v.at(i) / sum);
    }
    return result;
}

vector<double> Process::getFeatureVectorOfT(Data &data,vector<vector<int>> t){
    vector<double> fv(data.getNumberOfFeatures(),0.0);
    for(int i = 0 ; i < int(t.size()) ; i++){
        vector<int> state;
        state.push_back(t.at(i).at(0));
        state.push_back(t.at(i).at(1));
        int action = t.at(i).at(2);
        fv = add(fv, getFeatures(data,state,action));
    }
    return fv;
}

void Process::saveDeterministicObsModel(Data &data){
    vector<vector<int>> listOfStates = data.getListOfStates();
    vector<int> listOfActions = data.getListOfActions();

    for(int i = 0 ; i < int(listOfStates.size()) ; i++){
        vector<int> firstState = listOfStates.at(i);
        for(int j = 0 ; j < int(listOfActions.size()) ; j++){
            int action = listOfActions.at(j);
            vector<int> secondState = returnNextState(data,firstState,action);
            vector<double> firstPoint = getCenterPointInState(firstState.at(0));
            vector<double> secondPoint = getCenterPointInState(secondState.at(0));
            vector<double> obs = getObsSampleList(data,firstPoint, secondPoint);
            updateFlatObsList(data,firstState,action,obs,true);
        }
    }
    vector<Sample> flatObsList = data.getFlatObsList();
    ofstream myfile;
    string fileAddress = data.getObsModelAddress();
    myfile.open (fileAddress.c_str());
    std::ofstream ofs(fileAddress.c_str());
    boost::archive::text_oarchive oa(ofs);
    // write class instance to archive
    oa << flatObsList;
}

void Process::loadDeterministicObsModel(Data &data){
    vector<Sample> dObsModel;
    string fileAddress = data.getObsModelAddress();
    std::ifstream ifs(fileAddress);
    boost::archive::text_iarchive ia(ifs);
    ia >> dObsModel;
    data.setDeterministicObsModel(dObsModel);
}

void Process::saveClusteringObsModel(Data &data, int dataPerStateActionPair){
    vector<vector<int>> listOfStates = data.getListOfStates();
    vector<int> listOfActions = data.getListOfActions();
    int numberOfSamples = data.getNumberOfSamples();
    vector<Sample> clusteringObsModel;
    for(int i = 0 ; i < int(listOfStates.size()) ; i++){
        vector<int> firstState = listOfStates.at(i);
        for(int j = 0 ; j < int(listOfActions.size()) ; j++){
            int action = listOfActions.at(j);
            Sample centroid;
            centroid.p = 0; //I not using p in here so just assign zero
            double slopeSum = 0.0;
            double intercepSum = 0.0;

            for(int k = 0 ; k < dataPerStateActionPair ; k++){
                // I was using returnNextState(data,firstState,action)
                // I have to use transitionFunction(data,firstState,action);
                // to handle stochastic T in observation
                //                vector<int> secondState = returnNextState(data,firstState,action);
                vector<int> secondState = transitionFunction(data,firstState,action);
                vector<double> firstPoint = getCenterPointInState(firstState.at(0));
                vector<double> secondPoint = getCenterPointInState(secondState.at(0));
                vector<double> obs = getObsSampleList(data,firstPoint, secondPoint);
                vector<double> time_x = data.getTimeChunk();

                Maths::Regression::Linear linearRegression(numberOfSamples, time_x, obs); //Intensity_y = Obs
                double slope =linearRegression.getSlope();
                slopeSum = slopeSum + slope;
                double intercept =linearRegression.getIntercept();
                intercepSum = intercepSum + intercept;
            }
            centroid.values.push_back(slopeSum / dataPerStateActionPair);
            centroid.values.push_back(intercepSum / dataPerStateActionPair);
            centroid.values.push_back(firstState.at(0));
            centroid.values.push_back(firstState.at(1));
            centroid.values.push_back(action);
            clusteringObsModel.push_back(centroid);
        }
    }

    ofstream myfile;
    string fileAddress = data.getObsModelAddress();
    myfile.open (fileAddress.c_str());
    std::ofstream ofs(fileAddress.c_str());
    boost::archive::text_oarchive oa(ofs);
    // write class instance to archive
    oa << clusteringObsModel;
}

void Process::loadClusteringObsModel(Data &data){
    vector<Sample> cObsModel;
    string fileAddress = data.getObsModelAddress();
    std::ifstream ifs(fileAddress);
    boost::archive::text_iarchive ia(ifs);
    ia >> cObsModel;
    data.setClusteringObsModel(cObsModel);
}


#include <boost/numeric/ublas/matrix_proxy.hpp>
#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/vector.hpp>
#include <boost/numeric/ublas/vector_proxy.hpp>

#include <cmath>
#include <cstring>
#include <memory>
#include <vector>
#include <limits>
#include <algorithm>
#include <utility>
#include <iterator>

#include "statespace.hpp"

using namespace boost::numeric::ublas;

class ChangeEvent {
  /* struct (supporting comparison) representing
     a change-of-active-SSMs event. */
public:
  int t;
  int i_ssm;
  bool is_start;
  friend bool operator< (const ChangeEvent &c1, const ChangeEvent &c2);
};

bool operator< (const ChangeEvent &c1, const ChangeEvent &c2) {
    if (c1.t < c2.t) {
      return true;
    } else if (c1.t == c2.t) {
      // ends come before starts
      return c1.is_start > c2.is_start;
    }
    return false;
  }

int active_set_dimension(const std::vector<StateSpaceModel *> & ssms, vector<int> & active_set) {
  vector<int>::const_iterator it;
  int dimension = 0;
  for (it = active_set.begin(); it < active_set.end(); ++it) {
    if (*it < 0) continue;
    StateSpaceModel *ssm = ssms[*it];
    if (!ssm) continue;
    dimension += ssm->max_dimension;
  }
  return dimension;
}

TransientCombinedSSM::TransientCombinedSSM(\
       std::vector<StateSpaceModel *> & ssms, const vector<int> & start_idxs,
       const vector<int> & end_idxs, const std::vector<const double * > & scales,
       double obs_noise) : ssms(ssms), start_idxs(start_idxs),
			   end_idxs(end_idxs), scales(scales),
			   obs_noise(obs_noise), n_ssms(ssms.size()),
			   active_ssm_cache1_k(-1), active_ssm_cache2_k(-1) {

  this->ssms_tmp.resize(this->n_ssms);
  this->n_steps = *(max_element(end_idxs.begin(), end_idxs.end()));

  /* Compute a list of ChangeEvents, each representing the
   * activation or deactivation of a component SSM. The
   * list is sorted by the timestep at which the event
   * occurs (and secondarily, with deactivation events at
   * a timestep before activation events). */
  std::vector<ChangeEvent> events(start_idxs.size() + end_idxs.size());
  for (unsigned i=0; i < n_ssms; ++i) {
    events[2*i].t = start_idxs[i];
    events[2*i].i_ssm = i;
    events[2*i].is_start = true;

    events[2*i+1].t = end_idxs[i];
    events[2*i+1].i_ssm = i;
    events[2*i+1].is_start = false;
  }
  std::sort(events.begin(), events.end());

  /*
     Build a sorted list of changepoints (timesteps at which the set of active
     SSMs changes), along with the set of active SSMs at each changepoint.
     Also compute the max total dimension of the state space, by computing the
     dimension of each active set as we add it to the active_sets matrix.
   */
  // allocate the active_sets matrix by first computing the largest number of
  // SSMs that might be active at one time
  unsigned int n_active = 0;
  unsigned int max_active = 0;
  for(std::vector<ChangeEvent>::const_iterator idx = events.begin(); idx < events.end(); ++idx) {
    if (idx->is_start) n_active++; else n_active--;
    if (n_active > max_active) max_active = n_active;
  }
  active_sets.resize(2*n_ssms, max_active);

  vector<int> active_set(max_active);
  for (unsigned i=0; i < max_active; ++i) active_set(i) = -1;
  int t_prev = events[0].t, max_dimension = 0, i=0;
  std::vector<ChangeEvent>::const_iterator idx;
  for(i=0, idx = events.begin();
      idx < events.end(); ++idx) {

    // printf("event t %d i_ssm %d start %s\n", idx->t, idx->i_ssm, idx->is_start ? "true" : "false");

    // if this represents a change, add a new changepoint
    if (idx->t != t_prev) {
      changepoints.push_back(t_prev);

      // printf(" change detected (t_prev %d), copying to active_set %d\n", t_prev, i);

      // copy the current active set of SSMs
      // into the active_sets matrix (terminated with -1)
      for (unsigned k=0, j=0; k < max_active; ++k) {
	if (active_set(k) >= 0) {
	  active_sets(i, j++) = active_set(k);
	}
	if (j < max_active) {
	  active_sets(i, j) = -1;
	}
      }

      /*printf("matr active set(%d, ...) is", i);
      for(int tmpi=0; tmpi < max_active; ++tmpi) {
	printf(" %d", active_sets(i, tmpi));
      }
      printf("\n");*/

      t_prev = idx->t;
      i++;

      int current_dim = active_set_dimension(this->ssms, active_set);
      if (current_dim > max_dimension) max_dimension = current_dim;
    }

    // update the active set: add the new SSM if this
    // is a start event, otherwise remove it from
    // the active set.
    vector<int>::iterator it;
    if (idx->is_start) {
      it = std::find (active_set.begin(), active_set.end(), -1);
      if (it == active_set.end()) {
	printf("ERROR: need to start new SSM but no room in active_set vector!\n");
	exit(-1);
      } else {
	*it = idx->i_ssm;
      }
    } else {
      it = std::find (active_set.begin(), active_set.end(), idx->i_ssm);
      if (it == active_set.end()) {
	printf("ERROR: trying to remove SSM %d but it is not present in the active_set vector!\n", idx->i_ssm);
	exit(-1);
      } else {
	*it = -1;
      }
    }

  }

  this->max_dimension = max_dimension;
  this->is_cssm = false;
}

TransientCombinedSSM::~TransientCombinedSSM() {

  return;
};

int TransientCombinedSSM::active_set_idx(int k) {
  if (k == this->active_ssm_cache1_k)
    return this->active_ssm_cache1_v;
  else if (k == this->active_ssm_cache2_k)
    return this->active_ssm_cache2_v;

  std::vector<int>::iterator it = std::upper_bound(this->changepoints.begin(), this->changepoints.end(), k);
  int i = it - this->changepoints.begin() - 1;

  /*printf("active set at time k=%d is %d, changepoints %d %d %d %d\n", k, i,
	 i-1 >= 0 ? this->changepoints[i-1] : -1,
	 i >= 0 ? this->changepoints[i] : -1,
	 i+1 < this->changepoints.size() ? this->changepoints[i+1] : -1,
	 i+2 < this->changepoints.size() ? this->changepoints[i+2] : -1);*/

  this->active_ssm_cache2_k = this->active_ssm_cache1_k;
  this->active_ssm_cache2_v = this->active_ssm_cache1_v;
  this->active_ssm_cache1_k = k;
  this->active_ssm_cache1_v = i;
  return i;
}

unsigned int TransientCombinedSSM::state_size_at_timestep(unsigned int k) {

  unsigned int total_state_size = 0;

  int asidx = this->active_set_idx(k);
  if (asidx < 0) {
    return 0;
  }

  matrix_row < matrix<int> > ssm_indices = row(this->active_sets, asidx);
  for (matrix_row < matrix<int> >::const_iterator it = ssm_indices.begin();
       it < ssm_indices.end() && *it >= 0; ++it) {

    int i_ssm = *it;
    StateSpaceModel * ssm = this->ssms[i_ssm];
    if (!ssm) continue;

    unsigned int state_size = ssm->max_dimension;
    total_state_size += state_size;
  }
  return total_state_size;
}

int TransientCombinedSSM::apply_transition_matrix(const double * x, int k, double * result) {

  // first, loop over ssms active at the *previous*
  // timestep in order to cache the location of each
  // ssm in the previous state space.
  int j = 0;
  int asidx_prev = this->active_set_idx(k-1);
  int asidx = this->active_set_idx(k);
  bool same_active_set = (asidx==asidx_prev);

  if (k == 0) {
    printf("ERROR: requested transition INTO timestep 0 (invalid)!");
    exit(-1);
  }

  if (asidx < 0) {
    return 0;
  }

  if (!same_active_set) {
    matrix_row < matrix<int> > old_ssm_indices = row(this->active_sets, asidx_prev);
    for (matrix_row < matrix<int> >::const_iterator it = old_ssm_indices.begin();
	 it < old_ssm_indices.end() && *it >= 0; ++it) {

      // skip any null SSMs
      int i_ssm = *it;
      StateSpaceModel * ssm = this->ssms[i_ssm];
      if (!ssm) continue;

      this->ssms_tmp[i_ssm] = j;
      j += ssm->max_dimension;
    }
  }

  //printf("transition to time %d (asidx %d, %d):\n",  k, asidx_prev, asidx);

  // now apply the transition to the current time
  int i=0;
  matrix_row < matrix<int> > ssm_indices = row(this->active_sets, asidx);
  for (matrix_row < matrix<int> >::const_iterator it = ssm_indices.begin();
       it < ssm_indices.end() && *it >= 0; ++it) {

    // skip any null SSMs
    int i_ssm = *it;
    StateSpaceModel * ssm = this->ssms[i_ssm];
    if (!ssm) continue;

    unsigned int state_size = ssm->max_dimension;
    if (this->start_idxs[i_ssm] == k) {
      /* new ssms just get filled in as zero
         (prior means will be added by the
         transition_bias operator) */
      for (unsigned j=i; j < i+state_size; ++j) {
	result[j] = 0;
      }
      //printf("   new ssm %d active from %d to %d\n", i_ssm, i, i+state_size);

    } else {
      /* this ssm is persisting from the
       * previous timestep, so just run the
       * transition */
      unsigned int j = same_active_set ? i : this->ssms_tmp[i_ssm];

      ssm->apply_transition_matrix(x+j, k-this->start_idxs[i_ssm], result+i);
      //printf("   transitioning ssm %d, prev %d, in state %d to %d (sidx %d eidx %d)\n", i_ssm, j, i, i+state_size, this->start_idxs[i_ssm], this->end_idxs[i_ssm]);
    }
    i += state_size;
  }
  return i;
}

int TransientCombinedSSM::apply_transition_matrix( const matrix<double> &X,
						   unsigned int x_row_offset,
						   int k,
						   matrix<double> &result,
						   unsigned int r_row_offset,
						   unsigned int n) {
  int j = x_row_offset;
  int asidx_prev = this->active_set_idx(k-1);
  int asidx = this->active_set_idx(k);
  bool same_active_set = (asidx==asidx_prev);

  if (asidx < 0) {
    return 0;
  }

  if (!same_active_set) {
    matrix_row < matrix<int> > old_ssm_indices = row(this->active_sets, asidx_prev);
    for (matrix_row < matrix<int> >::const_iterator it = old_ssm_indices.begin();
	 it < old_ssm_indices.end() && *it >= 0; ++it) {

      // skip any null SSMs
      int i_ssm = *it;
      StateSpaceModel * ssm = this->ssms[i_ssm];
      if (!ssm) continue;

      this->ssms_tmp[i_ssm] = j;
      j += ssm->max_dimension;
    }
  }

  int i=x_row_offset;
  matrix_row < matrix<int> > ssm_indices = row(this->active_sets, asidx);
  for (matrix_row < matrix<int> >::const_iterator it = ssm_indices.begin();
       it < ssm_indices.end() && *it >= 0; ++it) {

    int i_ssm = *it;
    StateSpaceModel * ssm = this->ssms[i_ssm];
    if (!ssm) continue;

    unsigned int state_size = ssm->max_dimension;
    if (this->start_idxs[i_ssm] == k) {
      /* new ssms just get filled in as zero
         (prior means will be added by the
         transition_bias operator) */
      for (unsigned j=i; j < i+state_size; ++j) {
	for (unsigned jj=0; jj < n; ++jj) result(j, jj) = 0;
      }
    } else {
      unsigned int j = same_active_set ? i : this->ssms_tmp[i_ssm];
      //printf("MATR transitioning ssm %d, prev %d, in state %d to %d (sidx %d eidx %d)\n", i_ssm, j, i, i+state_size, this->start_idxs[i_ssm], this->end_idxs[i_ssm]);
      ssm->apply_transition_matrix(X, j, k-this->start_idxs[i_ssm], result, i, n);
    }
    i += state_size;
  }
  return i-x_row_offset;
}

void TransientCombinedSSM::transition_bias(int k, double *result) {

  int asidx = this->active_set_idx(k);
  if (asidx < 0) return;

  matrix_row < matrix<int> > ssm_indices = row(this->active_sets, asidx);
  for (matrix_row < matrix<int> >::const_iterator it = ssm_indices.begin();
       it < ssm_indices.end() && *it >= 0; ++it) {

    // skip any null SSMs
    int j = *it;
    StateSpaceModel * ssm = this->ssms[j];
    if (!ssm) continue;

    if (this->start_idxs[j] == k) {
      ssm->prior_mean(result);
    } else {
      ssm->transition_bias(k-this->start_idxs[j], result);
    }
    result += ssm->max_dimension;
  }
}

void TransientCombinedSSM::transition_noise_diag(int k, double *result) {
  int asidx = this->active_set_idx(k);
  if (asidx < 0) return;
  matrix_row < matrix<int> > ssm_indices = row(this->active_sets, asidx);
  for (matrix_row < matrix<int> >::const_iterator it = ssm_indices.begin();
       it < ssm_indices.end() && *it >= 0; ++it) {

    // skip any null SSMs
    int j = *it;
    StateSpaceModel * ssm = this->ssms[j];
    if (!ssm) continue;

    if (this->start_idxs[j] == k) {
      ssm->prior_vars(result);
    } else {
      ssm->transition_noise_diag(k-this->start_idxs[j], result);
    }
    result += ssm->max_dimension;
  }
}

double TransientCombinedSSM::apply_observation_matrix(const double *x, int k) {

  double r = 0;

  int asidx = this->active_set_idx(k);
  if (asidx < 0) return 0.0;
  matrix_row < matrix<int> > ssm_indices = row(this->active_sets, asidx);
  for (matrix_row < matrix<int> >::const_iterator it = ssm_indices.begin();
       it < ssm_indices.end() && *it >= 0; ++it) {

    int j = *it;
    StateSpaceModel * ssm = this->ssms[j];
    if (!ssm) continue;
    const double * scale = this->scales[j];
    double ri = ssm->apply_observation_matrix(x, k - this->start_idxs[j]);
    if (scale) ri *= scale[k - this->start_idxs[j]];
    r += ri;
    if (ssm) x += ssm->max_dimension;
  }
  return r;
}

void TransientCombinedSSM::apply_observation_matrix(const matrix<double> &X,
						    unsigned int row_offset, int k,
						    double *result, double *result_tmp, unsigned int n) {

  for (unsigned i=0; i < n; ++i) {
    result[i] = 0;
    result_tmp[i] = 0;
  }

  int asidx = this->active_set_idx(k);
  if (asidx < 0) return;

  matrix_row < matrix<int> > ssm_indices = row(this->active_sets, asidx);
  for (matrix_row < matrix<int> >::const_iterator it = ssm_indices.begin();
       it < ssm_indices.end() && *it >= 0; ++it) {
    int j = *it;
    StateSpaceModel * ssm = this->ssms[j];
    if (!ssm) continue;
    const double * scale = this->scales[j];

    unsigned int state_size = ssm->max_dimension;
    ssm->apply_observation_matrix(X, row_offset,
				  k-this->start_idxs[j],
				  result_tmp, NULL, n);
    // printf("TSSM step %d applying obs matrix on ssm %d state_size %d start_idx %d at row_offset %d n %d scale[%d] %f result[0] %f\n", k, j, state_size, this->start_idxs[j], row_offset, n, k-this->start_idxs[j], scale? scale[k-this->start_idxs[j]] : 1.0, result_tmp[0]);
    if (scale) {
      for (unsigned ii=0; ii < n; ++ii) {
	result[ii] += scale[k-this->start_idxs[j]] * result_tmp[ii];
      }
    } else {
      for (unsigned ii=0; ii < n; ++ii) {
	result[ii] += result_tmp[ii];
      }
    }
    row_offset += state_size;
  }

}

double TransientCombinedSSM::observation_bias(int k) {
  double bias = 0;

  int asidx = this->active_set_idx(k);
  if (asidx < 0) return 0.0;
  matrix_row < matrix<int> > ssm_indices = row(this->active_sets, asidx);
  for (matrix_row < matrix<int> >::const_iterator it = ssm_indices.begin();
       it < ssm_indices.end() && *it >= 0; ++it) {

    // skip any null SSMs
    int j = *it;
    StateSpaceModel * ssm = this->ssms[j];
    const double * scale = this->scales[j];
    int kk = k - this->start_idxs[j];

    double b = ssm ? ssm->observation_bias(kk) : 1.0;
    if (scale) b *= scale[kk];
    bias += b;
  }
  return bias;
}

double TransientCombinedSSM::observation_noise(int k) {
  return this->obs_noise;
}

bool TransientCombinedSSM::stationary(int k) {

  matrix_row < matrix<int> > s1 = row(this->active_sets, this->active_set_idx(k));
  if (k > 0) {
    matrix_row < matrix<int> > s2 = row(this->active_sets, this->active_set_idx(k-1));
    for (unsigned i=0; i < s1.size() && !(s1(i) == -1 && s2(i) == -1); ++i) {
      if (s1(i) != s2(i)) return false;
    }
  }

  for (matrix_row < matrix<int> >::const_iterator it = s1.begin();
       it < s1.end() && *it >= 0; ++it) {
    // skip any null SSMs
    int j = *it;
    if (this->scales[j]) return false;

    StateSpaceModel * ssm = this->ssms[j];
    if (ssm && !ssm->stationary(k-this->start_idxs[j])) return false;
  }
  return true;
}

int TransientCombinedSSM::prior_mean(double *result) {

  /*
     TODO: propagate through the forward model whenever
   */

  double * r1 = result;
  int asidx = this->active_set_idx(0);
  if (asidx < 0) return 0;
  matrix_row < matrix<int> > ssm_indices = row(this->active_sets, asidx);
  for (matrix_row < matrix<int> >::const_iterator it = ssm_indices.begin();
       it < ssm_indices.end() && *it >= 0; ++it) {

    int j = *it;
    StateSpaceModel * ssm = this->ssms[j];
    if (!ssm) continue;

    int state_size = ssm->max_dimension;
    for(int i=0; i < state_size; ++i) {
      result[i] = 0;
    }
    ssm->prior_mean(result);

    /* if this component starts at a negative time, push the
     * prior through the transition model to get the induced
     * distribution at time zero */
    if (this->start_idxs[j] < 0) {
      double * tmp_result = (double *) malloc(sizeof(double) * ssm->max_dimension);
      if (tmp_result==NULL) {
	printf("memory error, malloc failed!\n");
	exit(-1);
      }
      for (unsigned k=1; k <= -this->start_idxs[j]; ++k) {
	ssm->apply_transition_matrix(result, k, tmp_result);
	ssm->transition_bias(k, tmp_result);
	memcpy(result, tmp_result, ssm->max_dimension);
      }
      free(tmp_result);
    }

    result += state_size;
  }
  return result-r1;
}

int TransientCombinedSSM::prior_vars(double *result) {
  double * r1 = result;
  int asidx = this->active_set_idx(0);
  if (asidx < 0) return 0;
  matrix_row < matrix<int> > ssm_indices = row(this->active_sets, asidx);
  for (matrix_row < matrix<int> >::const_iterator it = ssm_indices.begin();
       it < ssm_indices.end() && *it >= 0; ++it) {

    int j = *it;
    StateSpaceModel * ssm = this->ssms[j];
    if (!ssm) continue;
    ssm->prior_vars(result);

    /*printf("ssm %d max dimension %d\n", j, ssm->max_dimension);
    for (double * q = result; q < result+ssm->max_dimension; ++q) {
      printf(" initial var %d = %f\n", q-result, *q);
      }*/

    if (this->start_idxs[j] < 0) {
      matrix<double> P (ssm->max_dimension, ssm->max_dimension);
      matrix<double> P2 (ssm->max_dimension, ssm->max_dimension);
      P.clear();
      for (int i=0; i < ssm->max_dimension; ++i) {
	P(i,i) = result[i];
      }

      for (unsigned k=1; k <= -this->start_idxs[j]; ++k) {
	ssm->apply_transition_matrix(P, 0, k, P2, 0, ssm->max_dimension); // P2 = FP
	noalias(P) = trans(P2); // P = PF'
	ssm->apply_transition_matrix(P, 0, k, P2, 0, ssm->max_dimension); // P2 = FPF'
	P = P2;
	ssm->transition_noise_diag(k, result);
	for (int i=0; i < ssm->max_dimension; ++i) P(i,i) += result[i];
      }
      for (int i=0; i < ssm->max_dimension; ++i) result[i] = P(i,i);
    }

    /*for (double * q = result; q < result+ssm->max_dimension; ++q) {
      printf(" post-prop var %d = %f\n", q-result, *q);
      }*/

    result += ssm->max_dimension;
  }

  /*for (double * q = r1; q < result; ++q) {
    printf("prior var %d = %f\n", q-r1, *q);
    }*/

  return result-r1;
}

void TransientCombinedSSM::init_coef_priors(std::vector<vector<double> > & cmeans,
					    std::vector<vector<double> > & cvars) {
  for (unsigned j=0; j < this->n_ssms; ++j) {
    StateSpaceModel * ssm = this->ssms[j];
    if (ssm && ssm->is_cssm) {
      CompactSupportSSM *cssm = (CompactSupportSSM *) ssm;
      vector<double> mean(cssm->coef_means);
      vector<double> var(cssm->coef_vars);
      cmeans.push_back(mean);
      cvars.push_back(var);
    } else {
      vector<double> mean;
      vector<double> var;
      cmeans.push_back(mean);
      cvars.push_back(var);
    }
  }
}

void TransientCombinedSSM::extract_all_coefs(FilterState &cache, int k,
					     std::vector<vector<double> > & cmeans,
					     std::vector<vector<double> > & cvars) {
  /*
    Assumes cache has a valid, current P matrix.
   */
  int asidx = this->active_set_idx(k);
  if (asidx < 0) return;
  matrix_row < matrix<int> > ssm_indices = row(this->active_sets, asidx);
  unsigned int state_offset = 0;
  for (matrix_row < matrix<int> >::const_iterator it = ssm_indices.begin();
       it < ssm_indices.end() && *it >= 0; ++it) {

    int j = *it;
    StateSpaceModel * ssm = this->ssms[j];
    if (!ssm) continue;
    if (ssm->is_cssm) {
      CompactSupportSSM *cssm = (CompactSupportSSM *) ssm;
      cssm->extract_coefs(cache.xk, cache.P,
			 state_offset, k - this->start_idxs[j],
			 cmeans[j],
			 cvars[j]);
    }
    state_offset += ssm->max_dimension;
  }
}

void TransientCombinedSSM::extract_component_means(double *xk, int k,
						   std::vector<vector<double> > & means) {
  int asidx = this->active_set_idx(k);
  if (asidx < 0) return;
  matrix_row < matrix<int> > ssm_indices = row(this->active_sets, asidx);
  unsigned int state_offset = 0;
  for (matrix_row < matrix<int> >::const_iterator it = ssm_indices.begin();
       it < ssm_indices.end() && *it >= 0; ++it) {

    int j = *it;
    StateSpaceModel * ssm = this->ssms[j];
    int kk = k - this->start_idxs[j];
    means[j](kk) = ssm ? ssm->apply_observation_matrix(xk + state_offset, kk) : 1.0;
    means[j](kk) += ssm ? ssm->observation_bias(kk) : 0.0;
    if (ssm) state_offset += ssm->max_dimension;
  }
}

void TransientCombinedSSM::extract_component_vars(matrix<double> &P, matrix<double> &P_tmp, int k,
						  std::vector<vector<double> > & vars) {
  int asidx = this->active_set_idx(k);
  if (asidx < 0) return;
  matrix_row < matrix<int> > ssm_indices = row(this->active_sets, asidx);
  unsigned int state_offset = 0;

  double * v_tmp = (double *) malloc(sizeof(double) * P.size1());
  double * v_tmp2 = (double *) malloc(sizeof(double) * P.size1());
  if (!v_tmp || !v_tmp2) {
    printf("memory allocation error in extract_component_vars!");
    exit(-1);
  }

  for (matrix_row < matrix<int> >::const_iterator it = ssm_indices.begin();
       it < ssm_indices.end() && *it >= 0; ++it) {

    int j = *it;
    StateSpaceModel * ssm = this->ssms[j];
    int kk = k - this->start_idxs[j];

    if (ssm) {
      // just be lazy and copy the submatrix for this component into its own matrix.
      subrange(P_tmp, 0, ssm->max_dimension, 0, ssm->max_dimension) = subrange(P, state_offset, state_offset+ssm->max_dimension, state_offset, state_offset+ssm->max_dimension);
      ssm->apply_observation_matrix(P_tmp, 0, kk, v_tmp, v_tmp2, ssm->max_dimension);
      vars[j](kk) = ssm->apply_observation_matrix(v_tmp, kk);
      state_offset += ssm->max_dimension;
      //printf("%d: set var at ssm %d kk %d = %f\n", k, j, kk, vars[j](kk));
    } else {
      vars[j](kk) = 0.0;
    }
  }

  free(v_tmp);
  free(v_tmp2);
}


/* pcmsolver_copyright_start */
/*
 *     PCMSolver, an API for the Polarizable Continuum Model
 *     Copyright (C) 2013-2016 Roberto Di Remigio, Luca Frediani and contributors
 *
 *     This file is part of PCMSolver.
 *
 *     PCMSolver 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 3 of the License, or
 *     (at your option) any later version.
 *
 *     PCMSolver 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 PCMSolver.  If not, see <http://www.gnu.org/licenses/>.
 *
 *     For information on the complete list of contributors to the
 *     PCMSolver API, see: <http://pcmsolver.readthedocs.io/>
 */
/* pcmsolver_copyright_end */

#ifndef GREENUTILS_HPP
#define GREENUTILS_HPP

#include "Config.hpp"

#include <Eigen/Core>

#include "DerivativeTypes.hpp"
#include "utils/Stencils.hpp"

namespace GreenUtils {
/*! Returns value of the directional derivative of the function passed for the pair of points p1, p2:
 *  \f$ \nabla_{\mathbf{p_2}}G(\mathbf{p}_1, \mathbf{p}_2)\cdot \mathbf{n}_{\mathbf{p}_2}\f$
 *  Notice that this method returns the directional derivative with respect
 *  to the probe point.
 *  \tparam DerivativeTraits evaluation strategy for the function and its derivatives
 *  \param[in] functor   function-object handle to the evaluation of the differentiable function
 *  \param[in] normal_p2 the normal vector to p2
 *  \param[in]        p1 first point
 *  \param[in]        p2 second point
 */
template <typename DerivativeTraits>
double derivativeProbe(const pcm::function<DerivativeTraits(DerivativeTraits *, DerivativeTraits *)> & functor,
                       const Eigen::Vector3d & normal_p2, const Eigen::Vector3d & p1, const Eigen::Vector3d & p2)
{
    DerivativeTraits t1[3], t2[3], der;
    t1[0] = p1(0); t1[1] = p1(1); t1[2] = p1(2);
    t2[0] = p2(0); t2[1] = p2(1); t2[2] = p2(2);
    t2[0][1] = normal_p2(0); t2[1][1] = normal_p2(1); t2[2][1] = normal_p2(2);
    der = functor(t1, t2);
    return der[1];
}

/*! Returns value of the directional derivative of the function passed for the pair of points p1, p2:
 *  \f$ \nabla_{\mathbf{p_1}}G(\mathbf{p}_1, \mathbf{p}_2)\cdot \mathbf{n}_{\mathbf{p}_1}\f$
 *  Notice that this method returns the directional derivative with respect
 *  to the source point.
 *  \tparam DerivativeTraits evaluation strategy for the function and its derivatives
 *  \param[in] functor   function-object handle to the evaluation of the differentiable function
 *  \param[in] normal_p1 the normal vector to p1
 *  \param[in]        p1 first point
 *  \param[in]        p2 second point
 */
template <typename DerivativeTraits>
double derivativeSource(const pcm::function<DerivativeTraits(DerivativeTraits *, DerivativeTraits *)> & functor,
                        const Eigen::Vector3d & normal_p1, const Eigen::Vector3d & p1, const Eigen::Vector3d & p2)
{
    DerivativeTraits t1[3], t2[3], der;
    t1[0] = p1(0); t1[1] = p1(1); t1[2] = p1(2);
    t1[0][1] = normal_p1(0); t1[1][1] = normal_p1(1); t1[2][1] = normal_p1(2);
    t2[0] = p2(0); t2[1] = p2(1); t2[2] = p2(2);
    der = functor(t1, t2);
    return der[1];
}

/*! Returns value of the directional derivative of the function passed for the pair of points p1, p2:
 *  \f$ \nabla_{\mathbf{p_2}}G(\mathbf{p}_1, \mathbf{p}_2)\cdot \mathbf{n}_{\mathbf{p}_2}\f$
 *  Notice that this method returns the directional derivative with respect
 *  to the probe point.
 *  \param[in] functor   function-object handle to the evaluation of the differentiable function
 *  \param[in] normal_p2 the normal vector to p2
 *  \param[in]        p1 first point
 *  \param[in]        p2 second point
 */
double derivativeProbe(const DifferentiableFunction & functor, const Eigen::Vector3d & normal_p2, const Eigen::Vector3d & p1, const Eigen::Vector3d & p2)
{
    return threePointStencil(pcm::bind(functor, pcm::_1, _2), p2, p1, normal_p2);
}

/*! Returns value of the directional derivative of the function passed for the pair of points p1, p2:
 *  \f$ \nabla_{\mathbf{p_1}}G(\mathbf{p}_1, \mathbf{p}_2)\cdot \mathbf{n}_{\mathbf{p}_1}\f$
 *  Notice that this method returns the directional derivative with respect
 *  to the source point.
 *  \param[in] functor   function-object handle to the evaluation of the differentiable function
 *  \param[in] normal_p1 the normal vector to p1
 *  \param[in]        p1 first point
 *  \param[in]        p2 second point
 */
double derivativeSource(const DifferentiableFunction & functor, const Eigen::Vector3d & normal_p1, const Eigen::Vector3d & p1, const Eigen::Vector3d & p2)
{
        return threePointStencil(pcm::bind(functor,_1, _2), p1, p2, normal_p1);
}

/*! Returns full gradient of the function passed for the pair of points p1, p2:
 *  \f$ \nabla_{\mathbf{p_1}}G(\mathbf{p}_1, \mathbf{p}_2)\f$
 *  Notice that this method returns the gradient with respect to the source point.
 *  \tparam DerivativeTraits evaluation strategy for the function and its derivatives
 *  \param[in] functor   function-object handle to the evaluation of the differentiable function
 *  \param[in] p1 first point
 *  \param[in] p2 second point
 */
template <typename DerivativeTraits>
Eigen::Vector3d gradientSource(const pcm::function<DerivativeTraits(DerivativeTraits *, DerivativeTraits *)> & functor,
                               const Eigen::Vector3d & p1, const Eigen::Vector3d & p2)
{
    return (Eigen::Vector3d() << derivativeSource(functor, Eigen::Vector3d::UnitX(), p1, p2),
                                 derivativeSource(functor, Eigen::Vector3d::UnitY(), p1, p2),
                                 derivativeSource(functor, Eigen::Vector3d::UnitZ(), p1, p2)).finished();
}

/*! Returns full gradient of the function passed for the pair of points p1, p2:
 *  \f$ \nabla_{\mathbf{p_2}}G(\mathbf{p}_1, \mathbf{p}_2)\f$
 *  Notice that this method returns the gradient with respect to the probe point.
 *  \tparam DerivativeTraits evaluation strategy for the function and its derivatives
 *  \param[in] functor   function-object handle to the evaluation of the differentiable function
 *  \param[in] p1 first point
 *  \param[in] p2 second point
 */
template <typename DerivativeTraits>
Eigen::Vector3d gradientProbe(const pcm::function<DerivativeTraits(DerivativeTraits *, DerivativeTraits *)> & functor,
                              const Eigen::Vector3d & p1, const Eigen::Vector3d & p2)
{
    return (Eigen::Vector3d() << derivativeProbe(functor, Eigen::Vector3d::UnitX(), p1, p2),
                                 derivativeProbe(functor, Eigen::Vector3d::UnitY(), p1, p2),
                                 derivativeProbe(functor, Eigen::Vector3d::UnitZ(), p1, p2)).finished();
}

/*! Returns full gradient of the function passed for the pair of points p1, p2:
 *  \f$ \nabla_{\mathbf{p_1}}G(\mathbf{p}_1, \mathbf{p}_2)\f$
 *  Notice that this method returns the gradient with respect to the source point.
 *  \param[in] functor   function-object handle to the evaluation of the differentiable function
 *  \param[in] p1 first point
 *  \param[in] p2 second point
 */
Eigen::Vector3d gradientSource(const DifferentiableFunction & functor, const Eigen::Vector3d & p1, const Eigen::Vector3d & p2)
{
    return (Eigen::Vector3d() << derivativeSource(functor, Eigen::Vector3d::UnitX(), p1, p2),
                                 derivativeSource(functor, Eigen::Vector3d::UnitY(), p1, p2),
                                 derivativeSource(functor, Eigen::Vector3d::UnitZ(), p1, p2)).finished();
}

/*! Returns full gradient of the function passed for the pair of points p1, p2:
 *  \f$ \nabla_{\mathbf{p_2}}G(\mathbf{p}_1, \mathbf{p}_2)\f$
 *  Notice that this method returns the gradient with respect to the probe point.
 *  \param[in] functor   function-object handle to the evaluation of the differentiable function
 *  \param[in] p1 first point
 *  \param[in] p2 second point
 */
Eigen::Vector3d gradientProbe(const DifferentiableFunction & functor, const Eigen::Vector3d & p1, const Eigen::Vector3d & p2)
{
    return (Eigen::Vector3d() << derivativeProbe(functor, Eigen::Vector3d::UnitX(), p1, p2),
                                 derivativeProbe(functor, Eigen::Vector3d::UnitY(), p1, p2),
                                 derivativeProbe(functor, Eigen::Vector3d::UnitZ(), p1, p2)).finished();
}
} // namespace GreenUtils

#endif // GREENUTILS_HPP


/*

   Copyright (c) 2006-2010, The Scripps Research Institute

   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.

   Author: Dr. Oleg Trott <ot14@columbia.edu>,
           The Olson Lab,
           The Scripps Research Institute

*/

#include "monte_carlo.h"
#include "coords.h"
#include "mutate.h"
#include "quasi_newton.h"

#include "commonMacros.h"
#include "wrapcl.h"
#include "random.h"
#include <iostream>

#include <fstream>
#include <boost/progress.hpp>
#include <thread>

//#define DISPLAY_ANALYSIS
//#define DATA_DISTRIBUTION_TEST
//#define BUILD_KERNEL_FROM_SOURCE
output_type monte_carlo::operator()(model& m, const precalculate& p, const igrid& ig, const precalculate& p_widened, const igrid& ig_widened, const vec& corner1, const vec& corner2, incrementable* increment_me, rng& generator) const {
	output_container tmp;
	this->operator()(m, tmp, p, ig, p_widened, ig_widened, corner1, corner2, increment_me, generator); // call the version that produces the whole container
	VINA_CHECK(!tmp.empty());
	return tmp.front();
}

bool metropolis_accept(fl old_f, fl new_f, fl temperature, rng& generator) {
	if(new_f < old_f) return true;
	const fl acceptance_probability = std::exp((old_f - new_f) / temperature);
	return random_fl(0, 1, generator) < acceptance_probability;
}

void monte_carlo::single_run(model& m, output_type& out, const precalculate& p, const igrid& ig, rng& generator) const {
	conf_size s = m.get_size();
	change g(s);
	vec authentic_v(1000, 1000, 1000);
	out.e = max_fl;
	output_type current(out);
	quasi_newton quasi_newton_par; quasi_newton_par.max_steps = ssd_par.evals;
	VINA_U_FOR(step, num_steps) {
		output_type candidate(current.c, max_fl);
		mutate_conf(candidate.c, m, mutation_amplitude, generator);
		quasi_newton_par(m, p, ig, candidate, g, hunt_cap);
		if(step == 0 || metropolis_accept(current.e, candidate.e, temperature, generator)) {
			quasi_newton_par(m, p, ig, candidate, g, authentic_v);
			current = candidate;
			if(current.e < out.e)
				out = current;
		}
	}
	quasi_newton_par(m, p, ig, out, g, authentic_v);
}

void monte_carlo::many_runs(model& m, output_container& out, const precalculate& p, const igrid& ig, const vec& corner1, const vec& corner2, sz num_runs, rng& generator) const {
	conf_size s = m.get_size();
	VINA_FOR(run, num_runs) {
		output_type tmp(s, 0);
		tmp.c.randomize(corner1, corner2, generator);
		single_run(m, tmp, p, ig, generator);
		out.push_back(new output_type(tmp));
	}
	out.sort();
}

std::vector<output_type> monte_carlo::cl_to_vina(output_type_cl result_ptr[], int exhaus) const {
	std::vector<output_type> results_vina;
	for (int i = 0; i < exhaus; i++) {
		output_type_cl tmp_cl = result_ptr[i];
		conf tmp_c;
		tmp_c.ligands.resize(1);
		// Position
		for (int j = 0; j < 3; j++)tmp_c.ligands[0].rigid.position[j] = tmp_cl.position[j];
		// Orientation
		qt q(tmp_cl.orientation[0], tmp_cl.orientation[1], tmp_cl.orientation[2], tmp_cl.orientation[3]);
		tmp_c.ligands[0].rigid.orientation = q;
		output_type tmp_vina(tmp_c, tmp_cl.e);
		// torsion
		for (int j = 0; j < tmp_cl.lig_torsion_size; j++)tmp_vina.c.ligands[0].torsions.push_back(tmp_cl.lig_torsion[j]);
		// coords
		for (int j = 0; j < MAX_NUM_OF_ATOMS; j++) {
			vec v_tmp(tmp_cl.coords[j][0], tmp_cl.coords[j][1], tmp_cl.coords[j][2]);
			if (v_tmp[0] * v_tmp[1] * v_tmp[2] != 0) tmp_vina.coords.push_back(v_tmp);
		}
		results_vina.push_back(tmp_vina);
	}
	return results_vina;
}

void monte_carlo::generate_uniform_position(const vec corner1, const vec corner2, std::vector<vec>& uniform_data, int exhaustiveness) const{
	assert(exhaustiveness == 125);
	int counter = 0;
	int n = 5;
	for (int i = 0; i < n; i++) {
		double p0 = (corner1.data[0] - corner2.data[0]) / n * i + corner2.data[0];
		for (int j = 0; j < n; j++) {
			double p1 = (corner1.data[1] - corner2.data[1]) / n * j + corner2.data[1];
			for (int k = 0; k < n; k++) {
				double p2 = (corner1.data[2] - corner2.data[2]) / n * k + corner2.data[2];
				uniform_data[counter] = vec(p0, p1, p2);
				counter++;
			}
		}
	}
	assert(counter == exhaustiveness);
}

#ifndef OPENCL_PART_2
// out is sorted
void monte_carlo::operator()(model& m, output_container& out, const precalculate& p, const igrid& ig, const precalculate& p_widened, const igrid& ig_widened, const vec& corner1, const vec& corner2, incrementable* increment_me, rng& generator) const {
	vec authentic_v(1000, 1000, 1000); // FIXME? this is here to avoid max_fl/max_fl
	conf_size s = m.get_size();
	change g(s);
	output_type tmp(s, 0);
	tmp.c.randomize(corner1, corner2, generator);
	fl best_e = max_fl;
	quasi_newton quasi_newton_par; quasi_newton_par.max_steps = ssd_par.evals;
	output_type origin = tmp;

	VINA_U_FOR(step, num_steps) {
		if (increment_me)
			++(*increment_me);
		output_type candidate = tmp;
		mutate_conf(candidate.c, m, mutation_amplitude, generator);
		quasi_newton_par(m, p, ig, candidate, g, hunt_cap);
		if (step == 0 || metropolis_accept(tmp.e, candidate.e, temperature, generator)) {
			tmp = candidate;

			m.set(tmp.c); // FIXME? useless?

			// FIXME only for very promising ones
			if (tmp.e < best_e || out.size() < num_saved_mins) {
				quasi_newton_par(m, p, ig, tmp, g, authentic_v);
				m.set(tmp.c); // FIXME? useless?
				tmp.coords = m.get_heavy_atom_movable_coords();
				add_to_output_container(out, tmp, min_rmsd, num_saved_mins); // 20 - max size
				if (tmp.e < best_e)
					best_e = tmp.e;
			}
		}
	}
	VINA_CHECK(!out.empty());
	VINA_CHECK(out.front().e <= out.back().e); // make sure the sorting worked in the correct order
}
#else

volatile bool finished = false; //20211119 Glinttsd
void print_process() {
	int count = 0;
	printf("\n");
	do
	{
#ifdef WIN32
		Sleep(100);
#else
		sleep(1);
#endif
		printf("\rPerform docking|");
		for (int i = 0; i < count; i++)printf(" ");
		printf("=======");
		for (int i = 0; i < 30 - count; i++)printf(" ");
		printf("|"); fflush(stdout);

		count++;
		count %= 30;
	} while (finished != true);
	printf("\rPerform docking|");
	for (int i = 0; i < 16; i++)printf("=");
	printf("done");
	for (int i = 0; i < 17; i++)printf("=");
	printf("|\n"); fflush(stdout);
}

void monte_carlo::operator()(model& m, output_container& out, const precalculate& p, const igrid& ig, const precalculate& p_widened, const igrid& ig_widened, const vec& corner1, const vec& corner2, incrementable* increment_me, rng& generator) const {
	/**************************************************************************/
	/***************************    OpenCL Init    ****************************/
	/**************************************************************************/

	cl_int err;
	cl_platform_id* platforms;
	cl_device_id* devices;
	cl_context context;
	cl_command_queue queue;
	cl_int gpu_platform_id;
	SetupPlatform(&platforms, &gpu_platform_id);
	SetupDevice(platforms, &devices, gpu_platform_id);
	SetupContext(platforms, devices, &context, 1, gpu_platform_id);
	SetupQueue(&queue, context, devices);
	char* program_file;
	cl_program program_cl;
	cl_program program;
	size_t program_size;

	//Read kernel source code
#ifdef BUILD_KERNEL_FROM_SOURCE
	printf("\nBuild kernels from source"); fflush(stdout);
	const std::string default_work_path = ".";
	const std::string include_path = default_work_path + "/OpenCL/inc"; //FIX it
	const std::string addtion = "";

	char* program_file_n[NUM_OF_FILES];
	size_t program_size_n[NUM_OF_FILES];
	std::string file_paths[NUM_OF_FILES] = {	default_work_path + "/OpenCL/src/kernels/code_head.cpp",
												default_work_path + "/OpenCL/src/kernels/mutate_conf.cpp",
												default_work_path + "/OpenCL/src/kernels/matrix.cpp",
												default_work_path + "/OpenCL/src/kernels/quasi_newton.cpp",
												default_work_path + "/OpenCL/src/kernels/kernel2.cl"}; // The order of files is important!

	read_n_file(program_file_n, program_size_n, file_paths, NUM_OF_FILES);
	std::string final_file;
	size_t final_size = NUM_OF_FILES - 1; // count '\n'
	for (int i = 0; i < NUM_OF_FILES; i++) {
		if (i == 0) final_file = program_file_n[0];
		else final_file = final_file + '\n' + (std::string)program_file_n[i];
		final_size += program_size_n[i];
	}
	const char* final_files_char = final_file.data();

	program_cl = clCreateProgramWithSource(context, 1, (const char**)&final_files_char, &final_size, &err); checkErr(err);
	SetupBuildProgramWithSource(program_cl, NULL, devices, include_path, addtion);
	SaveProgramToBinary(program_cl, "Kernel2_Opt.bin");
#endif

	//Display the progress
	printf("\nSearch depth is set to %d",search_depth);
	std::thread console_thread(print_process);

	program_cl = SetupBuildProgramWithBinary(context, devices, "Kernel2_Opt.bin");

	err = clUnloadPlatformCompiler(platforms[gpu_platform_id]); checkErr(err);
	//Set kernel arguments
	cl_kernel kernels[1];
	char kernel_name[][50] = { "kernel2" };
	SetupKernel(kernels, program_cl, 1, kernel_name);
	size_t max_wg_size; // max work item within one work group
	size_t max_wi_size[3]; // max work item within each dimension(global)
	err = clGetDeviceInfo(devices[0], CL_DEVICE_MAX_WORK_GROUP_SIZE, sizeof(size_t), &max_wg_size, NULL); checkErr(err);
	err = clGetDeviceInfo(devices[0], CL_DEVICE_MAX_WORK_ITEM_SIZES, 3 * sizeof(size_t), &max_wi_size, NULL); checkErr(err);

	/**************************************************************************/
	/************************    Original Vina code    ************************/
	/**************************************************************************/

	vec authentic_v(1000, 1000, 1000); // FIXME? this is here to avoid max_fl/max_fl
	conf_size s = m.get_size();
	change g(s);
	output_type tmp(s, 0);
	quasi_newton quasi_newton_par; const int quasi_newton_par_max_steps = ssd_par.evals;

	/**************************************************************************/
	/************************    Allocate CPU memory    ***********************/
	/**************************************************************************/

	// Preparing m related data
	m_cl m_cl;
	assert(m.atoms.size() < MAX_NUM_OF_ATOMS);

	for (int i = 0; i < m.atoms.size(); i++) {
		m_cl.atoms[i].types[0] = m.atoms[i].el;// To store 4 atoms types (el, ad, xs, sy)
		m_cl.atoms[i].types[1] = m.atoms[i].ad;
		m_cl.atoms[i].types[2] = m.atoms[i].xs;
		m_cl.atoms[i].types[3] = m.atoms[i].sy;
		for (int j = 0; j < 3; j++) {
			m_cl.atoms[i].coords[j] = m.atoms[i].coords[j];// To store atom coords
		}
	}

	// To store atoms coords
	for (int i = 0; i < m.coords.size(); i++) {
		for (int j = 0; j < 3; j++) {
			m_cl.m_coords.coords[i][j] = m.coords[i].data[j];
		}
	}

	//To store minus forces
	for (int i = 0; i < m.coords.size(); i++) {
		for (int j = 0; j < 3; j++) {
			m_cl.minus_forces.coords[i][j] = m.minus_forces[i].data[j];
		}
	}

	// Preparing ligand data
	assert(m.num_other_pairs() == 0); // m.other_paris is not supported!
	assert(m.ligands.size() == 1); // Only one ligand supported!
	m_cl.ligand.pairs.num_pairs = m.ligands[0].pairs.size();
	for (int i = 0; i < m_cl.ligand.pairs.num_pairs; i++) {
		m_cl.ligand.pairs.type_pair_index[i]	= m.ligands[0].pairs[i].type_pair_index;
		m_cl.ligand.pairs.a[i]					= m.ligands[0].pairs[i].a;
		m_cl.ligand.pairs.b[i]					= m.ligands[0].pairs[i].b;
	}
	m_cl.ligand.begin = m.ligands[0].begin; // 0
	m_cl.ligand.end = m.ligands[0].end; // 29
	ligand m_ligand = m.ligands[0]; // Only support one ligand
	assert(m_ligand.end < MAX_NUM_OF_ATOMS);

	// Store root node
	m_cl.ligand.rigid.atom_range[0][0] = m_ligand.node.begin;
	m_cl.ligand.rigid.atom_range[0][1] = m_ligand.node.end;
	for (int i = 0; i < 3; i++) m_cl.ligand.rigid.origin[0][i] = m_ligand.node.get_origin()[i];
	for (int i = 0; i < 9; i++) m_cl.ligand.rigid.orientation_m[0][i] = m_ligand.node.get_orientation_m().data[i];
	m_cl.ligand.rigid.orientation_q[0][0] = m_ligand.node.orientation().R_component_1();
	m_cl.ligand.rigid.orientation_q[0][1] = m_ligand.node.orientation().R_component_2();
	m_cl.ligand.rigid.orientation_q[0][2] = m_ligand.node.orientation().R_component_3();
	m_cl.ligand.rigid.orientation_q[0][3] = m_ligand.node.orientation().R_component_4();
	for (int i = 0; i < 3; i++) {m_cl.ligand.rigid.axis[0][i] = 0;m_cl.ligand.rigid.relative_axis[0][i] = 0;m_cl.ligand.rigid.relative_origin[0][i] = 0;}

	// Store children nodes (in depth-first order)
	struct tmp_struct {
		int start_index = 0;
		int parent_index = 0;
		void store_node(tree<segment>& child_ptr, rigid_cl& rigid) {
			start_index++; // start with index 1, index 0 is root node
			rigid.parent[start_index] = parent_index;
			rigid.atom_range[start_index][0] = child_ptr.node.begin;
			rigid.atom_range[start_index][1] = child_ptr.node.end;
			for (int i = 0; i < 9; i++) rigid.orientation_m[start_index][i] = child_ptr.node.get_orientation_m().data[i];
			rigid.orientation_q[start_index][0] = child_ptr.node.orientation().R_component_1();
			rigid.orientation_q[start_index][1] = child_ptr.node.orientation().R_component_2();
			rigid.orientation_q[start_index][2] = child_ptr.node.orientation().R_component_3();
			rigid.orientation_q[start_index][3] = child_ptr.node.orientation().R_component_4();
			for (int i = 0; i < 3; i++) {
				rigid.origin[start_index][i] = child_ptr.node.get_origin()[i];
				rigid.axis[start_index][i] = child_ptr.node.get_axis()[i];
				rigid.relative_axis[start_index][i] = child_ptr.node.relative_axis[i];
				rigid.relative_origin[start_index][i] = child_ptr.node.relative_origin[i];
			}
			if (child_ptr.children.size() == 0) return;
			else {
				assert(start_index < MAX_NUM_OF_RIGID);
				int parent_index_tmp = start_index;
				for (int i = 0; i < child_ptr.children.size(); i++) {
					this->parent_index = parent_index_tmp; // Update parent index
					this->store_node(child_ptr.children[i], rigid);
				}
			}
		}
	};
	tmp_struct ts;
	for (int i = 0; i < m_ligand.children.size(); i++) {
		ts.parent_index = 0; // Start a new branch, whose parent is 0
		ts.store_node(m_ligand.children[i], m_cl.ligand.rigid);
	}
	m_cl.ligand.rigid.num_children = ts.start_index;

	// set children_map
	for (int i = 0; i < MAX_NUM_OF_RIGID; i++)
		for (int j = 0; j < MAX_NUM_OF_RIGID; j++)
			m_cl.ligand.rigid.children_map[i][j] = false;
	for (int i = 1; i < m_cl.ligand.rigid.num_children + 1; i++) {
		int parent_index = m_cl.ligand.rigid.parent[i];
		m_cl.ligand.rigid.children_map[parent_index][i] = true;
	}
	m_cl.m_num_movable_atoms = m.num_movable_atoms();
	size_t m_cl_size = sizeof(m_cl);

	// Preparing ig related data
	ig_cl* ig_cl_ptr = (ig_cl*)malloc(sizeof(ig_cl));
	ig_cl_ptr->atu = ig.get_atu(); // atu
	ig_cl_ptr->slope = ig.get_slope(); // slope
	std::vector<grid> tmp_grids = ig.get_grids();
	int grid_size = tmp_grids.size();
	assert(GRIDS_SIZE == grid_size); // grid_size has to be 17

	for (int i = 0; i < grid_size; i++) {
		for (int j = 0; j < 3; j++) {
			ig_cl_ptr->grids[i].m_init[j] = tmp_grids[i].m_init[j];
			ig_cl_ptr->grids[i].m_factor[j] = tmp_grids[i].m_factor[j];
			ig_cl_ptr->grids[i].m_dim_fl_minus_1[j] = tmp_grids[i].m_dim_fl_minus_1[j];
			ig_cl_ptr->grids[i].m_factor_inv[j] = tmp_grids[i].m_factor_inv[j];
		}
		if (tmp_grids[i].m_data.dim0() != 0) {
			ig_cl_ptr->grids[i].m_i = tmp_grids[i].m_data.dim0(); assert(MAX_NUM_OF_GRID_MI >= ig_cl_ptr->grids[i].m_i);
			ig_cl_ptr->grids[i].m_j = tmp_grids[i].m_data.dim1(); assert(MAX_NUM_OF_GRID_MJ >= ig_cl_ptr->grids[i].m_j);
			ig_cl_ptr->grids[i].m_k = tmp_grids[i].m_data.dim2(); assert(MAX_NUM_OF_GRID_MK >= ig_cl_ptr->grids[i].m_k);

			//for (int j = 0; j < ig_cl_ptr->grids[i].m_i * ig_cl_ptr->grids[i].m_j * ig_cl_ptr->grids[i].m_k; j++) {
			//	ig_cl_ptr->grids[i].m_data[j] = tmp_grids[i].m_data.m_data[j];
			//}

			int m_i = tmp_grids[i].m_data.dim0();
			int m_j = tmp_grids[i].m_data.dim1();
			int m_k = tmp_grids[i].m_data.dim2();
			for (sz idx0 = 0; idx0 < m_i - 1; idx0++)
				for (sz idx1 = 0; idx1 < m_j - 1; idx1++)
					for (sz idx2 = 0; idx2 < m_k - 1; idx2++) {
						int base = (idx0 + m_i * (idx1 + m_j * idx2)) * 8;
						ig_cl_ptr->grids[i].m_data[base] = tmp_grids[i].m_data(idx0, idx1, idx2); // f000
						ig_cl_ptr->grids[i].m_data[base + 1] = tmp_grids[i].m_data(idx0 + 1, idx1, idx2); // f100
						ig_cl_ptr->grids[i].m_data[base + 2] = tmp_grids[i].m_data(idx0, idx1 + 1, idx2); // f010
						ig_cl_ptr->grids[i].m_data[base + 3] = tmp_grids[i].m_data(idx0 + 1, idx1 + 1, idx2); // f110
						ig_cl_ptr->grids[i].m_data[base + 4] = tmp_grids[i].m_data(idx0, idx1, idx2 + 1); // f001
						ig_cl_ptr->grids[i].m_data[base + 5] = tmp_grids[i].m_data(idx0 + 1, idx1, idx2 + 1); // f101
						ig_cl_ptr->grids[i].m_data[base + 6] = tmp_grids[i].m_data(idx0, idx1 + 1, idx2 + 1); // f011
						ig_cl_ptr->grids[i].m_data[base + 7] = tmp_grids[i].m_data(idx0 + 1, idx1 + 1, idx2 + 1); // f111
					}

		}
		else {
			ig_cl_ptr->grids[i].m_i = 0;
			ig_cl_ptr->grids[i].m_j = 0;
			ig_cl_ptr->grids[i].m_k = 0;
		}
	}
	size_t ig_cl_size = sizeof(*ig_cl_ptr);

	// Generating random ligand structures
	std::vector<output_type_cl*> rand_molec_struc_vec; rand_molec_struc_vec.resize(thread);
	for (int i = 0; i < thread; i++) {
		rand_molec_struc_vec[i] = (output_type_cl*)malloc(sizeof(output_type_cl));
	}

	//std::vector<output_type_cl> rand_molec_struc_vec;
	//rand_molec_struc_vec.resize(thread);
	int lig_torsion_size = tmp.c.ligands[0].torsions.size();
	int flex_torsion_size; if (tmp.c.flex.size() != 0) flex_torsion_size = tmp.c.flex[0].torsions.size(); else flex_torsion_size = 0;
	std::vector<vec> uniform_data;
	uniform_data.resize(thread);

	for (int i = 0; i < thread; i++) {
		tmp.c.randomize(corner1, corner2, generator); // generate a random structure
		//Generate positions with uniform probability
		//generate_uniform_position(corner1, corner2, uniform_data, exhaustiveness);
		//for (int j = 0; j < 3; j++) rand_molec_struc_vec[i].position[j] = uniform_data[i].data[j];
		for (int j = 0; j < 3; j++) rand_molec_struc_vec[i]->position[j] = tmp.c.ligands[0].rigid.position[j];
		assert(lig_torsion_size < MAX_NUM_OF_LIG_TORSION);
		for (int j = 0; j < lig_torsion_size; j++) rand_molec_struc_vec[i]->lig_torsion[j] = tmp.c.ligands[0].torsions[j];// Only support one ligand
		assert(flex_torsion_size < MAX_NUM_OF_FLEX_TORSION);
		for (int j = 0; j < flex_torsion_size; j++) rand_molec_struc_vec[i]->flex_torsion[j] = tmp.c.flex[0].torsions[j];// Only support one flex

		rand_molec_struc_vec[i]->orientation[0] = (float)tmp.c.ligands[0].rigid.orientation.R_component_1();
		rand_molec_struc_vec[i]->orientation[1] = (float)tmp.c.ligands[0].rigid.orientation.R_component_2();
		rand_molec_struc_vec[i]->orientation[2] = (float)tmp.c.ligands[0].rigid.orientation.R_component_3();
		rand_molec_struc_vec[i]->orientation[3] = (float)tmp.c.ligands[0].rigid.orientation.R_component_4();

		rand_molec_struc_vec[i]->lig_torsion_size = lig_torsion_size;
	}

	// Preaparing p related data
	p_cl p_cl;
	p_cl.m_cutoff_sqr = p.cutoff_sqr();
	p_cl.factor = p.factor;
	p_cl.n = p.n;
	assert(MAX_P_DATA_M_DATA_SIZE > p.data.m_data.size());
	for (int i = 0; i < p.data.m_data.size(); i++) {
		p_cl.m_data[i].factor = p.data.m_data[i].factor;
		assert(FAST_SIZE == p.data.m_data[i].fast.size());
		assert(SMOOTH_SIZE == p.data.m_data[i].smooth.size());
		for (int j = 0; j < FAST_SIZE; j++) {
			p_cl.m_data[i].fast[j] = p.data.m_data[i].fast[j];
		}
		for (int j = 0; j < SMOOTH_SIZE; j++) {
			p_cl.m_data[i].smooth[j][0] = p.data.m_data[i].smooth[j].first;
			p_cl.m_data[i].smooth[j][1] = p.data.m_data[i].smooth[j].second;
		}
	}
	size_t p_cl_size = sizeof(p_cl);

	// Generate random maps
	random_maps* rand_maps = (random_maps*)malloc(sizeof(random_maps));
	for (int i = 0; i < MAX_NUM_OF_RANDOM_MAP; i++) {
		rand_maps->int_map[i] = random_int(0, int(lig_torsion_size), generator);
		rand_maps->pi_map[i] = random_fl(-pi, pi, generator);
	}
	for (int i = 0; i < MAX_NUM_OF_RANDOM_MAP; i++) {
		vec rand_coords = random_inside_sphere(generator);
		for (int j = 0; j < 3 ; j ++) {
			rand_maps->sphere_map[i][j] = rand_coords[j];
		}
	}
	size_t rand_maps_size = sizeof(*rand_maps);

	float hunt_cap_float[3] = {hunt_cap[0], hunt_cap[1], hunt_cap[2]};
	float authentic_v_float[3] = { authentic_v[0],authentic_v[1], authentic_v[2] };
	float mutation_amplitude_float = mutation_amplitude;
	float epsilon_fl_float = epsilon_fl;
	int	total_wi = max_wi_size[0] * max_wi_size[1];
	/**************************************************************************/
	/************************    Allocate GPU memory    ***********************/
	/**************************************************************************/

	cl_mem rand_molec_struc_vec_gpu;
	CreateDeviceBuffer(&rand_molec_struc_vec_gpu, CL_MEM_READ_ONLY, thread* SIZE_OF_MOLEC_STRUC, context);
	for (int i = 0; i < thread; i++) {
		std::vector<float> pos(rand_molec_struc_vec[i]->position, rand_molec_struc_vec[i]->position + 3);
		std::vector<float> ori(rand_molec_struc_vec[i]->orientation, rand_molec_struc_vec[i]->orientation + 4);
		std::vector<float> lig_tor(rand_molec_struc_vec[i]->lig_torsion, rand_molec_struc_vec[i]->lig_torsion + MAX_NUM_OF_LIG_TORSION);
		std::vector<float> flex_tor(rand_molec_struc_vec[i]->flex_torsion, rand_molec_struc_vec[i]->flex_torsion + MAX_NUM_OF_FLEX_TORSION);
		float lig_tor_size = rand_molec_struc_vec[i]->lig_torsion_size;
		err = clEnqueueWriteBuffer(	queue, rand_molec_struc_vec_gpu, false, i * SIZE_OF_MOLEC_STRUC,
									pos.size() * sizeof(float), pos.data(), 0, NULL, NULL); checkErr(err);
		err = clEnqueueWriteBuffer(	queue, rand_molec_struc_vec_gpu, false, i * SIZE_OF_MOLEC_STRUC + pos.size() * sizeof(float),
									ori.size() * sizeof(float), ori.data(), 0, NULL, NULL); checkErr(err);
		err = clEnqueueWriteBuffer(	queue, rand_molec_struc_vec_gpu, false, i * SIZE_OF_MOLEC_STRUC + (pos.size() + ori.size() ) * sizeof(float),
									lig_tor.size() * sizeof(float), lig_tor.data(), 0, NULL, NULL); checkErr(err);
		err = clEnqueueWriteBuffer(queue, rand_molec_struc_vec_gpu, false, i * SIZE_OF_MOLEC_STRUC + (pos.size() + ori.size() + MAX_NUM_OF_LIG_TORSION) * sizeof(float),
									flex_tor.size() * sizeof(float), flex_tor.data(), 0, NULL, NULL); checkErr(err);
		err = clEnqueueWriteBuffer(queue, rand_molec_struc_vec_gpu, false, i * SIZE_OF_MOLEC_STRUC + (pos.size() + ori.size() + MAX_NUM_OF_LIG_TORSION + MAX_NUM_OF_FLEX_TORSION) * sizeof(float),
			sizeof(float), &lig_tor_size, 0, NULL, NULL); checkErr(err);
	}

	cl_mem best_e_gpu;
	CreateDeviceBuffer(&best_e_gpu, CL_MEM_READ_WRITE, thread * sizeof(float), context);
	err = clEnqueueFillBuffer(queue, best_e_gpu, &max_fl, sizeof(float), 0, thread * sizeof(float), 0, NULL, NULL); checkErr(err);

	cl_mem rand_maps_gpu;
	CreateDeviceBuffer(&rand_maps_gpu, CL_MEM_READ_ONLY, rand_maps_size, context);
	err = clEnqueueWriteBuffer(	queue, rand_maps_gpu, false, 0,	rand_maps_size, rand_maps, 0, NULL, NULL); checkErr(err);

	cl_mem hunt_cap_gpu;
	CreateDeviceBuffer(&hunt_cap_gpu, CL_MEM_READ_ONLY, 3 * sizeof(float), context);
	err = clEnqueueWriteBuffer(queue, hunt_cap_gpu, false, 0, 3 * sizeof(float), hunt_cap_float, 0, NULL, NULL); checkErr(err);

	// Preparing m related data
	cl_mem m_cl_gpu;
	CreateDeviceBuffer(&m_cl_gpu, CL_MEM_READ_WRITE, m_cl_size, context);
	err = clEnqueueWriteBuffer(queue, m_cl_gpu, false, 0, m_cl_size, &m_cl, 0, NULL, NULL); checkErr(err);

	// Preparing p related data
	cl_mem p_cl_gpu;
	CreateDeviceBuffer(&p_cl_gpu, CL_MEM_READ_ONLY, p_cl_size, context);
	err = clEnqueueWriteBuffer(queue, p_cl_gpu, false, 0, p_cl_size, &p_cl, 0, NULL, NULL); checkErr(err);

	// Preparing ig related data (cache related data)
	cl_mem ig_cl_gpu;
	CreateDeviceBuffer(&ig_cl_gpu, CL_MEM_READ_ONLY, ig_cl_size, context);
	err = clEnqueueWriteBuffer(queue, ig_cl_gpu, false, 0, ig_cl_size, ig_cl_ptr, 0, NULL, NULL); checkErr(err);

	cl_mem authentic_v_gpu;
	CreateDeviceBuffer(&authentic_v_gpu, CL_MEM_READ_ONLY, 3 * sizeof(float), context);
	err = clEnqueueWriteBuffer(queue, authentic_v_gpu, false, 0, 3 * sizeof(float), authentic_v_float, 0, NULL, NULL); checkErr(err);

	// Preparing result data
	cl_mem results;
	CreateDeviceBuffer(&results, CL_MEM_READ_WRITE | CL_MEM_ALLOC_HOST_PTR, thread * sizeof(output_type_cl), context);

	clFinish(queue);

	/**************************************************************************/
	/************************   Set kernel arguments    ***********************/
	/**************************************************************************/
	SetKernelArg(kernels[0], 0, sizeof(cl_mem),		&m_cl_gpu);
	SetKernelArg(kernels[0], 1, sizeof(cl_mem),		&ig_cl_gpu);
	SetKernelArg(kernels[0], 2, sizeof(cl_mem),		&p_cl_gpu);
	SetKernelArg(kernels[0], 3,	sizeof(cl_mem),		&rand_molec_struc_vec_gpu);
	SetKernelArg(kernels[0], 4, sizeof(cl_mem),		&best_e_gpu);
	SetKernelArg(kernels[0], 5, sizeof(int),		&quasi_newton_par_max_steps);
	SetKernelArg(kernels[0], 6, sizeof(unsigned int),	&num_steps);
	SetKernelArg(kernels[0], 7, sizeof(float),		&mutation_amplitude_float);
	SetKernelArg(kernels[0], 8, sizeof(cl_mem),		&rand_maps_gpu);
	SetKernelArg(kernels[0], 9, sizeof(float),		&epsilon_fl_float);
	SetKernelArg(kernels[0], 10, sizeof(cl_mem),	&hunt_cap_gpu);
	SetKernelArg(kernels[0], 11, sizeof(cl_mem),	&authentic_v_gpu);
	SetKernelArg(kernels[0], 12, sizeof(cl_mem),	&results);
	SetKernelArg(kernels[0], 13, sizeof(int),		&search_depth);
	SetKernelArg(kernels[0], 14, sizeof(int),		&thread);
	SetKernelArg(kernels[0], 15, sizeof(int),	&total_wi);
	/**************************************************************************/
	/****************************   Start kernel    ***************************/
	/**************************************************************************/
	size_t global_size[2] = {512, 32 };
	size_t local_size[2] = { 16,2 };
	cl_event monte_clarlo_cl;
	err = clEnqueueNDRangeKernel(queue, kernels[0], 2, 0, global_size, local_size, 0, NULL, &monte_clarlo_cl); checkErr(err);

	clWaitForEvents(1, &monte_clarlo_cl);

	finished = true;
	console_thread.join(); // wait the thread finish

	//getchar();

	// Maping result data
 	output_type_cl* result_ptr = (output_type_cl*)clEnqueueMapBuffer(queue, results, CL_TRUE, CL_MAP_READ,
																	0, thread * sizeof(output_type_cl),
																	0, NULL, NULL, &err); checkErr(err);

	std::vector<output_type> result_vina = cl_to_vina(result_ptr, thread);
	// if empty, something goes wrong in the device part
	if (result_vina.size() == 0) {
		printf("Error in the device part\n"); exit(-1);
	}

	// Unmaping result data
	err = clEnqueueUnmapMemObject(queue, results, result_ptr, 0, NULL, NULL); checkErr(err);

	// Write back to vina
	for (int i = 0; i < thread; i++) {
		add_to_output_container(out, result_vina[i], min_rmsd, num_saved_mins);
	}
	VINA_CHECK(!out.empty());
	VINA_CHECK(out.front().e <= out.back().e);

	/**************************************************************************/
	/*******************************   Finish    ******************************/
	/**************************************************************************/
	// Memory objects release
	err = clReleaseMemObject(m_cl_gpu); checkErr(err);
	err = clReleaseMemObject(ig_cl_gpu); checkErr(err);
	err = clReleaseMemObject(p_cl_gpu); checkErr(err);
	err = clReleaseMemObject(rand_molec_struc_vec_gpu); checkErr(err);
	err = clReleaseMemObject(best_e_gpu); checkErr(err);
	err = clReleaseMemObject(hunt_cap_gpu); checkErr(err);
	err = clReleaseMemObject(authentic_v_gpu); checkErr(err);
	err = clReleaseMemObject(results); checkErr(err);

	free(ig_cl_ptr);
	free(rand_maps);
	for (int i = 0; i < thread; i++)free(rand_molec_struc_vec[i]);

	// Output Analysis
	cl_ulong time_start, time_end;
	double total_time;
	err = clGetEventProfilingInfo(monte_clarlo_cl, CL_PROFILING_COMMAND_START, sizeof(time_start), &time_start, NULL); checkErr(err);
	err = clGetEventProfilingInfo(monte_clarlo_cl, CL_PROFILING_COMMAND_END, sizeof(time_end), &time_end, NULL); checkErr(err);
	total_time = time_end - time_start;
	printf("GPU monte carlo runtime: %0.3f s", (total_time / 1000000000.0));

#ifdef DISPLAY_ANALYSIS
	// Output Analysis
	cl_ulong time_start, time_end;
	double total_time;
	err = clGetEventProfilingInfo(monte_clarlo_cl, CL_PROFILING_COMMAND_START, sizeof(time_start), &time_start, NULL); checkErr(err);
	err = clGetEventProfilingInfo(monte_clarlo_cl, CL_PROFILING_COMMAND_END, sizeof(time_end), &time_end, NULL); checkErr(err);
	total_time = time_end - time_start;
	printf("\nGPU monte carlo runtime = %0.3f s\n", (total_time / 1000000000.0));

	std::ofstream file("gpu_runtime.txt");
	if (file.is_open())
	{
		file << "GPU monte carlo runtime = " << (double)(total_time / 1000000000.0) << " s" << std::endl;
		file.close();
	}

	cl_ulong private_mem_size;
	err = clGetKernelWorkGroupInfo(kernels[0], devices[0], CL_KERNEL_PRIVATE_MEM_SIZE, sizeof(cl_ulong), &private_mem_size, NULL); checkErr(err);
	printf("\nprivate mem used = %f KBytes", (double)private_mem_size / 1024);

	cl_ulong local_mem_size;
	err = clGetKernelWorkGroupInfo(kernels[0], devices[0], CL_KERNEL_LOCAL_MEM_SIZE, sizeof(cl_ulong), &local_mem_size, NULL); checkErr(err);
	printf("\nlocal mem used = %d Bytes", (int)local_mem_size);

	printf("\nconstant mem used = %0.3f MBytes \n", (double)(	p_cl_size + ig_cl_size + thread * SIZE_OF_MOLEC_STRUC) / (1024 * 1024));
#endif
	//getchar();
}
#endif // OPENCL_VERSION