#include <iostream>
#include "temporaltree.h"
#include "nlohmann/json.hpp"
#include "util.h"
#include "treejsonreader.h"
#include "treeordercomputationheuristic.h"
#include "treeordercomputationsaconstraints.h"
#include "treeordercomputationilp.h"
#include "treeordercomputationiterative.h"
#include "treewriter.h"
#include "gurobi_c++.h"

char *getCmdOption(char **begin, char **end, const std::string &option)
{
   char **itr = std::find(begin, end, option);
   if (itr != end && ++itr != end)
   {
      return *itr;
   }
   return 0;
}

bool cmdOptionExists(char **begin, char **end, const std::string &option)
{
   return std::find(begin, end, option) != end;
}

int main(int argc, char *argv[])
{
   // TemporalTreeJSONReaderNTG* tr = new TemporalTreeJSONReaderNTG();
   // auto treepointer = tr->readData("../data/trees/ViscousFingers.ntg");

   string filein(getCmdOption(argv, argv + argc, "-fin"));
   auto fout = getCmdOption(argv, argv + argc, "-fout");
   auto statsout = getCmdOption(argv, argv + argc, "-statsout");

   shared_ptr<TemporalTree> treepointer;

   if (hasEnding(filein, ".tree"))
   {
      TemporalTreeJSONReader *tr = new TemporalTreeJSONReader();
      treepointer = tr->readData(filein);
   }
   else
   {
      TemporalTreeJSONReaderNTG *tr = new TemporalTreeJSONReaderNTG();
      treepointer = tr->readData(filein);
   }
   json jAll;
   treepointer->computeReverseEdges();
   TemporalTree outTree = treepointer->getSelfDeaggregated();
   outTree.order.clear();
   jAll["nNodes"] = outTree.nodes.size();
   int nleaves = 0;
   int ntemporaledges = 0;
   set<int> times;
   map<int, int> time_leafcount;
   int maxleafcount = 0;
   for (size_t i = 0; i < outTree.nodes.size(); i++)
   {
      if (outTree.isLeaf(i))
         nleaves += 1;
      if (outTree.edgesTime.find(i) != outTree.edgesTime.end())
         ntemporaledges += outTree.edgesTime[i].size();
      times.insert(outTree.nodes[i].startTime());
      times.insert(outTree.nodes[i].endTime());
      time_leafcount[outTree.nodes[i].startTime()] += 1;
      maxleafcount = max(maxleafcount, time_leafcount[outTree.nodes[i].startTime()]);
   }
   jAll["nLeaves"] = nleaves;
   jAll["ntemporaledges"] = ntemporaledges;
   jAll["times"] = times.size();
   jAll["maxleafcount"] = maxleafcount;

   // auto optimalTree = computeTreeOrderILPAggregated(*treepointer);
   bool isBarycenter = cmdOptionExists(argv, argv + argc, "-Barycenter");
   bool isBarycenterLE = cmdOptionExists(argv, argv + argc, "-BarycenterLE");
   bool isILP = cmdOptionExists(argv, argv + argc, "-ILP");
   bool isILPQ = cmdOptionExists(argv, argv + argc, "-ILPQ");
   bool isILPLE = cmdOptionExists(argv, argv + argc, "-ILPLE");
   bool isSA = cmdOptionExists(argv, argv + argc, "-SA");
   bool isMedian = cmdOptionExists(argv, argv + argc, "-Median");
   bool isMedianLE = cmdOptionExists(argv, argv + argc, "-MedianLE");
   bool unweighted = cmdOptionExists(argv, argv + argc, "-unweighted");


   int sweeps = 10;
   if (cmdOptionExists(argv, argv+argc, "-sweeps")) {
      auto sweepsstring = getCmdOption(argv, argv + argc, "-sweeps");
      sweeps = stoi(sweepsstring);
   }

   mutex m;
   condition_variable cv;
   bool ismemlimit = false;
   GRBEnv env(true);
   env.set(GRB_IntParam_Threads, 1);
   env.set(GRB_IntParam_OutputFlag, 0);
   env.set(GRB_DoubleParam_MemLimit, 8);
   env.start();
   GRBModel model(env);
   bool istimeout = false;  
   auto t1 = high_resolution_clock::now();
   if (isBarycenter)
   {
      outTree = computeTreeOrderBarycenter(*treepointer, sweeps);
   }
   else if (isBarycenterLE)
   {
      outTree = computeTreeOrderBarycenterLeafEdges(*treepointer, sweeps);
   }
   else if (isMedian)
   {
      outTree = computeTreeOrderMedian(*treepointer, sweeps);
   }
   else if (isMedianLE)
   {
      outTree = computeTreeOrderMedianLeafEdges(*treepointer, sweeps);
   }
   else if (isILP)
   {
      try{
         outTree = computeTreeOrderILP(*treepointer, model);
      }
      catch(GRBException e) {
         if (e.getErrorCode() == GRB_ERROR_OUT_OF_MEMORY) {
            ismemlimit = true;
         }
      }      
   }
   else if (isILPLE)
   {
      try{
         outTree = computeTreeOrderLeafEdgesILP(*treepointer, model);
      }
      catch(GRBException e) {
         if (e.getErrorCode() == GRB_ERROR_OUT_OF_MEMORY) {
            ismemlimit = true;
         }
      }    
   }
   else if (isILPQ)
   {
      try{
         outTree = computeTreeOrderILPQuadr(*treepointer, model);
      }
      catch(GRBException e) {
         if (e.getErrorCode() == GRB_ERROR_OUT_OF_MEMORY) {
            ismemlimit = true;
         }
      }    
   }
   else if (isSA)
   {
      TemporalTree aggregated = treepointer->aggregate();
      treepointer = make_shared<TemporalTree>(aggregated);
      // treepointer->drawTree();
      TemporalTreeOrderComputationHeuristic *heuristic = new TemporalTreeOrderComputationHeuristic();
      heuristic->portInTree = shared_ptr<const TemporalTree>(treepointer.get());
      heuristic->initializeResources();
      heuristic->propOutputBestOrder = true;
      int numconstraints = heuristic->constraints.size();
      for (int i = 0; i < numconstraints*2; i++) {
         auto tnow = high_resolution_clock::now();
         auto durationnow = tnow-t1;
         if (durationnow > TIMELIMIT) {
            istimeout = true;
            break;
         }
         heuristic->singleStep();
         // heuristic->updateOutput();
         // outTree = *heuristic->portOutTree.get();
         // cout << "Number of edge crossings: " << treeorder::countCrossings(outTree) << endl;
         // cout<<numconstraints-heuristic->currentState.statistic.numFulfilledHierarchyConstraints()-heuristic->bestState.statistic.numFulFilledMergeSplitConstraints()<<endl;
      }
      heuristic->updateOutput();
      if (istimeout) {
         outTree = *heuristic->portOutTree.get();
      }
      else {
         const TemporalTree *outtree = heuristic->portOutTree.get();

         TemporalTreeOrderComputationSAConstraints *sa = new TemporalTreeOrderComputationSAConstraints();
         sa->propInitialTemperature = 25;
         sa->propIterationsMax = 10000000;
         sa->propIterationsPerTemp = 50;
         sa->propTemperatureDecay = 0.99;
         sa->propMinimumTemperature = 1e-6;
         sa->propOutputBestOrder = true;
         sa->portInTree = heuristic->portOutTree;
         sa->initializeResources();
         while (!sa->isConverged()) {
            auto tnow = high_resolution_clock::now();
            auto durationnow = tnow-t1;
            if (durationnow > TIMELIMIT) {
               istimeout = true;
               break;
            }
            sa->singleStep();
         }
         sa->updateOutput();
         outTree = *sa->portOutTree.get();
      }      
      outTree = outTree.getSelfDeaggregated();
   }

   auto t2 = high_resolution_clock::now();
   if ((isILP || isILPQ || isILPLE) && !ismemlimit) {
      istimeout = model.get(GRB_IntAttr_Status) != GRB_OPTIMAL;
   }

   if (outTree.order.size() > 0)
   {
      assert(treeorder::fitsWithTree(outTree, outTree.order));
      outTree.computeReverseEdges();
   }

   // Check that order is valid
   if (!isSA && outTree.order.size() > 0)
   {
      assert(outTree.getCntViolatedHierarchyConstraints() == 0);
      TemporalTree::TTreeOrderMap orderMap;
      treeorder::toOrderMap(orderMap, outTree.order);
      for (int nodeIndex = 0; nodeIndex < outTree.nodes.size(); nodeIndex++)
      {
         if (!outTree.isLeaf(nodeIndex))
         {
            vector<size_t> leaves;
            outTree.computeLeafDescendants(nodeIndex, leaves);
            // check that leaves are together
            int minindex = INT_MAX;
            int maxindex = 0;
            for (auto leaf : leaves)
            {
               minindex = min(minindex, (int)orderMap[leaf]);
               maxindex = max(maxindex, (int)orderMap[leaf]);
            }
            assert(maxindex - minindex + 1 == (int)leaves.size());
         }
      }
   }
   duration<double, std::milli> ms_double = t2 - t1;
   cout<<"nleaves: "<<nleaves<<endl;
   if (statsout == 0 && ((!istimeout && !ismemlimit) || outTree.order.size() > 0))
   {
      cout << "Number of edge crossings: " << treeorder::countCrossings(outTree) << endl;
      cout << "Number of leaf edge crossings: " << treeorder::countLeafEdgeCrossings(outTree) << endl;
      cout << "Violated hierarchy constraints: " << outTree.getCntViolatedHierarchyConstraints() << endl;
   }
   if (statsout == 0) {
      cout << "Runtime in ms: " << ms_double.count() << endl;
      cout << "is Timeout? "<< istimeout<<endl;
      cout << "is Memory Limit? "<<ismemlimit<<endl;
   }
      

   /* cout<<outTree.nodes.size()<<endl;
   cout<<"Order: ";
   for (auto nodeid: outTree.order){
      cout<<nodeid<<" ";
   }
   cout<<endl;*/

   if (fout != 0 && !istimeout && !ismemlimit)
   {
      string fileout(fout);
      if (unweighted) {
         outTree.setUniformWeights(0);    
      }      
      else {
         outTree.computeAccumulated(false);
      }
      TemporalTreeWriter writer;
      writer.propFilename = fileout;
      // writer.propNTGAddWeights = true;
      writer.propNTGAddOrder = true;
      // writer.propNTGScaleWeights = 0.01;
      writer.portInTree = make_shared<TemporalTree>(outTree);
      writer.process();
   }

   if (statsout != 0)
   {
      string fstatsout(statsout);

      if (!istimeout && !ismemlimit)
      {
         jAll["crossings"] = treeorder::countCrossings(outTree);
         jAll["crossingsle"] = treeorder::countLeafEdgeCrossings(outTree);
         jAll["violatedhierarchyconstr"] = outTree.getCntViolatedHierarchyConstraints();
         jAll["time"] = ms_double.count();
      }
      if (istimeout && outTree.order.size() > 0) {
         jAll["crossings"] = treeorder::countCrossings(outTree);
         jAll["crossingsle"] = treeorder::countLeafEdgeCrossings(outTree);
         jAll["violatedhierarchyconstr"] = outTree.getCntViolatedHierarchyConstraints();
      }
      jAll["istimeout"] = istimeout;
      jAll["ismemlimit"] = ismemlimit;


      if (isBarycenter)
         jAll["algorithm"] = "Barycenter";
      else if (isBarycenterLE)
         jAll["algorithm"] = "BarycenterLE";
      else if (isMedian)
         jAll["algorithm"] = "Median";
      else if (isMedianLE)
         jAll["algorithm"] = "MedianLE";
      else if (isILP)
         jAll["algorithm"] = "ILP";
      else if (isILPLE)
         jAll["algorithm"] = "ILPLE";
      else if (isILPQ)
         jAll["algorithm"] = "ILPQ";
      else if (isSA)
         jAll["algorithm"] = "SA";

      std::ofstream outfile;
      outfile.open(fstatsout);
      outfile << jAll << std::endl;
      outfile.close();
   }

   return 0;
}