/*
 * 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 "Verifier.h"

#include "galois/Galois.h"
#include "galois/Bag.h"
#include "galois/Timer.h"
#include "galois/graphs/SpatialTree.h"
#include "Lonestar/BoilerPlate.h"
#include "llvm/Support/CommandLine.h"

#include "galois/runtime/Profile.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>
    inputname(cll::Positional, cll::desc("<input file>"), cll::Required);
static cll::opt<std::string>
    doWriteMesh("writemesh",
                cll::desc("Write the mesh out to files with basename"),
                cll::value_desc("basename"));

using Tree = typename galois::graphs::SpatialTree2d<Point*>;

//! All Point* refer to elements in this bag
using basePointBag = typename galois::InsertBag<Point>;

//! [Define Insert Bag]
using ptrPointBag = typename galois::InsertBag<Point*>;

//! Our main functor
struct Process {
  Graph& graph;
  Tree& tree;
  ptrPointBag& ptrPoints;

  Process(Graph& g, Tree& t, ptrPointBag& p)
      : graph(g), tree(t), ptrPoints(p) {}

  typedef galois::PerIterAllocTy Alloc;

  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** rp = tree.find(p->t().x(), p->t().y());
    if (!rp)
      return false;

    (*rp)->get(galois::MethodFlag::WRITE);

    GNode someNode = (*rp)->someElement();

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

    return planarSearch(p, someNode, node);
  }

  void generateMesh() {
    typedef galois::worklists::PerThreadChunkLIFO<32> CA;
    galois::for_each(galois::iterate(ptrPoints),
                     [&, self = this](Point* p, auto& ctx) {
                       p->get(galois::MethodFlag::WRITE);
                       assert(!p->inMesh());

                       GNode node;
                       if (!self->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(self->graph.getData(node).inTriangle(p->t()));
                       assert(self->graph.containsNode(node));

                       Cavity<Alloc> cav(self->graph, ctx.getPerIterAlloc());
                       cav.init(node, p);
                       cav.build();
                       cav.update();
                       self->tree.insert(p->t().x(), p->t().y(), p);
                     },
                     galois::no_pushes(), galois::per_iter_alloc(),
                     galois::loopname("Main"), galois::wl<CA>());
  }
};

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 (const 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;
    }

    tree.init(minX, minY, maxX, maxY);

    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;
  Tree& tree;

public:
  ReadPoints(PointList& p, Tree& t) : points(p), tree(t) {}

  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);
    }
  }
};

struct ReadInput {
  Graph& graph;
  Tree& tree;
  basePointBag& basePoints;
  ptrPointBag& ptrPoints;

  ReadInput(Graph& g, Tree& t, basePointBag& b, ptrPointBag& p)
      : graph(g), tree(t), basePoints(b), ptrPoints(p) {}

  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);

    tree.insert(p1->t().x(), p1->t().y(), p1);

    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;
  }

  struct centerXCmp {
    template <typename T>
    bool operator()(const T& lhs, const T& rhs) const {
      return lhs.t().x() < rhs.t().x();
    }
  };

  struct centerYCmp {
    template <typename T>
    bool operator()(const T& lhs, const T& rhs) const {
      return lhs.t().y() < rhs.t().y();
    }
  };

  struct centerYCmpInv {
    template <typename T>
    bool operator()(const T& lhs, const T& rhs) const {
      return rhs.t().y() < lhs.t().y();
    }
  };

  template <typename Iter>
  void divide(const Iter& b, const Iter& e) {
    if (std::distance(b, e) > 64) {
      std::sort(b, e, centerXCmp());
      Iter m = galois::split_range(b, e);
      std::sort(b, m, centerYCmpInv());
      std::sort(m, e, centerYCmp());
      divide(b, galois::split_range(b, m));
      divide(galois::split_range(b, m), m);
      divide(m, galois::split_range(m, e));
      divide(galois::split_range(m, e), e);
    } else {
      std::random_shuffle(b, e);
    }
  }

  void layoutPoints(PointList& points) {
    divide(points.begin(), points.end() - 3);
    galois::do_all(galois::iterate(points.begin(), points.end() - 3),
                   [&](Point& p) {
                     Point* pr = &basePoints.push(p);
                     ptrPoints.push(pr);
                   });
    //! [Insert elements into InsertBag]
    Point* p1 = &basePoints.push(*(points.end() - 1));
    Point* p2 = &basePoints.push(*(points.end() - 2));
    Point* p3 = &basePoints.push(*(points.end() - 3));
    //! [Insert elements into InsertBag]
    addBoundaryNodes(p1, p2, p3);
  }

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

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

    galois::preAlloc(2 * 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::pagePoolSize()) // in pages
    );
    galois::reportPageAlloc("MeminfoPre");

    layoutPoints(points);
  }
};

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();
}

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();
}

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

  Graph graph;
  Tree tree;
  basePointBag basePoints;
  ptrPointBag ptrPoints;

  ReadInput(graph, tree, basePoints, ptrPoints)(inputname);

  galois::StatTimer T;
  T.start();
  galois::runtime::profileVtune(
      [&]() { Process(graph, tree, ptrPoints).generateMesh(); },
      "MeshGeneration");
  T.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, tree).from(inputname);
    writePoints(base.append(".node"), points);
  }

  return 0;
}


// This file is part of LatticeTester.
//
// LatticeTester
// Copyright (C) 2012-2021  The LatNet Builder author's, supervised by Pierre L'Ecuyer, Universite de Montreal.
//
// 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.

#include "netbuilder/Helpers/JoeKuo.h"
#include "netbuilder/Helpers/Path.h"
#include <cmath>

#include <string>
#include <fstream>
#include <boost/algorithm/string.hpp>
#include <boost/filesystem.hpp>

using namespace std;

namespace NetBuilder { namespace JoeKuo {

//===========================================================================

LatticeTester::Weight Weights::getWeight (const LatticeTester::Coordinates& projection) const
{
   if (projection.size()==2)
   {
         return std::pow((LatticeTester::Weight).9999, *projection.rbegin());
   }
   // fall back to zero
   return 0.0;
}

//===========================================================================

void Weights::format(ostream& os) const
{
   using LatticeTester::operator<<;
   os << "Weights()";
}

//===========================================================================

const char* ws = " \t\n\r\f\v";

// trim from end (right)
inline std::string& rtrim(std::string& s, const char* t = ws)
{
      s.erase(s.find_last_not_of(t) + 1);
      return s;
}

// trim from beginning (left)
inline std::string& ltrim(std::string& s, const char* t = ws)
{
      s.erase(0, s.find_first_not_of(t));
      return s;
}

// trim from both ends (left & right)
inline std::string& trim(std::string& s, const char* t = ws)
{
      return ltrim(rtrim(s, t), t);
}

std::vector<std::vector<uInteger>> readJoeKuoDirectionNumbers(Dimension dimension)
{
      assert(dimension >= 1 && dimension <= 21201);
      std::string path = PATH_TO_LATNETBUILDER_DIR + "/../share/latnetbuilder/data/JoeKuoSobolNets.csv";
      std::vector<std::vector<uInteger>> res(dimension);
      if (boost::filesystem::exists(path)){
            std::ifstream file(path);
            std::string sent;

            do
            {
            getline(file,sent);
            trim(sent);
            }
            while (sent != "###");

            getline(file,sent);

            for(unsigned int i = 1; i <= dimension; ++i)
            {
                  if(getline(file,sent))
                  {
                        std::vector<std::string> fields;
                        boost::split( fields, sent, boost::is_any_of( ";" ) );
                        for( const auto& token : fields)
                        {
                              res[i-1].push_back(std::stol(token));
                        }
                  }
                  else
                  {
                        break;
                  }
            }
      }
      else{
            throw runtime_error("Unable to locate data folder. The value of PATH_TO_LATNETBUILDER_DIR is probably incorrect. See netbuilder/Path.h.");
      }
      return res;
}

std::vector<DirectionNumbers> getJoeKuoDirectionNumbers(Dimension dimension)
{
      std::vector<std::vector<uInteger>> tmp = readJoeKuoDirectionNumbers(dimension);
      std::vector<DirectionNumbers> genVals(dimension);
      for(unsigned int j = 0; j < dimension; ++j)
      {
            genVals[j] = DirectionNumbers(j,tmp[j]);
      }
      return genVals;
}

DigitalNet<NetConstruction::SOBOL> createJoeKuoSobolNet(Dimension dimension, MatrixSize size)
{
      auto genVals = getJoeKuoDirectionNumbers(dimension);
      return DigitalNet<NetConstruction::SOBOL>(dimension, size, std::move(genVals));
}

std::unique_ptr<DigitalNet<NetConstruction::SOBOL>> createPtrToJoeKuoSobolNet(Dimension dimension, MatrixSize size)
{
      auto genVals = getJoeKuoDirectionNumbers(dimension);
      return std::make_unique<DigitalNet<NetConstruction::SOBOL>>(dimension, size, std::move(genVals));
}

}} // namespace


/* Copyright (c) 2015, Julian Straub <jstraub@csail.mit.edu>, Randi Cabezas <rcabezas@csail.mit.edu>
 * Licensed under the MIT license. See the license file LICENSE.
 */

#pragma once
#include <iostream>
#include <stdint.h>
#include <vector>
#include <Eigen/Dense>

#include <boost/shared_ptr.hpp>

#include <dpMM/dpMM.hpp>
#include <dpMM/cat.hpp>
#include <dpMM/dir.hpp>
#include <dpMM/niw.hpp>
#include <dpMM/sampler.hpp>
#include <dpMM/basemeasure.hpp>
#include <dpMM/niwBaseMeasure.hpp>
#include <dpMM/niwSphere.hpp>
#include <dpMM/dirBaseMeasure.hpp>
#include <mmf/mfBaseMeasure.hpp>

using namespace Eigen;
using std::cout;
using std::endl;
using boost::shared_ptr;
using std::vector;

template<typename T=double>
class DirMultiNaiveBayes : public DpMM<T>{

public:
  DirMultiNaiveBayes(std::ifstream &in, boost::mt19937 *rng);
  DirMultiNaiveBayes(const Dir<Cat<T>, T>& alpha, const vector<boost::shared_ptr<BaseMeasure<T> > >&thetas);
  DirMultiNaiveBayes(const Dir<Cat<T>, T>& alpha, const vector< vector<boost::shared_ptr<BaseMeasure<T> > > >&thetas);
  virtual ~DirMultiNaiveBayes();

  virtual void loadData(const vector<vector<Matrix<T,Dynamic,Dynamic> > > &x);//does nothing other than load data
  virtual void initialize(const vector<vector< Matrix<T,Dynamic,Dynamic> > >&x);
  virtual void initializeNoParamSampling(const vector<vector< Matrix<T,Dynamic,Dynamic> > >&x);
  virtual void initialize(const vector<vector< Matrix<T,Dynamic,Dynamic> > >&x, VectorXu &z);
  virtual void initialize(const boost::shared_ptr<ClGMMData<T> >&cld)
    {cout<<"not supported"<<endl; assert(false);};

  virtual void sampleLabels();
  virtual void MAPLabel();
  virtual void sampleParameters();

  virtual T logJoint(bool verbose=false);
  virtual const VectorXu& labels(){return z_;};
  virtual const VectorXu& getLabels(){return z_;};
  virtual uint32_t getK() const { return K_;};
  virtual uint32_t getM() const { return M_;};
  virtual uint32_t getN() const { return Nd_;};

//  virtual MatrixXu mostLikelyInds(uint32_t n);
  virtual MatrixXu mostLikelyInds(uint32_t n, Matrix<T,Dynamic,Dynamic>& logLikes);

