/* Author: Mincheul Kang */

#include <harmonious_sampling/HarmoniousLazyPRMstarMulti.h>
#include <ompl/base/objectives/PathLengthOptimizationObjective.h>
#include <ompl/base/goals/GoalSampleableRegion.h>
#include <ompl/util/GeometricEquations.h>
#include <ompl/geometric/planners/prm/ConnectionStrategy.h>
#include <ompl/tools/config/SelfConfig.h>
#include <ompl/base/spaces/RealVectorStateSpace.h>

#include <boost/lambda/bind.hpp>
#include <boost/graph/astar_search.hpp>
#include <boost/graph/lookup_edge.hpp>
#include <boost/foreach.hpp>
#include <fstream>
#include <queue>
#include <stack>
#include <utility>
#include <numeric>
#include <stdio.h>
#include <iostream>

#define foreach BOOST_FOREACH

namespace ompl {
    namespace magic {
/** \brief The number of nearest neighbors to consider by
      default in the construction of the PRM roadmap */
        static const unsigned int DEFAULT_NEAREST_NEIGHBORS_LAZY = 5;

/** \brief When optimizing solutions with lazy planners, this is the minimum
      number of path segments to add before attempting a new optimized solution
          extraction */
        static const unsigned int MIN_ADDED_SEGMENTS_FOR_LAZY_OPTIMIZATION = 5;
    }
}

ompl::geometric::HarmoniousLazyPRMstarMulti::HarmoniousLazyPRMstarMulti(const base::SpaceInformationPtr &si,
                                                                        base::HarmoniousSampler &hs,
                                                                        const std::vector<bool> &isContinuous,
                                                                        bool starStrategy) :
        base::Planner(si, "HarmoniousLazyPRMstarMulti"),
        starStrategy_(starStrategy),
        userSetConnectionStrategy_(false),
        maxDistance_(0.0),
        indexProperty_(boost::get(boost::vertex_index_t(), g_)),
        stateProperty_(boost::get(vertex_state_t(), g_)),
        radiusProperty_(boost::get(vertex_radius_t(), g_)),
        witnessProperty_(boost::get(vertex_witness_t(), g_)),
        costProperty_(boost::get(vertex_cost_t(), g_)),
        childrenProperty_(boost::get(vertex_children_t(), g_)),
        predecessorProperty_(boost::get(boost::vertex_predecessor_t(), g_)),
        colorProperty_(boost::get(boost::vertex_color_t(), g_)),
        weightProperty_(boost::get(boost::edge_weight_t(), g_)),
        vertexValidityProperty_(boost::get(vertex_flags_t(), g_)),
        edgeValidityProperty_(boost::get(edge_flags_t(), g_)),
        bestCost_(std::numeric_limits<double>::quiet_NaN()),
        iterations_(0),
        increaseIterations_(0),
        BisectionCC_(true),
        rewireFactor_(1.5),
        hs_(hs),
        isContinuous_(isContinuous)
{
    specs_.recognizedGoal = base::GOAL_SAMPLEABLE_REGION;
    specs_.approximateSolutions = false;
    specs_.optimizingPaths = true;

    Planner::declareParam<bool>("BisectionCC", this,
                                &HarmoniousLazyPRMstarMulti::setBisectionCC, std::string("."));
    Planner::declareParam<double>("RewireFactor", this,
                                  &HarmoniousLazyPRMstarMulti::setRewireFactor, std::string("."));

    addPlannerProgressProperty("iterations INTEGER",
                               std::bind(&HarmoniousLazyPRMstarMulti::getIterationCount, this));
    addPlannerProgressProperty("best cost REAL",
                               std::bind(&HarmoniousLazyPRMstarMulti::getBestCost, this));
    addPlannerProgressProperty("milestone count INTEGER",
                               std::bind(&HarmoniousLazyPRMstarMulti::getMilestoneCountString, this));
    addPlannerProgressProperty("edge count INTEGER",
                               std::bind(&HarmoniousLazyPRMstarMulti::getEdgeCountString, this));
}

ompl::geometric::HarmoniousLazyPRMstarMulti::~HarmoniousLazyPRMstarMulti() {
}

void ompl::geometric::HarmoniousLazyPRMstarMulti::setup() {
    Planner::setup();
    tools::SelfConfig sc(si_, getName());
    sc.configurePlannerRange(maxDistance_);

    if (!nn_) {
        nn_.reset(tools::SelfConfig::getDefaultNearestNeighbors<Vertex>(this));
        nn_->setDistanceFunction(std::bind(&HarmoniousLazyPRMstarMulti::distanceFunctionHarmonious, this, std::placeholders::_1, std::placeholders::_2));
    }

    if (!nnB_) {
        nnB_.reset(tools::SelfConfig::getDefaultNearestNeighbors<Vertex>(this));
        nnB_->setDistanceFunction(std::bind(&HarmoniousLazyPRMstarMulti::distanceFunctionHarmonious, this, std::placeholders::_1, std::placeholders::_2));
    }

    if (!nnM_) {
        nnM_.reset(tools::SelfConfig::getDefaultNearestNeighbors<Vertex>(this));
        nnM_->setDistanceFunction(std::bind(&HarmoniousLazyPRMstarMulti::distanceFunctionHarmonious, this, std::placeholders::_1, std::placeholders::_2));
    }

    if (!connectionStrategy_) {
//        if (starStrategy_)
//            connectionStrategy_ = KStarStrategy<Vertex>(std::bind(&HarmoniousLazyPRMstarMulti::milestoneCount, this), nn_, si_->getStateDimension());
//        else
//            connectionStrategy_ = KBoundedStrategy<Vertex>(magic::DEFAULT_NEAREST_NEIGHBORS_LAZY, maxDistance_, nn_);
    }

    double dimDbl = static_cast<double>(si_->getStateDimension());
    // double prunedMeasure_ = si_->getSpaceMeasure();

    k_rrgConstant_ = rewireFactor_ * boost::math::constants::e<double>() + (boost::math::constants::e<double>() / dimDbl);

    // Setup optimization objective
    //
    // If no optimization objective was specified, then default to
    // optimizing path length as computed by the distance() function
    // in the state space.
    if (pdef_) {
        if (pdef_->hasOptimizationObjective()) {
            opt_ = pdef_->getOptimizationObjective();
        }
        else {
            opt_.reset(new base::PathLengthOptimizationObjective(si_));

            if (!starStrategy_) {
                opt_->setCostThreshold(opt_->infiniteCost());
            }
        }
    }
    else {
        OMPL_ERROR("%s: problem definition is not set, deferring setup completion...\n", getName().c_str());
        setup_ = false;
    }
    sampler_ = si_->allocStateSampler();
}

void ompl::geometric::HarmoniousLazyPRMstarMulti::setRange(double distance) {
    maxDistance_ = distance;

    if (!userSetConnectionStrategy_) {
        connectionStrategy_.clear();
    }

    if (isSetup()) {
        setup();
    }
}

void ompl::geometric::HarmoniousLazyPRMstarMulti::setMaxNearestNeighbors(unsigned int k) {
    if (starStrategy_) {
        throw Exception("Cannot set the maximum nearest neighbors for " + getName());
    }

    if (!nn_) {
        nn_.reset(tools::SelfConfig::getDefaultNearestNeighbors<Vertex>(this));
        nn_->setDistanceFunction(std::bind(&HarmoniousLazyPRMstarMulti::distanceFunctionHarmonious, this,std::placeholders::_1, std::placeholders::_2));
    }

    if (!nnB_) {
        nnB_.reset(tools::SelfConfig::getDefaultNearestNeighbors<Vertex>(this));
        nnB_->setDistanceFunction(std::bind(&HarmoniousLazyPRMstarMulti::distanceFunctionHarmonious, this, std::placeholders::_1, std::placeholders::_2));
    }

    if (!nnM_) {
        nnM_.reset(tools::SelfConfig::getDefaultNearestNeighbors<Vertex>(this));
        nnM_->setDistanceFunction(std::bind(&HarmoniousLazyPRMstarMulti::distanceFunctionHarmonious, this, std::placeholders::_1, std::placeholders::_2));
    }

    if (!userSetConnectionStrategy_) {
        connectionStrategy_.clear();
    }

    if (isSetup()) {
        setup();
    }
}

void ompl::geometric::HarmoniousLazyPRMstarMulti::setProblemDefinition(const base::ProblemDefinitionPtr &pdef) {
    Planner::setProblemDefinition(pdef);
    clearQuery();
}

void ompl::geometric::HarmoniousLazyPRMstarMulti::clearQuery() {
    startM_.clear();
    goalM_.clear();
    pis_.restart();
}

void ompl::geometric::HarmoniousLazyPRMstarMulti::clear() {
    Planner::clear();
    freeMemory();

    if (nn_) {
        nn_->clear();
    }

    if (nnB_) {
        nnB_->clear();
    }

    if (nnM_) {
        nnM_->clear();
    }

    clearQuery();

    iterations_ = 0;
    bestCost_ = base::Cost(std::numeric_limits<double>::quiet_NaN());
}

void ompl::geometric::HarmoniousLazyPRMstarMulti::freeMemory() {
    foreach (Vertex v, boost::vertices(g_)) {
        si_->freeState(stateProperty_[v]);
    }

    g_.clear();
}

// Add newly sampled vertex and its adjancency edges connected to neigh neighbors.
ompl::geometric::HarmoniousLazyPRMstarMulti::Vertex ompl::geometric::HarmoniousLazyPRMstarMulti::addMilestone(base::State *state, bool isM, bool isChecked) {
    Vertex m = boost::add_vertex(g_);
    stateProperty_[m] = state;
    radiusProperty_[m] = std::numeric_limits<double>::infinity();
    costProperty_[m] = std::numeric_limits<double>::infinity();
    childrenProperty_[m] = new std::vector<Vertex>();
    witnessProperty_[m] = NULL;
    predecessorProperty_[m] = NULL;
    colorProperty_[m] = 0;
    vertexValidityProperty_[m] = (isChecked) ? VALIDITY_TRUE : VALIDITY_UNKNOWN;

    std::vector<Vertex> neighbors;
    std::vector<double> neighbors_costs;

    if(isM){
        // manipulation regions
        unsigned long int bSize = nnB_->size();
        unsigned int max_number_of_neighbors = std::ceil(k_rrgConstant_ * log(static_cast<double>(milestoneCount()-bSize) + 1u));
        neighbors.reserve(max_number_of_neighbors);
        neighbors_costs.reserve(max_number_of_neighbors);

        nnM_->nearestK(m, max_number_of_neighbors, neighbors);

        foreach (Vertex n, neighbors) {
            const double weight = distanceFunctionHarmonious(n, m);

            ompl::base::Cost cost_weight(weight);
            const Graph::edge_property_type properties(cost_weight);

            neighborProperty_[m].push_back(type_neighbor(n, weight));
            neighborProperty_[n].push_back(type_neighbor(m, weight));

            // If collision-free or optimized well,
            const Edge &e = boost::add_edge(n, m, properties, g_).first;
            edgeValidityProperty_[e] = VALIDITY_UNKNOWN;
        }

        // base regions
        max_number_of_neighbors = std::ceil(k_rrgConstant_ * log(static_cast<double>(bSize) + 1u));

        neighbors.reserve(max_number_of_neighbors);
        neighbors_costs.reserve(max_number_of_neighbors);

        nnB_->nearestK(m, max_number_of_neighbors, neighbors);

        foreach (Vertex n, neighbors) {
            const double weight = distanceFunctionHarmonious(n, m);

            ompl::base::Cost cost_weight(weight);
            const Graph::edge_property_type properties(cost_weight);

            neighborProperty_[m].push_back(type_neighbor(n, weight));
            neighborProperty_[n].push_back(type_neighbor(m, weight));

            // If collision-free or optimized well,
            const Edge &e = boost::add_edge(n, m, properties, g_).first;
            edgeValidityProperty_[e] = VALIDITY_UNKNOWN;
        }

        nnM_->add(m);
    }
    else{
        unsigned int max_number_of_neighbors = std::ceil(k_rrgConstant_ * log(static_cast<double>(milestoneCount()) + 1u));

        neighbors.reserve(max_number_of_neighbors);
        neighbors_costs.reserve(max_number_of_neighbors);

        nn_->nearestK(m, max_number_of_neighbors, neighbors);

        foreach (Vertex n, neighbors) {
            const double weight = distanceFunctionHarmonious(n, m);

            ompl::base::Cost cost_weight(weight);
            const Graph::edge_property_type properties(cost_weight);

            neighborProperty_[m].push_back(type_neighbor(n, weight));
            neighborProperty_[n].push_back(type_neighbor(m, weight));

            // If collision-free or optimized well,
            const Edge &e = boost::add_edge(n, m, properties, g_).first;
            edgeValidityProperty_[e] = VALIDITY_UNKNOWN;
        }

        nnB_->add(m);
    }

    nn_->add(m);

    return m;
}

