/*
 * This file belongs to the Galois project, a C++ library for exploiting
 * parallelism. The code is being released under the terms of the 3-Clause BSD
 * License (a copy is located in LICENSE.txt at the top-level directory).
 *
 * Copyright (C) 2018, The University of Texas at Austin. All rights reserved.
 * UNIVERSITY EXPRESSLY DISCLAIMS ANY AND ALL WARRANTIES CONCERNING THIS
 * SOFTWARE AND DOCUMENTATION, INCLUDING ANY WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR ANY PARTICULAR PURPOSE, NON-INFRINGEMENT AND WARRANTIES OF
 * PERFORMANCE, AND ANY WARRANTY THAT MIGHT OTHERWISE ARISE FROM COURSE OF
 * DEALING OR USAGE OF TRADE.  NO WARRANTY IS EITHER EXPRESS OR IMPLIED WITH
 * RESPECT TO THE USE OF THE SOFTWARE OR DOCUMENTATION. Under no circumstances
 * shall University be liable for incidental, special, indirect, direct or
 * consequential damages or loss of profits, interruption of business, or
 * related expenses which may arise from use of Software or Documentation,
 * including but not limited to those resulting from defects in Software and/or
 * Documentation, or loss or inaccuracy of data of any kind.
 */

#include "Point.h"
#include "Cavity.h"
#include "QuadTree.h"
#include "Verifier.h"

#include "galois/Galois.h"
#include "galois/Bag.h"
#include "galois/Timer.h"

#include "Lonestar/BoilerPlate.h"
#include "llvm/Support/CommandLine.h"

#include <boost/iterator/transform_iterator.hpp>
#include <boost/iterator/counting_iterator.hpp>

#include <algorithm>
#include <deque>
#include <fstream>
#include <iostream>
#include <limits>
#include <vector>

#include <string.h>
#include <unistd.h>

namespace cll = llvm::cl;

static const char* name = "Delaunay Triangulation";
static const char* desc =
    "Produces a Delaunay triangulation for a set of points";
static const char* url = "delaunay_triangulation";

static cll::opt<std::string>
    doWriteMesh("writemesh",
                cll::desc("Write the mesh out to files with basename"),
                cll::value_desc("basename"));
static cll::opt<std::string>
    doWritePoints("writepoints",
                  cll::desc("Write the (reordered) points to filename"),
                  cll::value_desc("filename"));
static cll::opt<bool>
    noReorderPoints("noreorder",
                    cll::desc("Don't reorder points to improve locality"),
                    cll::init(false));
static cll::opt<std::string>
    inputFile(cll::Positional, cll::desc("<input file>"), cll::Required);

enum DetAlgo { nondet, detBase, detPrefix, detDisjoint };

static cll::opt<DetAlgo>
    detAlgo(cll::desc("Deterministic algorithm:"),
            cll::values(clEnumVal(nondet, "Non-deterministic"),
                        clEnumVal(detBase, "Base execution"),
                        clEnumVal(detPrefix, "Prefix execution"),
                        clEnumVal(detDisjoint, "Disjoint execution")),
            cll::init(nondet));

//! Flag that forces user to be aware that they should be passing in a
//! mesh graph.
static cll::opt<bool>
    meshGraph("meshGraph", cll::desc("Specify that the input graph is a mesh"),
              cll::init(false));

struct GetPointer : public std::unary_function<Point&, Point*> {
  Point* operator()(Point& p) const { return &p; }
};

typedef std::vector<Point> PointList;

class ReadPoints {
  void addBoundaryPoints() {
    double minX, maxX, minY, maxY;

    minX = minY = std::numeric_limits<double>::max();
    maxX = maxY = std::numeric_limits<double>::min();

    for (auto& p : points) {
      double x = p.t().x();
      double y = p.t().y();
      if (x < minX)
        minX = x;
      else if (x > maxX)
        maxX = x;
      if (y < minY)
        minY = y;
      else if (y > maxY)
        maxY = y;
    }

    size_t size      = points.size();
    double width     = maxX - minX;
    double height    = maxY - minY;
    double maxLength = std::max(width, height);
    double centerX   = minX + width / 2.0;
    double centerY   = minY + height / 2.0;
    double radius =
        maxLength * 3.0; // radius of circle that should cover all points

    for (int i = 0; i < 3; ++i) {
      double dX = radius * cos(2 * M_PI * (i / 3.0));
      double dY = radius * sin(2 * M_PI * (i / 3.0));
      points.push_back(Point(centerX + dX, centerY + dY, size + i));
    }
  }

  void nextLine(std::ifstream& scanner) {
    scanner.ignore(std::numeric_limits<std::streamsize>::max(), '\n');
  }

  void fromTriangle(std::ifstream& scanner) {
    double x, y;
    long numPoints;

    scanner >> numPoints;

    int dim;
    scanner >> dim;
    assert(dim == 2);
    int k;
    scanner >> k; // number of attributes
    assert(k == 0);
    scanner >> k; // has boundary markers?

    for (long id = 0; id < numPoints; ++id) {
      scanner >> k; // point id
      scanner >> x >> y;
      nextLine(scanner);
      points.push_back(Point(x, y, id));
    }
  }

  void fromPointList(std::ifstream& scanner) {
    double x, y;

    // comment line
    nextLine(scanner);
    size_t id = 0;
    while (!scanner.eof()) {
      scanner >> x >> y;
      if (x == 0 && y == 0)
        break;
      points.push_back(Point(x, y, id++));
      x = y = 0;
      nextLine(scanner);
    }
  }

  PointList& points;

public:
  ReadPoints(PointList& p) : points(p) {}

  void from(const std::string& name) {
    std::ifstream scanner(name.c_str());
    if (!scanner.good()) {
      GALOIS_DIE("could not open file: ", name);
    }
    if (name.find(".node") == name.size() - 5) {
      fromTriangle(scanner);
    } else {
      fromPointList(scanner);
    }
    scanner.close();

    if (points.size())
      addBoundaryPoints();
    else {
      GALOIS_DIE("no points found in file: ", name);
    }
  }
};

static void writePoints(const std::string& filename, const PointList& points) {
  std::ofstream out(filename.c_str());
  // <num vertices> <dimension> <num attributes> <has boundary markers>
  out << points.size() << " 2 0 0\n";
  // out.setf(std::ios::fixed, std::ios::floatfield);
  out.setf(std::ios::scientific, std::ios::floatfield);
  out.precision(10);
  long id = 0;
  for (const auto& p : points) {
    const Tuple& t = p.t();
    out << id++ << " " << t.x() << " " << t.y() << " 0\n";
  }

  out.close();
}

using BasePoints = galois::InsertBag<Point>;
using PtrPoints  = galois::InsertBag<Point*>;
using Rounds     = std::vector<PtrPoints*>;

size_t maxRounds;
const int roundShift = 4; //! round sizes are portional to (1 << roundsShift)

static void copyPointsFromRounds(PointList& points, Rounds& rounds) {
  for (int i = maxRounds - 1; i >= 0; --i) {
    //! [Access elements of InsertBag]
    // PtrPoints expands to galois::InsertBag<Point*>
    // points is of type std::vector<Point>
    PtrPoints& pptrs = *rounds[i];
    for (auto ii : pptrs) {
      points.push_back(*ii);
    }
    //! [Access elements of InsertBag]
  }
}

struct ReadInput {
  Graph& graph;
  BasePoints& basePoints;
  Rounds& rounds;
  std::random_device rng;
  std::mt19937 urng;

  ReadInput(Graph& g, BasePoints& b, Rounds& r)
      : graph(g), basePoints(b), rounds(r), urng(rng()) {}

  void addBoundaryNodes(Point* p1, Point* p2, Point* p3) {
    Element large_triangle(p1, p2, p3);
    GNode large_node = graph.createNode(large_triangle);
    graph.addNode(large_node);

    p1->addElement(large_node);
    p2->addElement(large_node);
    p3->addElement(large_node);

    Element border_ele1(p1, p2);
    Element border_ele2(p2, p3);
    Element border_ele3(p3, p1);

    GNode border_node1 = graph.createNode(border_ele1);
    GNode border_node2 = graph.createNode(border_ele2);
    GNode border_node3 = graph.createNode(border_ele3);

    graph.addNode(border_node1);
    graph.addNode(border_node2);
    graph.addNode(border_node3);

    graph.getEdgeData(graph.addEdge(large_node, border_node1)) = 0;
    graph.getEdgeData(graph.addEdge(large_node, border_node2)) = 1;
    graph.getEdgeData(graph.addEdge(large_node, border_node3)) = 2;

    graph.getEdgeData(graph.addEdge(border_node1, large_node)) = 0;
    graph.getEdgeData(graph.addEdge(border_node2, large_node)) = 0;
    graph.getEdgeData(graph.addEdge(border_node3, large_node)) = 0;
  }

  template <typename L>
  void generateRoundsImpl(const L& loop, size_t size, PointList& points,
                          size_t log2) {
    loop(
        galois::iterate(size_t{0}, size),
        [&, this](size_t index) {
          const Point& p = points[index];

          Point* ptr = &(basePoints.push(p));
          int r      = 0;
          for (size_t i = 0; i < log2; ++i) {
            size_t mask = (1UL << (i + 1)) - 1;
            if ((index & mask) == (1UL << i)) {
              r = i;
              break;
            }
          }

          rounds[r / roundShift]->push(ptr);
        },
        galois::loopname("generateRoundsImpl"));
  }

  //! Blocked point distribution (exponentially increasing block size) with
  //! points randomized within a round
  void generateRoundsOld(PointList& points, bool randomize) {
    size_t counter = 0;
    size_t round   = 0;
    size_t next    = 1 << roundShift;
    std::vector<Point*> buf;

    PointList::iterator ii = points.begin(), ei = points.end();
    while (ii != ei) {
      Point* ptr = &(basePoints.push(*ii));
      buf.push_back(ptr);
      ++ii;
      if (ii == ei || counter > next) {
        next *= next;
        int r = maxRounds - 1 - round;
        if (randomize)
          std::shuffle(buf.begin(), buf.end(), urng);
        std::copy(buf.begin(), buf.end(), std::back_inserter(*rounds[r]));
        buf.clear();
        ++round;
      }
      ++counter;
    }
  }