  Matrix<T,Dynamic,1> getCounts();

  virtual void inferAll(uint32_t nIter, bool verbose=false);

  virtual void dump(std::ofstream& fOutMeans, std::ofstream& fOutCovs);
  virtual void dump_clean(std::ofstream &out);

  virtual vector<boost::shared_ptr<BaseMeasure<T> > > getThetas(uint32_t m) {
	  return(thetas_[m]);
  };
  virtual boost::shared_ptr<BaseMeasure<T> > getThetas(uint32_t m, uint32_t k) {
	  return(thetas_[m][k]);
  };

  virtual void setTheta(uint32_t m, uint32_t k, boost::shared_ptr<BaseMeasure<T> > newTheta ) {
	  thetas_[m][k] = newTheta;
  };

  virtual vector<T> evalLogLik(const vector<Matrix<T,Dynamic,1> > xnew, const vector<uint32_t> clusterInd,
							   const vector<uint32_t> comp2eval =vector<uint32_t>());

  virtual uint32_t sampleLabels(const vector<Matrix<T,Dynamic,1> > xnew,
								const vector<uint32_t> comp2eval =vector<uint32_t>());
  virtual vector<uint32_t> sampleLabels(const vector<vector<Matrix<T,Dynamic,1> > > xnew,
										const vector<uint32_t> comp2eval =vector<uint32_t>());
  virtual uint32_t MAPLabels(const vector<Matrix<T,Dynamic,1> > xnew,
							 const vector<uint32_t> comp2eval =vector<uint32_t>());
  virtual vector<uint32_t> MAPLabels(const vector<vector<Matrix<T,Dynamic,1> > > xnew,
									 const vector<uint32_t> comp2eval =vector<uint32_t>());
  virtual void updatePDF();

  vector<uint32_t> getLogEvalItersHist() {return logJointIterEval;}
  vector<T> getLogJointHist() {return logJointHist;}

protected:
  virtual T evalLogLik(const vector<Matrix<T,Dynamic,1> > xnew, const uint32_t clusterInd,
					   const vector<uint32_t> comp2eval);
  virtual uint32_t labels_sample_max(const vector<Matrix<T,Dynamic,1> > xnew,
									 const vector<uint32_t> comp2eval, const bool return_MAP_labels=false);

  uint32_t Nd_;
  uint32_t K_; //num cluseters
  uint32_t M_; //num data sources
  Dir<Cat<T>, T> dir_;
  Cat<T> pi_;
#ifdef CUDA
  SamplerGpu<T>* sampler_;
#else
  Sampler<T>* sampler_;
#endif
  virtual void initialize_sampler();
  Matrix<T,Dynamic,Dynamic> pdfs_;
//  Cat cat_;
  vector<vector<boost::shared_ptr<BaseMeasure<T> > > > thetas_;  // theta_[M][K]

  //suffiecient stats
  vector<vector<Matrix<T,Dynamic,Dynamic> > > x_; //x_[M][doc](:,word)
  VectorXu z_;

  virtual void helper_setDims();
  vector<VectorXu> dataDim;

  vector<uint32_t> logJointIterEval;
  vector<T> logJointHist;
};

// --------------------------------------- impl -------------------------------

template<typename T>
void DirMultiNaiveBayes<T>::initialize_sampler() {
	if (sampler_ != NULL) {
		delete sampler_;
		sampler_ = NULL;
	}
	//initialize sampler
	#ifdef CUDA
	  sampler_ = new SamplerGpu<T>(uint32_t(Nd_),K_,dir_.pRndGen_);
	#else
	  sampler_ = new Sampler<T>(dir_.pRndGen_);
	#endif
}

template<typename T>
DirMultiNaiveBayes<T>::DirMultiNaiveBayes(std::ifstream &in, boost::mt19937 *rng) :
 dir_(Matrix<T,2,1>::Ones(),rng), pi_(dir_.sample()),sampler_(NULL)
{
	//initialize the class from the file pointer given
	in >> M_;
	in >> K_;
	in >> Nd_;

	vector<uint32_t> dim;
	vector<baseMeasureType> type;
	Matrix<T,Dynamic,1> alpha(K_), pi(K_);

	for(uint32_t m = 0; m<M_; ++m){
		uint32_t temp;
		in >> temp;
		dim.push_back(temp);
	}

	for(uint32_t m = 0; m<M_; ++m){
		uint32_t temp;
		in >> temp;
		type.push_back(baseMeasureType(temp));
	}

	z_ = VectorXu(Nd_);
	for(uint32_t n=0; n<Nd_; ++n)
		in >> z_(n);

	for(uint32_t k=0; k<K_; ++k)
		in >> alpha(k);

	for(uint32_t k=0; k<K_; ++k)
		in >> pi(k);

	pdfs_ = Matrix<T,Dynamic,Dynamic>(Nd_,K_);
	//for(uint32_t n=0; n<Nd_*K_; ++n)
	//	in >> pdfs_(n%Nd_, (n-(n%Nd_))/Nd_);
	//	//in >> pdfs_((n-(n%Nd_))/Nd_, n%Nd_);
	////pdfs_ = pdfs_.transpose();

	//get parameters
	for(uint32_t m=0; m<M_; ++m) {
		baseMeasureType typeIter = baseMeasureType(type[m]);
		vector<boost::shared_ptr<BaseMeasure<T> > > thetaM;
		if(typeIter==NIW_SAMPLED) {
			uint32_t Diter = dim[m];
			T nu, kappa;
			Matrix<T,Dynamic,Dynamic> scatter(Diter, Diter), sigma(Diter,Diter);
			Matrix<T,Dynamic,1> theta(Diter), mu(Diter);
			for(uint32_t k=0; k<K_; ++k) {
				//get nu and kappa
					in>> nu;
					in>> kappa;
				//get theta
					for(uint32_t n=0; n<Diter; ++n)
						in >> theta(n);
					theta = theta.transpose();
				//get scatter
					for(uint32_t n=0; n<Diter*Diter; ++n)
						in >> scatter((n-(n%Diter))/Diter, n%Diter);
				//get mean
					for(uint32_t n=0; n<Diter; ++n)
						in >> mu(n);
					mu = mu.transpose();
				//get sigma
					for(uint32_t n=0; n<Diter*Diter; ++n)
						in >> sigma((n-(n%Diter))/Diter, n%Diter);

				//build theta[m][k]
				NIW<T> niw(scatter,theta,nu,kappa,rng);
				Normal<T> normal(mu,sigma,rng);
				boost::shared_ptr<NiwSampled<T> > baseIter( new NiwSampled<T>(niw, normal));

				//set
				thetaM.push_back(boost::shared_ptr<BaseMeasure<T> >(baseIter));
			}
		} else if(typeIter==NIW_SPHERE) {
			uint32_t Diter = dim[m]-1;
			T nu;
			T counts;
			Matrix<T,Dynamic,Dynamic> scatter(Diter, Diter), sigma(Diter,Diter),delta(Diter,Diter);
			Matrix<T,Dynamic,1> mu(Diter+1), mu_prior(Diter), north(Diter+1);
			for(uint32_t k=0; k<K_; ++k) {
				//get nu
					in>> nu;
					in>> counts;
				//get prior mean
					for(uint32_t n=0; n<Diter; ++n)
						in>> mu_prior(n);
				//get scatter
					for(uint32_t n=0; n<Diter*Diter; ++n)
						in >> scatter((n-(n%Diter))/Diter, n%Diter);
				//get delta
					for(uint32_t n=0; n<Diter*Diter; ++n)
						in >> delta((n-(n%Diter))/Diter, n%Diter);
				//get mean
					for(uint32_t n=0; n<(Diter+1); ++n)
						in >> mu(n);
					mu = mu.transpose();
				//get sigma
					for(uint32_t n=0; n<Diter*Diter; ++n)
						in >> sigma((n-(n%Diter))/Diter, n%Diter);
				//get north
					for(uint32_t n=0; n<(Diter+1); ++n)
						in >> north(n);
					north = north.transpose();

				//build theta[m][k]
				IW<T> iw(delta,nu, scatter, mu_prior, counts, rng);
				//IW<T> iw(delta,nu, rng);
				NiwSphere<T> niwSp(iw,rng);
				niwSp.S_ = Sphere<T>(north);
				niwSp.normalS_ = NormalSphere<T>(mu,sigma,rng);

				//set
				thetaM.push_back(boost::shared_ptr<BaseMeasure<T> >(niwSp.copy()));
				//boost::shared_ptr<NiwSphere<T> > baseIter( new NiwSphere<T>(iw,rng));
				//thetaM.push_back(boost::shared_ptr<BaseMeasure<T> >(baseIter));

			}

		} else if(typeIter==DIR_SAMPLED) {
			uint32_t Diter = dim[m];
			T localCount;
			Matrix<T,Dynamic,1> post_alpha(Diter), counts(Diter), pdf(Diter);

			for(uint32_t k=0; k<K_; ++k) {
				//get dir alpha
				for(uint32_t n=0; n<Diter; ++n)
					in >> post_alpha(n);

				//get dir counts
				for(uint32_t n=0; n<Diter; ++n)
					in >> counts(n);

				//get disc pdf
				for(uint32_t n=0; n<Diter; ++n)
					in >> pdf(n);

				in >> localCount;

				//build theta[m][k]
				Dir<Cat<T>,T> dirBase(post_alpha,counts, rng);
				DirSampled<Cat<T>,T> dirSamp(dirBase);
				Cat<T> disc(pdf,rng);
				dirSamp.disc_ = disc;
				dirSamp.count_	= localCount;

				//set
				thetaM.push_back(boost::shared_ptr<BaseMeasure<T> >(dirSamp.copy()));
			}

		} else if(typeIter==MF_T) {
			//			uint32_t Diter = dim[m];
			//			T localCount;
			Matrix<T,Dynamic,1> post_alpha(6), counts(6), pi_pdf(6);
			post_alpha.fill(1);
			counts.fill(0);
			pi_pdf.fill(0);

			Matrix<T,3,3>  R;
			Dir<Cat<T>, T> alpha(post_alpha,rng);
			Cat<T> pi(pi_pdf,rng);
			std::vector<shared_ptr<BaseMeasure<T> > > iwTs;
			std::vector<NormalSphere<T> > TGs;
			uint32_t nIter;

			for(uint32_t k=0; k<K_; ++k)
			{
        in >> nIter;
        //get dir alpha
        for(uint32_t n=0; n<6; ++n) in >> post_alpha(n);
        //get dir counts
        for(uint32_t n=0; n<6; ++n) in >> counts(n);
        //get pi pdf
        for(uint32_t n=0; n<6; ++n) in >> pi_pdf(n);
        alpha.alpha_ = post_alpha;
        alpha.setCounts(counts);
        pi.pdf(pi_pdf);
        // get rotation
        for(uint32_t n=0; n<9; ++n) in >> R(n/3,n%3);
        for(uint32_t j=0; j<6; ++j)
        {
          // load IW in tangent space
          Matrix<T,Dynamic,Dynamic> Delta(2,2);
          Matrix<T,Dynamic,Dynamic> Scatter(2,2);
          Matrix<T,Dynamic,1> mean(2);
          T nu,count;
          in >> nu;
          for(uint32_t n=0; n<4; ++n) in >> Delta(n/2,n%2);
          for(uint32_t n=0; n<4; ++n) in >> Scatter(n/2,n%2);
          for(uint32_t n=0; n<2; ++n) in >> mean(n);
          in >> count;
          IW<T> iw(Delta,nu,Scatter,mean,count,rng);
          iwTs.push_back(shared_ptr<IwTangent<T> >(
                new IwTangent<T>(iw,rng)));
          // load tangent space gaussians
          Matrix<T,Dynamic,Dynamic> Sigma(2,2);
          Matrix<T,Dynamic,1> mu(3);
          for(uint32_t n=0; n<4; ++n) in >> Sigma(n/2,n%2);
          for(uint32_t n=0; n<3; ++n) in >> mu(n);
          TGs.push_back(NormalSphere<T>(mu,Sigma,rng));
          reinterpret_cast<IwTangent<T>* >(
              iwTs[j].get())->normalS_=TGs[j];
        };
        // labels -- not set
        uint32_t Nk = 0;
        in >> Nk;
        VectorXu z(Nk);
        for(uint32_t n=0; n<Nk; ++n) in >> z(n);
        // build model
        DirMM<T> dirMM(alpha,iwTs);
        MfPrior<T> mfPrior(dirMM, nIter);
        MF<T> mf(R,pi,TGs);
				//set
				thetaM.push_back(boost::shared_ptr<BaseMeasure<T> >(
              new MfBase<T>(mfPrior, mf)));
			}

		} else {
				std::cerr << "[DirMultiNaiveBayes::dump_clean] error saving...returning" << endl;
				return;
		}
		thetas_.push_back(thetaM);
	}

	dir_ =  Dir<Cat<T>, T>(alpha,rng);
	pi_ = Cat<T>(pi,rng);

	this->initialize_sampler();

	//cout << "M:" << M_ << " K: " << K_ << " nd: " << Nd_ << endl;
	//for(uint32_t m = 0; m<M_; ++m){
	//	cout << m << ": d="<< dim[m] << ", type=" << type[m] << endl;
	//}
	//cout << "z_ " << z_.transpose() << endl;
	//cout << "alpha " << alpha.transpose() << endl;
}