ompl::base::PlannerStatus ompl::geometric::HarmoniousLazyPRMstarMulti::solve(const base::PlannerTerminationCondition &ptc) {
    // Initial checkup for start/goal configurations.
    checkValidity();

    // Add the valid start states as milestones
    while (const base::State *st = pis_.nextStart()) {
        Vertex st_vert = addMilestone(si_->cloneState(st), false);
        costProperty_[st_vert] = 0.0; // Initialize with 0 cost.
        startM_.push_back(st_vert);
    }

    if (startM_.size() == 0) {
        OMPL_ERROR("error-%s: There are no valid initial states!", getName().c_str());
        return base::PlannerStatus::INVALID_START;
    }

    unsigned long int nrStartStates = boost::num_vertices(g_);
    OMPL_INFORM("%s: Starting planning with %lu states already in datastructure", getName().c_str(), nrStartStates);

    bestCost_ = opt_->infiniteCost();
    base::State *workState = si_->allocState();
    bool fullyOptimized = false;
    base::PathPtr bestSolution;

    base::RealVectorStateSpace::StateType *rstate = static_cast<base::RealVectorStateSpace::StateType*>(workState);

    Vertex startV = startM_[0];

    do{
        if(hs_.makeGoal(si_, workState)){
            goalM_.push_back(addMilestone(si_->cloneState(workState), true));
        }
    } while(goalM_.size() == 0 && ptc == false);

    bool isM = false;
    while (ptc == false) {
        ++iterations_;

        int r_val = std::rand() % 100;
        if(r_val < hs_.params_.prob_goal_ && goalM_.size() < hs_.params_.max_num_goal_){
            if(hs_.makeGoal(si_, workState)){
                goalM_.push_back(addMilestone(si_->cloneState(workState), true));
            }
            else{
                continue;
            }
        }
        else if(r_val < hs_.params_.prob_uniform_){
            hs_.uniformSampling(workState);
            if(hs_.getReachable(workState)){
                isM = true;
            }
            else{
                isM = false;
            }
        }
        else{
            hs_.biasedSampling(workState);

            if(hs_.getReachable(workState)){
                hs_.jointSampling(workState);
                isM = true;
            }
            else{
                hs_.setBasicPose(workState);
                isM = false;
            }
        }

        if (!si_->isValid(workState)) {
            continue;
        }

        // Add collision-free vertices.
        Vertex addedVertex = addMilestone(si_->cloneState(workState), isM);

        // DSPT update.
        Decrease(addedVertex);

        // Only support a single pair of start and goal node.
        base::PathPtr solution;

        for (int i = 0; i < (int)goalM_.size(); i++) {
            Vertex goalV = goalM_[i];
            do {
                if (predecessorProperty_[goalV] == NULL || bestCost_.value() <= costProperty_[goalV])
                    break;
                solution = constructSolution(startV, goalV);
            } while (!solution);

            if (solution) {
                base::Cost c(costProperty_[goalV]);

                if (opt_->isCostBetterThan(c, bestCost_)) {
                    bestSolution = solution;

                    bestCost_ = c;
                }
            }
        }
    }
    if (goalM_.empty()) {
        OMPL_ERROR("%s: Unable to find any valid goal states", getName().c_str());
        return base::PlannerStatus::INVALID_GOAL;
    }

    si_->freeState(workState);

    if (bestSolution) {
        base::PlannerSolution psol(bestSolution);
        psol.setPlannerName(getName());
        // If the solution was optimized, we mark it as such
        psol.setOptimized(opt_, bestCost_, fullyOptimized);
        pdef_->addSolutionPath(psol);
    }

    OMPL_INFORM("%s: Created %lu vertices and %lu edges.", getName().c_str(), boost::num_vertices(g_) - 1, boost::num_edges(g_));
    OMPL_INFORM("%s: Goalstates %lu ", getName().c_str(), goalM_.size());
    OMPL_INFORM("Cost: %.6f", bestCost_.value());

    return bestSolution ? base::PlannerStatus::EXACT_SOLUTION :
           base::PlannerStatus::TIMEOUT;
}

// outedge, inedge? - doesn't matter, need to scan all the neighbors.
void ompl::geometric::HarmoniousLazyPRMstarMulti::Decrease(const Vertex &v) {
    typedef std::pair<double, Vertex> weight_vertex;
    std::priority_queue<weight_vertex> pq;

    // Initialize cost of v, i.e., finding best parent vertex in G(g_).
    BGL_FORALL_OUTEDGES(v, e, g_, Graph) {
        Vertex w = target(e, g_);
        double weight = weightProperty_[e].value();

        if (costProperty_[v] > costProperty_[w] + weight) {
            predecessorProperty_[v] = w;
            costProperty_[v] = costProperty_[w] + weight;
        }
    }

    // No need to invoke cancelAdoption since v is newly sampled.
    if (predecessorProperty_[v] != NULL) {
        childrenProperty_[predecessorProperty_[v]]->push_back(v);
    }

    // At this point, v has a best parent. From now on construct its subtree of descendants.

    pq.push(weight_vertex(-costProperty_[v], v)); // Invert the cost value for mimicking min-heap.

    while (!pq.empty()) {
        weight_vertex top = pq.top();
        pq.pop();

        double cost = -top.first; // Invert the cost value to be like min-heap.
        Vertex vert = top.second;

        if (cost > costProperty_[vert]) {
            continue;
        }

        BGL_FORALL_OUTEDGES(vert, e, g_, Graph) {
            Vertex w = target(e, g_);
            double weight = weightProperty_[e].value();
            double cost_w = costProperty_[w];

            if (cost_w > cost + weight) {
                costProperty_[w] = cost + weight;
                cancelAdoption(w);

                predecessorProperty_[w] = vert;
                childrenProperty_[vert]->push_back(w);
                pq.push(weight_vertex(-costProperty_[w], w));
            }
        }
    }

    // Now, DSPT is stable.
}

#define RED (increaseIterations_ + 1) // I know, it's bad #define.

void ompl::geometric::HarmoniousLazyPRMstarMulti::Increase(const Vertex vs) {
    // <Step 1. Preparation.
    // white is used for color of each vertex without initialization.
    // For each iteration, white is increased by 1, thus we can use it as
    // equal to or less than 'white' means 'white' color, 'red' otherwise.
    increaseIterations_ += 1;

    std::vector<Vertex> reds;
    typedef std::pair<double, Vertex> weight_vertex;
    std::priority_queue<weight_vertex> pq; //  Max-heap by default.

    pq.push(weight_vertex(-costProperty_[vs], vs));  // It works as if it is min-heap.

    // <Step 2. Coloring
    while (!pq.empty()) {
        weight_vertex top = pq.top();
        pq.pop();

        double cost = -top.first;
        Vertex vert = top.second;

        if (cost > costProperty_[vert]) {
            continue;  // Instead of heap-improve
        }

        // Probability to get a pink node? almost impossible!
        /*
        // If there exist a non-red neighbor q of z such that Dist(q) + w_(q, z) = D(z)
        // set pink, that means it can keep the current cost, thus it is not necessary to
        // iterate its children.
        // Otherwise, set red and enqueue all the children of z.
        bool pink_flag = false;
        BGL_FORALL_OUTEDGES(vert, e, g_, Graph) {
          Vertex w = target(e, g_);
          double weight = weightProperty_[e].value();

          if (colorProperty_[w] != RED && costProperty_[w] + weight == cost) {
            // Actually, '<' should not be happened all the time.
            // And even '==' would very rarely occur, but possible.
            pink_flag = true;

            cancelAdoption(vert);
            predecessorProperty_[vert] = w;
            childrenProperty_[w]->push_back(vert);
            break; // If there exsits, take anyone among them.
          }
        }

        if (pink_flag) {
          continue;
        }*/

        colorProperty_[vert] = RED; // Set to 'red'
        reds.push_back(vert);
        // Even with multiple starting red nodes, there will be no re-visit since each starting node is
        // a root node of sub'tree' in DSPT. That is, if statement within for loop might be useless.
        // Someone would be curious, e.g., then why do we have to use priority queue in step2 ?
        // Just for 'pink' case. Yeap. We need to identify all the other parent candidates are red or not
        // prior to checking current node.
        std::vector<Vertex> *children = childrenProperty_[vert];

        for (unsigned int i = 0; i < children->size(); i++) {
            pq.push(weight_vertex(-costProperty_[(*children)[i]], (*children)[i]));
        }
    }

    // 'pq' is empty at here.

    // <Step 3-a. Find best non-red parent for each red node.
    for (unsigned int i = 0; i < reds.size(); i++) {
        // TODO : need to be verified
        // Cost/predecessor initialization.
        costProperty_[reds[i]] = std::numeric_limits<double>::infinity();
        cancelAdoption(reds[i]);

        BGL_FORALL_OUTEDGES(reds[i], e, g_, Graph) {
            Vertex w = target(e, g_);
            double weight = weightProperty_[e].value();

            if (colorProperty_[w] == RED) {
                continue;  // If red, put aside for a while.
            }

            if (costProperty_[reds[i]] > costProperty_[w] + weight) {
                costProperty_[reds[i]] = costProperty_[w] + weight;
                predecessorProperty_[reds[i]] = w;
            }
        }

        if (predecessorProperty_[reds[i]] != NULL) {
            childrenProperty_[predecessorProperty_[reds[i]]]->push_back(reds[i]);
        }

        pq.push(weight_vertex(-costProperty_[reds[i]], reds[i]));
    }

    // <Step 3-b. Propagate the changes; rewiring for 'red' nodes whether it can replace
    // existing parent node of near neighbors.
    while (!pq.empty()) {
        weight_vertex top = pq.top();
        pq.pop();

        double cost = -top.first;
        Vertex vert = top.second;

        if (costProperty_[vert] < cost) {
            continue;
        }

        BGL_FORALL_OUTEDGES(vert, e, g_, Graph) {
            Vertex w = target(e, g_);
            double weight = weightProperty_[e].value();

            if (colorProperty_[w] != RED) {
                continue;  // If not red, then skip.
            }

            if (cost + weight < costProperty_[w]) {
                costProperty_[w] = cost + weight;

                cancelAdoption(w);

                predecessorProperty_[w] = vert;
                childrenProperty_[vert]->push_back(w);
                pq.push(weight_vertex(-costProperty_[w], w));
            }
        }
    }

    // The end!
    // colorProperty_ is not necessary to be cleansed out, just increase variable, 'RED'.
}