  void generateRounds(PointList& points, bool addBoundary) {
    size_t size = points.size() - 3;

    size_t log2 = std::max((size_t)floor(log(size) / log(2)), (size_t)1);
    maxRounds   = log2 / roundShift;
    for (size_t i = 0; i <= maxRounds;
         i++) { // rounds[maxRounds+1] for boundary points
      rounds.push_back(new galois::InsertBag<Point*>);
    }

    PointList ordered;
    // ordered.reserve(size);

    if (noReorderPoints) {
      std::copy(points.begin(), points.begin() + size,
                std::back_inserter(ordered));
      generateRoundsOld(ordered, false);
    } else {
      // Reorganize spatially
      QuadTree q(
          boost::make_transform_iterator(points.begin(), GetPointer()),
          boost::make_transform_iterator(points.begin() + size, GetPointer()));

      q.output(std::back_inserter(ordered));

      if (true) {
        if (detAlgo == nondet) {
          generateRoundsImpl(galois::DoAll(), size, ordered, log2);

        } else {
          generateRoundsImpl(galois::StdForEach(), size, ordered, log2);
        }
      } else {
        generateRoundsOld(ordered, true);
      }
    }

    if (!addBoundary)
      return;

    // Now, handle boundary points
    size_t last = points.size();
    //! [Insert elements into InsertBag]
    // basePoints is of type galois::InsertBag<Point>
    // points is of type std::vector<Point>
    Point* p1 = &(basePoints.push(points[last - 1]));
    Point* p2 = &(basePoints.push(points[last - 2]));
    Point* p3 = &(basePoints.push(points[last - 3]));
    //! [Insert elements into InsertBag]

    rounds[maxRounds]->push(p1);
    rounds[maxRounds]->push(p2);
    rounds[maxRounds]->push(p3);

    addBoundaryNodes(p1, p2, p3);
  }

  void operator()(const std::string& filename, bool addBoundary) {
    PointList points;
    ReadPoints(points).from(filename);

    std::cout << "configuration: " << points.size() << " points\n";

#if 1
    galois::preAlloc(
        32 * points.size() * sizeof(Element) *
        1.5 // mesh is about 2x number of points (for random points)
        / (galois::runtime::pagePoolSize()) // in pages
    );
#else
    galois::preAlloc(1 * numThreads // some per-thread state
                     + 2 * points.size() *
                           sizeof(Element) // mesh is about 2x number of points
                                           // (for random points)
                           * 32            // include graph node size
                           / (galois::runtime::hugePageSize) // in pages
    );
#endif
    galois::reportPageAlloc("MeminfoPre");

    galois::StatTimer T("generateRounds");
    T.start();
    generateRounds(points, addBoundary);
    T.stop();
  }
};

static void writeMesh(const std::string& filename, Graph& graph) {
  long numTriangles = 0;
  long numSegments  = 0;
  for (auto n : graph) {
    Element& e = graph.getData(n);
    if (e.boundary()) {
      numSegments++;
    } else {
      numTriangles++;
    }
  }

  long tid = 0;
  long sid = 0;
  std::string elementName(filename);
  std::string polyName(filename);

  elementName.append(".ele");
  polyName.append(".poly");

  std::ofstream eout(elementName.c_str());
  std::ofstream pout(polyName.c_str());
  // <num triangles> <nodes per triangle> <num attributes>
  eout << numTriangles << " 3 0\n";
  // <num vertices> <dimension> <num attributes> <has boundary markers>
  // ...
  // <num segments> <has boundary markers>
  pout << "0 2 0 0\n";
  pout << numSegments << " 1\n";
  for (auto n : graph) {
    const Element& e = graph.getData(n);
    if (e.boundary()) {
      // <segment id> <vertex> <vertex> <is boundary>
      pout << sid++ << " " << e.getPoint(0)->id() << " " << e.getPoint(1)->id()
           << " 1\n";
    } else {
      // <triangle id> <vertex> <vertex> <vertex> [in ccw order]
      eout << tid++ << " " << e.getPoint(0)->id() << " ";
      if (e.clockwise()) {
        eout << e.getPoint(2)->id() << " " << e.getPoint(1)->id() << "\n";
      } else {
        eout << e.getPoint(1)->id() << " " << e.getPoint(2)->id() << "\n";
      }
    }
  }

  eout.close();
  // <num holes>
  pout << "0\n";
  pout.close();
}

struct DelaunayTriangulation {

  QuadTree* tree;
  Graph& graph;

  struct ContainsTuple {
    const Graph& graph;
    Tuple tuple;
    ContainsTuple(const Graph& g, const Tuple& t) : graph(g), tuple(t) {}
    bool operator()(const GNode& n) const {
      assert(!graph.getData(n, galois::MethodFlag::UNPROTECTED).boundary());
      return graph.getData(n, galois::MethodFlag::UNPROTECTED)
          .inTriangle(tuple);
    }
  };

  void computeCenter(const Element& e, Tuple& t) const {
    for (int i = 0; i < 3; ++i) {
      const Tuple& o = e.getPoint(i)->t();
      for (int j = 0; j < 2; ++j) {
        t[j] += o[j];
      }
    }
    for (int j = 0; j < 2; ++j) {
      t[j] *= 1 / 3.0;
    }
  }

  void findBestNormal(const Element& element, const Point* p,
                      const Point*& bestP1, const Point*& bestP2) {
    Tuple center(0);
    computeCenter(element, center);
    int scale = element.clockwise() ? 1 : -1;

    Tuple origin = p->t() - center;
    //        double length2 = origin.x() * origin.x() + origin.y() *
    //        origin.y();
    bestP1 = bestP2 = NULL;
    double bestVal  = 0.0;
    for (int i = 0; i < 3; ++i) {
      int next = i + 1;
      if (next > 2)
        next -= 3;

      const Point* p1 = element.getPoint(i);
      const Point* p2 = element.getPoint(next);
      double dx       = p2->t().x() - p1->t().x();
      double dy       = p2->t().y() - p1->t().y();
      Tuple normal(scale * -dy, scale * dx);
      double val = normal.dot(origin); // / length2;
      if (bestP1 == NULL || val > bestVal) {
        bestVal = val;
        bestP1  = p1;
        bestP2  = p2;
      }
    }
    assert(bestP1 != NULL && bestP2 != NULL && bestVal > 0);
  }

  GNode findCorrespondingNode(GNode start, const Point* p1, const Point* p2) {
    for (auto ii : graph.edges(start)) {
      GNode dst  = graph.getEdgeDst(ii);
      Element& e = graph.getData(dst, galois::MethodFlag::UNPROTECTED);
      int count  = 0;
      for (int i = 0; i < e.dim(); ++i) {
        if (e.getPoint(i) == p1 || e.getPoint(i) == p2) {
          if (++count == 2)
            return dst;
        }
      }
    }
    GALOIS_DIE("unreachable");
    return start;
  }

  bool planarSearch(const Point* p, GNode start, GNode& node) {
    // Try simple hill climbing instead
    ContainsTuple contains(graph, p->t());
    while (!contains(start)) {
      Element& element = graph.getData(start, galois::MethodFlag::WRITE);
      if (element.boundary()) {
        // Should only happen when quad tree returns a boundary point which is
        // rare There's only one way to go from here
        assert(std::distance(graph.edge_begin(start), graph.edge_end(start)) ==
               1);
        start = graph.getEdgeDst(
            graph.edge_begin(start, galois::MethodFlag::WRITE));
      } else {
        // Find which neighbor will get us to point fastest by computing normal
        // vectors
        const Point *p1, *p2;
        findBestNormal(element, p, p1, p2);
        start = findCorrespondingNode(start, p1, p2);
      }
    }

    node = start;
    return true;
  }

  bool findContainingElement(const Point* p, GNode& node) {
    Point* result;
    if (!tree->find(p, result)) {
      return false;
    }

    result->get(galois::MethodFlag::WRITE);

    GNode someNode = result->someElement();

    // Not in mesh yet
    if (!someNode) {
      return false;
    }

    return planarSearch(p, someNode, node);
  }

  using Alloc = galois::PerIterAllocTy;

  struct LocalState {
    Cavity<Alloc> cav;
    LocalState(Graph& graph, Alloc& alloc) : cav(graph, alloc) {}
  };

  template <int Version, typename C>
  void processPoint(Point* p, C& ctx) {
    Cavity<Alloc>* cavp = NULL;

    if (Version == detDisjoint) {

      if (ctx.isFirstPass()) {
        LocalState* localState = ctx.template createLocalState<LocalState>(
            graph, ctx.getPerIterAlloc());
        cavp = &localState->cav;

      } else {

        LocalState* localState = ctx.template getLocalState<LocalState>();
        localState->cav.update();
        return;
      }
    }

    p->get(galois::MethodFlag::WRITE);
    assert(!p->inMesh());

    GNode node;
    if (!findContainingElement(p, node)) {
      // Someone updated an element while we were searching, producing
      // a semi-consistent state
      // ctx.push(p);
      // Current version is safe with locking so this shouldn't happen
      GALOIS_DIE("unreachable");
      return;
    }

    assert(graph.getData(node).inTriangle(p->t()));
    assert(graph.containsNode(node));

    if (Version == detDisjoint && ctx.isFirstPass()) {
      cavp->init(node, p);
      cavp->build();
    } else {
      Cavity<Alloc> cav(graph, ctx.getPerIterAlloc());
      cav.init(node, p);
      cav.build();
      if (Version == detPrefix)
        return;
      ctx.cautiousPoint();
      cav.update();
    }
  }

  template <int Version, typename WL, typename B, typename... Args>
  void generateMesh(B& pptrs, Args&&... args) {

    galois::for_each(
        galois::iterate(pptrs),
        [&, this](Point* p, auto& ctx) { this->processPoint<Version>(p, ctx); },
        galois::wl<WL>(), galois::loopname("generateMesh"),
        galois::local_state<LocalState>(), galois::per_iter_alloc(),
        galois::no_pushes(), std::forward<Args>(args)...);
  }
};

/*
template<int Version=detBase>
struct Process {

  //! Serial operator
  void operator()(Point* p) {
    p->get(galois::MethodFlag::WRITE);
    assert(!p->inMesh());

    GNode node;
    if (!findContainingElement(p, node)) {
      GALOIS_DIE("Could not find triangle containing point");
      return;
    }

    assert(graph.getData(node).inTriangle(p->t()));
    assert(graph.containsNode(node));

    Cavity<> cav(graph);
    cav.init(node, p);
    cav.build();
    cav.update();
  }
};
*/

static void run(Rounds& rounds, Graph& graph) {
  typedef galois::worklists::PerThreadChunkLIFO<32> Chunk;
  typedef galois::worklists::Deterministic<> DWL;

  for (int i = maxRounds - 1; i >= 0; --i) {

    galois::StatTimer BT("buildtree");
    BT.start();
    assert(rounds[i + 1]);
    PtrPoints& tptrs = *(rounds[i + 1]);
    QuadTree tree(tptrs.begin(), tptrs.end());
    BT.stop();

    galois::StatTimer PT("ParallelTime");
    PT.start();

    assert(rounds[i]);
    galois::InsertBag<Point*>& pptrs = *(rounds[i]);

    DelaunayTriangulation dt{&tree, graph};
    switch (detAlgo) {
    case nondet:
      dt.generateMesh<detBase, Chunk>(pptrs);
      break;
    case detBase:
      dt.generateMesh<detBase, DWL>(pptrs);
      break;
    case detPrefix: {
      auto nv = [&dt](Point* p, auto& ctx) {
        dt.processPoint<detPrefix>(p, ctx);
      };
      dt.generateMesh<detBase, DWL>(
          pptrs, galois::neighborhood_visitor<decltype(nv)>(nv));
      break;
    }
    case detDisjoint:
      dt.generateMesh<detDisjoint, DWL>(pptrs);
      break;
    default:
      GALOIS_DIE("unknown algorithm: ", detAlgo);
    }

    PT.stop();
  }
}