template<typename T>
DirMultiNaiveBayes<T>::DirMultiNaiveBayes(const Dir<Cat<T>,T>& alpha,
    const vector<boost::shared_ptr<BaseMeasure<T> > >& thetas) :
  sampler_(NULL), K_(alpha.K_), dir_(alpha), pi_(dir_.sample()), M_(uint32_t(thetas.size()))
{
	for (uint32_t m=0; m<M_; ++m)
	{
		vector<boost::shared_ptr<BaseMeasure<T> > > temp;
		for (uint32_t k=0; k<K_; ++k)
		{
      		temp.push_back(boost::shared_ptr<BaseMeasure<T> >(thetas[m]->copy()));
      	}
      thetas_.push_back(temp);
    }

#ifndef NDEBUG
	for(uint32_t m=0; m<int(M_); ++m) {
		for(int k=0; k<int(K_); ++k) {
				thetas_[m][k]->print();
		}
	}
#endif

};

template<typename T>
DirMultiNaiveBayes<T>::DirMultiNaiveBayes(const Dir<Cat<T>,T>& alpha,
    const vector< vector<boost::shared_ptr<BaseMeasure<T> > > >& theta) :
 K_(alpha.K_), M_(uint32_t(theta.size())),dir_(alpha),
  pi_(dir_.sample()), sampler_(NULL),thetas_(theta)
{ };

template<typename T>
DirMultiNaiveBayes<T>::~DirMultiNaiveBayes()
{
  if (sampler_ != NULL) {
	delete sampler_;
	sampler_ = NULL;
  }
};

template <typename T>
Matrix<T,Dynamic,1> DirMultiNaiveBayes<T>::getCounts()
{
  return counts<T,uint32_t>(z_,K_);
};

template<typename T>
void DirMultiNaiveBayes<T>::loadData(const vector<vector<Matrix<T,Dynamic,Dynamic> > > &x){
  x_ = x;
  this->helper_setDims();
}

template<typename T>
void DirMultiNaiveBayes<T>::initializeNoParamSampling(const vector< vector< Matrix<T,Dynamic,Dynamic> > > &x)
{
  // randomly init labels from prior
  Nd_= uint32_t(x.front().size());

  z_ = VectorXu::Zero(Nd_);
  //init data and labels from given
  pi_.sample(z_);

  x_ = x;

  pdfs_.setZero(Nd_,K_);

  this->initialize_sampler();
  this->helper_setDims();
};

template<typename T>
void DirMultiNaiveBayes<T>::initialize(const vector< vector< Matrix<T,Dynamic,Dynamic> > > &x)
{
  uint32_t Nd= uint32_t(x.front().size());

  // randomly init labels from prior
  VectorXu z;
  z.setZero(Nd);
  Cat<T> pi = dir_.sample();
  pi.sample(z);

  //delegate the initialization to the main intitialization function
  this->initialize(x,z);
};

template<typename T>
void DirMultiNaiveBayes<T>::initialize(const vector< vector< Matrix<T,Dynamic,Dynamic> > > &x, VectorXu &z)
{
  Nd_= uint32_t(x.front().size());

  //init data and labels from given
  x_ = x;
  z_ = z;

  pi_ = dir_.sample();

  pdfs_.setZero(Nd_,K_);

  this->initialize_sampler();
  this->helper_setDims();
  this->sampleParameters();
};

template<typename T>
void DirMultiNaiveBayes<T>::helper_setDims()
{
	dataDim.clear();
	dataDim.reserve(M_);
	for(uint32_t m=0; m<M_; ++m) {
		VectorXu temp(x_[m].size());
		for(uint32_t t=0; t<x_[m].size(); ++t)
			temp(t) = uint32_t(x_[m][t].cols());
		dataDim.push_back(temp);
	}
}

template<typename T>
void DirMultiNaiveBayes<T>::sampleLabels()
{
	// obtain posterior categorical under labels
	pi_ = dir_.posterior(z_).sample();
	//  cout<<pi_.pdf().transpose()<<endl;
	this->updatePDF();
	// sample z_i
	sampler_->sampleDiscPdf(pdfs_,z_);
};

template<typename T>
void DirMultiNaiveBayes<T>::updatePDF() {

// compute categorical distribution over label z_i
// no need to re-compute the array and log every iteration)
VectorXd logPdf_z_value = pi_.pdf().array().log();

#pragma omp parallel for
  for(int32_t d=0; d<int32_t(Nd_); ++d)
  {
	VectorXd logPdf_z = logPdf_z_value;
	for(uint32_t m=0; m<uint32_t(M_); ++m)
	{
		for(uint32_t k=0; k<K_; ++k)
		{
			//updated to SS
			logPdf_z[k] += thetas_[m][k]->logLikelihoodFromSS(x_[m][d]);
		}
	}
//    cout<<endl;
    // make pdf sum to 1. and exponentiate
    pdfs_.row(d) = (logPdf_z.array()-logSumExp(logPdf_z)).exp().matrix().transpose();
//    cout<<pi_.pdf().transpose()<<endl;
//    cout<<pdf.transpose()<<" |.|="<<pdf.sum();
//    cout<<" z_i="<<z_[d]<<endl;
  }

}

template<typename T>
void DirMultiNaiveBayes<T>::MAPLabel()
{
	/* it chooses the MAP label rather than sampling */
	this->updatePDF();

	#pragma omp parallel for
	 for(int32_t d=0; d<int32_t(Nd_); ++d) {
		 int r,c;
		 pdfs_.row(d).maxCoeff(&r, &c);
		 z_(d) = c;
	 }
};