// TODO : sync between children & edges.
void ompl::geometric::HarmoniousLazyPRMstarMulti::cancelAdoption(const Vertex &child) {
    if (predecessorProperty_[child] == NULL)
        return;

    std::vector<Vertex> *children = childrenProperty_[predecessorProperty_[child]];

    for (unsigned int i = 0; i < children->size(); i++) if ((*children)[i] == child) {
            std::swap((*children)[i], children->back());
            children->pop_back();
            break;
        }

    predecessorProperty_[child] = NULL;
}

// Vertex first.
ompl::base::PathPtr ompl::geometric::HarmoniousLazyPRMstarMulti::constructSolution(const Vertex &start, const Vertex &goal) {
    std::vector<Vertex> solution_path;

    // Construct a solution from DSPT.
    for (Vertex vert = goal; vert != NULL; vert = predecessorProperty_[vert])
        solution_path.push_back(vert);

    if (solution_path.empty() || solution_path.size() == 1)
        return base::PathPtr();

    // Goal = 0, Start = n - 1.
    // TODO : From goal or start ? which one is better?

    // auto from = solution_path.rbegin();
    std::vector<Vertex>::reverse_iterator from = solution_path.rbegin();
    for (std::vector<Vertex>::reverse_iterator to = from + 1; to != solution_path.rend(); ++to) {
        Edge e = boost::lookup_edge(*from, *to, g_).first; // Exhaustive search O(E) at worst case.
        unsigned int &evd = edgeValidityProperty_[e];

        if ((evd & VALIDITY_TRUE) == 0) { // Unknown
            bool result = true;
            // double weight = weightProperty_[e].value();

            result &= checkMotion(stateProperty_[*from], stateProperty_[*to]);

            if (result) {
                evd |= VALIDITY_TRUE;
            }
            else {
                boost::remove_edge(e, g_); // O(log(E/V)) time...

                cancelAdoption(*to);
                Increase(*to);
                return base::PathPtr();
            }
        }

        from = to;
    }

    PathGeometric *p = new PathGeometric(si_);

    // Feasible path is found, fetch optimized edges if possible.

    for (std::vector<Vertex>::const_reverse_iterator sol = solution_path.rbegin(); sol != solution_path.rend(); ++sol) {
        p->append(stateProperty_[*sol]);
    }

    return base::PathPtr(p);
}

void ompl::geometric::HarmoniousLazyPRMstarMulti::getPlannerData(base::PlannerData &data) const {
    Planner::getPlannerData(data);
    // Caution : it handles directional information regardless of the search graph setting
    //           which is undirectional graph.
}

double ompl::geometric::HarmoniousLazyPRMstarMulti::distanceFunction(const base::State *a, const base::State *b) const {
    return si_->distance(a, b);
}

double ompl::geometric::HarmoniousLazyPRMstarMulti::distanceFunctionHarmonious(const Vertex a, const Vertex b) const {
    double base = distanceFunctionBase(stateProperty_[a], stateProperty_[b]);
    double arm =  distanceFunctionJoints(stateProperty_[a], stateProperty_[b]);

    return base + arm;
}

double ompl::geometric::HarmoniousLazyPRMstarMulti::distanceFunctionBase(const base::State *a, const base::State *b) const {
    double dist = 0.0;
    int dim = si_->getStateDimension();

    std::vector<double> ca, cb;
    si_->getStateSpace()->copyToReals(ca, a);
    si_->getStateSpace()->copyToReals(cb, b);

    for (int i = 0; i < NUM_BASE_DOF; i++) {
        double dd = fabs(ca[i]-cb[i]);

        if(isContinuous_[i]) {
            if (dd > boost::math::constants::pi<double>()) {
                dd = 2.0 * boost::math::constants::pi<double>() - dd;
            }
        }

        dist += hs_.params_.weights_[i] * dd * hs_.params_.weights_[i] * dd;
    }

    return std::sqrt(dist);
}

double ompl::geometric::HarmoniousLazyPRMstarMulti::distanceFunctionJoints(const base::State *a, const base::State *b) const {
    double dist = 0.0;

    std::vector<double> ca, cb;
    si_->getStateSpace()->copyToReals(ca, a);
    si_->getStateSpace()->copyToReals(cb, b);

    for (int i = NUM_BASE_DOF; i < ca.size(); i++) {
        double dd = fabs(ca[i]-cb[i]);

        if(isContinuous_[i+3]){
            if (dd > boost::math::constants::pi<double>()) {
                dd = 2.0 * boost::math::constants::pi<double>() - dd;
            }
        }

        dist += hs_.params_.weights_[i] * dd * hs_.params_.weights_[i] * dd;
    }

    return std::sqrt(dist);
}

bool ompl::geometric::HarmoniousLazyPRMstarMulti::checkMotion(base::State *s1, base::State *s2) const {
//    /* Assume motion starts/ends in a valid configuration so v1/v2 are valid */
    bool result = true;
    int dim = si_->getStateDimension();
    int nd;

    int nd_j = validSegmentCount_compare(distanceFunctionJoints(s1, s2), hs_.params_.factor_joint_);
    int nd_b = validSegmentCount_compare(distanceFunctionBase(s1, s2), hs_.params_.factor_base_);
    if(nd_j > nd_b){
        nd = nd_j;
    }
    else{
        nd = nd_b;
    }

    if (nd > 1) {
        /* Temporary storage for the checked state */
        base::State *test = si_->allocState();
        std::queue<std::pair<unsigned int, unsigned int> > q;
        q.push(std::make_pair(1, nd - 1));

        while (!q.empty()) {
            std::pair<unsigned int, unsigned int> range = q.front();
            unsigned int mid;

            mid = (range.first + range.second) / 2;
            interpolate(s1, s2, (double)mid / (double)nd, test);
//            si_->getStateSpace()->interpolate(s1, s2, (double)mid / (double)nd, test);

            if (!si_->isValid(test)) {
                result = false;
                break;
            }

            q.pop();
            if (range.first < mid)
                q.push(std::make_pair(range.first, mid - 1));
            if (mid < range.second) {
                q.push(std::make_pair(mid + 1, range.second));
            } // if mid == first, no more recursion.
        }
    }

    return result;
}

unsigned int ompl::geometric::HarmoniousLazyPRMstarMulti::validSegmentCount_compare(const double dist,
                                                                                const double longestValidSegment) const {
    return si_->getStateSpace()->getValidSegmentCountFactor() * (unsigned int)ceil(dist / longestValidSegment);
}

void ompl::geometric::HarmoniousLazyPRMstarMulti::interpolate(const base::State *from, const base::State *to, const double t, base::State *state) const
{
    const base::RealVectorStateSpace::StateType *rfrom = static_cast<const base::RealVectorStateSpace::StateType*>(from);
    const base::RealVectorStateSpace::StateType *rto = static_cast<const base::RealVectorStateSpace::StateType*>(to);
    const base::RealVectorStateSpace::StateType *rstate = static_cast<base::RealVectorStateSpace::StateType*>(state);
    for (unsigned int i = 0 ; i < isContinuous_.size() ; ++i){
        if(isContinuous_[i]){
            double diff = rto->values[i] - rfrom->values[i];
            if (fabs(diff) <= boost::math::constants::pi<double>())
                rstate->values[i] = rfrom->values[i] + diff * t;
            else
            {
                double &v = rstate->values[i];
                if (diff > 0.0)
                    diff = 2.0 * boost::math::constants::pi<double>() - diff;
                else
                    diff = -2.0 * boost::math::constants::pi<double>() - diff;
                v = rfrom->values[i] - diff * t;
                // input states are within bounds, so the following check is sufficient
                if (v > boost::math::constants::pi<double>())
                    v -= 2.0 * boost::math::constants::pi<double>();
                else
                if (v < -boost::math::constants::pi<double>())
                    v += 2.0 * boost::math::constants::pi<double>();
            }
        }
        else{
            rstate->values[i] = rfrom->values[i] + (rto->values[i] - rfrom->values[i]) * t;
        }
    }
}


/*    Copyright (c) 2010-2019, Delft University of Technology
 *    All rigths reserved
 *
 *    This file is part of the Tudat. Redistribution and use in source and
 *    binary forms, with or without modification, are permitted exclusively
 *    under the terms of the Modified BSD license. You should have received
 *    a copy of the license with this file. If not, please or visit:
 *    http://tudat.tudelft.nl/LICENSE.
 */

#include <boost/lambda/lambda.hpp>

#include "tudat/interface/spice/spiceEphemeris.h"
#include "tudat/astro/ephemerides/customEphemeris.h"
#include "tudat/astro/ephemerides/keplerEphemeris.h"
#include "tudat/astro/ephemerides/multiArcEphemeris.h"
#include "tudat/astro/ephemerides/tabulatedEphemeris.h"
#include "tudat/astro/ephemerides/approximatePlanetPositions.h"
#include "tudat/astro/ephemerides/approximatePlanetPositionsCircularCoplanar.h"
#include "tudat/astro/ephemerides/constantEphemeris.h"
#include "tudat/math/interpolators/lagrangeInterpolator.h"
#include "tudat/simulation/environment_setup/createEphemeris.h"