void deleteRounds(Rounds& rounds) {
  for (auto r : rounds)
    delete r;
}

int main(int argc, char** argv) {
  galois::SharedMemSys G;
  LonestarStart(argc, argv, name, desc, url, &inputFile);

  galois::StatTimer totalTime("TimerTotal");
  totalTime.start();

  if (!meshGraph) {
    GALOIS_DIE("This application requires a mesh graph input;"
               " please use the -meshGraph flag "
               " to indicate the input is a mesh graph.");
  }

  Graph graph;

  //! All Point* refer to elements in this bag
  //! [Define InsertBag]
  // BasePoints expands to galois::InsertBag<Point>
  BasePoints basePoints;
  //! [Define InsertBag]

  Rounds rounds;

  bool writepoints = doWritePoints.size() > 0;
  ReadInput(graph, basePoints, rounds)(inputFile, !writepoints);
  if (writepoints) {
    std::cout << "Writing " << doWritePoints << "\n";
    PointList points;
    copyPointsFromRounds(points, rounds);
    writePoints(doWritePoints, points);
    deleteRounds(rounds);
    return 0;
  }

  const char* name = 0;
  switch (detAlgo) {
  case nondet:
    name = "nondet";
    break;
  case detBase:
    name = "detBase";
    break;
  case detPrefix:
    name = "detPrefix";
    break;
  case detDisjoint:
    name = "detDisjoint";
    break;
  default:
    name = "unknown";
    break;
  }
  galois::gInfo("Algorithm ", name);

  galois::StatTimer execTime("Timer_0");
  execTime.start();
  run(rounds, graph);
  execTime.stop();
  std::cout << "mesh size: " << graph.size() << "\n";

  galois::reportPageAlloc("MeminfoPost");

  if (!skipVerify) {
    Verifier verifier;
    if (!verifier.verify(&graph)) {
      GALOIS_DIE("triangulation failed");
    }
    std::cout << "Triangulation OK\n";
  }

  if (doWriteMesh.size()) {
    std::string base = doWriteMesh;
    std::cout << "Writing " << base << "\n";
    writeMesh(base.c_str(), graph);

    PointList points;
    // Reordering messes up connection between id and place in pointlist
    ReadPoints(points).from(inputFile);
    writePoints(base.append(".node"), points);
  }

  deleteRounds(rounds);

  totalTime.stop();

  return 0;
}


//---------------------------------------------------------------------------//
//!
//! \file   Utility_ElectronVoltUnit.hpp
//! \author Alex Robison
//! \brief  The electron volt unit declaration
//!
//---------------------------------------------------------------------------//

#ifndef UTILITY_ELECTRON_VOLT_UNIT_HPP
#define UTILITY_ELECTRON_VOLT_UNIT_HPP

// Boost Includes
#include <boost/units/systems/si/energy.hpp>
#include <boost/units/physical_dimensions/energy.hpp>
#include <boost/units/base_unit.hpp>
#include <boost/units/conversion.hpp>
#include <boost/units/make_scaled_unit.hpp>

namespace Utility{

namespace Units{

//! The electron volt base unit
struct ElectronVoltBaseUnit : public boost::units::base_unit<ElectronVoltBaseUnit,boost::units::energy_dimension,1>
{
  static const char* name() { return "electron volt"; }
  static const char* symbol() { return "eV"; }
};

//! The electron volt unit
typedef ElectronVoltBaseUnit::unit_type ElectronVolt;

//! The kilo-electron volt unit
typedef boost::units::make_scaled_unit<ElectronVolt, boost::units::scale<10,boost::units::static_rational<3> > >::type KiloElectronVolt;

//! The mega-electron volt unit
typedef boost::units::make_scaled_unit<ElectronVolt, boost::units::scale<10,boost::units::static_rational<6> > >::type MegaElectronVolt;

BOOST_UNITS_STATIC_CONSTANT( eV, ElectronVolt );
BOOST_UNITS_STATIC_CONSTANT( keV, KiloElectronVolt );
BOOST_UNITS_STATIC_CONSTANT( MeV, MegaElectronVolt );

} // end Units namespace

} // end Utility namespace

BOOST_UNITS_DEFINE_CONVERSION_FACTOR( Utility::Units::ElectronVoltBaseUnit, boost::units::si::energy, double, 1.602176565e-19 );
BOOST_UNITS_DEFINE_CONVERSION_FACTOR( boost::units::si::energy, Utility::Units::ElectronVoltBaseUnit, double, 1.0/1.602176565e-19 );

BOOST_UNITS_DEFAULT_CONVERSION( Utility::Units::ElectronVoltBaseUnit, boost::units::si::energy );

#endif // end UTILITY_ELECTRON_VOLT_UNIT_HPP

//---------------------------------------------------------------------------//
// end Utility_ElectronVoltUnit.hpp
//---------------------------------------------------------------------------//


/*****************************************************************************
*
* Copyright (c) 2003-2018 by The University of Queensland
* http://www.uq.edu.au
*
* Primary Business: Queensland, Australia
* Licensed under the Apache License, version 2.0
* http://www.apache.org/licenses/LICENSE-2.0
*
* Development until 2012 by Earth Systems Science Computational Center (ESSCC)
* Development 2012-2013 by School of Earth Sciences
* Development from 2014 by Centre for Geoscience Computing (GeoComp)
*
*****************************************************************************/

/****************************************************************************/

/* Paso: SparseMatrix */

/****************************************************************************/

/* Author: Lutz Gross, l.gross@uq.edu.au */

/****************************************************************************/

#include "SparseMatrix.h"
#include "BlockOps.h"
#include "MKL.h"
#include "Options.h"
#include "PasoUtil.h"
#include "Preconditioner.h"
#include "UMFPACK.h"
#include "mmio.h"

#include <boost/scoped_array.hpp>
#include <fstream>

/****************************************************************************/