template<typename T>
void DirMultiNaiveBayes<T>::sampleParameters()
{
//unpacks the contains here vector<vector<Matrix>> into what the posterior expects Matrix
	MatrixXu dim(M_,K_);
	for(uint32_t m=0; m<M_; ++m) {
		#pragma omp parallel for
		for(int32_t k=0; k<int32_t(K_); ++k) {
			VectorXu temp = (z_.array()==k).select(dataDim[m],0);
			dim(m,k) = temp.sum();
		}
	}

	for(int32_t m=0; m<M_; ++m) {
	#ifdef _WINDOWS
		//#pragma omp parallel for
		for(int32_t k=0; k<int32_t(K_); ++k) {
	#else
		//#pragma omp parallel for
		#pragma omp parallel for schedule(dynamic)
		for(int32_t k=0; k<int32_t(K_); ++k) {
	#endif

			if(dim(m,k)!=0) {

				//Matrix<T,Dynamic,1> ssIn = Matrix<T,Dynamic,1>::Zero(x_[m].front().rows());
				vector< Matrix<T,Dynamic,1> >dataIn;
				dataIn.reserve(dim(m,k));

				uint32_t count=0;

				for(int32_t d=0; d<Nd_; ++d) {
					if(z_[d]==k) {
						int add_size =int(x_[m][d].cols());
						//ssIn += x_[m][d]; //update iteration SS
						dataIn.push_back(x_[m][d]);
						count+=add_size;

						if(count==dim(m,k))
							break; //early out if you found them all
					}
				}
				//update values
				//thetas_[m][k]->posteriorFromSS(ssIn);

//        if(thetas_[m][k]->getBaseMeasureType() == MF_T)
//        {
//          cout<<"MF "<<k<<" ---------------- "<<endl;
//        }
				//always sends zeros and look for zeros
				thetas_[m][k]->posteriorFromSS(dataIn,VectorXu::Zero(dim(m,k)),0);

			} else {
				//Matrix<T,Dynamic,1> ssIn = Matrix<T,Dynamic,1>::Zero(x_[m].front().rows());
				//ssIn[0]=1; //set counts to 1 to avoid inf
				//the posterior needs to reset
				//thetas_[m][k]->posteriorFromSS(ssIn);
				//passing in one data point (all zeros, with 1 index value=0 and looking for 1)

				vector<Matrix<T,Dynamic,1> >dataIn;
				dataIn.push_back(Matrix<T,Dynamic,1>::Zero(x_[m].front().rows(),1));
				//the posterior needs to reset
        VectorXu zz = VectorXu::Zero(1);
				thetas_[m][k]->posteriorFromSS(dataIn,zz,1);
			}
		}
	}
};

template<typename T>
T DirMultiNaiveBayes<T>::logJoint(bool verbose)
{
  T logJoint = dir_.logPdf(pi_);
  if(verbose)
  	cout<<"\tlog p(pi)="<< logJoint << endl;

  for(int32_t m=0; m<int32_t(M_); ++m) {
	  T logPriorM = 0;
	  #pragma omp parallel for reduction(+:logPriorM)
	  for (int32_t k=0; k<int32_t(K_); ++k) {
		logPriorM = logPriorM + thetas_[m][k]->logPdfUnderPrior();
	  }
	  logJoint+=logPriorM;
	  if(verbose)
	  	cout<<"\tlog p(theta_" << m << ")="<< logPriorM << endl;

  }

	for (int32_t m=0; m<int32_t(M_); ++m) {
		T logThetaM=0;
		#pragma omp parallel for reduction(+:logThetaM)
		for (int32_t d=0; d<int32_t(Nd_); ++d) {
			logThetaM = logThetaM + thetas_[m][z_[d]]->logLikelihoodFromSS(x_[m][d]);
		}
		logJoint += logThetaM;
		if(verbose)
			cout<<"\tlog p(x|z,theta_" << m << ")=" << logThetaM << endl;

	}
    if(verbose)
		cout<<"log p(pi)*p(theta)*p(x|z,theta)=" << logJoint << endl;

  return logJoint;
};

template<typename T>
void DirMultiNaiveBayes<T>::inferAll(uint32_t nIter, bool verbose)
{
  if(verbose){
  	cout<<"[DirMultiNaiveBayes::inferALL] ------ inferingALL (nIter=" << nIter << ") ------"<<endl;
  	if(Nd_<=100) {
		cout <<"initial labels:"<< endl;
  		cout<<this->labels().transpose()<<endl;
	}
  }

  logJointIterEval.clear(); logJointHist.clear();
  if(verbose) {
	//all iterations stored
	logJointIterEval.reserve(nIter); logJointHist.reserve(nIter);
  } else {
	  //only mod 100 stored
	  logJointIterEval.reserve(int((nIter/100) + 1)); logJointHist.reserve(int((nIter/100) + 1));
  }

  for(uint32_t t=0; t<nIter; ++t)
  {
    this->sampleLabels();
    this->sampleParameters();
    if(verbose)
    {
      for(int m=0; m<int(M_); ++m) {
        for(int k=0; k<int(K_); ++k) {
          thetas_[m][k]->print();
        }
      }
    }
    if(verbose || t%int(ceil(nIter/100.))==0)
    {
      VectorXu Ns = counts<uint32_t,uint32_t>(
          this->labels(),K_).transpose();
      uint32_t K = K_;
      for(uint32_t k = 0; k<K_; ++k)
        if (Ns(k) == 0) --K;
      cout<<"@i "<<t<<": # "
        <<K<<" "<<std::setw(1)<<Ns.transpose() <<endl;

      T iterLogJoint = this->logJoint(true) ;
      //log iterJoint Prob
      logJointIterEval.push_back(t);
      logJointHist.push_back(iterLogJoint);

      if(Nd_<=10) {
        cout << "[" << std::setw(3)<< std::setfill('0')
          << t <<"] label: "
          << this->labels().transpose()
          << " [joint= " << std::setw(6) << iterLogJoint << "]"<< endl;
      } else {
        cout << "[" << std::setw(3)<< std::setfill('0')
          << t <<"] joint= "
          << std::setw(6) << iterLogJoint << endl;
      }
    }
//    if(verbose)
//    {
//      VectorXu Ns = counts<uint32_t,uint32_t>(this->labels(),K_).transpose();
//      uint32_t K = K_;
//      for(uint32_t k = 0; k<K_; ++k)
//        if (Ns(k) == 0) --K;
//      cout<<"@i "<<t<<": # "<<K<<" "<<std::setw(1) <<Ns.transpose() <<endl;
//    }
  }
  //keeps the MAP label in memory
  this->MAPLabel();
  this->sampleParameters();
}

template <typename T>
void DirMultiNaiveBayes<T>::dump(std::ofstream& fOutMeans, std::ofstream& fOutCovs)
{
	cout << "dumping MultiObs naiveBayes" << endl;
	cout << "doc index: " << endl;
	cout << this->labels().transpose() << endl;

	cout << "printing num components: " << endl;
	cout << M_ << endl;

	cout << "printing cluster params: " << endl;
	cout << K_ << endl;

	for(uint32_t m=0; m<M_; ++m) {
		cout << "component: " << m  << endl;
		for(uint32_t k=0; k<K_; ++k) {
			cout << "theta: " << k  << endl;
			thetas_[m][k]->print();
		}
	}

	cout << "printing mixture params: " << endl;
	pi_.print();
}

template <typename T>
void DirMultiNaiveBayes<T>::dump_clean(std::ofstream &out){
	//clean dump, only data with specific format
//FORMAT:
	//M 1x1
	//K	1x1
	//Nd 1x1
	//D[m] 1xM
	//Type[m] 1xM
	//labels 1xNd
	//Dir alpha 1xK
	//pi pdf 1xK
	//pdf KxNd
	// mixture parameters
	//params Loop over M then K each contains data type for specific type
	//for type 1 (NIWSampled)
		//---prior--- (NIW)
			//nu 1x1
			//kappa 1x1
			//theta 1xD
			//scatter DxD
		//---estimate (normal)
			//mu 1xD
			//Sigma DxD
	//for type 2 (NIWSphereFull)
		//----prior--- (IW)
			//nu 1x1
			//count 1x1
			//mean 1x(D-1)
			//scatter (D-1)x(D-1)
			//Delta	 (D-1)x(D-1)
		//--posterior (NormalSphere)
			//mean 1x(D-1)
			//Sigma (D-1)x(D-1)
		//--sphere--- (Sphere)
			//north 1x(D-1)
	//for type 3 (DirSampled)
		//--posterior (Dir)
			//alpha 1xK
			//counts 1xK
		//--distribution (Cat)
			//pdf 1xK
		//--counts (scalar)
			//counts 1x1
	//logJoint history
		// Niter 1x1
		// iterValue 1xNiter (iteration corresponding to the logValue)
		// logJoint	 1xNiter (logJoint)

	//this fixes issues with eigen matrices printing (eg, 00-0.7 )
	int curPres = int(out.precision());
	out.precision(10);
	IOFormat fullPresPrint(FullPrecision,DontAlignCols);

	//prints headers
	out << M_ << endl
		 << K_ << endl
		 << Nd_ << endl;

	//print dim
	for(uint32_t m=0; m<M_; ++m) {
		vector<boost::shared_ptr<BaseMeasure<T> > >  theta_base = this->getThetas(m);
		uint32_t temp = theta_base.front()->getDim();
		out << temp << " ";
		//out << x_[m].front().rows() << " ";
	}
	out << endl;
	//print type
	for(uint32_t m=0; m<M_; ++m) {
		out << thetas_[m].front()->getBaseMeasureType() << " ";
	}
	out << endl;

	//print labels
	out << this->labels().transpose() << endl;

	//print mixture parameters
	out << this->dir_.alpha_.transpose() << endl;
	out << this->pi_.pdf_.transpose().format(fullPresPrint) << endl;
	//out << this->pdfs_.transpose().format(fullPresPrint) << endl;

	//print parameters
	for(uint32_t m=0; m<M_; ++m) {
		vector<boost::shared_ptr<BaseMeasure<T> > >  theta_base = this->getThetas(m);
		for(uint32_t k=0; k<K_; ++k) {
			baseMeasureType type = theta_base[k]->getBaseMeasureType();
			if(type==NIW_SAMPLED) {
				boost::shared_ptr<NiwSampled<T> >  *theta_iter =
						reinterpret_cast<boost::shared_ptr<NiwSampled<T> >* >( &theta_base[k]);
					//printing prior
					NIW<T> prior = theta_iter->get()->niw0_;
					out << prior.nu_				 << endl <<
						    prior.kappa_			 << endl <<
						    prior.theta_.transpose() << endl <<
						    prior.Delta_.format(fullPresPrint)	<<	endl;

					//printing posterior
					Normal<T> norm = theta_iter->get()->normal_;
					out << norm.mu_.transpose() << endl;
					out << norm.Sigma().format(fullPresPrint) << endl;
			} else if(type==NIW_SPHERE) {
        boost::shared_ptr<NiwSphere<T> >  *theta_iter =
          reinterpret_cast<boost::shared_ptr<NiwSphere<T> >* >(
              &theta_base[k]);
				//prior
				IW<T> prior = theta_iter->get()->iw0_;
				out << prior.nu_ 				 << endl
					 << prior.count()			 << endl
				 	 << prior.mean().transpose() << endl
					 << prior.scatter().format(fullPresPrint)			 << endl
					 << prior.Delta_.format(fullPresPrint)			 << endl;

				//posterior
				NormalSphere<T> norm = theta_iter->get()->normalS_;
				out << norm.getMean().transpose().format(fullPresPrint) << endl;
				out << norm.Sigma().format(fullPresPrint) << endl;
				//sphere
				Sphere<T> sp = theta_iter->get()->S_;
				out << sp.north().transpose() << endl;
			} else if(type==DIR_SAMPLED) {
				boost::shared_ptr<DirSampled<Cat<T>,T> >  *theta_iter =
						reinterpret_cast<boost::shared_ptr<DirSampled<Cat<T>,T> >* >( &theta_base[k]);
				//posterior
				Dir<Catd,T>  post = theta_iter->get()->dir0_;
				out <<	post.alpha_.transpose()		<< endl;
				out <<	post.counts().transpose()	<< endl;

				//distribution
				Catd dist = theta_iter->get()->disc_;
				out <<	dist.pdf_.transpose() << endl;

				//counts
				T counts = theta_iter->get()->count_;
				out << counts << endl;

			} else if(type==MF_T) {
        boost::shared_ptr<MfBase<T> > *theta_iter =
          reinterpret_cast<boost::shared_ptr<MfBase<T> >* >(
              &theta_base[k]);

        out<< theta_iter->get()->mf0_.T_<<endl;

				//posterior
				Dir<Cat<T>, T>  post = theta_iter->get()->mf0_.dirMM().Alpha();
				out <<	post.alpha_.transpose()		<< endl;
				out <<	post.counts().transpose()	<< endl;
				Cat<T> pi = theta_iter->get()->mf0_.dirMM().Pi();
				out <<	pi.pdf().transpose() << endl;

        out << theta_iter->get()->mf_.R().format(fullPresPrint)<<endl;
        for(uint32_t j=0; j<6; ++j)
        {
          out<< theta_iter->get()->mf0_.theta(j)->iw0_.nu_ << endl;
          out<< theta_iter->get()->mf0_.theta(j)->iw0_.Delta_ << endl;
          out<< theta_iter->get()->mf0_.theta(j)->iw0_.scatter() << endl;
          out<< theta_iter->get()->mf0_.theta(j)->iw0_.mean() << endl;
          out<< theta_iter->get()->mf0_.theta(j)->iw0_.count() << endl;
          out<< theta_iter->get()->mf0_.theta(j)->normalS_.Sigma() << endl;
          out<< theta_iter->get()->mf0_.theta(j)->normalS_.getMean() << endl;
        }

        // output the labeling
        out<<theta_iter->get()->mf0_.dirMM().labels().size()<<endl;
	      out<<theta_iter->get()->mf0_.dirMM().labels().transpose()<<endl;

			} else {
					std::cerr << "[DirMultiNaiveBayes::dump_clean] error saving...returning" << endl;
					return;
			}
		}
	}

	//print logHistory
	out << int(logJointHist.size()) << endl;
	for(int i=0; i<logJointIterEval.size(); ++i)
		out << logJointIterEval[i] << " ";
	out << endl;

	for(int i=0; i<logJointHist.size(); ++i)
		out << logJointHist[i] << " ";
	out << endl;

	out.precision(curPres);

}

//template <typename T>
//void DirMultiNaiveBayes<T>::dump_clean(std::ofstream &out){
//	streambuf *coutbuf = std::cout.rdbuf(); //save old cout buffer
//	cout.rdbuf(out.rdbuf()); //redirect std::cout to fout1 buffer
//	this->dump_clean(); //write using cout to the specified buffer
//	std::cout.rdbuf(coutbuf); //reset to standard output again
//}

template <typename T>
T DirMultiNaiveBayes<T>::evalLogLik(const vector<Matrix<T,Dynamic,1> > xnew,
							  const uint32_t clusterInd, const vector<uint32_t> comp2eval)
{
	//T logJoint = pi_.pdf_(clusterInd);
	T logJoint  = 0;
	for (int32_t m=0; m<int32_t(comp2eval.size()); ++m)
	{
		logJoint += thetas_[comp2eval[m]][clusterInd]->logLikelihoodFromSS(xnew[m]);
	}

  return logJoint;
}

template <typename T>
vector<T> DirMultiNaiveBayes<T>::evalLogLik(const vector<Matrix<T,Dynamic,1> > xnew,
											const vector<uint32_t> clusterInd,
											const vector<uint32_t> comp2eval) {

	vector<uint32_t> comp2evalLocal = comp2eval;
	if(comp2evalLocal.empty()) {
		for(uint32_t m=0; m<M_; ++m)
			comp2evalLocal.push_back(m);
	}

	vector<T> out;
	for(uint32_t k=0; k<uint32_t(clusterInd.size()); ++k) {
		out.push_back(this->evalLogLik(xnew,clusterInd[k],comp2evalLocal));
	}
	return(out);
}

  template <typename T>
uint32_t DirMultiNaiveBayes<T>::labels_sample_max(const
    vector<Matrix<T,Dynamic,1> > xnew, const vector<uint32_t>
    comp2eval, const bool return_MAP_labels)
{
  /* xnew in the form x[docs][m][SS] */
  VectorXd logPdf_z = pi_.pdf().array().log();

  for(int32_t m=0; m<comp2eval.size(); ++m)
  {
    for(int32_t k=0; k<int32_t(K_); ++k)
    {
      logPdf_z[k] += thetas_[comp2eval[m]][k]->logLikelihoodFromSS(
          xnew[m]);
    }
  }

  // make pdf sum to 1. and exponentiate
  Matrix<T,Dynamic,Dynamic> pdfLocal =  Matrix<T,Dynamic,Dynamic>(1,K_);
  pdfLocal = (logPdf_z.array()-logSumExp(logPdf_z)).exp().matrix().transpose();

  VectorXu zout = VectorXu(1);

  if(return_MAP_labels) {
    // return MAP label
    int r,c;
    pdfLocal.maxCoeff(&r, &c);
    zout(0) = c;
  } else {
    // sample z_i
    sampler_->sampleDiscPdf(pdfLocal,zout);
  }

  return(zout(0));
};

template <typename T>
uint32_t DirMultiNaiveBayes<T>::sampleLabels(const vector<Matrix<T,Dynamic,1> > xnew,
											 const vector<uint32_t> comp2eval) {
	return(this->labels_sample_max(xnew, comp2eval, false));
}

template <typename T>
vector<uint32_t> DirMultiNaiveBayes<T>::sampleLabels(const vector<vector<Matrix<T,Dynamic,1> > > xnew,
													 const vector<uint32_t> comp2eval)
{
	/* xnew in the form x[doc][m][SS] */
	vector<uint32_t> out;
	for(uint32_t d=0; d<xnew.size(); ++d) {
		out.push_back(this->sampleLabels(xnew[d],comp2eval));
	}
	return(out);
}

template <typename T>
uint32_t DirMultiNaiveBayes<T>::MAPLabels(const vector<Matrix<T,Dynamic,1> > xnew,
										  const vector<uint32_t> comp2eval) {
	return(this->labels_sample_max(xnew, comp2eval, true));
}

template <typename T>
vector<uint32_t> DirMultiNaiveBayes<T>::MAPLabels(const vector<vector<Matrix<T,Dynamic,1> > > xnew,
												  const vector<uint32_t> comp2eval) {
	/* xnew in the form x[docs][m][SS] */
	vector<uint32_t> out;
	for(uint32_t d=0; d<xnew.size(); ++d) {
		out.push_back(this->MAPLabels(xnew[d],comp2eval));
	}
	return(out);
}

template<typename T>
MatrixXu DirMultiNaiveBayes<T>::mostLikelyInds(uint32_t n, Matrix<T,Dynamic,Dynamic>& logLikes)
{
  MatrixXu inds = MatrixXu::Zero(n,K_);
  logLikes = Matrix<T,Dynamic,Dynamic>::Ones(n,K_);

#pragma omp parallel for
  for (int32_t k=0; k<K_; ++k)
  {
    for (uint32_t i=0; i<z_.size(); ++i)
      if(z_(i) == k)
      {
        T logLike = 0.;
        // iterate over datasources and sum up their logLikes
        for(uint32_t m=0; m<uint32_t(M_); ++m)
        {
          logLike += thetas_[m][z_[i]]->logLikelihoodFromSS(x_[m][i]);
        }
        // keep only the top n and sorted
        for (uint32_t j=0; j<n; ++j)
          if(logLikes(j,k) < logLike)
          {
            for(uint32_t l=n-1; l>j; --l)
            {
              logLikes(l,k) = logLikes(l-1,k);
              inds(l,k) = inds(l-1,k);
            }
            logLikes(j,k) = logLike;
            inds(j,k) = i;
//            cout<<"after update "<<logLike<<endl;
//            Matrix<T,Dynamic,Dynamic> out(n,K_*2);
//            out<<logLikes.cast<T>(),inds.cast<T>();
//            cout<<out<<endl;
            break;
          }
      }
  }
  cout<<"::mostLikelyInds: logLikes"<<endl;
  cout<<logLikes<<endl;
  cout<<"::mostLikelyInds: inds"<<endl;
  cout<<inds<<endl;
  return inds;
};


// Copyright 2015-2022 The ALMA Project Developers
//
// 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.

#pragma once

/// @file
///
/// Code related to fully analytical 1D RTA solutions.

#include <constants.hpp>
#include <structures.hpp>
#include <qpoint_grid.hpp>
#include <Eigen/Dense>

namespace alma {

namespace analytic1D {

/// Class for computing basic thermal properties
/// (kappa, Cv, diffusivity) and cumulative functions
/// (resolved for MFP, energy, etc.) along a given transport direction.
/// Calculations involve appropriate thin film corrections
/// for both cross-plane and in-plane transport.

class BasicProperties_calculator {
public:
    /// Constructor: initialise internal variables
    BasicProperties_calculator(const alma::Crystal_structure* poscar,
                               const alma::Gamma_grid* grid,
                               const Eigen::ArrayXXd* w,
                               double T);

    /// Function setting the thermal transport axis
    void setDirection(const Eigen::Vector3d unitvector);

    /// Construct a logarithmic grid of MFP bins
    /// from MFPmin to MFPmax with Nbins elements
    void setLogMFPbins(double MFPmin, double MFPmax, int Nbins);

    /// Construct a logarithmic grid of MFP bins witn Nbins elements.
    /// MFPmax is assigned automatically to twice the largest MFP.
    /// MFPmin is assigned to 1e-6*largest MFP.
    void setAutoMFPbins(int Nbins);

    /// Construct a logarithmic grid of MFP bins witn Nbins elements.
    /// MFPmax is assigned automatically to the twice the largest projected MFP.
    /// MFPmin is assigned to 1e-6*largest projected MFP.
    void setAutoProjMFPbins(int Nbins);

    /// Construct a logarithmic grid of tau bins
    /// from taumin to taumax with Nbins elements
    void setLogRTbins(double taumin, double taumax, int Nbins);

    /// Construct a logarithmic grid of tau bins witn Nbins elements.
    /// taumax is assigned automatically to the twice the largest relaxation
    /// time.
    /// taumin is assigned to 1e-6*largest relaxation time.
    void setAutoRTbins(int Nbins);

    /// Construct a linear grid of omega bins
    /// from omegamin to omegamax with Nbins elements
    void setLinOmegabins(double omegamin, double omegamax, int Nbins);

    /// Construct a linear grid of omega bins witn Nbins elements.
    /// omegamax is assigned automatically to 1.1*largest angular frequency.
    /// omegamin is automatically set to 0.
    void setAutoOmegabins(int Nbins);

    /// Retrieve the MFP bins
    Eigen::VectorXd getMFPbins();

    /// Retrieve the relaxation time bins
    Eigen::VectorXd getRTbins();

    /// Retrieve the relaxation time bins
    Eigen::VectorXd getOmegabins();

    /// Specify that the medium should be treated as
    /// infinite bulk (default option)
    void setBulk();

    /// Specify that the medium should be treated as
    /// a thin film with provided thickness.
    /// This option corrects MFPs and relaxation times
    /// in conductivity calculations.
    void setInPlaneFilm(double filmthickness,
                        const Eigen::Vector3d normal,
                        double specularity = 0.0);
    void setCrossPlaneFilm(double filmthickness);

    /// Retrieve thermal conductivity
    double getConductivity();

    /// Obtain spectrally computed thermal conductivity
    double getSpectralConductivity();

    /// Compute the heat capacity and retrieve it
    double getCapacity();

    /// Compute the Fourier diffusivity and retrieve it
    double getDiffusivity();

    /// Retrieve the dominant projected MFP as defined below
    double getDominantProjMFP();

    /// Retrieve the dominant phonon relaxation time
    double getDominantRT();

    /// Compute the anisotropy index kappa_max/kappa_min and retrieve it.
    double getAnisotropyIndex();

    /// Choose resolving cumulative curves by MFP
    void resolveByMFP();

    /// Choose resolving cumulative curves by projected MFP
    void resolveByProjMFP();

    /// Choose resolving cumulative curves by relaxation time
    void resolveByRT();

    /// Choose resolving cumulative curves by angular frequency
    void resolveByOmega();

    /// Compute the cumulative thermal conductivity curve
    /// with respect to the MFP bins and retrieve it
    Eigen::VectorXd getCumulativeConductivity();

    /// Compute the cumulative heat capacity curve
    /// with respect to the MFP bins and retrieve it
    Eigen::VectorXd getCumulativeCapacity();

    /// Compute the cumulative hydrodinamic l2
    /// using RTA and assuming isotropy
    Eigen::VectorXd getCumulativel2RTAiso();

    /// Set the number of bins to be used in DOS evaluation
    void setDOSgridsize(int nbins);

    /// Obtain the DOS energy grid (in meV)
    Eigen::VectorXd getDOSgrid();

    /// Calculate and obtain the DOS
    Eigen::VectorXd getDOS();

private:
    /// Pointer to description of the unit cell
    const alma::Crystal_structure* poscar;
    /// Pointer to phonon spectrum on a regular q-point grid
    const alma::Gamma_grid* grid;
    /// Pointer to scattering rates for all modes irreducible points in the grid
    const Eigen::ArrayXXd* w;
    /// Temperature
    double T;

    /// Unitvector describing 1D thermal transport axis
    Eigen::Vector3d unitvector;
    /// Unitvector describing the film normal for in-plane transport
    Eigen::Vector3d normalvector;
    /// MFP bins for cumulative curves
    Eigen::VectorXd MFPbins;
    /// relaxation time bins for cumulative curves
    Eigen::VectorXd taubins;
    /// omega bins for cumulative curves
    Eigen::VectorXd omegabins;
    /// Max MFP value present in medium
    double internal_MFPmax;
    /// Max proj MFP value present in medium
    double internal_MFPprojmax;
    /// Max relaxation time present in medium
    double internal_taumax;
    /// Max phonon frequency present in medium
    double internal_omegamax;
    /// Thin film specifiers
    bool thinfilm;
    bool crossplane;
    double filmthickness;
    double specularity; // for in-plane films
    /// Basic properties
    double kappa;
    double Cv;
    /// "Dominant" projected MFP, defined as kappa-weighted average
    /// <Lambda_proj> = sum[(kappamode/kappabulk)*Lambda_proj]
    double dominantProjMFP;
    /// "Dominant" relaxation time, defined as kappa-weighted average
    /// <RT> = sum[(kappamode/kappabulk)*RT_mode]
    double dominantRT;
    /// 3D rotation matrix that rotates the cartesian coordinate system such
    /// that
    /// the transport axis becomes (1,0,0) and film normal becomes (0,0,1)
    Eigen::Matrix3d FuchsRotation;
    /// Precalculate the 3D rotation matrix used for in-plane Fuch corrections.
    void initFuchsRotation();
    /// Precalculate thermal properties
    void updateMe();
    /// Calculated cumulative conductivity function
    Eigen::VectorXd kappacumul;
    /// Calculated cumulative capacity function
    Eigen::VectorXd Cvcumul;
    /// Calculated cumulative l**2
    Eigen::VectorXd l2cumul;

    /// Specifier for kappacumul calculation
    enum kappacumulID {
        resolve_by_MFP,
        resolve_by_ProjMFP,
        resolve_by_RT,
        resolve_by_omega
    };
    int kappacumulIdentifier;
    /// Number of bins to be used for DOS evaluation
    int DOS_Nbins;
};

/// Class for computing the RTA propagator function psi(xi) of a medium.
/// The psi function fully determines the analytical single pulse
/// response of the infinite bulk in weakly quasi-ballistic regime
/// (time scales exceeding phonon relaxation times):
/// Energy density in Fourier-Laplace domain = 1/[s + psi(xi)]
/// Energy density in Fourier-time domain = exp[-psi(xi)*t]

class psi_calculator {
public:
    /// Constructor: initialise internal variables
    psi_calculator(const alma::Crystal_structure* poscar,
                   const alma::Gamma_grid* grid,
                   const Eigen::ArrayXXd* w,
                   double T);

    /// Function setting the thermal transport axis
    void setDirection(const Eigen::Vector3d unitvector);

    /// Construct a linear grid of spatial frequencies
    /// from ximin to ximax with Nxi elements
    void setLinGrid(double ximin, double ximax, int Nxi);

    /// Construct a logarithmic grid of spatial frequencies
    /// from ximin to ximax with Nxi elements
    void setLogGrid(double ximin, double ximax, int Nxi);

    /// Manually set a grid of spatial frequencies
    void setXiGrid(const Eigen::Ref<const Eigen::VectorXd> xigrid);

    /// Retrieve the spatial frequency grid
    Eigen::VectorXd getSpatialFrequencies();

    /// Determine whether computation output should be
    /// normalised by the Fourier solution Dbulk*xi^2
    void normaliseOutput(bool norm);

    /// Obtain Fourier diffusivity of the medium
    double getDiffusivity();

    /// Compute psi function and retrieve it
    Eigen::VectorXd getPsi();

private:
    /// Pointer to description of the unit cell
    const alma::Crystal_structure* poscar;
    /// Pointer to phonon spectrum on a regular q-point grid
    const alma::Gamma_grid* grid;
    /// Pointer to scattering rates for all modes irreducible points in the grid
    const Eigen::ArrayXXd* w;
    /// Temperature
    double T;

    /// Unitvector describing 1D thermal transport axis
    Eigen::Vector3d unitvector;
    /// Spatial frequency grid
    Eigen::VectorXd xi;
    /// Normalisation specifier
    bool scaleoutput;
    /// Fourier diffusivity
    double Dbulk;
    void updateDiffusivity();
    /// Calculated psi function
    Eigen::VectorXd psi;
};

/////////////////// SPR_calculator_FourierLaplace ///////////////////

/// Class for computing the exact analytical RTA single pulse
/// energy density response of the infinite bulk medium
/// in Fourier-Laplace domain. [see PRB 91 085202 (2015)]

class SPR_calculator_FourierLaplace {
public:
    /// Constructor: initialise internal variables
    SPR_calculator_FourierLaplace(const alma::Crystal_structure* poscar,
                                  const alma::Gamma_grid* grid,
                                  const Eigen::ArrayXXd* w,
                                  double T);

    /// Function setting the thermal transport axis
    void setDirection(const Eigen::Vector3d unitvector);

    /// Construct a linear grid of spatial frequencies
    /// from ximin to ximax with Nxi elements
    void setLinSpatialGrid(double ximin, double ximax, int Nxi);

    /// Construct a logarithmic grid of spatial frequencies
    /// from ximin to ximax with Nxi elements
    void setLogSpatialGrid(double ximin, double ximax, int Nxi);

    /// Construct a linear grid of temporal frequencies
    /// from fmin to fmax with Nf elements.
    /// From this a Laplace grid s = 2*pi*1i*f is created.
    void setLinTemporalGrid(double fmin, double fmax, int Nf);

    /// Construct a logarithmic grid of temporal frequencies
    /// from fmin to fmax with Nf elements.
    /// From this a Laplace grid s = 2*pi*1i*f is created.
    void setLogTemporalGrid(double fmin, double fmax, int Nf);

    /// Retrieve the spatial frequency grid
    Eigen::VectorXd getSpatialFrequencies();

    /// Retrieve the temporal frequency grid
    Eigen::VectorXd getTemporalFrequencies();

    /// Compute SPR and retrieve it
    Eigen::MatrixXcd getSPR();

private:
    /// Pointer to description of the unit cell
    const alma::Crystal_structure* poscar;
    /// Pointer to phonon spectrum on a regular q-point grid
    const alma::Gamma_grid* grid;
    /// Pointer to scattering rates for all modes irreducible points in the grid
    const Eigen::ArrayXXd* w;
    /// Temperature
    double T;

    /// Unitvector describing 1D thermal transport axis
    Eigen::Vector3d unitvector;
    /// Spatial frequency grid
    Eigen::VectorXd xi;
    /// Temporal frequency grids
    Eigen::VectorXd f;
    Eigen::VectorXcd s;
    /// Calculated SPR
    Eigen::MatrixXcd Pxis;
};

/////////////////// SPR_calculator_RealSpace ///////////////////

/// Class for computing the approximate analytical RTA single pulse
/// energy density response in real space at a given time.
///
/// The approach is only valid in weakly quasiballistic regime
/// (time scales exceeding phonon relaxation times) so that
/// P(xi,t) \approx exp[-psi(xi)*t].
///
/// Calculations are performed by evaluating the Fourier inversion
/// (1/pi)*Integral(exp[-psi(xi)*t]*cos(xi*x),xi=0..infinity)
/// semi-analytically with 2nd order Filon-type quadrature.

class SPR_calculator_RealSpace {
public:
    /// Constructor: initialise internal variables
    SPR_calculator_RealSpace(const alma::Crystal_structure* poscar,
                             const alma::Gamma_grid* grid,
                             const Eigen::ArrayXXd* w,
                             double T);

    /// Function setting the thermal transport axis
    void setDirection(const Eigen::Vector3d unitvector);

    /// Construct a linear space grid
    /// from xmin to xmax with Nx elements
    void setLinGrid(double xmin, double xmax, int Nx);

    /// Construct a logarithmic space grid
    /// from xmin to xmax with Nx elements
    void setLogGrid(double xmin, double xmax, int Nx);

    /// Declare whether the space grid is normalised.
    /// If yes, space grid values are counted in
    /// Fourier diffusion lengths, i.e.
    /// x_actual = x_grid*sqrt(2*Dbulk*t)
    void declareGridNormalised(bool norm);

    // Optional normalisation of calculation output
    // by sqrt(4*pi*Dbulk*t)
    void normaliseOutput(bool norm);

    /// Set time value to be used for calculations
    void setTime(double t);

    /// Set MFP bins to be used for resolveSPRbyMFP()
    void setLogMFPbins(double MFPmin, double MFPmax, int Nbins);

    /// Retrieve the spatial grid
    Eigen::VectorXd getGrid();
    Eigen::VectorXd getNormalisedGrid();

    /// Retrieve time value
    double getTime();

    /// Retrieve MFP bins
    Eigen::VectorXd getMFPbins();

    /// Retrieve Fourier diffusivity along thermal transport axis
    double getDiffusivity();

    /// Compute source transient P(x=0) at the provided times
    Eigen::VectorXd getSourceTransient(
        const Eigen::Ref<Eigen::VectorXd> timegrid);

    /// Compute SPR and retrieve it
    Eigen::VectorXd getSPR();

    /// Compute SPR of individual modes resolved by MFP
    Eigen::MatrixXd resolveSPRbyMFP();

private:
    /// Pointer to description of the unit cell
    const alma::Crystal_structure* poscar;
    /// Pointer to phonon spectrum on a regular q-point grid
    const alma::Gamma_grid* grid;
    /// Pointer to scattering rates for all modes irreducible points in the grid
    const Eigen::ArrayXXd* w;
    /// Temperature
    double T;

    /// Unitvector describing 1D thermal transport axis
    Eigen::Vector3d unitvector;
    /// Unitvector describing the film normal for inplane transport
    Eigen::Vector3d filmnormal;
    /// Spatial grid
    Eigen::VectorXd x;
    bool gridIsNormalised;
    /// Output specifier
    bool normaliseoutput;
    /// MFP bins
    Eigen::VectorXd MFPbins;
    /// time value
    double t;
    /// diffusivity of the medium along the transport axis
    double Dbulk;
    void updateDiffusivity();
    /// Calculated SPR
    Eigen::VectorXd Pxt;
    /// Calculated SPR resolved by MFP
    Eigen::MatrixXd Pmodes;
};

/////////////////// MSD_calculator_Laplace ///////////////////

/// Class for computing the exact analytical RTA solution for
/// mean square thermal energy displacement. By definition,
/// the MSD is the variance of the macroscopic energy density:
/// MSD = Integral(x^2*P(x,s),x=-infinity..infinity).
/// Results are valid across all regimes
/// (from fully ballistic to fully diffusive transport).
/// Calculations are performed fully analytically by using
/// the moment generating properties of the energy density
/// in spatial frequency domain:
/// MSD = minus second derivative of P(xi,s) at xi = 0.

class MSD_calculator_Laplace {
public:
    /// Constructor: initialise internal variables
    MSD_calculator_Laplace(const alma::Crystal_structure* poscar,
                           const alma::Gamma_grid* grid,
                           const Eigen::ArrayXXd* w,
                           double T);

    /// Function setting the thermal transport axis
    void setDirection(const Eigen::Vector3d unitvector);

    /// Construct a linear temporal frequency grid
    /// from fmin to fmax with Nf elements
    void setLinGrid(double fmin, double fmax, int Nf);

    /// Construct a logarithmic temporal frequency grid
    /// from fmin to fmax with Nf elements
    void setLogGrid(double fmin, double fmax, int Nf);

    /// Manually set a Laplace grid
    void setLaplaceGrid(const Eigen::Ref<const Eigen::VectorXcd> sgrid);

    /// Retrieve the temporal frequency grid
    Eigen::VectorXd getGrid();
    Eigen::VectorXcd getLaplaceGrid();

    /// Compute MSD and retrieve it
    Eigen::VectorXcd getMSD();

private:
    /// Pointer to description of the unit cell
    const alma::Crystal_structure* poscar;
    /// Pointer to phonon spectrum on a regular q-point grid
    const alma::Gamma_grid* grid;
    /// Pointer to scattering rates for all modes irreducible points in the grid
    const Eigen::ArrayXXd* w;
    /// Temperature
    double T;

    /// Unitvector describing 1D thermal transport axis
    Eigen::Vector3d unitvector;
    /// Temporal frequency grid
    Eigen::VectorXd f;
    /// Laplace grid
    Eigen::VectorXcd s;
    /// Calculated MSD function
    Eigen::VectorXcd MSD;
};

/////////////////// MSD_calculator_RealTime ///////////////////

/// Class for computing the exact analytical RTA solution for
/// mean square thermal energy displacement in time domain.
///
/// Calculations are performed by numerically inverting
/// output from MSD_calculator_Laplace to time domain
/// with a Gaver-Stehfest scheme.

class MSD_calculator_RealTime {
public:
    /// Constructor: initialise internal variables
    MSD_calculator_RealTime(const alma::Crystal_structure* poscar,
                            const alma::Gamma_grid* grid,
                            const Eigen::ArrayXXd* w,
                            double T);

    /// Function setting the thermal transport axis
    void setDirection(const Eigen::Vector3d unitvector);

    /// Construct a linear time grid
    /// from tmin to tmax with Nt elements
    void setLinGrid(double tmin, double tmax, int Nt);

    /// Construct a logarithmic time grid
    /// from tmin to tmax with Nt elements
    void setLogGrid(double tmin, double tmax, int Nt);

    /// Manually set a time grid
    void setTimeGrid(const Eigen::Ref<const Eigen::VectorXd> tgrid);

    /// Retrieve the timegrid
    Eigen::VectorXd getGrid();

    /// Optional choice to normalise calculated results
    /// by Fourier solution 2*D*t.
    void normaliseOutput(bool norm);

    /// Retrieve Fourier diffusivity
    double getDiffusivity();

    /// Compute MSD and retrieve it
    Eigen::VectorXd getMSD();

private:
    /// Pointer to description of the unit cell
    const alma::Crystal_structure* poscar;
    /// Pointer to phonon spectrum on a regular q-point grid
    const alma::Gamma_grid* grid;
    /// Pointer to scattering rates for all modes irreducible points in the grid
    const Eigen::ArrayXXd* w;
    /// Temperature
    double T;

    /// Unitvector describing 1D thermal transport axis
    Eigen::Vector3d unitvector;
    /// Unitvector describing the film normal for inplane transport
    Eigen::Vector3d filmnormal;
    /// Time grid
    Eigen::VectorXd t;
    /// Fourier diffusivity
    double Dbulk;
    void updateDiffusivity();
    /// output specifier
    bool normaliseoutput;
    /// Calculated MSD function
    Eigen::VectorXd MSD;

    /// Gaver-Stehfest Laplace inversion coefficients
    int GS_depth;
    Eigen::VectorXd GS_coeffs;
    /// Helper function for Gaver-Stehfest Laplace inversion
    double factorial(int n);
};
} // end namespace analytic1D
} // end namespace alma


//
// Created by markus on 23.04.21.
//

#include <random>
#include <boost/math/distributions/exponential.hpp>
#include "ParticleFilter.h"
#include "../Simulator/world/world.h"
#include "../Simulator/interactive_objects/DistanceSensor.h"
#include "helpers.h"
#include "DummyRobot.h"

ParticleFilter::ParticleFilter(RobotControlInterface *robot, std::string map_filename, int initialN, int targetN) {
    this->particles_world = new World(map_filename, "Particle filter");
    this->particles_world->addMapObject(this->mapLine);
    this->robot = robot;

    this->N = initialN;
    this->TARGET_N = targetN;

    // create initial particles distribution and show the simulation map
    this->particles = ParticleFilter::create_uniform_particles(this->particles_world->get_map_bounds().x,
                                                               this->particles_world->get_map_bounds().y,
                                                               this->N);
    this->estimatedRobot = new MapRobot(robot->get_radius());
    this->estimatedRobot->setColor(CV_RGB(255, 255, 255));
    this->particles_world->addMapObject(this->estimatedRobot);
    ParticleFilter::updateParticleSimulation(this->particles, SHOW_WHATS_GOING_ON);
}

std::vector<std::array<double, 3>> ParticleFilter::create_uniform_particles(double x_range, double y_range, int N) {
    std::vector<std::array<double, 3>> particles;
    std::default_random_engine generator;
    generator.seed(std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::system_clock::now().time_since_epoch()).count());

    std::uniform_real_distribution<double> x_distribution(0, x_range);
    std::uniform_real_distribution<double> y_distribution(0, y_range);
    std::uniform_real_distribution<double> angle_distribution(-M_PI, M_PI);

    for (int i = 0; i < N; i++) {
        particles.push_back(std::array<double, 3>({x_distribution(generator),
                                                   y_distribution(generator),
                                                   angle_distribution(generator)}));
    }
    return particles;
}