namespace tudat
{

namespace simulation_setup
{

using namespace ephemerides;

//! Function to create an ephemeris model.
std::shared_ptr< ephemerides::Ephemeris > createBodyEphemeris(
        const std::shared_ptr< EphemerisSettings > ephemerisSettings,
        const std::string& bodyName )
{
    // Declare return object.
    std::shared_ptr< ephemerides::Ephemeris > ephemeris;

    if( ephemerisSettings->getMakeMultiArcEphemeris( ) )
    {
        std::map< double, std::shared_ptr< Ephemeris > > singleArcEphemerides;
        ephemerisSettings->resetMakeMultiArcEphemeris( false );

        singleArcEphemerides[ -std::numeric_limits< double >::lowest( ) ] = createBodyEphemeris(
                    ephemerisSettings, bodyName );

        ephemeris = std::make_shared< MultiArcEphemeris >(
                    singleArcEphemerides, ephemerisSettings->getFrameOrigin( ), ephemerisSettings->getFrameOrientation( ) );
    }
    else
    {

        // Check which type of ephemeris model is to be created.
        switch( ephemerisSettings->getEphemerisType( ) )
        {
        case direct_spice_ephemeris:
        {
            // Check consistency of type and class.
            std::shared_ptr< DirectSpiceEphemerisSettings > directEphemerisSettings =
                    std::dynamic_pointer_cast< DirectSpiceEphemerisSettings >( ephemerisSettings );
            if( directEphemerisSettings == nullptr )
            {
                throw std::runtime_error(
                            "Error, expected direct spice ephemeris settings for body " + bodyName );
            }
            else
            {
                std::string inputName = ( directEphemerisSettings->getBodyNameOverride( ) == "" ) ?
                            bodyName : directEphemerisSettings->getBodyNameOverride( );

                // Create corresponding ephemeris object.
                ephemeris = std::make_shared< SpiceEphemeris >(
                            inputName,
                            directEphemerisSettings->getFrameOrigin( ),
                            directEphemerisSettings->getCorrectForStellarAberration( ),
                            directEphemerisSettings->getCorrectForLightTimeAberration( ),
                            directEphemerisSettings->getConvergeLighTimeAberration( ),
                            directEphemerisSettings->getFrameOrientation( ) );
            }
            break;
        }
        case interpolated_spice:
        {
            // Check consistency of type and class.
            std::shared_ptr< InterpolatedSpiceEphemerisSettings > interpolatedEphemerisSettings =
                    std::dynamic_pointer_cast< InterpolatedSpiceEphemerisSettings >(
                        ephemerisSettings );
            if( interpolatedEphemerisSettings == nullptr )
            {
                throw std::runtime_error(
                            "Error, expected tabulated spice ephemeris settings for body " + bodyName );
            }
            else
            {
                // Since only the barycenters of planetary systems are included in the standard DE
                // ephemerides, append 'Barycenter' to body name.
                std::string inputName = ( interpolatedEphemerisSettings->getBodyNameOverride( ) == "" ) ?
                            bodyName : interpolatedEphemerisSettings->getBodyNameOverride( );
                // Create corresponding ephemeris object.
                if( !interpolatedEphemerisSettings->getUseLongDoubleStates( ) )
                {
                    ephemeris = createTabulatedEphemerisFromSpice< double, double >(
                                inputName,
                                interpolatedEphemerisSettings->getInitialTime( ),
                                interpolatedEphemerisSettings->getFinalTime( ),
                                interpolatedEphemerisSettings->getTimeStep( ),
                                interpolatedEphemerisSettings->getFrameOrigin( ),
                                interpolatedEphemerisSettings->getFrameOrientation( ),
                                interpolatedEphemerisSettings->getInterpolatorSettings( ) );
                }
                else
                {
#if( TUDAT_BUILD_WITH_EXTENDED_PRECISION_PROPAGATION_TOOLS )

                    ephemeris = createTabulatedEphemerisFromSpice< long double, double >(
                                inputName,
                                static_cast< long double >( interpolatedEphemerisSettings->getInitialTime( ) ),
                                static_cast< long double >( interpolatedEphemerisSettings->getFinalTime( ) ),
                                static_cast< long double >( interpolatedEphemerisSettings->getTimeStep( ) ),
                                interpolatedEphemerisSettings->getFrameOrigin( ),
                                interpolatedEphemerisSettings->getFrameOrientation( ),
                                interpolatedEphemerisSettings->getInterpolatorSettings( ) );
#else
                    throw std::runtime_error( "Error, long double compilation is turned off; requested long double tabulated ephemeris" );
#endif
                }
            }
            break;
        }
        case tabulated_ephemeris:
        {
            // Check consistency of type and class.
            std::shared_ptr< TabulatedEphemerisSettings > tabulatedEphemerisSettings =
                    std::dynamic_pointer_cast< TabulatedEphemerisSettings >( ephemerisSettings );
            if( tabulatedEphemerisSettings == nullptr )
            {
                throw std::runtime_error(
                            "Error, expected tabulated ephemeris settings for body " + bodyName );
            }
            else
            {
                // Create corresponding ephemeris object.
                if( !tabulatedEphemerisSettings->getUseLongDoubleStates( ) )
                {
                    if( tabulatedEphemerisSettings->getBodyStateHistory( ).size( ) != 0 )
                    {
                        ephemeris = std::make_shared< TabulatedCartesianEphemeris< > >(
                                    std::make_shared<
                                    interpolators::LagrangeInterpolator< double, Eigen::Vector6d > >
                                    ( tabulatedEphemerisSettings->getBodyStateHistory( ), 6,
                                      interpolators::huntingAlgorithm,
                                      interpolators::lagrange_cubic_spline_boundary_interpolation ),
                                    tabulatedEphemerisSettings->getFrameOrigin( ),
                                    tabulatedEphemerisSettings->getFrameOrientation( ) );
                    }
                    else
                    {
                        ephemeris = std::make_shared< TabulatedCartesianEphemeris< > >(
                                    std::shared_ptr< interpolators::OneDimensionalInterpolator< double, Eigen::Vector6d > >( ),
                                      tabulatedEphemerisSettings->getFrameOrigin( ),
                                      tabulatedEphemerisSettings->getFrameOrientation( ) );
                    }
                }
                else
                {
#if( TUDAT_BUILD_WITH_EXTENDED_PRECISION_PROPAGATION_TOOLS )

                    // Cast input history to required type.
                    if( tabulatedEphemerisSettings->getBodyStateHistory( ).size( ) != 0 )
                    {
                        std::map< double, Eigen::Vector6d > originalStateHistory =
                                tabulatedEphemerisSettings->getBodyStateHistory( );
                        std::map< double, Eigen::Matrix< long double, 6, 1 > > longStateHistory;

                        for( std::map< double, Eigen::Vector6d >::const_iterator stateIterator =
                             originalStateHistory.begin( ); stateIterator != originalStateHistory.end( ); stateIterator++ )
                        {
                            longStateHistory[ stateIterator->first ] = stateIterator->second.cast< long double >( );
                            ephemeris =
                                    std::make_shared< TabulatedCartesianEphemeris< long double, double > >(
                                        std::make_shared< interpolators::LagrangeInterpolator<
                                        double, Eigen::Matrix< long double, 6, 1 > > >
                                        ( longStateHistory, 6,
                                          interpolators::huntingAlgorithm,
                                          interpolators::lagrange_cubic_spline_boundary_interpolation ),
                                        tabulatedEphemerisSettings->getFrameOrigin( ),
                                        tabulatedEphemerisSettings->getFrameOrientation( ) );
                        }
                    }
                    else
                    {
                        ephemeris = std::make_shared< TabulatedCartesianEphemeris< long double, double > >(
                                    std::shared_ptr< interpolators::OneDimensionalInterpolator<
                                    double, Eigen::Matrix< long double, 6, 1 > > >( ),
                                      tabulatedEphemerisSettings->getFrameOrigin( ),
                                      tabulatedEphemerisSettings->getFrameOrientation( ) );
                    }
#else
                    throw std::runtime_error( "Error, long double compilation is turned off; requested long double tabulated ephemeris" );
#endif
                }
            }
            break;
        }
        case auto_generated_tabulated_ephemeris:
        {
            // Check consistency of type and class.
            std::shared_ptr< AutoGeneratedTabulatedEphemerisSettings > tabulatedEphemerisSettings =
                    std::dynamic_pointer_cast< AutoGeneratedTabulatedEphemerisSettings >( ephemerisSettings );
            if( tabulatedEphemerisSettings == nullptr )
            {
                throw std::runtime_error( "Error, expected auto-generate tabulated ephemeris settings for " + bodyName );
            }
            else
            {
                // Create ephemeris
                ephemeris = getTabulatedEphemeris(
                            createBodyEphemeris( tabulatedEphemerisSettings->getEphemerisSettings( ), bodyName ),
                            tabulatedEphemerisSettings->getStartTime( ), tabulatedEphemerisSettings->getEndTime( ),
                            tabulatedEphemerisSettings->getTimeStep( ), tabulatedEphemerisSettings->getInterpolatorSettings( ) );
            }
            break;
        }
        case constant_ephemeris:
        {
            // Check consistency of type and class.
            std::shared_ptr< ConstantEphemerisSettings > constantEphemerisSettings =
                    std::dynamic_pointer_cast< ConstantEphemerisSettings >( ephemerisSettings );
            if( constantEphemerisSettings == nullptr )
            {
                throw std::runtime_error( "Error, expected constant ephemeris settings for " + bodyName );
            }
            else
            {
                // Create ephemeris
                ephemeris = std::make_shared< ConstantEphemeris >(
                            [ = ]( ){ return constantEphemerisSettings->getConstantState( ); },
                            constantEphemerisSettings->getFrameOrigin( ),
                            constantEphemerisSettings->getFrameOrientation( ) );
            }
            break;
        }
        case custom_ephemeris:
        {
            // Check consistency of type and class.
            std::shared_ptr< CustomEphemerisSettings > customEphemerisSettings =
                    std::dynamic_pointer_cast< CustomEphemerisSettings >( ephemerisSettings );
            if( customEphemerisSettings == nullptr )
            {
                throw std::runtime_error( "Error, expected custom ephemeris settings for " + bodyName );
            }
            else
            {
                // Create ephemeris
                ephemeris = std::make_shared< CustomEphemeris >(
                            customEphemerisSettings->getCustomStateFunction( ),
                            customEphemerisSettings->getFrameOrigin( ),
                            customEphemerisSettings->getFrameOrientation( ) );
            }
            break;
        }

        case kepler_ephemeris:
        {
            // Check consistency of type and class.
            std::shared_ptr< KeplerEphemerisSettings > keplerEphemerisSettings =
                    std::dynamic_pointer_cast< KeplerEphemerisSettings >( ephemerisSettings );
            if( keplerEphemerisSettings == nullptr )
            {
                throw std::runtime_error( "Error, expected Kepler ephemeris settings for " + bodyName );
            }
            else
            {
                // Create ephemeris
                ephemeris = std::make_shared< KeplerEphemeris >(
                            keplerEphemerisSettings->getInitialStateInKeplerianElements( ),
                            keplerEphemerisSettings->getEpochOfInitialState( ),
                            keplerEphemerisSettings->getCentralBodyGravitationalParameter( ),
                            keplerEphemerisSettings->getFrameOrigin( ),
                            keplerEphemerisSettings->getFrameOrientation( ),
                            keplerEphemerisSettings->getRootFinderAbsoluteTolerance( ),
                            keplerEphemerisSettings->getRootFinderMaximumNumberOfIterations( ) );
            }
            break;
        }
        case approximate_planet_positions:
        {
            // Check consistency of type and class.
            std::shared_ptr< ApproximateJplEphemerisSettings > approximateEphemerisSettings =
                    std::dynamic_pointer_cast< ApproximateJplEphemerisSettings >(
                        ephemerisSettings );
            if( approximateEphemerisSettings == nullptr )
            {
                throw std::runtime_error(
                            "Error, expected approximate ephemeris settings for body " + bodyName );
            }
            else
            {

//                ephemerides::ApproximatePlanetPositionsBase::BodiesWithEphemerisData bodyToUse;
//                if( approximateEphemerisSettings->getBodyIdentifier( ) ==
//                        ephemerides::ApproximatePlanetPositionsBase::BodiesWithEphemerisData::undefined )
//                {
//                    try
//                    {
//                        bodyToUse = ephemerides::ApproximatePlanetPositionsBase::getBodiesWithEphemerisDataId( bodyName );
//                    }
//                    catch( std::runtime_error const& )

//                    {
//                        throw std::runtime_error( "Error, approximate ephemeris not available for body: " + bodyName + " when creating ephemeris." );
//                    }
//                }
//                else
//                {
//                    bodyToUse = approximateEphemerisSettings->getBodyIdentifier( );
//                }
//=======
//>>>>>>> origin/feature/mga_estimation_refactor_merge

                // Create corresponding ephemeris object.
                if( approximateEphemerisSettings->getUseCircularCoplanarApproximation( ) )
                {
                    ephemeris = std::make_shared< ApproximateJplCircularCoplanarEphemeris >(
                                approximateEphemerisSettings->getBodyName( ) );
                }
                else
                {
                    ephemeris = std::make_shared< ApproximateJplEphemeris >(
                                approximateEphemerisSettings->getBodyName( ) );
                }
            }
            break;
        }
		case direct_tle_ephemeris:
		{
			// Check consistency of type and class.
			std::shared_ptr< DirectTleEphemerisSettings > directTleEphemerisSettings =
					std::dynamic_pointer_cast< DirectTleEphemerisSettings >( ephemerisSettings );
			if( directTleEphemerisSettings == nullptr )
			{
				throw std::runtime_error(
						"Error, expected direct TLE ephemeris settings for body " + bodyName );
			}
			// Check if the Earth is present in the simulation
			else
			{
				//std::string inputName = bodyName;

				// Check period of the satellite for correct SDP setting

				// Create corresponding ephemeris object.
				ephemeris = std::make_shared< TleEphemeris >(
						directTleEphemerisSettings->getFrameOrigin(),
						directTleEphemerisSettings->getFrameOrientation(),
						directTleEphemerisSettings->getTle()
						);
			}
			break;
		}
		case interpolated_tle_ephemeris:
		{
			// Check consistency of type and class.
			std::shared_ptr< InterpolatedTleEphemerisSettings > interpolatedTleEphemerisSettings =
					std::dynamic_pointer_cast< InterpolatedTleEphemerisSettings >(
							ephemerisSettings );
			if( interpolatedTleEphemerisSettings == nullptr )
			{
				throw std::runtime_error(
						"Error, expected interpolated TLE ephemeris settings for body " + bodyName );
			}
			else
			{

				ephemeris = createTabulatedEphemerisFromTLE< double, double >(
						bodyName,
						interpolatedTleEphemerisSettings->getInitialTime( ),
						interpolatedTleEphemerisSettings->getFinalTime( ),
						interpolatedTleEphemerisSettings->getTimeStep( ),
						interpolatedTleEphemerisSettings->getFrameOrigin( ),
						interpolatedTleEphemerisSettings->getFrameOrientation( ),
						interpolatedTleEphemerisSettings->getTle( ),
						interpolatedTleEphemerisSettings->getInterpolatorSettings( )
						);
			}
			break;
		}
        case scaled_ephemeris:
        {
            // Check consistency of type and class.
            std::shared_ptr< ScaledEphemerisSettings > scaledEphemeriSettings =
                    std::dynamic_pointer_cast< ScaledEphemerisSettings >(
                        ephemerisSettings );
            if( scaledEphemeriSettings == nullptr )
            {
                throw std::runtime_error(
                            "Error, expected scaled ephemeris settings for body " + bodyName );
            }
            else
            {
                std::shared_ptr< Ephemeris > baseEphemeris = createBodyEphemeris(
                            scaledEphemeriSettings->getBaseSettings( ), bodyName );
                ephemeris = std::make_shared< ScaledEphemeris >(
                            baseEphemeris, scaledEphemeriSettings->getScaling( ), scaledEphemeriSettings->getIsScalingAbsolute( ) );
            }
            break;
        }
        default:
        {
            throw std::runtime_error(
                        "Error, did not recognize ephemeris model settings type " +
                        std::to_string( ephemerisSettings->getEphemerisType( ) ) );
        }
        }
    }
    return ephemeris;

}

//! Function that retrieves the time interval at which an ephemeris can be safely interrogated
std::pair< double, double > getSafeInterpolationInterval( const std::shared_ptr< ephemerides::Ephemeris > ephemerisModel )
{
    // Make default output pair
    std::pair< double, double > safeInterval = std::make_pair(
                std::numeric_limits< double >::lowest( ),  std::numeric_limits< double >::max( ) );

    // Check if model is tabulated, and retrieve safe interval from model
    if( isTabulatedEphemeris( ephemerisModel ) )
    {
        safeInterval = getTabulatedEphemerisSafeInterval( ephemerisModel );
    }
    // Check if model is multi-arc, and retrieve safe intervals from first and last arc.
    else if( std::dynamic_pointer_cast< ephemerides::MultiArcEphemeris >( ephemerisModel ) != nullptr )
    {
        std::shared_ptr< ephemerides::MultiArcEphemeris > multiArcEphemerisModel  =
                std::dynamic_pointer_cast< ephemerides::MultiArcEphemeris >( ephemerisModel );
        safeInterval.first = getSafeInterpolationInterval( multiArcEphemerisModel->getSingleArcEphemerides( ).at( 0 ) ).first;
        safeInterval.second = getSafeInterpolationInterval(
                    multiArcEphemerisModel->getSingleArcEphemerides( ).at(
                        multiArcEphemerisModel->getSingleArcEphemerides( ).size( ) - 1 ) ).second;
    }
    return safeInterval;
}

} // namespace simulation_setup

} // namespace tudat