namespace paso {

using escript::IndexList;

/* debug: print the entries */
/*
void print_entries(index_t *r, index_t *c, double *v, int nz)
{
    for(int i=0; i<nz; i++)
        printf("(%ld, %ld) == %e\n", (long)r[i], (long)c[i], v[i]);
}
*/

/* swap function */
void swap(index_t *r, index_t *c, double *v, int left, int right)
{
    double v_temp;
    index_t temp;

    temp = r[left];
    r[left] = r[right];
    r[right] = temp;

    temp = c[left];
    c[left] = c[right];
    c[right] = temp;

    v_temp = v[left];
    v[left] = v[right];
    v[right] = v_temp;
}

void q_sort(index_t *row, index_t *col, double *val, int begin, int end, int N)
{
    int l, r;
    index_t pivot, lval;

    if (end > begin) {
        pivot = N * row[begin] + col[begin];
        l = begin + 1;
        r = end;

        while (l < r) {
            lval = N * row[l] + col[l];
            if (lval < pivot)
                l++;
            else {
                r--;
                swap( row, col, val, l, r );
            }
        }
        l--;
        swap(row, col, val, begin, l);
        q_sort(row, col, val, begin, l, N);
        q_sort(row, col, val, r, end, N);
    }
}

/* Allocates a SparseMatrix of given type using the given matrix pattern.
   Values are initialized with zero.
   If patternIsUnrolled and type & MATRIX_FORMAT_BLK1, it is assumed that the
   pattern is already unrolled to match the requested block size
   and offsets. Otherwise unrolling and offset adjustment will be performed.
*/
SparseMatrix::SparseMatrix(SparseMatrixType ntype, Pattern_ptr npattern,
                           dim_t rowBlockSize, dim_t colBlockSize,
                           bool patternIsUnrolled) :
    type(ntype),
    val(NULL),
    solver_package(PASO_PASO),
    solver_p(NULL)
{
    if (patternIsUnrolled) {
        if ((ntype & MATRIX_FORMAT_OFFSET1) != (npattern->type & MATRIX_FORMAT_OFFSET1)) {
            throw PasoException("SparseMatrix: requested offset and pattern offset do not match.");
        }
    }
    // do we need to apply unrolling?
    bool unroll
          // we don't like non-square blocks
        = (rowBlockSize != colBlockSize)
#ifndef ESYS_HAVE_LAPACK
          // or any block size bigger than 3
          || (colBlockSize > 3)
#endif
          // or if block size one requested and the block size is not 1
          || ((ntype & MATRIX_FORMAT_BLK1) && (colBlockSize > 1))
          // or if offsets don't match
          || ((ntype & MATRIX_FORMAT_OFFSET1) != (npattern->type & MATRIX_FORMAT_OFFSET1));

    SparseMatrixType pattern_format_out = (ntype & MATRIX_FORMAT_OFFSET1)
                             ? MATRIX_FORMAT_OFFSET1 : MATRIX_FORMAT_DEFAULT;

    // === compressed sparse columns ===
    if (ntype & MATRIX_FORMAT_CSC) {
        if (unroll) {
            if (patternIsUnrolled) {
                pattern = npattern;
            } else {
                pattern = npattern->unrollBlocks(pattern_format_out,
                                                 colBlockSize, rowBlockSize);
            }
            row_block_size = 1;
            col_block_size = 1;
        } else {
            pattern = npattern->unrollBlocks(pattern_format_out, 1, 1);
            row_block_size = rowBlockSize;
            col_block_size = colBlockSize;
        }
        numRows = pattern->numInput;
        numCols = pattern->numOutput;
    } else {
    // === compressed sparse row ===
        if (unroll) {
            if (patternIsUnrolled) {
                pattern = npattern;
            } else {
                pattern = npattern->unrollBlocks(pattern_format_out,
                                                 rowBlockSize, colBlockSize);
            }
            row_block_size = 1;
            col_block_size = 1;
        } else {
            pattern = npattern->unrollBlocks(pattern_format_out, 1, 1);
            row_block_size = rowBlockSize;
            col_block_size = colBlockSize;
        }
        numRows = pattern->numOutput;
        numCols = pattern->numInput;
    }
    if (ntype & MATRIX_FORMAT_DIAGONAL_BLOCK) {
        block_size = std::min(row_block_size, col_block_size);
    } else {
        block_size = row_block_size*col_block_size;
    }
    len = (size_t)(pattern->len)*(size_t)(block_size);

    val=new double[len];
    setValues(0.);
}

SparseMatrix::~SparseMatrix()
{
    switch (solver_package) {
        case PASO_SMOOTHER:
            Preconditioner_LocalSmoother_free((Preconditioner_LocalSmoother*) solver_p);
            break;

        case PASO_MKL:
            MKL_free(this);
            break;

        case PASO_UMFPACK:
            UMFPACK_free(this);
            break;
    }
    delete[] val;
}

SparseMatrix_ptr SparseMatrix::loadMM_toCSR(const char* filename)
{
    SparseMatrix_ptr out;
    int i;
    MM_typecode matrixCode;

    // open the file
    std::ifstream f(filename);
    if (f.fail()) {
        throw PasoException("SparseMatrix::loadMM_toCSR: Cannot open file for reading.");
    }

    // process banner
    if (mm_read_banner(f, &matrixCode) != 0) {
        f.close();
        throw PasoException("SparseMatrix::loadMM_toCSR: Error processing MM banner.");
    }
    if (!(mm_is_real(matrixCode) && mm_is_sparse(matrixCode) && mm_is_general(matrixCode))) {
        f.close();
        throw PasoException("SparseMatrix::loadMM_toCSR: found Matrix Market type is not supported.");
    }

    // get matrix size
    int M, N, nz;

    if (mm_read_mtx_crd_size(f, &M, &N, &nz) != 0) {
        f.close();
        throw PasoException("SparseMatrix::loadMM_toCSR: Could not parse matrix size.");
    }

    // prepare storage
    index_t* col_ind = new index_t[nz];
    index_t* row_ind = new index_t[nz];
    index_t* row_ptr = new index_t[M+1];
    double* val = new double[nz];

    // perform actual read of elements
    for (i=0; i<nz; i++) {
        f >> row_ind[i] >> col_ind[i] >> val[i];
        //scan_ret = fscanf(fileHandle_p, "%d %d %le\n", &row_ind[i], &col_ind[i], &val[i]);
        if (!f.good()) {
            delete[] val;
            delete[] row_ind;
            delete[] col_ind;
            delete[] row_ptr;
            f.close();
            return out;
        }
        row_ind[i]--;
        col_ind[i]--;
    }
    f.close();

    // sort the entries
    q_sort(row_ind, col_ind, val, 0, nz, N);

    // setup row_ptr
    int curr_row = 0;
    for(i=0; (i<nz && curr_row<M); curr_row++) {
        while(row_ind[i] != curr_row)
            i++;
        row_ptr[curr_row] = i;
    }
    row_ptr[M] = nz;

    Pattern_ptr mainPattern(new Pattern(MATRIX_FORMAT_DEFAULT, M, N,
                                        row_ptr, col_ind));
    out.reset(new SparseMatrix(MATRIX_FORMAT_DEFAULT, mainPattern, 1, 1, true));

    // copy values
    for (i=0; i<nz; i++)
        out->val[i] = val[i];

    delete[] val;
    delete[] row_ind;
    return out;
}

void SparseMatrix::saveMM(const char* filename) const
{
    if (col_block_size != row_block_size) {
        throw PasoException("SparseMatrix::saveMM: currently only square blocks are supported.");
    }

    // open the file
    std::ofstream f(filename);
    if (f.fail()) {
        throw PasoException("SparseMatrix::saveMM: File could not be opened for writing");
    }
    if (type & MATRIX_FORMAT_CSC) {
        throw PasoException("SparseMatrix::saveMM does not support CSC.");
    } else {
        MM_typecode matcode;
        mm_initialize_typecode(&matcode);
        mm_set_matrix(&matcode);
        mm_set_coordinate(&matcode);
        mm_set_real(&matcode);

        const dim_t N = getNumRows();
        const dim_t M = getNumCols();
        mm_write_banner(f, matcode);
        mm_write_mtx_crd_size(f, N*row_block_size,
                              M*col_block_size, pattern->ptr[N]*block_size);

        const index_t offset=(type & MATRIX_FORMAT_OFFSET1 ? 1:0);

        f.precision(15);

        if (type & MATRIX_FORMAT_DIAGONAL_BLOCK) {
            for (dim_t i=0; i<N; i++) {
                for (dim_t iptr = pattern->ptr[i]-offset; iptr<pattern->ptr[i+1]-offset; ++iptr) {
                    const dim_t j=pattern->index[iptr]-offset;
                    for (dim_t ib=0; ib<block_size; ib++) {
                        const dim_t irow=ib+row_block_size*i;
                        const dim_t icol=ib+col_block_size*j;
                        f << irow+1 << " " << icol+1 << " "
                          << val[iptr*block_size+ib] << std::endl;
                    }
                }
            }
        } else {
            for (dim_t i=0; i<N; i++) {
                for (dim_t iptr = pattern->ptr[i]-offset; iptr<pattern->ptr[i+1]-offset; ++iptr) {
                    const dim_t j=pattern->index[iptr]-offset;
                    for (dim_t irb=0; irb<row_block_size; irb++) {
                        const dim_t irow=irb+row_block_size*i;
                        for (dim_t icb=0; icb<col_block_size; icb++) {
                            const dim_t icol=icb+col_block_size*j;
                            f << irow+1 << " " << icol+1 << " "
                                << val[iptr*block_size+irb+row_block_size*icb]
                                << std::endl;
                        }
                    }
                }
            }
        }
    }
    // close the file
    f.close();
}

void SparseMatrix::addAbsRow_CSR_OFFSET0(double* array) const
{
    const dim_t nOut = pattern->numOutput;
#pragma omp parallel for
    for (dim_t ir=0; ir < nOut; ir++) {
        for (dim_t irb=0; irb < row_block_size; irb++) {
            const dim_t irow = irb+row_block_size*ir;
            double fac=0.;
            for (index_t iptr=pattern->ptr[ir]; iptr < pattern->ptr[ir+1]; iptr++) {
                for (dim_t icb=0; icb < col_block_size; icb++) {
                    const index_t idx = iptr*block_size+irb+row_block_size*icb;
                    fac += std::abs(val[idx]);
                }
            }
            array[irow]+=fac;
        }
    }
}

void SparseMatrix::maxAbsRow_CSR_OFFSET0(double* array) const
{
    const dim_t nOut = pattern->numOutput;
#pragma omp parallel for
    for (dim_t ir=0; ir < nOut; ir++) {
        for (dim_t irb=0; irb < row_block_size; irb++) {
            const dim_t irow = irb+row_block_size*ir;
            double fac=0.;
            for (index_t iptr=pattern->ptr[ir]; iptr < pattern->ptr[ir+1]; iptr++) {
                for (dim_t icb=0; icb < col_block_size; icb++) {
                    const index_t idx = iptr*block_size+irb+row_block_size*icb;
                    fac=std::max(fac, std::abs(val[idx]));
                }
            }
            array[irow]=std::max(array[irow], fac);
        }
    }
}

void SparseMatrix::addRow_CSR_OFFSET0(double* array) const
{
    const dim_t nOut = pattern->numOutput;
#pragma omp parallel for
    for (dim_t ir=0; ir < nOut; ir++) {
        for (dim_t irb=0; irb < row_block_size; irb++) {
            dim_t irow=irb+row_block_size*ir;
            double fac=0.;
            for (index_t iptr=pattern->ptr[ir]; iptr<pattern->ptr[ir+1]; iptr++) {
                for (dim_t icb=0; icb < col_block_size; icb++)
                    fac += val[iptr*block_size+irb+row_block_size*icb];

            }
            array[irow]+=fac;
        }
    }
}

void SparseMatrix::copyBlockToMainDiagonal(const double* in)
{
    const dim_t n = pattern->numOutput;
    const dim_t nblk = block_size;
    const size_t nblk_size = sizeof(double)*nblk;
    const index_t* main_ptr = borrowMainDiagonalPointer();
#pragma omp parallel for
    for (index_t ir=0; ir < n;ir++) {
        memcpy((void*)&val[main_ptr[ir]*nblk], (void*)&in[nblk*ir], nblk_size);
    }
}

void SparseMatrix::copyBlockFromMainDiagonal(double* out) const
{
    const dim_t n = pattern->numOutput;
    const dim_t nblk = block_size;
    const size_t nblk_size = sizeof(double)*nblk;
    const index_t* main_ptr = borrowMainDiagonalPointer();
#pragma omp parallel for
    for (index_t ir=0; ir < n; ir++) {
        memcpy((void*)&out[nblk*ir], (void*)&val[main_ptr[ir]*nblk], nblk_size);
    }
}

void SparseMatrix::copyFromMainDiagonal(double* out) const
{
    const dim_t n = pattern->numOutput;
    const dim_t nblk = block_size;
    const dim_t blk = std::min(row_block_size, col_block_size);
    const index_t* main_ptr = borrowMainDiagonalPointer();
#pragma omp parallel for
    for (index_t ir=0; ir < n; ir++) {
        for (index_t ib=0; ib < blk; ib++) {
            out[ir*blk+ib] = val[main_ptr[ir]*nblk+ib+row_block_size*ib];
        }
    }
}

void SparseMatrix::copyToMainDiagonal(const double* in)
{
    const dim_t n = pattern->numOutput;
    const dim_t nblk = block_size;
    const dim_t blk = std::min(row_block_size, col_block_size);
    const index_t* main_ptr = borrowMainDiagonalPointer();
#pragma omp parallel for
    for (index_t ir=0; ir < n; ir++) {
        for (index_t ib=0; ib < blk; ib++) {
            val[main_ptr[ir]*nblk+ib+row_block_size*ib] = in[ir*blk+ib];
        }
    }
}

void SparseMatrix::applyDiagonal_CSR_OFFSET0(const double* left,
                                             const double* right)
{
    const dim_t row_block = row_block_size;
    const dim_t col_block = col_block_size;
    const dim_t n_block = row_block*col_block;
    const dim_t nOut = pattern->numOutput;

#pragma omp parallel for
    for (index_t ir=0; ir < nOut; ir++) {
        for (index_t irb=0; irb < row_block; irb++) {
            const index_t irow = irb+row_block*ir;
            const double rtmp = left[irow];
            for (index_t iptr=pattern->ptr[ir]; iptr < pattern->ptr[ir+1]; iptr++) {
                #pragma ivdep
                for (index_t icb=0; icb < col_block_size; icb++) {
                    const index_t icol = icb+col_block*pattern->index[iptr];
                    const index_t l = iptr*n_block + irb+row_block*icb;
                    val[l] *= rtmp*right[icol];
                }
            }
        }
    }
}

void SparseMatrix::setValues(double value)
{
    const index_t index_offset=(type & MATRIX_FORMAT_OFFSET1 ? 1:0);
    if (!pattern->isEmpty()) {
        const dim_t nOut = pattern->numOutput;
#pragma omp parallel for
        for (dim_t i=0; i < nOut; ++i) {
            for (index_t iptr=pattern->ptr[i]-index_offset; iptr < pattern->ptr[i+1]-index_offset; ++iptr) {
                for (dim_t j=0; j<block_size; ++j)
                    val[iptr*block_size+j] = value;
            }
        }
    }
}

void SparseMatrix::invMain(double* inv_diag, index_t* pivot) const
{
    int failed = 0;
    double A11;
    const dim_t n=numRows;
    const dim_t n_block=row_block_size;
    const dim_t m_block=col_block_size;
    dim_t i;
    index_t iPtr;
    index_t* main_ptr=pattern->borrowMainDiagonalPointer();
    // check matrix is square
    if (m_block != n_block) {
        throw PasoException("SparseMatrix::invMain: square block size expected.");
    }
    if (n_block == 1) {
#pragma omp parallel for private(i, iPtr, A11) schedule(static)
        for (i = 0; i < n; i++) {
            iPtr = main_ptr[i];
            A11 = val[iPtr];
            if (std::abs(A11) > 0.) {
                inv_diag[i]=1./A11;
            } else {
                failed=1;
            }
        }
    } else if (n_block==2) {
#pragma omp parallel for private(i, iPtr) schedule(static)
        for (i = 0; i < n; i++) {
            iPtr = main_ptr[i];
            BlockOps_invM_2(&inv_diag[i*4], &val[iPtr*4], &failed);
        }
    } else if (n_block==3) {
#pragma omp parallel for private(i, iPtr) schedule(static)
        for (i = 0; i < n; i++) {
            iPtr = main_ptr[i];
            BlockOps_invM_3(&inv_diag[i*9], &val[iPtr*9], &failed);
        }
    } else {
#pragma omp parallel for private(i, iPtr) schedule(static)
        for (i = 0; i < n; i++) {
            iPtr = main_ptr[i];
            BlockOps_Cpy_N(block_size, &inv_diag[i*block_size], &val[iPtr*block_size]);
            BlockOps_invM_N(n_block, &inv_diag[i*block_size], &pivot[i*n_block], &failed);
        }
    }
    if (failed > 0) {
        throw PasoException("SparseMatrix::invMain: non-regular main diagonal block.");
    }
}

void SparseMatrix::applyBlockMatrix(double* block_diag, index_t* pivot,
                                    double* x, const double *b) const
{
    const dim_t n = numRows;
    const dim_t n_block = row_block_size;
    util::copy(n_block*n, x, b);
    BlockOps_solveAll(n_block, n, block_diag, pivot, x);
}

SparseMatrix_ptr SparseMatrix::getTranspose() const
{
    const dim_t m = numCols;
    const dim_t n = numRows;
    boost::scoped_array<IndexList> index_list(new IndexList[m]);

    for (dim_t i=0; i<n; ++i) {
        for (index_t iptr2=pattern->ptr[i]; iptr2<pattern->ptr[i+1]; ++iptr2) {
            const index_t j = pattern->index[iptr2];
            index_list[j].insertIndex(i);
        }
    }

    Pattern_ptr ATpattern(Pattern::fromIndexListArray(0,m,index_list.get(),0,n,0));
    SparseMatrix_ptr AT(new SparseMatrix(type, ATpattern, col_block_size, row_block_size, false));

    if ( ((type & MATRIX_FORMAT_DIAGONAL_BLOCK) && (block_size == 1)) ||
         (row_block_size == 1 && col_block_size == 1)) {
#pragma omp parallel for
        for (dim_t i=0; i<m; ++i) {
            for (index_t iptr_AT=AT->pattern->ptr[i]; iptr_AT<AT->pattern->ptr[i+1]; ++iptr_AT) {
                const index_t j = AT->pattern->index[iptr_AT];
                index_t jptr_A = pattern->ptr[j];
                const index_t* start_p = &pattern->index[jptr_A];
                const index_t* where_p=(index_t*)bsearch(&i, start_p,
                                          pattern->ptr[j+1]-jptr_A,
                                          sizeof(index_t), util::comparIndex);
                if (where_p != NULL) { // this should always be the case
                    jptr_A += (index_t)(where_p-start_p);
                    AT->val[iptr_AT] = val[jptr_A];
                }
            }
        }
    } else {
        if (type & MATRIX_FORMAT_DIAGONAL_BLOCK) {
#pragma omp parallel for
            for (dim_t i=0; i<m; ++i) {
                for (index_t iptr_AT=AT->pattern->ptr[i]; iptr_AT<AT->pattern->ptr[i+1]; ++iptr_AT) {
                    const index_t j = AT->pattern->index[iptr_AT];
                    index_t jptr_A = pattern->ptr[j];
                    const index_t* start_p = &pattern->index[jptr_A];
                    const index_t* where_p = (index_t*)bsearch(&i, start_p,
                                         pattern->ptr[j+1]-jptr_A,
                                         sizeof(index_t), util::comparIndex);
                    if (where_p != NULL) { // this should always be the case
                        jptr_A += (index_t)(where_p-start_p);
                        for (dim_t ib=0; ib < block_size; ++ib)
                            AT->val[iptr_AT*block_size+ib] = val[jptr_A*block_size+ib];
                    }
                }
            }
        } else {
#pragma omp parallel for
            for (dim_t i=0; i<m; ++i) {
                for (index_t iptr_AT=AT->pattern->ptr[i]; iptr_AT<AT->pattern->ptr[i+1]; ++iptr_AT) {
                    const index_t j = AT->pattern->index[iptr_AT];
                    index_t jptr_A = pattern->ptr[j];
                    const index_t* start_p = &pattern->index[jptr_A];
                    const index_t* where_p=(index_t*)bsearch(&i, start_p,
                                       pattern->ptr[j + 1]-jptr_A,
                                       sizeof(index_t), util::comparIndex);
                    if (where_p != NULL) { // this should always be the case
                        jptr_A += (index_t)(where_p-start_p);
                        for (index_t irb=0; irb < row_block_size; ++irb) {
                            for (index_t icb=0 ; icb < col_block_size; ++icb) {
                                AT->val[iptr_AT*block_size+icb+col_block_size*irb] = val[jptr_A*block_size+irb+row_block_size*icb];
                            }
                        }
                    }
                }
            }
        }
    }
    return AT;
}

SparseMatrix_ptr SparseMatrix::unroll(SparseMatrixType newType) const
{
    const index_t out_type = (newType & MATRIX_FORMAT_BLK1) ? newType : newType + MATRIX_FORMAT_BLK1;
    SparseMatrix_ptr out(new SparseMatrix(out_type, pattern, row_block_size, col_block_size, false));

    const dim_t n = numRows;
    const index_t A_offset = (type & MATRIX_FORMAT_OFFSET1 ? 1 : 0);
    const index_t out_offset = (out_type & MATRIX_FORMAT_OFFSET1 ? 1 : 0);

    if (out->type & MATRIX_FORMAT_CSC) {
#pragma omp parallel for
        for (dim_t i=0; i<n; ++i) {
            for (index_t iptr=pattern->ptr[i]-A_offset; iptr<pattern->ptr[i+1]-A_offset; ++iptr) {
                const index_t j = pattern->index[iptr]-A_offset;
                for (dim_t icb=0; icb<col_block_size; ++icb) {
                    const index_t icol=j*col_block_size+icb;
                    const index_t* start_p=&out->pattern->index[out->pattern->ptr[icol]-out_offset];
                    const index_t l_col=out->pattern->ptr[icol+1]-out->pattern->ptr[icol];
                    for (dim_t irb=0; irb<row_block_size; ++irb) {
                        const index_t irow=row_block_size*i+irb+out_offset;
                        const index_t* where_p = (index_t*)bsearch(&irow,
                                    start_p, l_col, sizeof(index_t),
                                    util::comparIndex);
                        if (where_p != NULL)
                            out->val[out->pattern->ptr[icol]-out_offset+(index_t)(where_p-start_p)] =
                                val[block_size*iptr+irb+row_block_size*icb];
                    }
                }
            }
        }
    } else {
#pragma omp parallel for
        for (dim_t i=0; i<n; ++i) {
            for (index_t iptr=pattern->ptr[i]-A_offset; iptr<pattern->ptr[i+1]-A_offset; ++iptr) {
                const index_t j = pattern->index[iptr]-A_offset;
                for (dim_t irb=0; irb<row_block_size; ++irb) {
                    const index_t irow=row_block_size*i+irb;
                    const index_t* start_p = &out->pattern->index[out->pattern->ptr[irow]-out_offset];
                    const index_t l_row=out->pattern->ptr[irow+1]-out->pattern->ptr[irow];
                    for (dim_t icb=0; icb<col_block_size; ++icb) {
                        const index_t icol=j*col_block_size+icb+out_offset;
                        const index_t* where_p = (index_t*)bsearch(&icol,
                                    start_p, l_row, sizeof(index_t),
                                    util::comparIndex);
                        if (where_p != NULL)
                            out->val[out->pattern->ptr[irow]-out_offset+(index_t)(where_p-start_p)] =
                                val[block_size*iptr+irb+row_block_size*icb];
                    }
                }
            }
        }
    }
    return out;
}

} // namespace paso