void ParticleFilter::updateParticleSimulation(std::vector<std::array<double, 3>> particles, bool showMap, std::vector<double> weights) {
    // get sensor config to recreate (clone) the same sensors as on the robot
    std::vector<std::array<double, 2>> sensor_config;
    for (DistanceSensor *sensor : DistanceSensor::filter_for_distance_sensor(this->robot->get_sensors())) {
        sensor_config.push_back(std::array<double, 2>({sensor->get_sensor_angle(), sensor->get_sensor_max_distance()}));
    }

    // If weights are provided they will be visualized on the map. For the color scaling we have to know what the maximum probability of one particle is
    double maxWeight = 0;
    bool useCustomWeightColors = false;
    if (!weights.empty()) {
        useCustomWeightColors = true;
        for (double weight : weights) {
            maxWeight = maxWeight > weight ? maxWeight : weight;
        }
    }

    // add or delete robots from the world if the amount does not match
    if (this->particles_world->getRobotsList().size() > particles.size()) {
        // there are too many particles -> remove some
        this->particles_world->deleteRobotByIndex(particles.size() - 1, this->particles_world->getRobotsList().size() - particles.size() + 1, true);
    } else if (this->particles_world->getRobotsList().size() < particles.size()) {
        // there arent enough particles -> add some
        for (long i = this->particles_world->getRobotsList().size(); i<particles.size(); ++i) {
            // create robot
            auto* dummyRobot = new DummyRobot("particle", this->robot->get_radius(), this->particles_world);
            // set default visualization options
            dummyRobot->setDrawOptions(CV_RGB(255, 0, 0), true, 1);
            // add robot to world
            this->particles_world->addRobot(dummyRobot);

            // create sensors
            for (std::array<double, 2> s_config : sensor_config) {
                DistanceSensor *sensor = new DistanceSensor(this->particles_world, dummyRobot, s_config[0], s_config[1], true);
                dummyRobot->add_sensor(sensor);
            }
        }
    }

    // update robots
    std::vector<Robot*> dummyRobotsList = this->particles_world->getRobotsList();
    for (int i = 0; i < particles.size(); ++i) {
        std::array<double, 3> particle = particles[i];
        DummyRobot* dummyRobot = (DummyRobot*) dummyRobotsList[i];
        dummyRobot->reposition(particle[0], particle[1], particle[2]);

        // set custom visualization options
        if (useCustomWeightColors) {
            dummyRobot->setDrawOptions(getColor(weights[i] / maxWeight), true, 1);
        }

        for (SensorInterface* sensor : dummyRobot->get_sensors()) {
            sensor->update_sensor_data(true);
        }
    }

    // show map if enabled
    if (showMap) {
        this->particles_world->show_map(true);
    }
}