#include "diffuclass.h"
#include <iostream>
#include <math.h>
#include <boost/filesystem.hpp>

//------------------------------------------------------------------------------
//                           Declare functions
//------------------------------------------------------------------------------

extern PetscErrorCode ComputeMatrix(KSP, Mat, Mat, void*);

PetscErrorCode MyMatMult(Mat,Vec,Vec);

//------------------------------------------------------------------------------
//							  Public Methods
//------------------------------------------------------------------------------
//TODO: include all boundary condition elements and individually parallelized grids
void diffusionPETSc::initDiffusion(eQ::diffusionSolver::params &initParams)
//void diffusionPETSc::initDiffusion(MPI_Comm comm, std::vector<std::string> filePaths, int argc, char* argv[])
{
    auto argc = initParams.argc;
    auto argv = initParams.argv;
    int 	 i, j;

    /*
    //Split input communicator into subcommunicators for each desired diffusion grid
    MPI_Comm subComm;
    int 	 myRankMPI, commSize, i, j;

    MPI_Comm_size(comm, &commSize);
    MPI_Comm_rank(comm, &myRankMPI);

    MPI_Comm_split(comm, myRankMPI, 0, &subComm);

    initData.subCommunicator = subComm;
    //Record vector of diffusion constants and choose the appropriate constant for each processor
    std::vector<double>		d_vector;
    d_vector = std::vector<double>(eQ::data::parameters["D_HSL"].get<std::vector<double>>());
    initData.diffusionConstant = d_vector.at(myRankMPI % d_vector.size());

    //Set domain size based on parameters
    initData.xLengthMicrons = int(eQ::data::parameters["simulationTrapWidthMicrons"]);
    initData.yLengthMicrons = int(eQ::data::parameters["simulationTrapHeightMicrons"]);

    //Set time step
    initData.dt = eQ::data::parameters["dt"];

    //Set grid spacing in domain
    initData.h = 1 / double(eQ::data::parameters["nodesPerMicronSignaling"]);

    //Set the appropriate filepath to each grid for data recording
    initData.directoryName = filePaths.at(myRankMPI % d_vector.size());
*/
    initData.subCommunicator = initParams.comm;
    initData.diffusionConstant = initParams.D_HSL;//direct read from init params
    initData.xLengthMicrons = initParams.trapWidthMicrons;
    initData.yLengthMicrons = initParams.trapHeightMicrons;
    initData.dt = initParams.dt;
    initData.h = 1.0 / initParams.nodesPerMicron;

    initData.directoryName = initParams.filePath + "petsc";
     boost::filesystem::path dstFolder = initData.directoryName;
     boost::filesystem::create_directory(dstFolder);

     initData.objectName = "grid";

    //Set initial boundary conditions
    if("DIRICHLET_0" == eQ::data::parameters["boundaryType"]){
        initData.homogeneousDirichlet = PETSC_TRUE;

        initData.topDirichletCoefficient = 1;
        initData.bottomDirichletCoefficient = 1;
        initData.leftDirichletCoefficient = 1;
        initData.rightDirichletCoefficient = 1;

        initData.topNeumannCoefficient = 0;
        initData.bottomNeumannCoefficient = 0;
        initData.leftNeumannCoefficient = 0;
        initData.rightNeumannCoefficient = 0;

        initData.topBoundaryValue = 0;
        initData.bottomBoundaryValue = 0;
        initData.leftBoundaryValue = 0;
        initData.rightBoundaryValue = 0;
    }

    //Initialize PETSc element of diffusion
    InitializeDiffusion(&initData, argc, argv);

    //Resize and populate vectors for data transfer
    solution_vector.resize(gridNodesX * gridNodesY);
    allXCoordinates.resize(gridNodesX * gridNodesY);
    allYCoordinates.resize(gridNodesX * gridNodesY);

    for (j=0; j < gridNodesY; j++){
        for (i=0; i < gridNodesX; i++){
            allXCoordinates.at(i + j * gridNodesX) = i*initData.h;
            allYCoordinates.at(i + j * gridNodesX) = j*initData.h;
            solution_vector.at(i + j * gridNodesX) = 0;
        }
    }

    //Copy initial conditions to solution vector
    ReadGridValues(allXCoordinates, allYCoordinates, &solution_vector);
}

void diffusionPETSc::stepDiffusion(void)
{
    //Copy modified solution vector to u0
    //solution_vector is our public vector of all solution data - copy that over.

    WriteGridValues(allXCoordinates, allYCoordinates, solution_vector);

    TimeStep();

    ReadGridValues(allXCoordinates, allYCoordinates, &solution_vector);
}

//TODO: implement each wall separately
void diffusionPETSc::setBoundaryValues(const eQ::data::parametersType &bvals)
{
    //JW: over-ride for now: set boundaries explicitly by writing the data structures
    return;

    if(bool(bvals["allBoundaries"]) == true){
        gridData->topBoundaryValue = double(bvals["boundaryValue"]);
        gridData->bottomBoundaryValue = double(bvals["boundaryValue"]);
        gridData->leftBoundaryValue = double(bvals["boundaryValue"]);
        gridData->rightBoundaryValue = double(bvals["boundaryValue"]);
    }
}

//TODO: verify accuracy, check about choosing walls
eQ::data::parametersType diffusionPETSc::getBoundaryFlux(void)
{
    //Compute and return boundary flux across all walls
    eQ::data::parametersType 	fluxData;
    double 			   	totalBoundaryFlux, boundarySlope;
    int 				i, j;

    for (i = 0; i < gridNodesX; i++){
        for (j = 0; j < gridNodesY; j++){
            if (i == 0){
                boundarySlope = boundarySlope + (solution_vector.at(2 + (gridNodesX) * j) - \
                        solution_vector.at((gridNodesX)*j))/(2*gridData->h);
            }
            else if (i == gridNodesX - 1){
                boundarySlope = boundarySlope + (solution_vector.at(i - 2 + (gridNodesX-1)*j) - \
                        solution_vector.at(i + (gridNodesX-1)*j))/(2*gridData->h);
            }
            else if (j == 0){
                boundarySlope = boundarySlope + (solution_vector.at(i + (gridNodesX-1) * 2) - \
                        solution_vector.at(i))/(2*gridData->h);
            }
            else if (j == gridNodesY - 1){
                boundarySlope = boundarySlope + (solution_vector.at(i + (gridNodesX-1) * (j - 2)) - \
                        solution_vector.at(i + (gridNodesX-1)*j))/(2*gridData->h);
            }
        }
    }

    totalBoundaryFlux = gridData->diffusionConstant * gridData->dt * boundarySlope;
    fluxData["totalFlux"] = totalBoundaryFlux;
    return fluxData;
}