/**
 * @file GlobalLearning.hpp
 * @brief Implementation of the classes for global learning algorithms.
 * @author Ankit Srivastava <asrivast@gatech.edu>
 *
 * Copyright 2020 Georgia Institute of Technology
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#ifndef DETAIL_GLOBALLEARNING_HPP_
#define DETAIL_GLOBALLEARNING_HPP_

#include "common/SetUtils.hpp"
#include "utils/Logging.hpp"

#include <boost/math/special_functions/binomial.hpp>
#include <boost/iterator/counting_iterator.hpp>
#include <boost/iterator/zip_iterator.hpp>
#include <boost/tuple/tuple.hpp>

template <typename Data, typename Var, typename Set>
/**
 * @brief Constructs the object with the given data.
 */
GlobalLearning<Data, Var, Set>::GlobalLearning(
  const mxx::comm& comm,
  const Data& data,
  const double alpha,
  const Var maxConditioning
) : ConstraintBasedLearning<Data, Var, Set>(comm, data, alpha, maxConditioning),
    m_cachedNeighbors(),
    m_removedEdges()
{
  TIMER_RESET(m_tSync);
  TIMER_RESET(m_tDist);
}

template <typename Data, typename Var, typename Set>
/**
 * @brief Default destructor that prints out timing information.
 */
GlobalLearning<Data, Var, Set>::~GlobalLearning(
)
{
  if (this->m_comm.is_first()) {
    TIMER_ELAPSED_NONZERO("Time taken in redistributing: ", m_tDist);
    TIMER_ELAPSED_NONZERO("Time taken in synchronizing neighbors: ", m_tSync);
  }
}

template <typename Data, typename Var, typename Set>
const Set&
GlobalLearning<Data, Var, Set>::getPC(
  const Var target
) const
{
  return m_cachedNeighbors.at(target);
}

template <typename Data, typename Var, typename Set>
const Set&
GlobalLearning<Data, Var, Set>::getMB(
  const Var target
) const
{
  return m_cachedNeighbors.at(target);
}