std::vector<std::array<double, 3>> ParticleFilter::particles_predict(std::vector<std::array<double, 3>> *oldParticles,
                                                                     double move_distance,
                                                                     double move_angle,
                                                                     double standard_deviation) {
    std::vector<std::array<double, 3>> updated_particles;

    // init normal_distribution
    std::random_device rd{};
    std::mt19937 gen{rd()};
    std::normal_distribution<> d{0, standard_deviation};

    // update oldParticles
    for (std::array<double, 3> particle : *oldParticles) {
        // update moved distance and angle with a norm distribution value (x * ((100 + <norm_dist>) / 100))
        double distance_with_distribution = move_distance * ((100 + d(gen)) / 100);
        double angle_with_distribution = move_angle * ((100 + d(gen)) / 100);

        double newRobotAngle = angle_with_distribution + particle[2];

        // update oldParticles with distance and angle
        updated_particles.push_back(std::array<double, 3>(
                {
                        particle[0] + cos(newRobotAngle) * distance_with_distribution,
                        particle[1] + sin(newRobotAngle) * distance_with_distribution,
                        newRobotAngle
                }));
    }

    return updated_particles;
}

std::vector<double> ParticleFilter::particles_update(std::vector<double> robotSensorValues, std::vector<Robot *> simRobots, double lambda) {
    std::vector<double> updated_weights = std::vector<double>(simRobots.size(), 1.0);
    auto d = boost::math::exponential_distribution<>{lambda};

    double sumOfWeights = 0;

    for (int i = 0; i < updated_weights.size(); ++i) {
        // collect sensor values of simulated particles
        std::vector<DistanceSensor *> sensors = DistanceSensor::filter_for_distance_sensor(simRobots.at(i)->get_sensors());
        int j = 0;
        for (DistanceSensor *sensor : sensors) {
            // calculate correlation
            updated_weights[i] *= boost::math::pdf(d, abs(sensor->get_simplified_sensor_value() - robotSensorValues[j]));
            updated_weights[i] += 1.e-300; // prevent rounding to 0
            j++;
        }

        sumOfWeights += updated_weights[i];
    }

    // normalize weights: sum of weights should equal 1
    for (double &updated_weight : updated_weights) {
        updated_weight /= sumOfWeights;
    }

    return updated_weights;
}