void diffusionPETSc::writeDiffusionFiles(double timestamp)
{
    //Work in timestamp
    RecordData();
}

/*
void diffusionPETSc::writeDataFiles(double dt)
{
    //Need to write full function
}
*/

void diffusionPETSc::finalize(void)
{
    DiffusionFinalize(PETSC_TRUE);
}

//------------------------------------------------------------------------------
//                            Private Methods
//------------------------------------------------------------------------------

PetscErrorCode diffusionPETSc::ApplyBoundaryConditions()
{
  //Applys Neumann/Robin Boundary Conditions to RHS vector
  PetscInt       i,j,xm,ym,xs,ys;
  PetscScalar    **RHSarray;

  //Get subgrid information for each processor and get arrays from vector
  ierr = DMDAGetCorners(distributedArray,&xs,&ys,0,&xm,&ym,0);CHKERRQ(ierr);
  ierr = DMDAVecGetArray(distributedArray, globalVector, &RHSarray);CHKERRQ(ierr);

  //JW NOTE:  each node can check its range whether it has boundary points
  //  then, traverse the boundary explicitly, rather than traversing the entire i,j space
    const double twoFh = 2 * gridData->fourierNumber * gridData->h;//
//    double topD = gridData->topDirichletCoefficient;
//    double bottomD = gridData->bottomDirichletCoefficient;
//    double leftD = gridData->leftDirichletCoefficient;
//    double rightD = gridData->rightDirichletCoefficient;
        const double topN = gridData->topNeumannCoefficient;
        const double bottomN = gridData->bottomNeumannCoefficient;
        const double leftN = gridData->leftNeumannCoefficient;
        const double rightN = gridData->rightNeumannCoefficient;

  //Alter values in arrays to match Boundary Conditions
  for (j = ys; j < ys+ym; j++){
    for (i = xs; i < xs+xm; i++){
//      if (i == 0 || i == gridNodesX-1 || j == 0 || j == gridNodesY-1){

          //TOP WALL
          if(j == gridNodesY-1)
          {
              const double topBV = topBoundaryValue[i];
              //Set non-Dirichlet for top wall
              if (topN != 0)
              {//if TL (TR) corner, only set to N if left (right) wall is N
                  if ((i != 0 || leftN != 0) && (i != gridNodesX-1 || rightN != 0))
                        RHSarray[j][i] += (twoFh * topBV) / topN;
              }
              else
//                  RHSarray[j][i] = topBV / topD;
                  RHSarray[j][i] = topBV;//if N==0, then D must be 1
          }
          //BOTTOM WALL
          else if(j == 0)
          {
              const double bottomBV = bottomBoundaryValue[i];
              //Set non-Dirichlet for bottom wall
              if (bottomN != 0)
              {
                  if ((i != 0 || leftN != 0) && (i != gridNodesX-1 || rightN != 0))
                        RHSarray[j][i] +=  (twoFh * bottomBV) / bottomN;
              }
              else
//                  RHSarray[j][i] = bottomBV / bottomD;
                  RHSarray[j][i] = bottomBV;//if N==0, then D must be 1
          }
          //RIGHT WALL
          if(i == gridNodesX-1)
          {
              const double rightBV = gridData->rightBoundaryValue;
              //Set non-Dirichlet for right wall
              if (rightN != 0)
              {
                  if ((j != 0 || bottomN != 0) && (j != gridNodesY-1 || topN != 0))
                      RHSarray[j][i] += (twoFh * rightBV) / rightN;
              }
              else
//                  RHSarray[j][i] = rightBV / rightD;
                  RHSarray[j][i] = rightBV;//if N==0, then D must be 1
          }
          //LEFT WALL
          else if(i == 0)
          {
              const double leftBV = gridData->leftBoundaryValue;
              //Set non-Dirichlet for left wall
              if (leftN != 0)
              {
                  if ((j != 0 || bottomN != 0) && (j != gridNodesY-1 || topN != 0))
                      RHSarray[j][i] += (twoFh * leftBV) / leftN;
              }
              else
//                  RHSarray[j][i] = leftBV / leftD;
                  RHSarray[j][i] = leftBV;//if N==0, then D must be 1
          }
//      }
    }
  }
/*
  //Alter values in arrays to match Boundary Conditions
  for (j = ys; j < ys+ym; j++){
    for (i = xs; i < xs+xm; i++){
      if (i == 0 || i == gridNodesX-1 || j == 0 || j == gridNodesY-1){
        //Set non-Dirichlet for top wall
        if (gridData->topNeumannCoefficient != 0 && j == gridNodesY-1){
          //Only change corners if they are also Neumann Boundary Conditions
          if ((i != 0 || gridData->leftNeumannCoefficient != 0) &&
              (i != gridNodesX-1 || gridData->rightNeumannCoefficient != 0)){
                RHSarray[j][i] = RHSarray[j][i] + (2 * gridData->fourierNumber * gridData->h * \
                  gridData->topBoundaryValue) / gridData->topNeumannCoefficient;
          }
        }
        //Set non-Dirichlet for bottom wall
        if (gridData->bottomNeumannCoefficient != 0 && j == 0){
          //Only change corners if they are also Neumann Boundary Conditions
          if ((i != 0 || gridData->leftNeumannCoefficient != 0) &&
              (i != gridNodesX-1 || gridData->rightNeumannCoefficient != 0)){
            RHSarray[j][i] = RHSarray[j][i] + (2 * gridData->fourierNumber * gridData->h * \
              gridData->bottomBoundaryValue) / gridData->bottomNeumannCoefficient;
          }
        }
        //Set non-Dirichlet for left wall
        if (gridData->leftNeumannCoefficient != 0 && i == 0){
          //Only change corners if they are also Neumann Boundary Conditions
          if ((j != 0 || gridData->bottomNeumannCoefficient != 0) &&
              (j != gridNodesY-1 || gridData->topNeumannCoefficient != 0)){
            RHSarray[j][i] = RHSarray[j][i] + (2 * gridData->fourierNumber * gridData->h * \
              gridData->leftBoundaryValue) / gridData->leftNeumannCoefficient;
          }
        }
        //Set non-Dirichlet for right wall
        if (gridData->rightNeumannCoefficient != 0 && i == gridNodesX-1){
          //Only change corners if they are also Neumann Boundary Conditions
          if ((j != 0 || gridData->bottomNeumannCoefficient != 0) &&
              (j != gridNodesY-1 || gridData->topNeumannCoefficient != 0)){
            RHSarray[j][i] = RHSarray[j][i] + (2 * gridData->fourierNumber * gridData->h * \
              gridData->rightBoundaryValue) / gridData->rightNeumannCoefficient;
          }
        }

        //Otherwise, set Dirichlet conditions
        if (gridData->topNeumannCoefficient == 0 && j == gridNodesY-1){
          RHSarray[j][i] = gridData->topBoundaryValue / gridData->topDirichletCoefficient;
        }
        else if (gridData->bottomNeumannCoefficient == 0 && j == 0){
          RHSarray[j][i] = gridData->bottomBoundaryValue / gridData->bottomDirichletCoefficient;
        }
        else if (gridData->leftNeumannCoefficient == 0 && i == 0){
          RHSarray[j][i] = gridData->leftBoundaryValue / gridData->leftDirichletCoefficient;
        }
        else if (gridData->rightNeumannCoefficient == 0 && i == gridNodesX-1){
          RHSarray[j][i] = gridData->rightBoundaryValue / gridData->rightDirichletCoefficient;
        }
      }
    }
  }
*/
  ierr = DMDAVecRestoreArray(distributedArray, globalVector, &RHSarray);CHKERRQ(ierr);

  return ierr;
}

//Figure out argc and argv
PetscErrorCode diffusionPETSc::InitializeDiffusion (DiffusionData *dataStruct, int argc, char **argv)
{
  //Initializes sets data values, inializes PETSc, and otherwise takes care of
  //all preparation needed in order to iterate. It should be noted that this
  //function does NOT set any initial conditions.

  gridData = dataStruct;
  //Initialize PETSc and assign passed communicator to DIFFU_COMM
  DIFFU_COMM = gridData->subCommunicator;
  ierr = PetscInitialize(&argc,&argv,(char*)0,NULL);
  MPI_Comm_size(DIFFU_COMM, &commSize);

  //Set up viewer for print statements
  printViewer = PETSC_VIEWER_STDOUT_(DIFFU_COMM);

  //Calculate crucial values for grid
  gridNodesX = (gridData->xLengthMicrons/gridData->h) + 1;
  gridNodesY = (gridData->yLengthMicrons/gridData->h) + 1;
  gridData->totalNodes = gridNodesX * gridNodesY;
  gridData->fourierNumber = (gridData->diffusionConstant*gridData->dt) / (gridData->h*gridData->h);

  //Set boundary and stencil types for distributed array
  boundary = DM_BOUNDARY_NONE;
  stencilType = DMDA_STENCIL_STAR;

  //Create Krylov Subspace method(KSP) context and Distributed Array(DMDA)
  ierr = KSPCreate(DIFFU_COMM, &krylovSolver);CHKERRQ(ierr);
  ierr = DMDACreate2d(DIFFU_COMM, boundary, boundary, stencilType, gridNodesX, gridNodesY, PETSC_DECIDE, \
    PETSC_DECIDE, 1, 1, NULL, NULL, &distributedArray);CHKERRQ(ierr);

  //Configure DMDA
  ierr = DMSetFromOptions(distributedArray);CHKERRQ(ierr);
  ierr = DMSetUp(distributedArray);CHKERRQ(ierr);
  ierr = DMDASetUniformCoordinates(distributedArray, 0, gridData->xLengthMicrons, 0, \
    gridData->yLengthMicrons, 0, 0);CHKERRQ(ierr);
  ierr = DMView(distributedArray, printViewer);CHKERRQ(ierr);

  //Create solution vector and local helper vector for scattering
  ierr = DMCreateGlobalVector(distributedArray, &globalVector);CHKERRQ(ierr);
  ierr = DMCreateLocalVector(distributedArray, &gridData->localVector);CHKERRQ(ierr);

  //Obtain additional needed values for data structure
  ierr = VecGetOwnershipRange(globalVector, &gridData->localMinimum, &gridData->localMaximum);CHKERRQ(ierr);

  //Configure KSP
  ierr = KSPSetDM(krylovSolver, distributedArray);CHKERRQ(ierr);
  ierr = KSPSetComputeOperators(krylovSolver, ComputeMatrix, gridData);CHKERRQ(ierr);
  ierr = KSPGetPC(krylovSolver, &preconditioner);CHKERRQ(ierr);
  ierr = PCSetType(preconditioner, PCNONE);CHKERRQ(ierr);
  ierr = KSPSetType(krylovSolver, KSPFBCGSR);
  ierr = KSPSetFromOptions(krylovSolver);CHKERRQ(ierr);
  ierr = KSPSetUp(krylovSolver);CHKERRQ(ierr);

  //Prepare Application Ordering (AO)
  ierr = DMDAGetAO(distributedArray, &appOrder);CHKERRQ(ierr);

  ierr = PetscPrintf(DIFFU_COMM, "\nPETSC CLASS InitializeDiffusion(): %s \n", gridData->directoryName.c_str());CHKERRQ(ierr);

  //Set initial timestep number
  step = 0;
  return ierr;
}