template <typename Data, typename Var, typename Set>
/**
 * @brief Function for initializing the data structures used for
 *        learning the skeleton.
 *
 * @param myEdges The array to be initialized with the edges
 *                to be handled by this processor.
 * @param allNeighbors The map containing the neighborhood set
 *                     for all the variables in the network.
 * @param allNeighbors The map used for tracking the removals for
 *                     for all the variables in the network.
 */
void
GlobalLearning<Data, Var, Set>::initializeLearning(
  std::vector<std::tuple<Var, Var, double>>& myEdges,
  std::unordered_map<Var, Set>& allNeighbors,
  std::unordered_map<Var, Set>& removedNeighbors
) const
{
  // First, block decompose all the edges on all the processors
  auto n = this->m_allVars.size();
  mxx::blk_dist dist((n * (n - 1)) / 2, this->m_comm);
  auto myOffset = dist.eprefix_size();
  auto mySize = dist.local_size();

  auto vars = std::vector<Var>(this->m_allVars.begin(), this->m_allVars.end());
  auto k = 0u;
  auto currOffset = 0u;
  while (currOffset + (n - (k + 1)) <= static_cast<uint32_t>(myOffset)) {
    currOffset += (n - (k + 1));
    ++k;
  }
  auto primary = vars.begin() + k;
  auto secondary = primary + 1 + (myOffset - currOffset);

  myEdges.resize(mySize);
  for (auto i = 0u; i < mySize; ++i, ++secondary) {
    if (secondary == vars.end()) {
      // Start a new primary variable if the secondary variable is exhausted
      ++primary;
      secondary = primary + 1;
    }
    // Initialize the p-values
    myEdges[i] = std::make_tuple(*primary, *secondary, 0.0);
  }
  // Then, initialize all the neighbor sets
  for (const auto v : vars) {
    allNeighbors.insert(std::make_pair(v, set_init(this->getCandidates(v), this->m_data.numVars())));
    removedNeighbors.insert(std::make_pair(v, set_init(Set(), this->m_data.numVars())));
  }
}

template <typename Data, typename Var, typename Set>
bool
GlobalLearning<Data, Var, Set>::fixWeightedImbalance(
  std::vector<std::tuple<Var, Var, double>>& myEdges,
  const std::vector<double>& myWeights,
  const double imbalanceThreshold
) const
{
  bool fixed = false;
  double myTotalWeight = std::accumulate(myWeights.cbegin(), myWeights.cend(), 0.0);
  TIMER_START(this->m_tMxx);
  double totalWeight = mxx::allreduce(myTotalWeight, this->m_comm);
  TIMER_PAUSE(this->m_tMxx);
  auto imbalance = 0.0;
  if (totalWeight > 0) {
    TIMER_START(this->m_tMxx);
    auto maxWeight = mxx::allreduce(myTotalWeight, mxx::max<double>(), this->m_comm);
    TIMER_PAUSE(this->m_tMxx);
    auto avgWeight = totalWeight / this->m_comm.size();
    imbalance = (maxWeight / avgWeight) - 1.0;
    //if (this->m_comm.is_first()) {
      //std::cout << "Imbalance: " << imbalance << std::endl;
    //}
  }
  if (std::isgreater(imbalance, imbalanceThreshold)) {
    TIMER_START(this->m_tMxx);
    // Get the weight on previous processors
    auto globalPrefix = mxx::exscan(myTotalWeight, this->m_comm);
    TIMER_PAUSE(this->m_tMxx);
    std::vector<uint64_t> sendCounts(this->m_comm.size(), 0);
    if (std::isgreater(myTotalWeight, 0)) {
      double div = totalWeight / this->m_comm.size();
      auto proc = static_cast<uint32_t>(std::floor(globalPrefix / div));
      auto procFirst = 0u;
      double localPrefix = 0;
      double leftWeight = myTotalWeight;
      std::vector<double> myWeightsPrefix(myWeights.size());
      std::partial_sum(myWeights.cbegin(), myWeights.cend(), myWeightsPrefix.begin());
      for (; std::isgreater(leftWeight, 0) && (proc < this->m_comm.size()); ++proc) {
        double sendWeight = std::min<double>((div * (proc + 1)) - globalPrefix, leftWeight);
        auto procLast = std::distance(myWeightsPrefix.cbegin(),
                                      std::lower_bound(myWeightsPrefix.cbegin(),
                                                       myWeightsPrefix.cend(),
                                                       localPrefix + sendWeight));

        if (std::isless(myWeightsPrefix[procLast], localPrefix + sendWeight)) {
          ++procLast;
        }
        sendCounts[proc] = (procLast - procFirst) + 1;
        leftWeight -= (myWeightsPrefix[procLast] - localPrefix);
        globalPrefix += (myWeightsPrefix[procLast] - localPrefix);
        localPrefix = myWeightsPrefix[procLast];
        procFirst = procLast + 1;
      }
    }
    TIMER_START(this->m_tMxx);
    auto recvCounts = mxx::all2all(sendCounts, this->m_comm);
    TIMER_PAUSE(this->m_tMxx);
    auto myRecv = std::accumulate(recvCounts.begin(), recvCounts.end(), 0u);
    std::vector<std::tuple<Var, Var, double>> newEdges(myRecv);
    const std::tuple<Var, Var, double>* sendBuf = (myEdges.size() > 0) ? &myEdges[0] : nullptr;
    std::tuple<Var, Var, double>* recvBuf = (newEdges.size() > 0) ? &newEdges[0] : nullptr;
    TIMER_START(this->m_tMxx);
    mxx::all2allv(sendBuf, sendCounts, recvBuf, recvCounts, this->m_comm);
    TIMER_PAUSE(this->m_tMxx);
    myEdges = std::move(newEdges);
    fixed = true;
  }
  return fixed;
}

template <typename Data, typename Var, typename Set>
/**
 * @brief Function for storing removed edges if they might be relevant
 *        for directing edges later.
 *
 * @param myRemoved A vector of all the removed edges.
 * @param myDSepSets Vector containing d-separating sets for all the removed edges.
 * @param allNeighbors Neighbor sets for all the variables.
 */
void
GlobalLearning<Data, Var, Set>::storeRemovedEdges(
  std::vector<std::tuple<Var, Var, double>>&& myRemoved,
  std::vector<Set>&& myDSepSets,
  const std::unordered_map<Var, Set>& allNeighbors,
  const bool duplicateEdges
) const
{
  LOG_MESSAGE_IF(myRemoved.size() != myDSepSets.size(),
                 error, "Mismatch between number of edges and d-separating sets.");
  // Only retain information which can be used for directing edges later
  auto newEnd = std::remove_if(boost::make_zip_iterator(boost::make_tuple(myRemoved.begin(), myDSepSets.begin())),
                               boost::make_zip_iterator(boost::make_tuple(myRemoved.end(), myDSepSets.end())),
                               [&allNeighbors] (const boost::tuple<std::tuple<Var, Var, double>, Set>& e)
                                               { return set_intersection(allNeighbors.at(std::get<0>(boost::get<0>(e))),
                                                                         allNeighbors.at(std::get<1>(boost::get<0>(e)))).empty(); });
  myRemoved.erase(boost::get<0>(newEnd.get_iterator_tuple()), myRemoved.end());
  myDSepSets.erase(boost::get<1>(newEnd.get_iterator_tuple()), myDSepSets.end());
  TIMER_START(this->m_tMxx);
  // Get the prefix size of all the d-separating sets
  auto removedSizes = mxx::allgather(myRemoved.size(), this->m_comm);
  TIMER_PAUSE(this->m_tMxx);
  auto removedPrefix = std::accumulate(removedSizes.begin(), removedSizes.begin() + this->m_comm.rank(), 0u);
  std::vector<std::tuple<Var, Var, double, uint32_t>> myRemovedEdges(myRemoved.size());
  std::transform(std::make_move_iterator(myRemoved.begin()), std::make_move_iterator(myRemoved.end()),
                 boost::counting_iterator<uint32_t>(removedPrefix), myRemovedEdges.begin(),
                 [] (const std::tuple<Var, Var, double>&& e, const uint32_t idx)
                    { return std::make_tuple(std::get<0>(e), std::get<1>(e), std::get<2>(e), idx); });
  TIMER_START(this->m_tMxx);
  // Block redistribute all the removed edges
  mxx::stable_distribute_inplace(myRemovedEdges, this->m_comm);
  // Sort the removed edges in parallel
  mxx::comm nonzero_comm(static_cast<MPI_Comm>(this->m_comm));
  if (mxx::any_of(myRemovedEdges.size() == 0, this->m_comm)) {
    nonzero_comm = this->m_comm.split(myRemovedEdges.size() > 0);
  }
  if (myRemovedEdges.size() > 0) {
    auto sortEdges = [] (const std::tuple<Var, Var, double, uint32_t>& a, const std::tuple<Var, Var, double, uint32_t>& b)
                        { return std::make_pair(std::get<0>(a), std::get<1>(a)) < std::make_pair(std::get<0>(b), std::get<1>(b)); };
    if (!duplicateEdges) {
      mxx::sort(myRemovedEdges.begin(), myRemovedEdges.end(), sortEdges, nonzero_comm);
    }
    else {
      mxx::stable_sort(myRemovedEdges.begin(), myRemovedEdges.end(), sortEdges, nonzero_comm);
      auto newEnd = mxx::unique(myRemovedEdges.begin(), myRemovedEdges.end(),
                                [] (const std::tuple<Var, Var, double, uint32_t>& a, const std::tuple<Var, Var, double, uint32_t>& b)
                                   { return std::make_pair(std::get<0>(a), std::get<1>(a)) == std::make_pair(std::get<0>(b), std::get<1>(b)); },
                                nonzero_comm);
      myRemovedEdges.erase(newEnd, myRemovedEdges.end());
    }
  }
  // Gather all the removed edges on all the processors
  auto allRemovedEdges = mxx::allgatherv(std::move(myRemovedEdges), this->m_comm);
  // Also gather all the d-separating sets on all the processors
  // XXX: We can not use Set as part of a std::tuple during communication
  //      Therefore, we gather all the sets separately
  auto allDSepSets = set_allgatherv(std::move(myDSepSets), removedSizes, this->m_data.numVars(), this->m_comm);
  TIMER_PAUSE(this->m_tMxx);
  // Finally, store the removed edges with the corresponding d-separating sets
  m_removedEdges.resize(allRemovedEdges.size());
  std::transform(std::make_move_iterator(allRemovedEdges.begin()), std::make_move_iterator(allRemovedEdges.end()),
                 m_removedEdges.begin(),
                 [&allDSepSets] (const std::tuple<Var, Var, double, uint32_t>&& e)
                                { return std::make_tuple(std::get<0>(e), std::get<1>(e), std::get<2>(e), allDSepSets.at(std::get<3>(e))); });
}

template <typename Data, typename Var, typename Set>
/**
 * @brief Function for constructing skeleton from the neighbor sets.
 *
 * @param allNeighbors Neighborhood sets for all the variables.
 */