std::tuple<std::vector<std::array<double, 3>>, std::vector<double>> ParticleFilter::particles_resample(
        std::vector<std::array<double, 3>> *oldParticles,
        std::vector<double> *weights,
        int N,
        bool enableRandomParticles,
        double noise,
        cv::Point2d mapBounds) {
    // grouping algorithm for resampling is explained here: https://robotics.stackexchange.com/a/481
    // creates an index list of entries to keep
    std::vector<double> positions = {};
    for (int i = 0; i < N; ++i) {
        positions.push_back((i + ((double) rand() / RAND_MAX)) / N);
    }

    std::vector<int> indexes(N, 0);
    std::vector<double> cumulativeSum = {};
    cumulativeSum.push_back(weights->at(0));
    for (int i = 1; i < N; ++i) {
        cumulativeSum.push_back(cumulativeSum.at(i - 1) + weights->at(i));
    }

    int i = 0;
    int j = 0;
    while (i < N && j < N) {
        if (positions[i] < cumulativeSum[j]) {
            indexes[i] = j;
            i += 1;
        } else {
            j += 1;
        }
    }

    // create a new particles list. It will be based on the previous one. Based on the index list created above some values will be kept, others
    // will be deleted. There will be many duplicates, which is expected. Each value will be modified by a noise value.
    std::default_random_engine generator;
    std::uniform_real_distribution<double> uniformNoise(-noise, noise);

    std::vector<std::array<double, 3>> updatedParticles = {};
    std::vector<double> updatedWeights = {};
    double weightSum = 0;
    for (int i = 0; i < N; ++i) {
        updatedParticles.push_back(std::array<double, 3>({oldParticles->at(indexes[i])[0] + uniformNoise(generator),
                                                          oldParticles->at(indexes[i])[1] + uniformNoise(generator),
                                                          oldParticles->at(indexes[i])[2] + uniformNoise(generator) / 10
                                                         }));
        updatedWeights.push_back(weights->at(indexes[i]));
        weightSum += weights->at(indexes[i]);
    }

    // normalize weights
    for (int i = 0; i < N; ++i) {
        weights->at(indexes[i]) /= weightSum;
    }

    return std::tuple<std::vector<std::array<double, 3>>, std::vector<double>>(updatedParticles, updatedWeights);
}

ParticleEvaluationData ParticleFilter::update() {
    return this->update(this->robot->get_last_tick_movement_distance(), this->robot->get_last_tick_movement_angle());
}