PetscErrorCode diffusionPETSc::TimeStep()
{
  //Takes a step forward in time in the simulation
  step++;
  ierr = ApplyBoundaryConditions();CHKERRQ(ierr);
  ierr = KSPSolve(krylovSolver, globalVector, globalVector);CHKERRQ(ierr);
  ierr = PetscPrintf(DIFFU_COMM,"\r%d steps complete", step);CHKERRQ(ierr);
  return ierr;
}

//Ensure vector references work as desired
PetscErrorCode diffusionPETSc::WriteGridValues(vector<double> xCoordinates, vector<double> \
  yCoordinates, vector<double> values)
{
  //Takes values and their coordinates in vector form, converts coordinates to
  //a discretized form, and then inserts the values into the grid based on these
  //coordinates.

  //Declare temporary variables for use in method
  PetscInt      xm, ym, xs, ys, xindex, yindex;
  unsigned int  size, i;
  PetscScalar   **globalarray;

  //Record size variable and check for size consistency across vectors
  size = xCoordinates.size();
  if (size != yCoordinates.size() || size != values.size()){
    ierr = PetscPrintf(PETSC_COMM_SELF, \
        "Error: vector sizes for WriteGridValues must match\n");CHKERRQ(ierr);
  }

  //Get coordinate range and array for all data on processor
  ierr = DMDAGetCorners(distributedArray,&xs,&ys,0,&xm,&ym,0);CHKERRQ(ierr);
  ierr = DMDAVecGetArray(distributedArray, globalVector, &globalarray);CHKERRQ(ierr);

  //Iterate through all coordinates, filter out those not on processor, and
  //insert values into grid
  for (i=0; i<size; i++){
    xindex = round(xCoordinates.at(i)/gridData->h);
    yindex = round(yCoordinates.at(i)/gridData->h);
    if (xindex >= xs && xindex < xs+xm && yindex >= ys && yindex < ys+ym){
        globalarray[yindex][xindex] = values.at(i);
    }
  }

  //Restore array to vector, finalizing the insert of these new values.
  ierr = DMDAVecRestoreArray(distributedArray, globalVector, &globalarray);CHKERRQ(ierr);

  return ierr;
}

//TODO: Ensure vector references work as desired
//Extensively test commSize > 1 case - need MPI to properly share data?
PetscErrorCode diffusionPETSc::ReadGridValues(vector<double> xCoordinates, vector<double> \
  yCoordinates, vector<double> *values)
{
  //Given some coordinates and a vector for values, discretizes the coordinates
  //and finds the values associated with them in the grid. Puts these values
  //into the given vector. Note that this only does so for on-process values.

  //Record size variable and check for size consistency across vectors
  unsigned int  readSize;
  readSize = xCoordinates.size();
  if (readSize != yCoordinates.size() || readSize != values->size()){
    ierr = PetscPrintf(PETSC_COMM_SELF, \
    "Error: vector sizes for ReadGridValues must match\n");CHKERRQ(ierr);
    return ierr;
  }

  //Read values from grid - use alternative algorithm if grid is parallelized
  if (commSize == 1){
    //Declare temporary variables for use in method
    PetscInt            xm, ym, xs, ys, xindex, yindex;
    unsigned int        i;
    PetscScalar         **readarray;

    //Get coordinate range and array for all data on processor
    ierr = DMDAGetCorners(distributedArray,&xs,&ys,0,&xm,&ym,0);CHKERRQ(ierr);
    ierr = DMDAVecGetArrayRead(distributedArray, globalVector, &readarray);CHKERRQ(ierr);

    //Get list of natural grid coordinates
    for (i=0; i < readSize; i++){
        xindex = round(xCoordinates.at(i)/gridData->h);
        yindex = round(yCoordinates.at(i)/gridData->h);
        if (xindex >= xs && xindex < xs+xm && yindex >= ys && yindex < ys+ym){
            values->at(i) = readarray[yindex][xindex];
        }
    }

    //Restore read array to vector
    ierr = DMDAVecRestoreArrayRead(distributedArray, globalVector, &readarray);CHKERRQ(ierr);
  }

  else{
    //Declare temporary variables for use in method
    PetscInt 		   	xindex, yindex, globalIndex[readSize],
                        localIndex[readSize];
    unsigned int		i;
    const PetscScalar	*readarray;
    IS					from, to;
    Vec					sequentialVector;
    VecScatter			scatter;

    //Create index vector for later PETSc vector scattering
    for (i=0; i < readSize; i++){
        localIndex[i] = i;
    }

    //Get list of natural grid coordinates
    for (i=0; i < readSize; i++){
        xindex = round(xCoordinates.at(i)/gridData->h);
        yindex = round(yCoordinates.at(i)/gridData->h);
        globalIndex[i] = yindex*gridNodesX + xindex;
    }

    //Replace natural grid coordinates with PETSc grid coordinates
    ierr = AOApplicationToPetsc(appOrder, readSize, globalIndex);CHKERRQ(ierr);

    //Prepares index sets for PETSc vector scattering
    ierr = ISCreateGeneral(DIFFU_COMM, readSize, globalIndex, PETSC_COPY_VALUES, &from);
        CHKERRQ(ierr);
    ierr = ISCreateGeneral(PETSC_COMM_SELF, readSize, localIndex, PETSC_COPY_VALUES, &to);
        CHKERRQ(ierr);

    //Create PETSc vector to be scattered to
    ierr = VecCreateSeq(PETSC_COMM_SELF, readSize, &sequentialVector);CHKERRQ(ierr);

    //Create PETSc vector scatter and execute it
    ierr = VecScatterCreate(globalVector, from, sequentialVector, to, &scatter);CHKERRQ(ierr);
    ierr = VecScatterBegin(scatter, globalVector, sequentialVector, INSERT_VALUES, SCATTER_FORWARD);
        CHKERRQ(ierr);
    ierr = VecScatterEnd(scatter, globalVector, sequentialVector, INSERT_VALUES, SCATTER_FORWARD);
        CHKERRQ(ierr);

    //Obtain data array from scattered PETSc vector
    ierr = VecGetArrayRead(sequentialVector, &readarray);CHKERRQ(ierr);

    //Copy data array into vector (non-PETSc vector)
    for (i=0; i < readSize; i++){
        values->at(i) = readarray[i];
    }

    //Restore array to vector
    ierr = VecRestoreArrayRead(sequentialVector, &readarray);CHKERRQ(ierr);

    //Destroy unnecessary objects
    ierr = ISDestroy(&from);CHKERRQ(ierr);
    ierr = ISDestroy(&to);CHKERRQ(ierr);
    ierr = VecScatterDestroy(&scatter);CHKERRQ(ierr);
    ierr = VecDestroy(&sequentialVector);CHKERRQ(ierr);
  }

  return ierr;
}

PetscErrorCode diffusionPETSc::RecordData()
{
  //Writes grid data into a .vtr file to be read by ParaView

  //Declare temporary variables for use in method
  char filename[100];
  PetscViewer vtrviewer;

  //Get name of file to be written
  sprintf(filename, "%s/%s%04d.vtr", gridData->directoryName.c_str(), \
          gridData->objectName.c_str(), step);

    //    sprintf(filename, "%s/%s%d.vtr", gridData->directoryName.c_str(),

  //Open viewer and write data to file
  ierr = PetscViewerVTKOpen(DIFFU_COMM, filename, FILE_MODE_WRITE, \
    &vtrviewer);CHKERRQ(ierr);
  ierr = PetscViewerPushFormat(vtrviewer, PETSC_VIEWER_VTK_VTR);CHKERRQ(ierr);
  ierr = VecView(globalVector, vtrviewer);CHKERRQ(ierr);

  //Destroy viewer and return ierr
  ierr = PetscViewerDestroy(&vtrviewer);CHKERRQ(ierr);
  return ierr;
}

PetscErrorCode diffusionPETSc::DiffusionFinalize(PetscBool last)
{
    //Destroys all remaining objects in class and finalizes PETSc if "last" is set
    //to true.

    //Destroy objects created for an instance of this class
    ierr = DMDestroy(&distributedArray);CHKERRQ(ierr);
    ierr = KSPDestroy(&krylovSolver);CHKERRQ(ierr);
    ierr = VecDestroy(&globalVector);CHKERRQ(ierr);
    ierr = VecDestroy(&gridData->localVector);CHKERRQ(ierr);

    ierr = PetscPrintf(DIFFU_COMM, "\nDone\n");CHKERRQ(ierr);

    //If this is the last instance of this class to be finished, finalize PETSc.
    if (last){
        ierr = PetscFinalize();
    }

    return ierr;
}

//------------------------------------------------------------------------------
//                    Helper functions (external to class)
//------------------------------------------------------------------------------

PetscErrorCode ComputeMatrix(KSP krylovSolver, Mat A, Mat jac, void *user){
    //Establishes shell matrix for KSP operation

    //Initialize function---------------------------------------------------------
    //Declare variables
    PetscErrorCode ierr;
    DiffusionData  *gridData = (DiffusionData*)user;

    //Begin function
    PetscFunctionBegin;

    //Create matrices-------------------------------------------------------------
    //Form matrix for KSP process
    ierr = MatSetSizes(A, gridData->localMaximum-gridData->localMinimum, gridData->localMaximum-\
        gridData->localMinimum, gridData->totalNodes, gridData->totalNodes);CHKERRQ(ierr);
    ierr = MatSetType(A,MATSHELL);CHKERRQ(ierr);
    ierr = MatShellSetContext(A,gridData);CHKERRQ(ierr);
    ierr = MatSetUp(A);CHKERRQ(ierr);

    //Set multiplication function for KSP process matrix
    ierr = MatShellSetOperation(A,MATOP_MULT,(void(*)(void))MyMatMult);
        CHKERRQ(ierr);

    //Return error code-----------------------------------------------------------
    PetscFunctionReturn(ierr);
}