BayesianNetwork<Var>
GlobalLearning<Data, Var, Set>::constructSkeleton(
  const std::unordered_map<Var, Set>&& allNeighbors
) const
{
  BayesianNetwork<Var> bn(this->m_data.varNames(this->m_allVars));
  for (const auto& neighbors : allNeighbors) {
    const auto x = neighbors.first;
    for (const auto y : neighbors.second) {
      if (x < y) {
        LOG_MESSAGE(info, "+ Adding the edge %s <-> %s",
                          this->m_data.varName(x), this->m_data.varName(y));
        bn.addEdge(x, y, true);
      }
    }
  }
  m_cachedNeighbors = std::move(allNeighbors);
  return bn;
}

template <typename Data, typename Var, typename Set>
/**
 * @brief Checks if the given v-structure (y-x-z) is a collider.
 *        Also returns the p-value for the collider.
 *
 * This function checks if x is part of the set that rendered
 * y and z independent. If not, y-x-z is a collider.
 */
std::pair<bool, double>
GlobalLearning<Data, Var, Set>::checkCollider(
  const Var y,
  const Var x,
  const Var z
) const
{
  static auto findPair = [] (const std::tuple<Var, Var, double, Set>& t, const std::pair<Var, Var>& e)
                            { return std::make_pair(std::get<0>(t), std::get<1>(t)) < e; };
  // Always try to find the forward edge
  auto edge = (y < z) ? std::make_pair(y, z) : std::make_pair(z, y);
  auto eIt = std::lower_bound(m_removedEdges.begin(), m_removedEdges.end(), edge, findPair);
  if ((eIt == m_removedEdges.end()) ||
      (std::make_pair(std::get<0>(*eIt), std::get<1>(*eIt)) != edge)) {
    auto pv = this->m_data.pValue(y, z);
    LOG_MESSAGE(debug, "Computed p-value for edge %s - %s is %g",
                       this->m_data.varName(y), this->m_data.varName(z), pv);
    return std::make_pair(true, pv);
  }
  else {
    auto collider = !(std::get<3>(*eIt).contains(x));
    auto pv = std::get<2>(*eIt);
    LOG_MESSAGE(debug, "Stored p-value for edge %s - %s is %g",
                       this->m_data.varName(y), this->m_data.varName(z), std::get<2>(*eIt));
    return std::make_pair(collider, pv);
  }
}

template <typename Data, typename Var, typename Set>
PCStableCommon<Data, Var, Set>::PCStableCommon(
  const mxx::comm& comm,
  const Data& data,
  const double alpha,
  const Var maxConditioning
) : GlobalLearning<Data, Var, Set>(comm, data, alpha, maxConditioning)
{
}

template <typename Data, typename Var, typename Set>
std::pair<double, Set>
PCStableCommon<Data, Var, Set>::checkEdge(
  const std::tuple<Var, Var, double>& edge,
  const std::unordered_map<Var, Set>& allNeighbors,
  std::unordered_map<Var, Set>& removedNeighbors,
  const uint32_t setSize,
  const bool checkBackward
) const
{
  auto pv = 0.0;
  auto dsep = set_init(Set(), this->m_data.numVars());
  const auto x = std::get<0>(edge);
  const auto y = std::get<1>(edge);
  auto xNeighbors = allNeighbors.at(x);
  auto yNeighbors = allNeighbors.at(y);
  LOG_MESSAGE(debug, "Checking the edge %s <-> %s, d-separating set of size %u",
                     this->m_data.varName(x), this->m_data.varName(y), setSize);
  auto remove = false;
  // First, check the edge using neighbors of x
  if (xNeighbors.size() > setSize) {
    xNeighbors.erase(y);
    std::tie(pv, dsep) = this->m_data.maxPValueSubset(this->m_alpha, x, y, xNeighbors, setSize, setSize);
    remove = this->m_data.isIndependent(this->m_alpha, pv);
  }
  // Then, check the edge using neighbors of y if all of the following hold:
  // 1. The conditioning set used for checking is not empty
  // 2. The neighbors of x did not remove it already
  // 3. The size of the neighborhood of y is greater than the conditioning set size
  if (checkBackward && !remove && (yNeighbors.size() > setSize)) {
    yNeighbors.erase(x);
    // Further, only check if neighborhood of y has some elements which are
    // not present in the neighborhood of x
    if (!set_difference(yNeighbors, xNeighbors).empty()) {
      std::tie(pv, dsep) = this->m_data.maxPValueSubset(this->m_alpha, x, y, yNeighbors, setSize, setSize);
      remove = this->m_data.isIndependent(this->m_alpha, pv);
    }
  }
  LOG_MESSAGE(debug, "%s and %s are " + std::string(this->m_data.isIndependent(this->m_alpha, pv) ? "independent" : "dependent") +
                     " (p-value = %g)",
                     this->m_data.varName(x), this->m_data.varName(y), pv);
  LOG_MESSAGE_IF(remove, debug, "- Removing the edge %s <-> %s",
                                this->m_data.varName(x), this->m_data.varName(y));
  if (remove) {
    // Mark y for removal from neighborhood of x
    removedNeighbors.at(x).insert(y);
    // Mark x for removal from neighborhood of y
    removedNeighbors.at(y).insert(x);
  }
  return std::make_pair(pv, dsep);
}

template <typename Data, typename Var, typename Set>
BayesianNetwork<Var>
PCStableCommon<Data, Var, Set>::getSkeleton_sequential(
  const bool directEdges
) const
{
  std::vector<std::tuple<Var, Var, double>> allEdges;
  std::unordered_map<Var, Set> allNeighbors;
  std::unordered_map<Var, Set> removedNeighbors;
  this->initializeLearning(allEdges, allNeighbors, removedNeighbors);
  auto maxSize = std::min(this->m_maxConditioning, static_cast<Var>(this->m_allVars.size() - 2));
  for (auto s = 0u; (s <= maxSize) && !allEdges.empty(); ++s) {
    LOG_MESSAGE(debug, "Testing %u edges using sets of size %u", allEdges.size(), s);
    TIMER_DECLARE(tIter);
    for (auto& e : allEdges) {
      auto result = this->checkEdge(e, allNeighbors, removedNeighbors, s, (s > 0));
      std::get<2>(e) = result.first;
      // We need to store the removed edges since the d-separating
      // set may be required for directing edges later
      if (directEdges && (s > 0) &&
          this->m_data.isIndependent(this->m_alpha, result.first)) {
        TIMER_START(this->m_tDirect);
        Var x, y;
        std::tie(x, y, std::ignore) = e;
        // We only use this information for directing colliders of type x-z-y
        // However, there is no need to store it if there are no candidates for z
        // i.e., no common neighbors between x and y
        if (!set_intersection(allNeighbors.at(x), allNeighbors.at(y)).empty()) {
          this->m_removedEdges.push_back(std::make_tuple(x, y, result.first, result.second));
        }
        TIMER_PAUSE(this->m_tDirect);
      }
    }
    for (auto& rn : removedNeighbors) {
      if (!rn.second.empty()) {
        allNeighbors[rn.first] = set_difference(allNeighbors.at(rn.first), rn.second);
        rn.second.clear();
      }
    }
    auto newEnd = std::remove_if(allEdges.begin(), allEdges.end(),
                                 [this, &allNeighbors, &s] (const std::tuple<Var, Var, double>& e)
                                                           { return this->m_data.isIndependent(this->m_alpha, std::get<2>(e)) ||
                                                                    (allNeighbors.at(std::get<0>(e)).size() <= (s + 1) &&
                                                                     allNeighbors.at(std::get<1>(e)).size() <= (s + 1)); });
    allEdges.erase(newEnd, allEdges.end());
    if (this->m_comm.is_first()) {
      TIMER_ELAPSED("Time taken in testing all sets of size " + std::to_string(s) + ": ", tIter);
    }
    if (directEdges) {
      TIMER_START(this->m_tDirect);
      auto newEnd = std::remove_if(this->m_removedEdges.begin(), this->m_removedEdges.end(),
                                   [&allNeighbors] (const std::tuple<Var, Var, double, Set>& t)
                                                   { return set_intersection(allNeighbors.at(std::get<0>(t)),
                                                                             allNeighbors.at(std::get<1>(t))).empty(); });
      this->m_removedEdges.erase(newEnd, this->m_removedEdges.end());
      TIMER_PAUSE(this->m_tDirect);
    }
  }
  if (directEdges) {
    TIMER_START(this->m_tDirect);
    std::sort(this->m_removedEdges.begin(), this->m_removedEdges.end(),
              [] (const std::tuple<Var, Var, double, Set>& a, const std::tuple<Var, Var, double, Set>& b)
                 { return std::make_pair(std::get<0>(a), std::get<1>(a)) < std::make_pair(std::get<0>(b), std::get<1>(b)); });
    TIMER_PAUSE(this->m_tDirect);
  }
  return this->constructSkeleton(std::move(allNeighbors));
}

template <typename Data, typename Var, typename Set>
PCStable<Data, Var, Set>::PCStable(
  const mxx::comm& comm,
  const Data& data,
  const double alpha,
  const Var maxConditioning
) : PCStableCommon<Data, Var, Set>(comm, data, alpha, maxConditioning)
{
}

template <typename Data, typename Var, typename Set>
/**
 * @brief Function that synchronizes the candidate sets by taking an
 *        intersection of the sets across all the processors.
 *
 * @param mySets A map with the neighbors of the variables
 *               remaining on this processor.
 */
void
PCStable<Data, Var, Set>::syncSets(
  std::unordered_map<Var, Set>& mySets
) const
{
  TIMER_START(this->m_tMxx);
  set_allintersect_indexed(mySets, this->m_allVars, this->m_data.numVars(), this->m_comm);
  TIMER_PAUSE(this->m_tMxx);
}