ParticleEvaluationData ParticleFilter::update(double distance, double angle) {
    // check if in location finding phase
    int removeElementsCount = N * INITIAL_PHASE_REPLACE_SHARE;;
    if (!this->initialLocationFinished && this->iterationsCounter > 0) {
        // check if location is found
        if (this->locationCertaintyEstimation > CERTAINTY_ESTIMATION_THRESHOLD) {
            // check uncertainty metric if location seems to be already successful
            std::cout << "Found location" << std::endl;
            this->initialLocationFinished = true;
            this->N = TARGET_N;
            this->useRandomParticles = false;
        }

        // if location not yet found: select the worst rated particles and replace them with random
        if (!this->initialLocationFinished) {
            // bubble sort
            // could be replaced with a more efficient algorithm if its computationally too expensive
            bool finished = false;
            for (int i = 0; !finished && i < this->particles.size(); ++i) {
                finished = true;
                for (int j = i + 1; j < this->particles.size(); ++j) {
                    if (this->weights[i] > this->weights[j]) {
                        finished = false;
                        double tmpWeight = this->weights[i];
                        std::array<double, 3> tmpParticle = this->particles[i];
                        this->weights[i] = this->weights[j];
                        this->weights[j] = tmpWeight;
                        this->particles[i] = this->particles[j];
                        this->particles[j] = tmpParticle;
                    }
                }
            }

            // drop the first particles (lowest probability)
            this->weights.erase(this->weights.begin(), this->weights.begin() + removeElementsCount);
            this->particles.erase(this->particles.begin(), this->particles.begin() + removeElementsCount);

            // generate new particles to replace the removed ones
            std::vector<std::array<double, 3>> randomParticles = create_uniform_particles(this->particles_world->get_map_bounds().x,
                                                                                this->particles_world->get_map_bounds().y,
                                                                                removeElementsCount);
            this->particles.insert(this->particles.end(), randomParticles.begin(), randomParticles.end());
        }
    }
    this->iterationsCounter++;

    // Begin of actual particle filter steps

    // Particle Filter Step 1) particle prediction: apply estimated movement distance / angle (currently in this simulator there is no noise for those values)
    std::vector<std::array<double, 3>> updated_particles = particles_predict(&this->particles,
                                                                             distance,
                                                                             angle,
                                                                             4);

    // Particle Filter Step 2) update weights: update the weight of each particle based on how much the simulated sensor values correspond to the ones of our "real" robot
    // generate vector containing the current sensor values of the "real" robot
    std::vector<double> robotSensorValues = {};
    for (DistanceSensor *sensor : DistanceSensor::filter_for_distance_sensor(this->robot->get_sensors())) {
        robotSensorValues.push_back(sensor->get_simplified_sensor_value());
    }

    // this allows to read the sensor values of all simulated particles in particles_update()
    this->updateParticleSimulation(updated_particles, false);

    // this will do the actual weights update
    std::vector<double> updated_weights = this->particles_update(robotSensorValues, this->particles_world->getRobotsList(), 0.01);

    // Particle Filter Step 3) resample: create new particles based on the weights. They will be mostly around the most previous particles with the highest weight.
    std::tuple<std::vector<std::array<double, 3>>, std::vector<double>> resampledTuple = this->particles_resample(&updated_particles,
                                                                                &updated_weights,
                                                                                this->N,
                                                                                this->useRandomParticles,
                                                                                5,
                                                                                this->particles_world->get_map_bounds());

    // END of actual particle filter steps

    // calculate estimated position
    std::array<double, 3> estimatedParticle{};
    if (!this->initialLocationFinished) {
        // if in the initial phase: only the first x particles are relevant, the others are random
        int relevantParticlesCount = N - removeElementsCount;
        estimatedParticle = weightedAverageParticle(std::vector<std::array<double,3>>(updated_particles.begin(), updated_particles.begin()+relevantParticlesCount), std::vector<double>(updated_weights.begin(), updated_weights.begin()+relevantParticlesCount));
    } else {
        estimatedParticle = weightedAverageParticle(updated_particles, updated_weights);
    }

    // estimate whether the prediction is probably already accurate and if yes, add to history, if not: clear history
    // also visualize it on the map as a line
    if (!this->estimationHistory.empty()) {
        // TODO: replace this algorithm with a better one looking further in the past and calculate an uncertainty value
        double maxAcceptableMovementDistance = pow(this->robot->get_max_move_speed() / GAME_TPS * 2.5, 2); // theoretically maximum possible movement distance * X as square to prevent requirement of sqrt
        double distanceFromLastEstimation = pow(this->estimationHistory.back()[0] - estimatedParticle[0], 2) + pow(this->estimationHistory.back()[1] - estimatedParticle[1], 2);
        if (distanceFromLastEstimation >= maxAcceptableMovementDistance) {
            this->estimationHistory.clear();
            this->locationCertaintyEstimation = this->locationCertaintyEstimation > 0 ? this->locationCertaintyEstimation - 5 / GAME_TPS : 0;
            if (SHOW_WHATS_GOING_ON && !this->initialLocationFinished) this->mapLine->clearPoints();
            std::cout << "Estimated location jumped too far, this is an indication for a wrong location estimation" << std::endl;
        } else {
            this->locationCertaintyEstimation = this->locationCertaintyEstimation < 10 ? this->locationCertaintyEstimation + 1 / GAME_TPS : 10;
        }
    }
    if (SHOW_WHATS_GOING_ON) this->estimationHistory.push_back(estimatedParticle);
    this->mapLine->addPoint(cv::Point2i((int) estimatedParticle[0], (int) estimatedParticle[1]));

    // show current particles with their weights and the estimated robot location
    // nice for visualization but will impact performance!
    if (SHOW_WHATS_GOING_ON) {
        this->updateParticleSimulation(updated_particles, false, updated_weights);
        this->estimatedRobot->reposition(cv::Point2d(estimatedParticle[0], estimatedParticle[1]), estimatedParticle[2]);
        this->particles_world->show_map(true);
    }

    // save particles and weights
    this->particles = std::get<0>(resampledTuple);
    this->weights = std::get<1>(resampledTuple);

    return ParticleEvaluationData({cv::Point2d(estimatedParticle[0], estimatedParticle[1]), estimatedParticle[2], locationCertaintyEstimation >= CERTAINTY_ESTIMATION_THRESHOLD});
}


//    Copyright 2019 Jij Inc.

//    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 OPENJIJ_UTILITY_EIGEN_HPP__
#define OPENJIJ_UTILITY_EIGEN_HPP__

#include <graph/all.hpp>
#include <Eigen/Dense>
#include <Eigen/Sparse>

namespace openjij {
    namespace utility {

        /**
         * @brief get Eigen Matrix type from Graph Type
         *
         * @tparam GraphType
         * @tparam Options Eigen Options (RowMajor or ColMajor)
         */
        template<typename GraphType, int Options=Eigen::ColMajor>
            struct get_eigen_matrix_type{};

        /**
         * @brief get Eigen Matrix type from Graph Type
         *
         * @tparam GraphType
         * @tparam Options Eigen Options (RowMajor or ColMajor)
         */
        template<typename FloatType, int Options>
            struct get_eigen_matrix_type<graph::Dense<FloatType>, Options>{
                using type = Eigen::Matrix<FloatType, Eigen::Dynamic, Eigen::Dynamic, Options>;
            };

        /**
         * @brief get Eigen Matrix type from Graph Type
         *
         * @tparam GraphType
         * @tparam Options Eigen Options (RowMajor or ColMajor)
         */
        template<typename FloatType, int Options>
            struct get_eigen_matrix_type<graph::Sparse<FloatType>, Options>{
                using type = Eigen::SparseMatrix<FloatType, Options>;
            };

        /**
         * @brief generate Eigen Vector from std::vector
         *
         * @tparam FloatType
         * @tparam Options Eigen Options (RowMajor or ColMajor)
         * @param init_spin
         *
         * @return generated Eigen Vector (init_spin.size()+1 x 1)
         */
        template<typename FloatType, int Options=Eigen::ColMajor>
            inline static Eigen::Matrix<FloatType, Eigen::Dynamic, 1, Options>
            gen_vector_from_std_vector(const graph::Spins& init_spin){
                Eigen::Matrix<FloatType, Eigen::Dynamic, 1, Options> ret_vec(init_spin.size()+1);

                //initialize spin
                for(size_t i=0; i<init_spin.size(); i++){
                    ret_vec(i) = init_spin[i];
                }

                //for local field
                ret_vec[init_spin.size()] = 1;

                return ret_vec;
            }

        /**
         * @brief generate Eigen Matrix from TrotterSpins
         *
         * @tparam FloatType
         * @tparam Options Eigen Options (RowMajor or ColMajor)
         * @param trotter_spins
         *
         * @return generated Eigen Matrix (trotter_spins[0].size()+1 x trotter_spins.size())
         */
        template<typename FloatType, int Options=Eigen::ColMajor>
            inline static Eigen::Matrix<FloatType, Eigen::Dynamic, Eigen::Dynamic, Options>
            gen_matrix_from_trotter_spins(const std::vector<graph::Spins>& trotter_spins){
                Eigen::Matrix<FloatType, Eigen::Dynamic, Eigen::Dynamic, Options> ret_mat(trotter_spins[0].size()+1, trotter_spins.size());

                //initialize spin
                for(size_t j=0; j<trotter_spins.size(); j++){
                    for(size_t i=0; i<trotter_spins[j].size(); i++){
                        ret_mat(i,j) = trotter_spins[j][i];
                    }
                }

                //dummy spins
                for(size_t j=0; j<trotter_spins.size(); j++){
                    ret_mat(trotter_spins[0].size(),j) = 1;
                }

                return ret_mat;
            }

        /**
         * @brief generate Eigen Dense Matrix from Dense graph
         *
         * @tparam Options Eigen Options (RowMajor or ColMajor)
         * @tparam FloatType
         * @param graph
         *
         * @return generated Eigen Dense Matrix (graph.get_num_spins()+1 x graph.get_num_spins()+1)
         */
        template<int Options=Eigen::ColMajor, typename FloatType>
            inline static Eigen::Matrix<FloatType, Eigen::Dynamic, Eigen::Dynamic, Options>
            gen_matrix_from_graph(const graph::Dense<FloatType>& graph){
                //initialize interaction
                Eigen::Matrix<FloatType, Eigen::Dynamic, Eigen::Dynamic, Options> ret_mat(graph.get_num_spins()+1, graph.get_num_spins()+1);

                ret_mat.setZero();

                for(size_t i=0; i<graph.get_num_spins(); i++){
                    for(size_t j=i+1; j<graph.get_num_spins(); j++){
                        ret_mat(i,j) = graph.J(i,j);
                        ret_mat(j,i) = graph.J(i,j);
                    }
                }

                //for local field
                for(size_t i=0; i<graph.get_num_spins(); i++){
                    ret_mat(i,graph.get_num_spins()) = graph.h(i);
                    ret_mat(graph.get_num_spins(),i) = graph.h(i);
                }

                //for local field
                ret_mat(graph.get_num_spins(),graph.get_num_spins()) = 1;

                return ret_mat;
            }

        /**
         * @brief generate Eigen Sparse Matrix from Sparse graph
         *
         * @tparam Options Eigen Options (RowMajor or ColMajor)
         * @tparam FloatType
         * @param graph
         *
         * @return generated Eigen Sparse Matrix (graph.get_num_spins()+1 x graph.get_num_spins()+1)
         */
        template<int Options=Eigen::ColMajor, typename FloatType>
            inline static Eigen::SparseMatrix<FloatType, Options>
            gen_matrix_from_graph(const graph::Sparse<FloatType>& graph){
                //initialize interaction
                Eigen::SparseMatrix<FloatType, Options> ret_mat(graph.get_num_spins()+1, graph.get_num_spins()+1);

                ret_mat.setZero();

                //make triplet list
                using T = std::vector<Eigen::Triplet<FloatType>>;
                T t_list;

                for(size_t ind=0; ind<graph.get_num_spins(); ind++){
                    for(size_t adj_ind : graph.adj_nodes(ind)){
                        if(ind != adj_ind){
                            t_list.emplace_back(ind, adj_ind, graph.J(ind, adj_ind));
                        }
                        else{
                            t_list.emplace_back(ind, graph.get_num_spins(), graph.h(ind));
                            t_list.emplace_back(graph.get_num_spins(), ind, graph.h(ind));
                        }
                    }
                }

                t_list.emplace_back(graph.get_num_spins(), graph.get_num_spins(), 1);

                ret_mat.setFromTriplets(t_list.begin(), t_list.end());

                return ret_mat;
            }

    } // namespace utility
} // namespace openjij

#endif