PetscErrorCode MyMatMult(Mat A, Vec X, Vec Y){
    //Defines the multiplication function for shell matrix

    //Initialize function---------------------------------------------------------
    //Declare variables
    PetscErrorCode ierr;
    void           *ptr;
    DiffusionData  *user;
    PetscScalar    **yarray, **xarray;
    DM             distributedArray;
    PetscInt       i,j,xm,ym,xs,ys,gxm,gym,gxs,gys;

    //Begin function
    PetscFunctionBegin;

    //Retrieve context and DM
    ierr = MatShellGetContext(A,&ptr);CHKERRQ(ierr);
    ierr = VecGetDM(X,&distributedArray);CHKERRQ(ierr);
    user = (DiffusionData*)ptr;

    //Prepare for computation-----------------------------------------------------
    //Begin scattering
    ierr = DMGlobalToLocalBegin(distributedArray, X, INSERT_VALUES, user->localVector);CHKERRQ(ierr);

    //Get local domain in grid
    ierr = DMDAGetCorners(distributedArray,&xs,&ys,0,&xm,&ym,0);CHKERRQ(ierr);
    ierr = DMDAGetGhostCorners(distributedArray,&gxs,&gys,0,&gxm,&gym,0);CHKERRQ(ierr);

    //Get local array for product vector
    ierr = DMDAVecGetArray(distributedArray, Y, &yarray);CHKERRQ(ierr);

    //Finish scattering and get local array from scattered vector
    ierr = DMGlobalToLocalEnd(distributedArray, X, INSERT_VALUES, user->localVector);CHKERRQ(ierr);
    ierr = DMDAVecGetArrayRead(distributedArray, user->localVector, &xarray);CHKERRQ(ierr);

    //Compute product-------------------------------------------------------------
/*
    const double h = user->h;
    const double F = user->fourierNumber;
    const double twoF = 2*user->fourierNumber;
    const double topD = user->topDirichletCoefficient;
    const double bottomD = user->bottomDirichletCoefficient;
    const double leftD = user->leftDirichletCoefficient;
    const double rightD = user->rightDirichletCoefficient;
        const double topN = user->topNeumannCoefficient;
        const double bottomN = user->bottomNeumannCoefficient;
        const double leftN = user->leftNeumannCoefficient;
        const double rightN = user->rightNeumannCoefficient;

    for (j=ys; j<ys+ym; j++) {
        for (i=xs; i<xs+xm; i++) {
            //Sets value for edge points
//            if (i == gxs || j == gys || i == gxs+gxm-1 || j == gys+gym-1){

            //TOP WALL
            if(j == gys+gym-1)
            {
                if (topN != 0)
                {
                    //Evaluate left corner when left wall is Neumann/Robin
                    if (i == gxs && leftN != 0)
                    {
                        yarray[j][i] =
                                -twoF * (xarray[j][i+1] + xarray[j-1][i])
                            + (1 + (2 + (h * topD/topN) + (h * leftD/leftN)) * twoF) * xarray[j][i];
                    }
                    //Evaluate right corner when right wall is Neumann/Robin
                    else if (i == gxs+gxm-1 && rightN != 0)
                    {
                        yarray[j][i] =
                                -twoF * (xarray[j][i-1] + xarray[j-1][i])
                            + (1 + (2 + (h * topD/topN) + (h * rightD/rightN)) * twoF) * xarray[j][i];
                    }
                    //Evaluate all other edge points on top wall
                    else{
                        yarray[j][i] =
                                -F * (2*xarray[j-1][i] + xarray[j][i-1] + xarray[j][i+1])
                                + (1 + (2 + h*topD/topN) * twoF) * xarray[j][i];
                    }
                }
                //Evaluate wall in case of Dirichlet condition
                else
                    yarray[j][i] = xarray[j][i];
            }
            //BOTTOM WALL
            else if(j == gys)
            {
                if (bottomN != 0)
                {
                    //Evaluate left corner when left wall is Neumann/Robin
                    if (i == gxs && leftN != 0){
                        yarray[j][i] =
                                -twoF * (xarray[j][i+1] + xarray[j+1][i])
                            + (1 + (2 + (h * leftD/leftN) +  (h * bottomD/bottomN)) * twoF) * xarray[j][i];
                    }
                    //Evaluate right corner when right wall is Neumann/Robin
                    else if (i == gxs+gxm-1 && rightN != 0){
                        yarray[j][i] =
                                -twoF * (xarray[j][i-1] + xarray[j+1][i])
                            + ( 1 + (2 + (h * rightD/rightN) + (h * bottomD/bottomN)) * twoF) * xarray[j][i];
                    }
                    //Evaluate all other edge points on bottom wall
                    else{
                        yarray[j][i] =
                                -F * (2*xarray[j+1][i] + xarray[j][i-1] + xarray[j][i+1])
                                + (1 + (2 + h*bottomD/bottomN) * twoF) * xarray[j][i];
                    }
                }
                //Evaluate wall in case of Dirichlet condition
                else
                    yarray[j][i] = xarray[j][i];
            }
            //LEFT WALL
            else if(i == gxs)
            {
//                else if ((user->leftNeumannCoefficient != 0 && i == gxs) && (j != gys && j != gys+gym-1)){
                if (leftN != 0)
                    yarray[j][i] =
                            -F * (xarray[j-1][i] + xarray[j+1][i] + 2*xarray[j][i+1])
                            + (1 + (2 + (user->h*leftD/leftN)) * twoF) * xarray[j][i];
                //Evaluate wall in case of Dirichlet condition
                else
                    yarray[j][i] = xarray[j][i];
            }
            //RIGHT WALL
            else if(i == gxs+gxm-1)
            {
//                else if ((user->rightNeumannCoefficient != 0 && i == gxs+gxm-1) && (j != gys && j != gys+gym-1)){
                if (rightN != 0)
                    yarray[j][i] =
                            -F * (xarray[j-1][i] + xarray[j+1][i] -  2*xarray[j][i-1])
                            + (1 + (2 + (h*rightD/rightN)) * twoF) * xarray[j][i];
                //Evaluate wall in case of Dirichlet condition
                else
                    yarray[j][i] = xarray[j][i];
            }

            //Sets value for non-edge points
            else
                yarray[j][i] = 1.0 - F*(
                        xarray[j-1][i] + xarray[j+1][i] + xarray[j][i-1] + xarray[j][i+1] - 4*xarray[j][i]);
                //        yarray[j][i] = -user->fourierNumber*xarray[j-1][i]-user->fourierNumber*xarray[j+1][i]-\
                //            user->fourierNumber*xarray[j][i-1] -user->fourierNumber*xarray[j][i+1]+\
                //            (1+4*user->fourierNumber)*xarray[j][i];
        }
    }
*/

///*
    for (j=ys; j<ys+ym; j++) {
        for (i=xs; i<xs+xm; i++) {
        //Sets value for edge points
          if (i == gxs || j == gys || i == gxs+gxm-1 || j == gys+gym-1){
                //Evaluate in case where top wall is Neumann/Robin
                if (user->topNeumannCoefficient != 0 && j == gys+gym-1){
                    //Evaluate left corner when left wall is Neumann/Robin
                    if (i == gxs && user->leftNeumannCoefficient != 0){
                        yarray[j][i] = -2*user->fourierNumber*xarray[j][i+1]-2*user->fourierNumber*xarray[j-1][i]\
                            +(1+(4 + (2*user->h*user->topDirichletCoefficient/user->topNeumannCoefficient)\
                            +(2*user->h*user->leftDirichletCoefficient/user->leftNeumannCoefficient))*\
                            user->fourierNumber)*xarray[j][i];
                    }
                    //Evaluate right corner when right wall is Neumann/Robin
                    else if (i == gxs+gxm-1 && user->rightNeumannCoefficient != 0){
                        yarray[j][i] = -2*user->fourierNumber*xarray[j][i-1]-2*user->fourierNumber*xarray[j-1][i]\
                            +(1+(4 + (2*user->h*user->topDirichletCoefficient/user->topNeumannCoefficient)\
                            +(2*user->h*user->rightDirichletCoefficient/user->rightNeumannCoefficient))*\
                            user->fourierNumber)*xarray[j][i];
                    }
                    //Evaluate all other edge points on top wall
                    else{
                        yarray[j][i] = -2*user->fourierNumber*xarray[j-1][i]-user->fourierNumber*xarray[j][i-1]\
                            -user->fourierNumber*xarray[j][i+1]+(1 + (4 + 2*user->h*user->topDirichletCoefficient\
                            /user->topNeumannCoefficient)*user->fourierNumber)*xarray[j][i];
                    }
                }
                //Evaluate in case where bottom wall is Neumann/Robin
                else if (user->bottomNeumannCoefficient != 0 && j == gys){
                    //Evaluate left corner when left wall is Neumann/Robin
                    if (i == gxs && user->leftNeumannCoefficient != 0){
                        yarray[j][i] = -2*user->fourierNumber*xarray[j][i+1]-2*user->fourierNumber*xarray[j+1][i]\
                            +(1+(4 + (2*user->h*user->leftDirichletCoefficient/user->leftNeumannCoefficient) + \
                            (2*user->h*user->bottomDirichletCoefficient/user->bottomNeumannCoefficient))*\
                            user->fourierNumber)*xarray[j][i];
                    }
                    //Evaluate right corner when right wall is Neumann/Robin
                    else if (i == gxs+gxm-1 && user->rightNeumannCoefficient != 0){
                        yarray[j][i] = -2*user->fourierNumber*xarray[j][i-1]-2*user->fourierNumber*xarray[j+1][i]\
                            +(1+(4 + (2*user->h*user->rightDirichletCoefficient/user->rightNeumannCoefficient) + \
                            (2*user->h*user->bottomDirichletCoefficient/user->bottomNeumannCoefficient))*\
                            user->fourierNumber)*xarray[j][i];
                    }
                    //Evaluate all other edge points on bottom wall
                    else{
                        yarray[j][i] =
                                -2*user->fourierNumber*xarray[j+1][i]
                                -user->fourierNumber*xarray[j][i-1]
                            -user->fourierNumber*xarray[j][i+1]
                                +(1+(4 + (2*user->h*user->bottomDirichletCoefficient\
                            /user->bottomNeumannCoefficient))*user->fourierNumber)*xarray[j][i];
                    }
                }
                //Evaluate in case where left wall is Neumann/Robin
                else if ((user->leftNeumannCoefficient != 0 && i == gxs) && (j != gys && j != gys+gym-1)){
                    yarray[j][i] = -user->fourierNumber*xarray[j-1][i]-user->fourierNumber*xarray[j+1][i]-\
                        2*user->fourierNumber*xarray[j][i+1]+(1+(4 + (2*user->h*user->leftDirichletCoefficient\
                        /user->leftNeumannCoefficient))*user->fourierNumber)*xarray[j][i];
                }
                //Evaluate in case where right wall is Neumann/Robin
                else if ((user->rightNeumannCoefficient != 0 && i == gxs+gxm-1) && (j != gys && j != gys+gym-1)){
                    yarray[j][i] = -user->fourierNumber*xarray[j-1][i]-user->fourierNumber*xarray[j+1][i]-\
                        2*user->fourierNumber*xarray[j][i-1]+(1+(4 + (2*user->h*user->rightDirichletCoefficient\
                        /user->rightNeumannCoefficient))*user->fourierNumber)*xarray[j][i];
                }
                //Evaluate wall in case of Dirichlet condition
                else{
                    yarray[j][i] = xarray[j][i];
                }
            }
            //Sets value for non-edge points
            else
                yarray[j][i] = -user->fourierNumber*xarray[j-1][i]-user->fourierNumber*xarray[j+1][i]-\
                    user->fourierNumber*xarray[j][i-1] -user->fourierNumber*xarray[j][i+1]+\
                    (1+4*user->fourierNumber)*xarray[j][i];
        }
    }
//*/

    //Finish function-------------------------------------------------------------
    //Restore arrays
    ierr = DMDAVecRestoreArray(distributedArray, Y, &yarray);CHKERRQ(ierr);
    ierr = DMDAVecRestoreArrayRead(distributedArray, user->localVector, &xarray);CHKERRQ(ierr);

    //Return error code
    PetscFunctionReturn(ierr);
}