template <typename Data, typename Var, typename Set>
BayesianNetwork<Var>
PCStable<Data, Var, Set>::getSkeleton_parallel(
  const bool directEdges,
  const double imbalanceThreshold
) const
{
  // Initialize the learning in a similar fashion as done sequentially
  // Create nC2 edges - one for each unordered variable pair
  std::vector<std::tuple<Var, Var, double>> myEdges;
  std::unordered_map<Var, Set> allNeighbors;
  std::unordered_map<Var, Set> removedNeighbors;
  this->initializeLearning(myEdges, allNeighbors, removedNeighbors);
  // The following two data structures are used for getting
  // d-separating set and p-value for all the removed edges
  std::vector<std::tuple<Var, Var, double>> myRemoved;
  std::vector<Set> myDSepSets;
  auto maxSize = std::min(this->m_maxConditioning, static_cast<Var>(this->m_allVars.size() - 2));
  for (auto s = 0u; (s <= maxSize) && mxx::any_of(myEdges.size() > 0, this->m_comm); ++s) {
    LOG_MESSAGE(debug, "Testing %u edges using sets of size %u", myEdges.size(), s);
    TIMER_DECLARE(tIter);
    for (auto& e : myEdges) {
      auto result = this->checkEdge(e, allNeighbors, removedNeighbors, s, (s > 0));
      std::get<2>(e) = result.first;
      // We need to store the removed edges since the d-separating
      // set may be required for directing edges later
      if (directEdges && (s > 0) &&
          this->m_data.isIndependent(this->m_alpha, result.first)) {
        TIMER_START(this->m_tDirect);
        Var x, y;
        std::tie(x, y, std::ignore) = e;
        // We only use this information for directing colliders of type x-z-y
        // However, there is no need to store it if there are no candidates for z
        // i.e., no common neighbors between x and y
        if (!set_intersection(allNeighbors.at(x), allNeighbors.at(y)).empty()) {
          myRemoved.push_back(e);
          myDSepSets.push_back(result.second);
        }
        TIMER_PAUSE(this->m_tDirect);
      }
    }
    auto newEnd = std::remove_if(myEdges.begin(), myEdges.end(),
                                 [this] (const std::tuple<Var, Var, double>& e)
                                        { return this->m_data.isIndependent(this->m_alpha, std::get<2>(e)); });
    myEdges.erase(newEnd, myEdges.end());
    for (auto& rn : removedNeighbors) {
      if (!rn.second.empty()) {
        allNeighbors[rn.first] = set_difference(allNeighbors.at(rn.first), rn.second);
        rn.second.clear();
      }
    }
    this->m_comm.barrier();
    TIMER_START(this->m_tSync);
    this->syncSets(allNeighbors);
    TIMER_PAUSE(this->m_tSync);
    newEnd = std::remove_if(myEdges.begin(), myEdges.end(),
                            [&allNeighbors, &s] (const std::tuple<Var, Var, double>& e)
                                                { return allNeighbors.at(std::get<0>(e)).size() <= (s + 1) &&
                                                         allNeighbors.at(std::get<1>(e)).size() <= (s + 1); });
    myEdges.erase(newEnd, myEdges.end());
    if (this->m_comm.is_first()) {
      TIMER_ELAPSED("Time taken in testing all sets of size " + std::to_string(s) + ": ", tIter);
    }
    if (std::isgreaterequal(imbalanceThreshold, 0.0)) {
      std::vector<double> myWeights(myEdges.size(), 0.0);
      for (auto e = 0u; e < myEdges.size(); ++e) {
        auto n1 = allNeighbors.at(std::get<0>(myEdges[e])).size();
        auto n2 = allNeighbors.at(std::get<1>(myEdges[e])).size();
        if (n1 > s + 1) {
          myWeights[e] = boost::math::binomial_coefficient<double>(n1 - 1, s + 1);
        }
        if (n2 > s + 1) {
          myWeights[e] += boost::math::binomial_coefficient<double>(n2 - 1, s + 1);
        }
      }
      TIMER_START(this->m_tDist);
      this->fixWeightedImbalance(myEdges, myWeights, imbalanceThreshold);
      TIMER_PAUSE(this->m_tDist);
    }
  }
  if (directEdges) {
    TIMER_START(this->m_tDirect);
    this->storeRemovedEdges(std::move(myRemoved), std::move(myDSepSets), allNeighbors);
    TIMER_PAUSE(this->m_tDirect);
  }
  return this->constructSkeleton(std::move(allNeighbors));
}

template <typename Data, typename Var, typename Set>
PCStable2<Data, Var, Set>::PCStable2(
  const mxx::comm& comm,
  const Data& data,
  const double alpha,
  const Var maxConditioning
) : PCStableCommon<Data, Var, Set>(comm, data, alpha, maxConditioning)
{
}

template <typename Data, typename Var, typename Set>
BayesianNetwork<Var>
PCStable2<Data, Var, Set>::getSkeleton_parallel(
  const bool directEdges,
  const double imbalanceThreshold
) const
{
  // Initialize the learning in a similar fashion as done sequentially
  // Create nC2 edges - one for each unordered variable pair
  std::vector<std::tuple<Var, Var, double>> myEdges;
  std::unordered_map<Var, Set> allNeighbors;
  std::unordered_map<Var, Set> removedNeighbors;
  this->initializeLearning(myEdges, allNeighbors, removedNeighbors);
  // The following two data structures are used for getting
  // d-separating set and p-value for all the removed edges
  std::vector<std::tuple<Var, Var, double>> myRemoved;
  std::vector<Set> myDSepSets;
  std::set<std::tuple<Var, Var, double>> myBackwardRemoved;
  auto maxSize = std::min(this->m_maxConditioning, static_cast<Var>(this->m_allVars.size() - 2));
  for (auto s = 0u; (s <= maxSize) && mxx::any_of(myEdges.size() > 0, this->m_comm); ++s) {
    LOG_MESSAGE(debug, "Testing %u edges using sets of size %u", myEdges.size(), s);
    TIMER_DECLARE(tIter);
    for (auto& e : myEdges) {
      auto result = this->checkEdge(e, allNeighbors, removedNeighbors, s, false);
      std::get<2>(e) = result.first;
      // We need to store the removed edges since the d-separating
      // set may be required for directing edges later
      if (directEdges && (s > 0) &&
          this->m_data.isIndependent(this->m_alpha, result.first)) {
        TIMER_START(this->m_tDirect);
        Var x, y;
        double pv;
        std::tie(x, y, pv) = e;
        // We only use this information for directing colliders of type x-z-y
        // However, there is no need to store it if there are no candidates for z
        // i.e., no common neighbors between x and y
        if (!set_intersection(allNeighbors.at(x), allNeighbors.at(y)).empty()) {
          myRemoved.push_back((x < y) ? e : std::make_tuple(y, x, pv));
          myDSepSets.push_back(result.second);
        }
        TIMER_PAUSE(this->m_tDirect);
      }
    }
    this->m_comm.barrier();
    if (this->m_comm.is_first()) {
      TIMER_ELAPSED("Time taken in testing all sets of size " + std::to_string(s) + ": ", tIter);
    }
    TIMER_START(this->m_tSync);
    this->syncSets(removedNeighbors);
    TIMER_PAUSE(this->m_tSync);
    // Remove all the edges found to be independent on this processor
    // Also remove all the edges found to be independent on any other processor
    auto newEnd = std::remove_if(myEdges.begin(), myEdges.end(),
                                 [this, &removedNeighbors] (const std::tuple<Var, Var, double>& e)
                                                           { return this->m_data.isIndependent(this->m_alpha, std::get<2>(e)) ||
                                                                    removedNeighbors.at(std::get<0>(e)).contains(std::get<1>(e)); });
    myEdges.erase(newEnd, myEdges.end());
    for (auto& rn : removedNeighbors) {
      if (!rn.second.empty()) {
        allNeighbors[rn.first] = set_difference(allNeighbors.at(rn.first), rn.second);
        rn.second.clear();
      }
    }
    if (s == 0) {
      // Both direction of edges can be tested simultaneously when s > 0
      // Create a reverse edge for all the remaining edges
      auto origSize = myEdges.size();
      myEdges.resize(origSize * 2);
      Var x, y;
      for (auto f = 0u, b = origSize; f < origSize; ++f, ++b) {
        std::tie(x, y, std::ignore) = myEdges[f];
        myEdges[b] = std::make_tuple(y, x, 0.0);
      }
    }
    // Remove the edges that can no longer be tested using the first variable's neighborhood
    newEnd = std::remove_if(myEdges.begin(), myEdges.end(),
                            [&allNeighbors, &s] (const std::tuple<Var, Var, double>& e)
                                                { return allNeighbors.at(std::get<0>(e)).size() <= (s + 1); });
    myEdges.erase(newEnd, myEdges.end());
    std::set<std::tuple<Var, Var, double>> remove;
    // Now, copy all the backward edges, i.e., (x, y) such that x > y
    // and the neighborhood of x is the same as that of y
    std::copy_if(myEdges.begin(), myEdges.end(), std::inserter(remove, remove.begin()),
                 [&allNeighbors] (const std::tuple<Var, Var, double>& e)
                                 { return (std::get<0>(e) > std::get<1>(e)) &&
                                          (set_difference(allNeighbors.at(std::get<0>(e)),
                                                          allNeighbors.at(std::get<1>(e))) == Set({std::get<1>(e)})); });
    newEnd = std::remove_if(myEdges.begin(), myEdges.end(), [&remove] (const std::tuple<Var, Var, double>& e)
                                                                      { return remove.find(e) != remove.end(); });
    myEdges.erase(newEnd, myEdges.end());
    // Check if any of the backward edges removed in previous iterations
    // can be added back to the list of edges
    std::set<std::tuple<Var, Var, double>> addBackwardRemoved;
    Var x, y;
    for (const auto& e : myBackwardRemoved) {
      std::tie(x, y, std::ignore) = e;
      if ((allNeighbors.at(x).size() > (s + 1)) &&
          (set_difference(allNeighbors.at(x), allNeighbors.at(y)) != Set({y}))) {
        addBackwardRemoved.insert(e);
      }
    }
    // Remove the edges to be added back from the set of removed edges
    std::set<std::tuple<Var, Var, double>> remaining;
    std::set_difference(std::make_move_iterator(myBackwardRemoved.begin()),
                        std::make_move_iterator(myBackwardRemoved.end()),
                        addBackwardRemoved.begin(), addBackwardRemoved.end(),
                        std::inserter(remaining, remaining.begin()));
    myBackwardRemoved.clear();
    // Also add the edges removed in this iteration
    std::set_union(std::make_move_iterator(remaining.begin()),
                   std::make_move_iterator(remaining.end()),
                   remove.begin(), remove.end(),
                   std::inserter(myBackwardRemoved, myBackwardRemoved.begin()));
    // Re-add the backward edges which now need to be tested
    myEdges.insert(myEdges.end(), addBackwardRemoved.begin(), addBackwardRemoved.end());
    if (std::isgreaterequal(imbalanceThreshold, 0.0)) {
      std::vector<double> myWeights(myEdges.size());
      for (auto e = 0u; e < myEdges.size(); ++e) {
        auto n = allNeighbors.at(std::get<0>(myEdges[e])).size();
        myWeights[e] = boost::math::binomial_coefficient<double>(n - 1, s + 1);
      }
      TIMER_START(this->m_tDist);
      this->fixWeightedImbalance(myEdges, myWeights, imbalanceThreshold);
      TIMER_PAUSE(this->m_tDist);
    }
  }
  if (directEdges) {
    TIMER_START(this->m_tDirect);
    this->storeRemovedEdges(std::move(myRemoved), std::move(myDSepSets), allNeighbors, true);
    TIMER_PAUSE(this->m_tDirect);
  }
  return this->constructSkeleton(std::move(allNeighbors));
}

#endif // DETAIL_GLOBALLEARNING_HPP_