#include "LHCOpticsApproximator.h"
#include <vector>
#include <iostream>
#include "TROOT.h"
#include "TMath.h"
#include <boost/shared_ptr.hpp>
#include "CurrentMemoryUsage.h"

ClassImp(LHCOpticsApproximator);
ClassImp(LHCApertureApproximator);

void LHCOpticsApproximator::Init()
{
  out_polynomials.clear();
  apertures_.clear();
  out_polynomials.push_back(&x_parametrisation);
  out_polynomials.push_back(&theta_x_parametrisation);
  out_polynomials.push_back(&y_parametrisation);
  out_polynomials.push_back(&theta_y_parametrisation);

  coord_names.clear();
  coord_names.push_back("x");
  coord_names.push_back("theta_x");
  coord_names.push_back("y");
  coord_names.push_back("theta_y");
  coord_names.push_back("ksi");

  s_begin_ = 0.0;
  s_end_ = 0.0;
  trained_ = false;
}

LHCOpticsApproximator::LHCOpticsApproximator(std::string name, std::string title, TMultiDimFet::EMDFPolyType polynom_type, std::string beam_direction, double nominal_beam_energy)
: x_parametrisation(5, polynom_type, "k"),
  theta_x_parametrisation(5, polynom_type, "k"),
  y_parametrisation(5, polynom_type, "k"),
  theta_y_parametrisation(5, polynom_type, "k")
{
  this->SetName(name.c_str());
  this->SetTitle(title.c_str());
  Init();

  if(beam_direction == "lhcb1")
    beam = lhcb1;
  else if(beam_direction == "lhcb2")
    beam = lhcb2;
  else
    beam = lhcb1;

  nominal_beam_energy_ = nominal_beam_energy;
  nominal_beam_momentum_ = TMath::Sqrt(nominal_beam_energy_*nominal_beam_energy_ - 0.938272029*0.938272029);
}

LHCOpticsApproximator::LHCOpticsApproximator(std::string name, std::string title, TMultiDimFet::EMDFPolyType polynom_type, std::string beam_direction, double nominal_beam_energy,
  TMultiDimFet *given_x_parametrisation, TMultiDimFet *given_theta_x_parametrisation, TMultiDimFet *given_y_parametrisation, TMultiDimFet *given_theta_y_parametrisation)
{
  this->x_parametrisation = *given_x_parametrisation;
  this->theta_x_parametrisation = *given_theta_x_parametrisation;
  this->y_parametrisation = *given_y_parametrisation;
  this->theta_y_parametrisation = *given_theta_y_parametrisation;
  this->SetName(name.c_str());
  this->SetTitle(title.c_str());
  Init();

  if(beam_direction == "lhcb1")
    beam = lhcb1;
  else if(beam_direction == "lhcb2")
    beam = lhcb2;
  else
    beam = lhcb1;

  nominal_beam_energy_ = nominal_beam_energy;
  nominal_beam_momentum_ = TMath::Sqrt(nominal_beam_energy_*nominal_beam_energy_ - 0.938272029*0.938272029);
  trained_ = true;
}

LHCOpticsApproximator::LHCOpticsApproximator()
{
  Init();
  beam = lhcb1;
  nominal_beam_energy_ = 7000;
  nominal_beam_momentum_ = TMath::Sqrt(nominal_beam_energy_*nominal_beam_energy_ - 0.938272029*0.938272029);
}

bool LHCOpticsApproximator::Transport(double *in, double *out, bool check_apertures)
{
  if(in==NULL || out==NULL || !trained_)
    return false;

  bool res = CheckInputRange(in);

  out[0] = x_parametrisation.Eval(in);
  out[1] = theta_x_parametrisation.Eval(in);
  out[2] = y_parametrisation.Eval(in);
  out[3] = theta_y_parametrisation.Eval(in);
  out[4] = in[4];

  if(check_apertures)
  {
    for(int i=0; i<apertures_.size(); i++)
    {
      res = res && apertures_[i].CheckAperture(in);
    }
  }
  return res;
}

bool LHCOpticsApproximator::Transport_m_GeV(double in_pos[3], double in_momentum[3], double out_pos[3], double out_momentum[3],
        bool check_apertures, double z2_z1_dist)
{
  double in[5];
  double out[5];
  double part_mom = 0.0;
  for(int i=0; i<3; i++)
    part_mom += in_momentum[i]*in_momentum[i];

  part_mom = TMath::Sqrt(part_mom);

  in[0] = in_pos[0];
  in[1] = in_momentum[0]/nominal_beam_momentum_;
  in[2] = in_pos[1];
  in[3] = in_momentum[1]/nominal_beam_momentum_;
  in[4] = (part_mom-nominal_beam_momentum_)/nominal_beam_momentum_;

  bool res = Transport(in, out, check_apertures);

  out_pos[0] = out[0];
  out_pos[1] = out[2];
  out_pos[2] = in_pos[2] + z2_z1_dist;

  out_momentum[0] = out[1]*nominal_beam_momentum_;
  out_momentum[1] = out[3]*nominal_beam_momentum_;
  double part_out_total_mom = (out[4]+1)*nominal_beam_momentum_;
  out_momentum[2] = TMath::Sqrt(part_out_total_mom*part_out_total_mom - out_momentum[0]*out_momentum[0] - out_momentum[1]*out_momentum[1]);
  out_momentum[2] = TMath::Sign(out_momentum[2], in_momentum[2]);

  return res;
}

bool LHCOpticsApproximator::Transport(const MadKinematicDescriptor *in, MadKinematicDescriptor *out, bool check_apertures)
{
  if(in==NULL || out==NULL || !trained_)
    return false;

  Double_t input[5];
  Double_t output[5];
  input[0] = in->x;
  input[1] = in->theta_x;
  input[2] = in->y;
  input[3] = in->theta_y;
  input[4] = in->ksi;

  bool res = Transport(input, output, check_apertures);

  out->x = output[0];
  out->theta_x = output[1];
  out->y = output[2];
  out->theta_y = output[3];
  out->ksi = output[4];

  return res;
}

LHCOpticsApproximator::LHCOpticsApproximator(const LHCOpticsApproximator &org) : TNamed(org), x_parametrisation(org.x_parametrisation), theta_x_parametrisation(org.theta_x_parametrisation), y_parametrisation(org.y_parametrisation), theta_y_parametrisation(org.theta_y_parametrisation)
{
//  std::cout<<"LHCOpticsApproximator::LHCOpticsApproximator(const LHCOpticsApproximator &org) entered"<<std::endl;
    void Init();
    s_begin_ = org.s_begin_;
    s_end_ = org.s_end_;
    trained_ = org.trained_;
    apertures_ = org.apertures_;
    beam = org.beam;
    nominal_beam_energy_ = org.nominal_beam_energy_;
    nominal_beam_momentum_ = org.nominal_beam_momentum_;
//  std::cout<<"LHCOpticsApproximator::LHCOpticsApproximator(const LHCOpticsApproximator &org) left"<<std::endl;
}

double LHCOpticsApproximator::GetBegin(){
  return s_begin_;
}

double LHCOpticsApproximator::GetEnd(){
  return s_end_;
}

LHCOpticsApproximator & LHCOpticsApproximator::operator=(const LHCOpticsApproximator &org)
{
  if(this!=&org)
  {
    void Init();
    TNamed::operator=(org);
    s_begin_ = org.s_begin_;
    s_end_ = org.s_end_;
    trained_ = org.trained_;

    x_parametrisation = org.x_parametrisation;
    theta_x_parametrisation = org.theta_x_parametrisation;
    y_parametrisation = org.y_parametrisation;
    theta_y_parametrisation = org.theta_y_parametrisation;
    apertures_ = org.apertures_;
    beam = org.beam;
    nominal_beam_energy_ = org.nominal_beam_energy_;
    nominal_beam_momentum_ = org.nominal_beam_momentum_;
  }
}

void LHCOpticsApproximator::Train(TTree *inp_tree, std::string data_prefix, polynomials_selection mode, int max_degree_x, int max_degree_tx, int max_degree_y, int max_degree_ty, bool common_terms, double *prec)
{
  if(inp_tree==NULL)
    return;

  PrintCurrentMemoryUsage("Train, begin");

  InitializeApproximators(mode, max_degree_x, max_degree_tx, max_degree_y, max_degree_ty, common_terms);
  std::cout<<this->GetName()<<" is being trained..."<<std::endl;

  //in-variables
  //x_in, theta_x_in, y_in, theta_y_in, ksi_in, s_in
  double in_var[6];

  //out-variables
  //x_out, theta_x_out, y_out, theta_y_out, ksi_out, s_out, valid_out;
  double out_var[7];

  //in- out-lables
  std::string x_in_lab = "x_in";
  std::string theta_x_in_lab = "theta_x_in";
  std::string y_in_lab = "y_in";
  std::string theta_y_in_lab = "theta_y_in";
  std::string ksi_in_lab = "ksi_in";
  std::string s_in_lab = "s_in";

  std::string x_out_lab = data_prefix + "_x_out";
  std::string theta_x_out_lab = data_prefix + "_theta_x_out";
  std::string y_out_lab = data_prefix + "_y_out";
  std::string theta_y_out_lab = data_prefix + "_theta_y_out";
  std::string ksi_out_lab = data_prefix + "_ksi_out";
  std::string s_out_lab = data_prefix + "_s_out";
  std::string valid_out_lab = data_prefix + "_valid_out";

  //disable not needed branches to speed up the readin
  inp_tree->SetBranchStatus("*",0);  //disable all branches
  inp_tree->SetBranchStatus(x_in_lab.c_str(),1);
  inp_tree->SetBranchStatus(theta_x_in_lab.c_str(),1);
  inp_tree->SetBranchStatus(y_in_lab.c_str(),1);
  inp_tree->SetBranchStatus(theta_y_in_lab.c_str(),1);
  inp_tree->SetBranchStatus(ksi_in_lab.c_str(),1);
  inp_tree->SetBranchStatus(x_out_lab.c_str(),1);
  inp_tree->SetBranchStatus(theta_x_out_lab.c_str(),1);
  inp_tree->SetBranchStatus(y_out_lab.c_str(),1);
  inp_tree->SetBranchStatus(theta_y_out_lab.c_str(),1);
  inp_tree->SetBranchStatus(s_out_lab.c_str(),1);
  inp_tree->SetBranchStatus(ksi_out_lab.c_str(),1);
  inp_tree->SetBranchStatus(valid_out_lab.c_str(),1);

  //set input data adresses
  inp_tree->SetBranchAddress(x_in_lab.c_str(), &(in_var[0]) );
  inp_tree->SetBranchAddress(theta_x_in_lab.c_str(), &(in_var[1]) );
  inp_tree->SetBranchAddress(y_in_lab.c_str(), &(in_var[2]) );
  inp_tree->SetBranchAddress(theta_y_in_lab.c_str(), &(in_var[3]) );
  inp_tree->SetBranchAddress(ksi_in_lab.c_str(), &(in_var[4]) );
  inp_tree->SetBranchAddress(s_in_lab.c_str(), &(in_var[5]) );

  //set output data adresses
  inp_tree->SetBranchAddress(x_out_lab.c_str(), &(out_var[0]) );
  inp_tree->SetBranchAddress(theta_x_out_lab.c_str(), &(out_var[1]) );
  inp_tree->SetBranchAddress(y_out_lab.c_str(), &(out_var[2]) );
  inp_tree->SetBranchAddress(theta_y_out_lab.c_str(), &(out_var[3]) );
  inp_tree->SetBranchAddress(ksi_out_lab.c_str(), &(out_var[4]) );
  inp_tree->SetBranchAddress(s_out_lab.c_str(), &(out_var[5]) );
  inp_tree->SetBranchAddress(valid_out_lab.c_str(), &(out_var[6]) );

  Long64_t entries = inp_tree->GetEntries();
  if(entries>0)
  {
    inp_tree->SetBranchStatus(s_in_lab.c_str(),1);
    inp_tree->SetBranchStatus(s_out_lab.c_str(),1);
    inp_tree->GetEntry(0);
    s_begin_ = in_var[5];
    s_end_ = out_var[5];
    inp_tree->SetBranchStatus(s_in_lab.c_str(),0);
    inp_tree->SetBranchStatus(s_out_lab.c_str(),0);
  }

  PrintCurrentMemoryUsage("Train, before setting data");
  //set input and output variables for fitting
  for(Long64_t i=0; i<entries; ++i)
  {
    inp_tree->GetEntry(i);
    if(out_var[6] != 0)  //if out data valid
    {
      x_parametrisation.AddRow(in_var, out_var[0], 0);
      theta_x_parametrisation.AddRow(in_var, out_var[1], 0);
      y_parametrisation.AddRow(in_var, out_var[2], 0);
      theta_y_parametrisation.AddRow(in_var, out_var[3], 0);
    }
  }

  std::cout<<"Optical functions parametrizations from "<<s_begin_<<" to "<<s_end_<<std::endl;
  PrintInputRange();
  PrintCurrentMemoryUsage("Train, before fitting");
  for(int i=0; i<4; i++)
  {
    double best_precision=0.0;
    if(prec)
      best_precision = prec[i];
    out_polynomials[i]->FindParameterization(best_precision);
    std::cout<<"Out variable "<<coord_names[i]<<" polynomial"<<std::endl;
    out_polynomials[i]->PrintPolynomialsSpecial("M");
    std::cout<<std::endl;
  }

  trained_ = true;
  PrintCurrentMemoryUsage("Train, function end");
}

void LHCOpticsApproximator::InitializeApproximators(polynomials_selection mode, int max_degree_x, int max_degree_tx, int max_degree_y, int max_degree_ty, bool common_terms)
{
  SetDefaultAproximatorSettings(x_parametrisation, X, max_degree_x);
  SetDefaultAproximatorSettings(theta_x_parametrisation, THETA_X, max_degree_tx);
  SetDefaultAproximatorSettings(y_parametrisation, Y, max_degree_y);
  SetDefaultAproximatorSettings(theta_y_parametrisation, THETA_Y, max_degree_ty);

  if(mode == PREDEFINED)
  {
    SetTermsManually(x_parametrisation, X, max_degree_x, common_terms);
    SetTermsManually(theta_x_parametrisation, THETA_X, max_degree_tx, common_terms);
    SetTermsManually(y_parametrisation, Y, max_degree_y, common_terms);
    SetTermsManually(theta_y_parametrisation, THETA_Y, max_degree_ty, common_terms);
  }
}

void LHCOpticsApproximator::SetDefaultAproximatorSettings(TMultiDimFet &approximator, variable_type var_type, int max_degree)
{
  if(max_degree<1 || max_degree>20)
    max_degree = 10;

  if(var_type == X || var_type == THETA_X)
  {
    Int_t mPowers[] = { 2, 4, 2, 4, max_degree };
    approximator.SetMaxPowers (mPowers);
    approximator.SetMaxFunctions (3000);
    approximator.SetMaxStudy (3000);
    approximator.SetMaxTerms (3000);
    approximator.SetPowerLimit (1.6);
//    approximator.SetMinAngle (2e-4);
//    approximator.SetMaxAngle (10);
    approximator.SetMinRelativeError (1e-13);
  }

  if(var_type == Y || var_type == THETA_Y)
  {
    Int_t mPowers[] = { 2, 4, 2, 4, max_degree };
    approximator.SetMaxPowers (mPowers);
    approximator.SetMaxFunctions (3000);
    approximator.SetMaxStudy (3000);
    approximator.SetMaxTerms (3000);
    approximator.SetPowerLimit (1.6);
//    approximator.SetMinAngle (2e-4);
//    approximator.SetMaxAngle (10);
    approximator.SetMinRelativeError (1e-13);
  }
}

void LHCOpticsApproximator::SetTermsManually(TMultiDimFet &approximator, variable_type variable, int max_degree, bool common_terms)
{
  if(max_degree<1 || max_degree>20)
    max_degree = 10;

  //put terms of shape:
  //1,0,0,0,t    0,1,0,0,t    0,2,0,0,t    0,3,0,0,t    0,0,0,0,t
  //t: 0,1,...,max_degree
//  int total_terms = 5*(max_degree+1);
//  int table_size = total_terms*5;
//  Int_t powers[table_size];

  std::vector<Int_t> term_literals;
  term_literals.reserve(5000);

  if(variable == X || variable == THETA_X)
  {
    //1,0,0,0,t
    for(int i=0; i<=max_degree; ++i)
    {
      term_literals.push_back(1);
      term_literals.push_back(0);
      term_literals.push_back(0);
      term_literals.push_back(0);
      term_literals.push_back(i);
    }
    //0,1,0,0,t
    for(int i=0; i<=max_degree; ++i)
    {
      term_literals.push_back(0);
      term_literals.push_back(1);
      term_literals.push_back(0);
      term_literals.push_back(0);
      term_literals.push_back(i);
    }
    //0,2,0,0,t
    for(int i=0; i<=max_degree; ++i)
    {
      term_literals.push_back(0);
      term_literals.push_back(2);
      term_literals.push_back(0);
      term_literals.push_back(0);
      term_literals.push_back(i);
    }
    //0,3,0,0,t
    for(int i=0; i<=max_degree; ++i)
    {
      term_literals.push_back(0);
      term_literals.push_back(3);
      term_literals.push_back(0);
      term_literals.push_back(0);
      term_literals.push_back(i);
    }
    //0,0,0,0,t
    for(int i=0; i<=max_degree; ++i)
    {
      term_literals.push_back(0);
      term_literals.push_back(0);
      term_literals.push_back(0);
      term_literals.push_back(0);
      term_literals.push_back(i);
    }
  }

  if(variable == Y || variable == THETA_Y)
  {
    //0,0,1,0,t
    for(int i=0; i<=max_degree; ++i)
    {
      term_literals.push_back(0);
      term_literals.push_back(0);
      term_literals.push_back(1);
      term_literals.push_back(0);
      term_literals.push_back(i);
    }
    //0,0,0,1,t
    for(int i=0; i<=max_degree; ++i)
    {
      term_literals.push_back(0);
      term_literals.push_back(0);
      term_literals.push_back(0);
      term_literals.push_back(1);
      term_literals.push_back(i);
    }
    //0,0,0,2,t
    for(int i=0; i<=max_degree; ++i)
    {
      term_literals.push_back(0);
      term_literals.push_back(0);
      term_literals.push_back(0);
      term_literals.push_back(2);
      term_literals.push_back(i);
    }
    //0,0,0,3,t
    for(int i=0; i<=max_degree; ++i)
    {
      term_literals.push_back(0);
      term_literals.push_back(0);
      term_literals.push_back(0);
      term_literals.push_back(3);
      term_literals.push_back(i);
    }
    //0,0,0,0,t
    for(int i=0; i<=max_degree; ++i)
    {
      term_literals.push_back(0);
      term_literals.push_back(0);
      term_literals.push_back(0);
      term_literals.push_back(0);
      term_literals.push_back(i);
    }
  }

  //push common terms
  if(common_terms)
  {
    term_literals.push_back(1), term_literals.push_back(1), term_literals.push_back(0), term_literals.push_back(0), term_literals.push_back(0);
    term_literals.push_back(1), term_literals.push_back(0), term_literals.push_back(1), term_literals.push_back(0), term_literals.push_back(0);
    term_literals.push_back(1), term_literals.push_back(0), term_literals.push_back(0), term_literals.push_back(1), term_literals.push_back(0);
    term_literals.push_back(0), term_literals.push_back(1), term_literals.push_back(1), term_literals.push_back(0), term_literals.push_back(0);
    term_literals.push_back(0), term_literals.push_back(1), term_literals.push_back(0), term_literals.push_back(1), term_literals.push_back(0);
    term_literals.push_back(0), term_literals.push_back(0), term_literals.push_back(1), term_literals.push_back(1), term_literals.push_back(0);

    term_literals.push_back(1), term_literals.push_back(1), term_literals.push_back(0), term_literals.push_back(0), term_literals.push_back(1);
    term_literals.push_back(1), term_literals.push_back(0), term_literals.push_back(1), term_literals.push_back(0), term_literals.push_back(1);
    term_literals.push_back(1), term_literals.push_back(0), term_literals.push_back(0), term_literals.push_back(1), term_literals.push_back(1);
    term_literals.push_back(0), term_literals.push_back(1), term_literals.push_back(1), term_literals.push_back(0), term_literals.push_back(1);
    term_literals.push_back(0), term_literals.push_back(1), term_literals.push_back(0), term_literals.push_back(1), term_literals.push_back(1);
    term_literals.push_back(0), term_literals.push_back(0), term_literals.push_back(1), term_literals.push_back(1), term_literals.push_back(1);

    term_literals.push_back(1), term_literals.push_back(2), term_literals.push_back(0), term_literals.push_back(0), term_literals.push_back(0);
    term_literals.push_back(1), term_literals.push_back(0), term_literals.push_back(2), term_literals.push_back(0), term_literals.push_back(0);
    term_literals.push_back(1), term_literals.push_back(0), term_literals.push_back(0), term_literals.push_back(2), term_literals.push_back(0);
    term_literals.push_back(2), term_literals.push_back(1), term_literals.push_back(0), term_literals.push_back(0), term_literals.push_back(0);
    term_literals.push_back(0), term_literals.push_back(1), term_literals.push_back(2), term_literals.push_back(0), term_literals.push_back(0);
    term_literals.push_back(0), term_literals.push_back(1), term_literals.push_back(0), term_literals.push_back(2), term_literals.push_back(0);
    term_literals.push_back(2), term_literals.push_back(0), term_literals.push_back(1), term_literals.push_back(0), term_literals.push_back(0);
    term_literals.push_back(0), term_literals.push_back(2), term_literals.push_back(1), term_literals.push_back(0), term_literals.push_back(0);
    term_literals.push_back(0), term_literals.push_back(0), term_literals.push_back(1), term_literals.push_back(2), term_literals.push_back(0);
    term_literals.push_back(2), term_literals.push_back(0), term_literals.push_back(0), term_literals.push_back(1), term_literals.push_back(0);
    term_literals.push_back(0), term_literals.push_back(2), term_literals.push_back(0), term_literals.push_back(1), term_literals.push_back(0);
    term_literals.push_back(0), term_literals.push_back(0), term_literals.push_back(2), term_literals.push_back(1), term_literals.push_back(0);

    term_literals.push_back(1), term_literals.push_back(2), term_literals.push_back(0), term_literals.push_back(0), term_literals.push_back(1);
    term_literals.push_back(1), term_literals.push_back(0), term_literals.push_back(2), term_literals.push_back(0), term_literals.push_back(1);
    term_literals.push_back(1), term_literals.push_back(0), term_literals.push_back(0), term_literals.push_back(2), term_literals.push_back(1);
    term_literals.push_back(2), term_literals.push_back(1), term_literals.push_back(0), term_literals.push_back(0), term_literals.push_back(1);
    term_literals.push_back(0), term_literals.push_back(1), term_literals.push_back(2), term_literals.push_back(0), term_literals.push_back(1);
    term_literals.push_back(0), term_literals.push_back(1), term_literals.push_back(0), term_literals.push_back(2), term_literals.push_back(1);
    term_literals.push_back(2), term_literals.push_back(0), term_literals.push_back(1), term_literals.push_back(0), term_literals.push_back(1);
    term_literals.push_back(0), term_literals.push_back(2), term_literals.push_back(1), term_literals.push_back(0), term_literals.push_back(1);
    term_literals.push_back(0), term_literals.push_back(0), term_literals.push_back(1), term_literals.push_back(2), term_literals.push_back(1);
    term_literals.push_back(2), term_literals.push_back(0), term_literals.push_back(0), term_literals.push_back(1), term_literals.push_back(1);
    term_literals.push_back(0), term_literals.push_back(2), term_literals.push_back(0), term_literals.push_back(1), term_literals.push_back(1);
    term_literals.push_back(0), term_literals.push_back(0), term_literals.push_back(2), term_literals.push_back(1), term_literals.push_back(1);
  }

  Int_t powers[term_literals.size()];
  for(int i=0; i<term_literals.size(); ++i)
  {
    powers[i] = term_literals[i];
  }
  approximator.SetPowers(powers, term_literals.size()/5);
}

void LHCOpticsApproximator::Test(TTree *inp_tree, TFile *f_out, std::string data_prefix, std::string base_out_dir)
{
  PrintCurrentMemoryUsage("Test, begin");
  if(inp_tree==NULL || f_out==NULL)
    return;

  std::cout<<this->GetName()<<" is being tested..."<<std::endl;
  //in-variables
  //x_in, theta_x_in, y_in, theta_y_in, ksi_in, s_in
  double in_var[6];

  //out-variables
  //x_out, theta_x_out, y_out, theta_y_out, ksi_out, s_out, valid_out;
  double out_var[7];

  //in- out-lables
  std::string x_in_lab = "x_in";
  std::string theta_x_in_lab = "theta_x_in";
  std::string y_in_lab = "y_in";
  std::string theta_y_in_lab = "theta_y_in";
  std::string ksi_in_lab = "ksi_in";
  std::string s_in_lab = "s_in";

  std::string x_out_lab = data_prefix + "_x_out";
  std::string theta_x_out_lab = data_prefix + "_theta_x_out";
  std::string y_out_lab = data_prefix + "_y_out";
  std::string theta_y_out_lab = data_prefix + "_theta_y_out";
  std::string ksi_out_lab = data_prefix + "_ksi_out";
  std::string s_out_lab = data_prefix + "_s_out";
  std::string valid_out_lab = data_prefix + "_valid_out";

  //disable not needed branches to speed up the readin
  inp_tree->SetBranchStatus("*",0);  //disable all branches
  inp_tree->SetBranchStatus(x_in_lab.c_str(),1);
  inp_tree->SetBranchStatus(theta_x_in_lab.c_str(),1);
  inp_tree->SetBranchStatus(y_in_lab.c_str(),1);
  inp_tree->SetBranchStatus(theta_y_in_lab.c_str(),1);
  inp_tree->SetBranchStatus(ksi_in_lab.c_str(),1);
  inp_tree->SetBranchStatus(x_out_lab.c_str(),1);
  inp_tree->SetBranchStatus(theta_x_out_lab.c_str(),1);
  inp_tree->SetBranchStatus(y_out_lab.c_str(),1);
  inp_tree->SetBranchStatus(theta_y_out_lab.c_str(),1);
  inp_tree->SetBranchStatus(ksi_out_lab.c_str(),1);
  inp_tree->SetBranchStatus(s_out_lab.c_str(),1);
  inp_tree->SetBranchStatus(valid_out_lab.c_str(),1);

  //set input data adresses
  inp_tree->SetBranchAddress(x_in_lab.c_str(), &(in_var[0]) );
  inp_tree->SetBranchAddress(theta_x_in_lab.c_str(), &(in_var[1]) );
  inp_tree->SetBranchAddress(y_in_lab.c_str(), &(in_var[2]) );
  inp_tree->SetBranchAddress(theta_y_in_lab.c_str(), &(in_var[3]) );
  inp_tree->SetBranchAddress(ksi_in_lab.c_str(), &(in_var[4]) );
  inp_tree->SetBranchAddress(s_in_lab.c_str(), &(in_var[5]) );

  //set output data adresses
  inp_tree->SetBranchAddress(x_out_lab.c_str(), &(out_var[0]) );
  inp_tree->SetBranchAddress(theta_x_out_lab.c_str(), &(out_var[1]) );
  inp_tree->SetBranchAddress(y_out_lab.c_str(), &(out_var[2]) );
  inp_tree->SetBranchAddress(theta_y_out_lab.c_str(), &(out_var[3]) );
  inp_tree->SetBranchAddress(ksi_out_lab.c_str(), &(out_var[4]) );
  inp_tree->SetBranchAddress(s_out_lab.c_str(), &(out_var[5]) );
  inp_tree->SetBranchAddress(valid_out_lab.c_str(), &(out_var[6]) );

  PrintCurrentMemoryUsage("Test, before allocating the histograms");
  //test histogramms
  TH1D *err_hists[4];
  TH2D *err_inp_cor_hists[4][5];
  TH2D *err_out_cor_hists[4][5];

  AllocateErrorHists(err_hists);
  AllocateErrorInputCorHists(err_inp_cor_hists);
  AllocateErrorOutputCorHists(err_out_cor_hists);

  Long64_t entries = inp_tree->GetEntries();
  //set input and output variables for fitting
  for(Long64_t i=0; i<entries; ++i)
  {
    double errors[4];
    inp_tree->GetEntry(i);

    errors[0] = out_var[0] - x_parametrisation.Eval(in_var);
    errors[1] = out_var[1] - theta_x_parametrisation.Eval(in_var);
    errors[2] = out_var[2] - y_parametrisation.Eval(in_var);
    errors[3] = out_var[3] - theta_y_parametrisation.Eval(in_var);

    FillErrorHistograms(errors, err_hists);
    FillErrorDataCorHistograms(errors, in_var, err_inp_cor_hists);
    FillErrorDataCorHistograms(errors, out_var, err_out_cor_hists);
  }

  PrintCurrentMemoryUsage("Test, before writing the histograms");
  WriteHistograms(err_hists, err_inp_cor_hists, err_out_cor_hists, f_out, base_out_dir);
  std::cout<<"Histograms have been written."<<std::endl;

  DeleteErrorHists(err_hists);
  DeleteErrorCorHistograms(err_inp_cor_hists);
  DeleteErrorCorHistograms(err_out_cor_hists);
  PrintCurrentMemoryUsage("Test, end of function");
}

void LHCOpticsApproximator::AllocateErrorHists(TH1D *err_hists[4])
{
  std::vector<std::string> error_labels;
  error_labels.push_back("x error");
  error_labels.push_back("theta_x error");
  error_labels.push_back("y error");
  error_labels.push_back("theta_y error");

  for(int i=0; i<4; ++i)
  {
    err_hists[i] = new TH1D(error_labels[i].c_str(), error_labels[i].c_str(), 100, -0.0000000001, 0.0000000001);
    err_hists[i]->SetXTitle(error_labels[i].c_str());
    err_hists[i]->SetYTitle("counts");
    err_hists[i]->SetDirectory(0);
    err_hists[i]->SetCanExtend(TH1::kAllAxes);
  }
}

void LHCOpticsApproximator::TestAperture(TTree *inp_tree, TTree *out_tree)  //x, theta_x, y, theta_y, ksi, mad_accepted, parametriz_accepted
{
  if(inp_tree==NULL || out_tree==NULL)
    return;

  Long64_t entries = inp_tree->GetEntries();
  double entry[7];
  double parametrization_out[5];

  inp_tree->SetBranchAddress("x", &(entry[0]) );
  inp_tree->SetBranchAddress("theta_x", &(entry[1]) );
  inp_tree->SetBranchAddress("y", &(entry[2]) );
  inp_tree->SetBranchAddress("theta_y", &(entry[3]) );
  inp_tree->SetBranchAddress("ksi", &(entry[4]) );
  inp_tree->SetBranchAddress("mad_accept", &(entry[5]) );
  inp_tree->SetBranchAddress("par_accept", &(entry[6]) );

  out_tree->SetBranchAddress("x", &(entry[0]) );
  out_tree->SetBranchAddress("theta_x", &(entry[1]) );
  out_tree->SetBranchAddress("y", &(entry[2]) );
  out_tree->SetBranchAddress("theta_y", &(entry[3]) );
  out_tree->SetBranchAddress("ksi", &(entry[4]) );
  out_tree->SetBranchAddress("mad_accept", &(entry[5]) );
  out_tree->SetBranchAddress("par_accept", &(entry[6]) );

//  int ind=0;
  for(Long64_t i=0; i<entries; i++)
  {
    inp_tree->GetEntry(i);
//    for(int j=0; j<7; j++)
//      std::cout<<entry[j]<<" ";
//    std::cout<<std::endl;

    bool res = Transport(entry, parametrization_out, true);
//    for(int j=0; j<5; j++)
//      std::cout<<parametrization_out[j]<<" ";
//    std::cout<<"TestAperture "<<res<<std::endl;

    if( res )
      entry[6] = 1.0;
    else
      entry[6] = 0.0;

    out_tree->Fill();
//    if(ind++>300)
//      exit(0);
  }
}

void LHCOpticsApproximator::AllocateErrorInputCorHists(TH2D *err_inp_cor_hists[4][5])
{
  std::vector<std::string> error_labels;
  std::vector<std::string> data_labels;

  error_labels.push_back("x error");
  error_labels.push_back("theta_x error");
  error_labels.push_back("y error");
  error_labels.push_back("theta_y error");

  data_labels.push_back("x input");
  data_labels.push_back("theta_x input");
  data_labels.push_back("y input");
  data_labels.push_back("theta_y input");
  data_labels.push_back("ksi input");

  for(int eri=0; eri<4; ++eri)
  {
    for(int dati=0; dati<5; ++dati)
    {
      std::string name = error_labels[eri] + " vs. " + data_labels[dati];
      std::string title = name;
      err_inp_cor_hists[eri][dati] = new TH2D(name.c_str(), title.c_str(), 100,-0.0000000001,0.0000000001,100,-0.0000000001,0.0000000001);
      err_inp_cor_hists[eri][dati]->SetXTitle(error_labels[eri].c_str());
      err_inp_cor_hists[eri][dati]->SetYTitle(data_labels[dati].c_str());
      err_inp_cor_hists[eri][dati]->SetDirectory(0);
      err_inp_cor_hists[eri][dati]->SetCanExtend(TH2::kAllAxes);
    }
  }
}

void LHCOpticsApproximator::AllocateErrorOutputCorHists(TH2D *err_out_cor_hists[4][5])
{
  std::vector<std::string> error_labels;
  std::vector<std::string> data_labels;

  error_labels.push_back("x error");
  error_labels.push_back("theta_x error");
  error_labels.push_back("y error");
  error_labels.push_back("theta_y error");

  data_labels.push_back("x output");
  data_labels.push_back("theta_x output");
  data_labels.push_back("y output");
  data_labels.push_back("theta_y output");
  data_labels.push_back("ksi output");

  for(int eri=0; eri<4; ++eri)
  {
    for(int dati=0; dati<5; ++dati)
    {
      std::string name = error_labels[eri] + " vs. " + data_labels[dati];
      std::string title = name;
      err_out_cor_hists[eri][dati] = new TH2D(name.c_str(), title.c_str(), 100,-0.0000000001,0.0000000001,100,-0.0000000001,0.0000000001);
      err_out_cor_hists[eri][dati]->SetXTitle(error_labels[eri].c_str());
      err_out_cor_hists[eri][dati]->SetYTitle(data_labels[dati].c_str());
      err_out_cor_hists[eri][dati]->SetDirectory(0);
      err_out_cor_hists[eri][dati]->SetCanExtend(TH2::kAllAxes);
    }
  }
}

void LHCOpticsApproximator::FillErrorHistograms(double errors[4], TH1D *err_hists[4])
{
  for(int i=0; i<4; ++i)
  {
    err_hists[i]->Fill(errors[i]);
  }
}

void LHCOpticsApproximator::FillErrorDataCorHistograms(double errors[4], double var[5], TH2D *err_cor_hists[4][5])
{
  for(int eri=0; eri<4; ++eri)
  {
    for(int dati=0; dati<5; ++dati)
    {
      err_cor_hists[eri][dati]->Fill(errors[eri], var[dati]);
    }
  }
}

void LHCOpticsApproximator::DeleteErrorHists(TH1D *err_hists[4])
{
  for(int i=0; i<4; ++i)
  {
    delete err_hists[i];
  }
}

void LHCOpticsApproximator::DeleteErrorCorHistograms(TH2D *err_cor_hists[4][5])
{
  for(int eri=0; eri<4; ++eri)
  {
    for(int dati=0; dati<5; ++dati)
    {
      delete err_cor_hists[eri][dati];
    }
  }
}

void LHCOpticsApproximator::WriteHistograms(TH1D *err_hists[4], TH2D *err_inp_cor_hists[4][5], TH2D *err_out_cor_hists[4][5], TFile *f_out, std::string base_out_dir)
{
  if(f_out==NULL)
    return;

  f_out->cd();
  if(!gDirectory->cd(base_out_dir.c_str()))
    gDirectory->mkdir(base_out_dir.c_str());

  gDirectory->cd(base_out_dir.c_str());
  gDirectory->mkdir(this->GetName());
  gDirectory->cd(this->GetName());
  gDirectory->mkdir("x");
  gDirectory->mkdir("theta_x");
  gDirectory->mkdir("y");
  gDirectory->mkdir("theta_y");

  gDirectory->cd("x");
  err_hists[0]->Write("", TObject::kWriteDelete);
  for(int i=0; i<5; i++)
  {
    err_inp_cor_hists[0][i]->Write("", TObject::kWriteDelete);
    err_out_cor_hists[0][i]->Write("", TObject::kWriteDelete);
  }

  gDirectory->cd("..");
  gDirectory->cd("theta_x");
  err_hists[1]->Write("", TObject::kWriteDelete);
  for(int i=0; i<5; i++)
  {
    err_inp_cor_hists[1][i]->Write("", TObject::kWriteDelete);
    err_out_cor_hists[1][i]->Write("", TObject::kWriteDelete);
  }

  gDirectory->cd("..");
  gDirectory->cd("y");
  err_hists[2]->Write("", TObject::kWriteDelete);
  for(int i=0; i<5; i++)
  {
    err_inp_cor_hists[2][i]->Write("", TObject::kWriteDelete);
    err_out_cor_hists[2][i]->Write("", TObject::kWriteDelete);
  }

  gDirectory->cd("..");
  gDirectory->cd("theta_y");
  err_hists[3]->Write("", TObject::kWriteDelete);
  for(int i=0; i<5; i++)
  {
    err_inp_cor_hists[3][i]->Write("", TObject::kWriteDelete);
    err_out_cor_hists[3][i]->Write("", TObject::kWriteDelete);
  }
  gDirectory->cd("..");
  gDirectory->cd("..");
}

void LHCOpticsApproximator::PrintInputRange()
{
  const TVectorD* min_var = x_parametrisation.GetMinVariables();
  const TVectorD* max_var = x_parametrisation.GetMaxVariables();

  std::cout<<"Covered input parameters range:"<<std::endl;
  for(int i=0; i<5; i++)
  {
    std::cout<<(*min_var)(i)<<" < "<<coord_names[i]<<" < "<<(*max_var)(i)<<std::endl;
  }
  std::cout<<std::endl;
}

bool LHCOpticsApproximator::CheckInputRange(double *in)  //x, thx, y, thy, ksi
{
  const TVectorD* min_var = x_parametrisation.GetMinVariables();
  const TVectorD* max_var = x_parametrisation.GetMaxVariables();
  bool res = true;

  for(int i=0; i<5; i++)
  {
    res = res && in[i]>=(*min_var)(i) && in[i]<=(*max_var)(i);
//    std::cout<<"test ranges: "<<in[i]<<" "<<(*min_var)(i)<<" "<<in[i]<<" "<<(*max_var)(i)<<std::endl;
  }
//  std::cout<<"result dupa:"<<res<<std::endl;
  return res;
}

void LHCOpticsApproximator::AddRectEllipseAperture(const LHCOpticsApproximator &in, double rect_x, double rect_y, double r_el_x, double r_el_y)
{
  apertures_.push_back(LHCApertureApproximator(in, rect_x, rect_y, r_el_x, r_el_y, LHCApertureApproximator::RECTELLIPSE));
}

//////////////////////////////////////////////////////////////////

LHCApertureApproximator::LHCApertureApproximator()
{
  rect_x_ = rect_y_ = r_el_x_ = r_el_y_ = 0.0;
  ap_type_ = NO_APERTURE;
}

LHCApertureApproximator::LHCApertureApproximator(const LHCOpticsApproximator &in, double rect_x, double rect_y, double r_el_x, double r_el_y,
        aperture_type type) : LHCOpticsApproximator(in)
{
  rect_x_ = rect_x;
  rect_y_ = rect_y;
  r_el_x_ = r_el_x;
  r_el_y_ = r_el_y;
  ap_type_ = type;
}

bool LHCApertureApproximator::CheckAperture(double *in)  //x, thx. y, thy, ksi
{
  double out[5];
  bool result = Transport(in, out);
//  std::cout<<"in apreture: "<<out[0]<<" "<<out[1]<<" "<<out[2]<<" "<<out[3]<<" "<<out[4]<<std::endl;
//  std::cout<<rect_x_<<", "<<rect_y_<<", "<<r_el_x_<<", "<<r_el_y_<<", "<<ap_type_<<std::endl;

//  std::cout<<"LHCApertureApproximator::CheckAperture(double *in)  //x, thx. y, thy, ksi"<<std::endl;
  if(ap_type_==RECTELLIPSE)
  {
//    std::cout<<"before cond. checked: "<<result<<std::endl;
    result = result && out[0]<rect_x_ && out[0]>-rect_x_ && out[2]<rect_y_ && out[2]>-rect_y_ &&
        ( out[0]*out[0]/(r_el_x_*r_el_x_) + out[2]*out[2]/(r_el_y_*r_el_y_) < 1 );
//    std::cout<<"LHCApertureApproximator::CheckAperture(double *in)  //x, thx. y, thy, ksi"<<std::endl;
//    std::cout<<"after cond. checked: "<<result<<std::endl;
  }
  return result;
}

/*
bool LHCApertureApproximator::CheckAperture(MadKinematicDescriptor *in)  //x, thx. y, thy, ksi
{
  MadKinematicDescriptor out;
  bool result = Transport(in, &out);
  if(ap_type_==RECTELLIPSE)
  {

    result = result && out.x<rect_x_ && out.x>-rect_x_ && out.y<rect_y_ && out.y>-rect_y_ &&
        ( out.x*out.x/(r_el_x_*r_el_x_) + out.y*out.y/(r_el_y_*r_el_y_) < 1 );
  }
  return result;
}
*/

void LHCOpticsApproximator::PrintOpticalFunctions()
{
  std::cout<<std::endl<<"Linear terms of optical functions:"<<std::endl;
  for(int i=0; i<4; i++)
  {
    PrintCoordinateOpticalFunctions(*out_polynomials[i], coord_names[i], coord_names);
  }
}

void LHCOpticsApproximator::PrintCoordinateOpticalFunctions(TMultiDimFet &parametrization, const std::string &coord_name, const std::vector<std::string> &input_vars)
{
  double in[5];
  double out;
  double d_out_d_in[5];
  double d_par = 1e-5;
  double bias = 0;

  for(int j=0; j<5; j++)
      in[j]=0.0;

  const TVectorD* min_var = x_parametrisation.GetMinVariables();
  const TVectorD* max_var = x_parametrisation.GetMaxVariables();

  bias = parametrization.Eval(in);

  for(int i=0; i<5; i++)
  {
    for(int j=0; j<5; j++)
      in[j]=0.0;

    d_par = -((*max_var)[i]-(*min_var)[i])/10.0;
    in[i] = d_par;
    d_out_d_in[i] = parametrization.Eval(in);
    in[i] = 0.0;
    d_out_d_in[i] = d_out_d_in[i] - parametrization.Eval(in);
    d_out_d_in[i] = d_out_d_in[i]/d_par;
  }
  std::cout<<coord_name<<" = "<<bias;
  for(int i=0; i<5; i++)
  {
    std::cout<<" + "<<d_out_d_in[i]<<"*"<<input_vars[i];
  }
  std::cout<<std::endl;
}

//real angles in the matrix, MADX convention used only for input
void LHCOpticsApproximator::GetLineariasedTransportMatrixX(
    double mad_init_x, double mad_init_thx, double mad_init_y, double mad_init_thy,
    double mad_init_xi, TMatrixD &transp_matrix, double d_mad_x, double d_mad_thx)
{
  double MADX_momentum_correction_factor = 1.0 + mad_init_xi;
  transp_matrix.ResizeTo(2,2);
  double in[5];
  in[0] = mad_init_x;
  in[1] = mad_init_thx;
  in[2] = mad_init_y;
  in[3] = mad_init_thy;
  in[4] = mad_init_xi;

  double out[5];

  Transport(in, out);
  double x1 = out[0];
  double thx1 = out[1];

  in[0] = mad_init_x + d_mad_x;
  Transport(in, out);
  double x2_dx = out[0];
  double thx2_dx = out[1];

  in[0] = mad_init_x;
  in[1] = mad_init_thx + d_mad_thx;  //?
  Transport(in, out);
  double x2_dthx = out[0];
  double thx2_dthx = out[1];

//  | dx/dx,   dx/dthx    |
//  | dthx/dx, dtchx/dthx |

  transp_matrix(0,0) = (x2_dx-x1)/d_mad_x;
  transp_matrix(1,0) = (thx2_dx-thx1)/(d_mad_x*MADX_momentum_correction_factor);
  transp_matrix(0,1) = MADX_momentum_correction_factor*(x2_dthx-x1)/d_mad_thx;
  transp_matrix(1,1) = (thx2_dthx-thx1)/d_mad_thx;
}

//real angles in the matrix, MADX convention used only for input
void LHCOpticsApproximator::GetLineariasedTransportMatrixY(
    double mad_init_x, double mad_init_thx, double mad_init_y, double mad_init_thy,
    double mad_init_xi, TMatrixD &transp_matrix, double d_mad_y, double d_mad_thy)
{
  double MADX_momentum_correction_factor = 1.0 + mad_init_xi;
  transp_matrix.ResizeTo(2,2);
  double in[5];
  in[0] = mad_init_x;
  in[1] = mad_init_thx;
  in[2] = mad_init_y;
  in[3] = mad_init_thy;
  in[4] = mad_init_xi;

  double out[5];

  Transport(in, out);
  double y1 = out[2];
  double thy1 = out[3];

  in[2] = mad_init_y + d_mad_y;
  Transport(in, out);
  double y2_dy = out[2];
  double thy2_dy = out[3];

  in[2] = mad_init_y;
  in[3] = mad_init_thy + d_mad_thy;  //?
  Transport(in, out);
  double y2_dthy = out[2];
  double thy2_dthy = out[3];

//  | dy/dy,   dy/dthy    |
//  | dthy/dy, dtchy/dthy |

  transp_matrix(0,0) = (y2_dy-y1)/d_mad_y;
  transp_matrix(1,0) = (thy2_dy-thy1)/(d_mad_y*MADX_momentum_correction_factor);
  transp_matrix(0,1) = MADX_momentum_correction_factor*(y2_dthy-y1)/d_mad_thy;
  transp_matrix(1,1) = (thy2_dthy-thy1)/d_mad_thy;
}

//MADX convention used only for input
double LHCOpticsApproximator::GetDx(
    double mad_init_x, double mad_init_thx, double mad_init_y, double mad_init_thy,
    double mad_init_xi, double d_mad_xi)
{
  double in[5];
  in[0] = mad_init_x;
  in[1] = mad_init_thx;
  in[2] = mad_init_y;
  in[3] = mad_init_thy;
  in[4] = mad_init_xi;

  double out[5];

  Transport(in, out);
  double x1 = out[0];

  in[4] = mad_init_xi + d_mad_xi;
  Transport(in, out);
  double x2_dxi = out[0];
  double dispersion = (x2_dxi-x1)/d_mad_xi;

  return dispersion;
}

//MADX convention used only for input
//angular dispersion
double LHCOpticsApproximator::GetDxds(
    double mad_init_x, double mad_init_thx, double mad_init_y, double mad_init_thy,
    double mad_init_xi, double d_mad_xi)
{
  double MADX_momentum_correction_factor = 1.0 + mad_init_xi;
  double in[5];
  in[0] = mad_init_x;
  in[1] = mad_init_thx;
  in[2] = mad_init_y;
  in[3] = mad_init_thy;
  in[4] = mad_init_xi;

  double out[5];

  Transport(in, out);
  double thx1 = out[1]/MADX_momentum_correction_factor;

  in[4] = mad_init_xi + d_mad_xi;
  Transport(in, out);
  double thx2_dxi = out[1]/MADX_momentum_correction_factor;
  double dispersion = (thx2_dxi-thx1)/d_mad_xi;

  return dispersion;
}


/// @file  triangle.hpp
/// @brief Declarations for methods based on triangle constraints

#pragma once
#ifndef OGT_EMBED_TRIANGLE_HPP
#define OGT_EMBED_TRIANGLE_HPP

#include <ogt/config.hpp>
#include <ogt/core/collection.hpp>
#include <ogt/embed/embed.hpp>
#include <Eigen/Dense>
#include <iostream>
#include <json.hpp>
#include <memory>
#include <vector>

using OGT_NAMESPACE::core::Collection;
using OGT_NAMESPACE::core::Oracle;
using OGT_NAMESPACE::core::Traversal;

namespace OGT_NAMESPACE {
namespace embed {

struct CmpConstraint; // Forward declaration

/// A triangle with two edges' lengths constrained to be scaled to within some
/// interval of the length of the third edge.
struct TriangleConstraint {

	/// Initializing constructor
	TriangleConstraint(size_t p, size_t x, size_t y, double xmin,
		double xmax, double ymin, double ymax)
		: p(p), x(x), y(y), xmin(xmin), xmax(xmax), ymin(ymin), ymax(ymax)
	{
	}

	/// Initialize from a vector, e.g. loaded from .csv
	TriangleConstraint(const Eigen::VectorXd& vec)
		: p(vec(0)), x(vec(1)), y(vec(2)), xmin(vec(3)), xmax(vec(4))
		, ymin(vec(5)), ymax(vec(6))
	{
	}

	/// Elementwise equality test
	bool operator==(const TriangleConstraint& other) const {
		return p == other.p && x == other.x && y == other.y
			&& xmin == other.xmin && xmax == other.xmax && ymin == other.ymin
			&& ymax == other.ymax;
	}

	/// Tighten the constraint using the triangle inequality.
	///
	/// We can update xmin and ymin as follows.
	/// ~~~~
	/// xy <= px + py <= px + xy * ymax
	///    --> px >= xy(1 - ymax)
	///    --> px >= xy * max(xmin, 1 - ymax)
	///    --> py >= xy * max(ymin, 1 - xmax) (by symmetry)
	/// ~~~~
	void tighten();

	/// Rotate the constraint to use a different edge as the reference distance.
	/// If toX is true, the output's p will be the input's x point. Otherwise,
	/// the output's p will be the input's y point.
	///
	/// The update rule is derived as follows.
	/// ~~~~
	/// xmin <= px/xy <= xmax
	///     --> 1/xmax <= xy/px <= 1/xmin
	/// py <= ymax * xy
	///    <= ymax * px / xmin
	///    --> ymin / xmax <= py/px <= ymax / xmin
	/// ~~~~
	TriangleConstraint rotate(bool toX, double maxRatio = INFINITY) const;

	/// The point incident to the two constrained edges.
	size_t p;

	/// One of the points incident to the reference edge.
	size_t x;

	/// One of the points incident to the reference edge.
	size_t y;

	/// A lower bound on the ratio dist(p,x) / dist(x,y)
	double xmin;

	/// An upper bound on the ratio dist(p,x) / dist(x,y)
	double xmax;

	/// A lower bound on the ratio dist(p,y) / dist(x,y)
	double ymin;

	/// An upper bound on the ratio dist(p,y) / dist(x,y)
	double ymax;
};

/// Manages distance intervals from a set of constraints.
struct DistIntervals {

	/// Initialize to NaN matrices.
	DistIntervals(size_t nObj);

	/// Set the specified distance and interval width
	void set(Eigen::Index ii, Eigen::Index jj, double dmin, double dmax);

	/// Set the appropriate entries from a constraint, if the reference distance
	/// is already set. Returns true if successful.
	bool setFromCon(const TriangleConstraint& con);

	/// The center of each interval. Can be used as distance estimates.
	/// Uninitialized elements have the value NaN.
	Eigen::MatrixXd D;

	/// The width of each interval.
	/// Uninitialized elements have the value NaN.
	Eigen::MatrixXd W;
};

/// Stream a constraint to output
std::ostream& operator<<(std::ostream& os, const TriangleConstraint& con);

/// Sorts constraints into the order needed for embedTrianglesWithSpheres().
/// Throws std::invalid_argument on failure.
void sortTriangleConstraints(std::vector<TriangleConstraint>& cons, bool prune);

/// Validate the triangle constraints to determine whether they fully specify
/// an embedding in some dimensionality.
bool validateTriangleConstraints(const std::vector<TriangleConstraint>& cons,
	nlohmann::json* json);

/// The algorithm to use when selecting reference edges for triangle constraints
enum EdgeMethod {

	/// Pair farthest edges together
	EdgeFarthest,

	/// Select edges by density
	EdgeByDensity,

	/// Select all edges
	EdgeByAll,

	/// Recurse into local regions based on density rules
	EdgeByLocality
};

/// Find triangular constraints for a collection, using some traversal order to
/// identify reference edges. The first point on an edge will be selected by the
/// order, and the second point will be the farthest point from it.
/// If visitBoth = true, the second point will be visited for each edge.
/// This is useful for Collection::frftTraversal() and
/// Collection::stopWhenAffDep(), but not for Collection::stopAfter().
std::vector<TriangleConstraint> findTriangleConstraints(
	std::shared_ptr<Collection> coll, Traversal order, size_t nDim,
	bool visitBoth, bool basicIntervals, EdgeMethod edges);

/// Find triangular constraints for a collection with reference to some pair.
std::vector<TriangleConstraint> findTriangleConstraints(
	std::shared_ptr<Collection> coll, size_t pt1, size_t pt2,
	bool basicIntervals);

/// Find triangular constraints using the specified distance matrix.
/// The constraints will be the correct distances plus/minus eps.
/// Useful for testing.
std::vector<TriangleConstraint> findTriangleConstraints(
	const Eigen::MatrixXd& dists, double eps, size_t nRef);

/// Update triangle constraints using distances from an embedding of a subset
/// of the objects.
std::vector<TriangleConstraint> updateTriangleConstraints(
	const std::vector<TriangleConstraint>& inCons,
	std::shared_ptr<Oracle> ranker, const Eigen::MatrixXd& pos,
	const Eigen::VectorXd& objects);

/// Load triangle constraints from a data file.
std::vector<TriangleConstraint> loadTriangleConstraints(std::string path);

/// Count the number of violated triangle constraints.
size_t numBadTriangles(const std::vector<TriangleConstraint>& cons,
	const Eigen::MatrixXd& dists);

/// Calculate the loss for a single constraint
double triangleLoss(const TriangleConstraint& con,
	const Eigen::MatrixXd& dists);

/// Calculate the loss for a set of triangle constraints and distances.
template<class Container>
double triangleLoss(const Container& cons, const Eigen::MatrixXd& dists) {
	double loss = 0;
	for (const TriangleConstraint& con : cons) {
		loss += triangleLoss(con, dists);
	}
	return loss;
}

/// Calculate the loss for a set of triangle constraints and distances.
template<class ForwardIterator>
double triangleLoss(ForwardIterator begin, ForwardIterator end,
	const Eigen::MatrixXd& dists) {
	double loss = 0;
	for (auto it = begin; it != end; it++) {
		loss += triangleLoss(*it, dists);
	}
	return loss;
}

/// Sorts constraints into the order needed for embedTrianglesWithSpheres()
/// and embedTrianglesWithShells().
/// Throws std::invalid_argument on failure.
void sortConstraints(std::vector<TriangleConstraint>& cons, bool addRotations,
	double maxRotationRatio);

/// Reduce constraint bounds based on embedded coordinates.
/// Used for embedTrianglesWithSpheres() and embedTrianglesWithShells().
std::vector<TriangleConstraint> tightenConstraints(
	const std::map<size_t,std::vector<const TriangleConstraint*>> conMap,
	const Eigen::MatrixXd& X);

/// Find an embedding in R^nDim consistent with the triangle ratio constraints.
/// This method may throw EmbedErr on some errors.
///
/// This method works by first positioning the reference distances, and then
/// placing all the other points with respect to them.
/// It therefore requires that there be an ordering of reference distances such
/// that the points in each distance are constrained by the prior distances.
/// The points of the first reference distance are placed, arbitrarily, at the
/// origin (0, 0, 0, 0, ...) and the first elementary vector (1, 0, 0, 0, ...).
/// The method will attempt to find an ordering of constraints to satisfy its
/// needs, and will throw std::invalid_argument when it can't.
/// The constraints returned by findTriangleConstraints() when it is passed some
/// traversal order are generally valid input.
EmbedResult embedTrianglesWithSpheres(
	std::vector<TriangleConstraint> cons, EmbedConfig config);

/// Find an embedding in R^nDim consistent with the triangle ratio constraints.
/// This method may throw EmbedErr on some errors.
///
/// This method works similarly to embedTrianglesWithSpheres(), but takes the
/// spherical shells into account instead of simply aiming for the center of
/// each shell.
EmbedResult embedTrianglesWithShells(
	std::vector<TriangleConstraint> cons, EmbedConfig config);

/// Find an embedding in R^d (d = X0.cols()) consistent with the triangle
/// ratio constraints.
/// This approach uses a quasi-newton method (L-BFGS) to minimize a loss
/// function defined in terms of the constraints.
/// Returns the embedding and the loss achieved.
EmbedResult embedTrianglesWithOpt(
	std::vector<TriangleConstraint> cons, const Eigen::MatrixXd& X0,
	EmbedConfig config);

/// Find an embedding in R^d consistent with the triangle ratio constraints.
/// Distances are found to a set of d+1 reference vertices using the triangle
/// constraints, and then the set is embedded from the distances.
EmbedResult embedTrianglesWithCM(std::vector<TriangleConstraint> cons,
	EmbedConfig config);

/// Embeds a dataset using a mixture of triangles and comparisons.
EmbedResult embedTrianglesWithMixture(std::vector<TriangleConstraint> tris,
	std::vector<CmpConstraint> cmps, double lambda, const Eigen::MatrixXd& X0,
	EmbedConfig config);

} // end namespace embed
} // end namespace OGT_NAMESPACE
#endif /* OGT_EMBED_TRIANGLE_HPP */


#include <Eigen/Dense>
#include "mex.h"
#include "math.h"
#include <iostream>
#include <omp.h>
#include <vector>
#include "MinIndexedPQ.h"

using namespace std;

using Eigen::Matrix2d;
using Eigen::Vector2d;
using Eigen::VectorXd;
using Eigen::SelfAdjointEigenSolver;
using Eigen::Map;

// To cache the dot products that we need to compute between the
// query image features and the predictor image features.  Each
// cachecell is associated with one locaton in the query image
struct CacheCell{
  // stores the locations in the predictor image where the
  // cache is valid
  int xmin;
  int ymin;
  int xmax;
  int ymax;
  // The actual dot products.
  double* data;
  CacheCell():xmin(0),ymin(0),xmax(0),ymax(0),data(NULL){}
};

vector<vector<CacheCell> > cache;

//extern bool mxUnshareArray(const mxArray *pr, const bool noDeepCopy);

mxArray* getfield(mxArray* str, const string& fnam){
  int field_num = mxGetFieldNumber(str, fnam.c_str());
  return mxGetFieldByNumber(str, 0, field_num);
}

void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[]) {
  // Note that you should disable the parallel loop
  // before you enable debugging, or you will get segfaults!
  const bool DEBUG=false;
  // TODO: The first thing this mex function does is make a copy of its
  // input, which is really stupid because the structure of corresp does
  // not change, and the old values don't need to be kept.  Supposedly
  // these lines will let you avoid the copy, but I haven't tested it.
  //mxUnshareArray(const_cast<mxarray*> prhs[1],false);
  //mxArray* corresp=const_cast<mxarray*>(prhs[1]);
  mxArray* corresp=mxDuplicateArray(prhs[1]);
  int maxupdates = (int) mxGetScalar(prhs[3]);
  const mxArray* hogim = prhs[0];
  const int* hogimdims=mxGetDimensions(hogim);
  double* hogdata = mxGetPr(hogim);
  int ndims=hogimdims[0];
  int hogrows=hogimdims[1];
  int hogcols=hogimdims[2];
  int npyrs=* mxGetDimensions(prhs[2]);
  int* correspidx=(int *) mxGetData(prhs[4]);
  bool* inferred=(bool*)mxGetData(prhs[5]);
  int* inferinds=(int *)mxGetData(prhs[6]);
  int ntoinfer=*mxGetDimensions(prhs[6]);
  double* confidences = mxGetPr(prhs[7]);
  double lambda = mxGetScalar(prhs[8]);
  bool deleteCache = mxGetScalar(prhs[9])!=0;
  double* numneighbors = mxGetPr(prhs[10]);
  double lambdaprime = mxGetScalar(prhs[11]);
  if(DEBUG){mexPrintf("lam2 %f\n",lambda);mexEvalString("drawnow;");}
  if(DEBUG){mexPrintf("%d %d\n",ntoinfer,npyrs);}
  double** transf_out = new double*[npyrs];
  int ntransfs=0;
  // if deleteCache is 1, we're starting a new image, so clear out all
  // the cached dot products.
  if(deleteCache){
    for(int i = 0; i<cache.size(); ++i){
      for(int j=0; j<cache[i].size(); ++j){
        delete cache[i][j].data;
      }
    }
    cache.clear();
  }
  if(cache.size()==0){
    for(int i = 0; i<npyrs; ++i){
      cache.push_back(vector<CacheCell>(hogrows*hogcols));
    }
  }
  int cachehit=0;
  int cachemiss=0;
  string errmsg="";
  bool error=false;
  if(DEBUG){mexPrintf("start main loop over pyramids\n");}

  // Each pyramid's correspondence can be inferred in parallel.
  #pragma omp parallel for reduction(+:cachehit,cachemiss) shared(error,errmsg)
  for(int pyridx=0; pyridx<npyrs; ++pyridx){
    // These help us estimate when we need to return,
    // first by counting the number of times we've updated a
    // mu and sigma pair, and second by seeing how long it's
    // been since something moved a lot.
    int nupdates=0;
    int nsincebigmove=0;

    // transf==alpha in the paper.  Figure out which cells need to have an
    // alpha value inferred before we can start updating mu's and sigmas.
    // alpha's aren't kept around between calls to optimizecorresp, so
    // when we begin we need to infer an alpha for every cell in the
    // condition region.
    //
    // Likewise, cellheap associates with each cell an estimate of how
    // much that cell will change if it's updated.  We attempt to update
    // first any cell that seems like it will move a lot; this lets us
    // focus our computation on the parts of f that make the biggest
    // difference.
    vector<int> mustcomptransf;
    MinIndexedPQ cellheap(hogrows*hogcols);
    for(int ti = 0; ti<ntoinfer; ++ti){
      mustcomptransf.push_back(inferinds[ti]);
      cellheap.insert(inferinds[ti],-1000000);
    }
    vector<Matrix2d> transfs(hogrows*hogcols);
    if(DEBUG){cout<<"start main optimization loop\n";}
    // On each iteration of this loop, we update one mu and one sigma.
    while(!error){
      ++nupdates;
      ++nsincebigmove;

      // However, we may need to compute updates for many alphas, especially
      // at the beginning.
      for(int ti=0; ti<mustcomptransf.size(); ++ti){
        int infidx=mustcomptransf[ti];
        int xpos=infidx/hogrows;
        int ypos=infidx-xpos*hogrows;
        Matrix2d xsigma=Matrix2d::Zero();
        Matrix2d ysigma=Matrix2d::Zero();
        Vector2d xb=Vector2d::Zero();
        Vector2d yb=Vector2d::Zero();
        // Iterate over the edges of the lattice that this \alpha participates
        // in.  We aggregate statistics for both vertical and horizontal edges
        // simultaneously.  xsigma/ysigma aggregate the quadratic term
        // of equation 15, and xb/yb aggregate the linear term.
        for(int edgex=-2; edgex<=2; ++edgex){
          for(int edgey=-2; edgey<=2; ++edgey){
            if(edgex>-2&&edgex+xpos-1>=0&&edgex+xpos<hogcols&&edgey+ypos>=0&&edgey+ypos<hogrows &&
                  inferred[edgey+ypos+hogrows*(edgex+xpos-1)]&&inferred[edgey+ypos+hogrows*(edgex+xpos)]){
              mxArray* othcorresp=mxGetCell(corresp,edgey+ypos+(edgex+xpos)*hogrows);
              mxArray* othcorresp2=mxGetCell(corresp,edgey+ypos+(edgex+xpos-1)*hogrows);
              Map<Vector2d> othmu(mxGetPr(mxGetFieldByNumber(othcorresp, 0, 0))+2*pyridx);
              Map<Matrix2d> othcovar(mxGetPr(mxGetFieldByNumber(othcorresp, 0, 1))+4*pyridx);
              Map<Vector2d> othmu2(mxGetPr(mxGetFieldByNumber(othcorresp2, 0, 0))+2*pyridx);
              Map<Matrix2d> othcovar2(mxGetPr(mxGetFieldByNumber(othcorresp2, 0, 1))+4*pyridx);
              xsigma=xsigma+(othcovar.inverse()+othcovar2.inverse());
              xb=xb+(othcovar.inverse()+othcovar2.inverse())*(othmu-othmu2);
            }
            if(edgey>-2&&edgex+xpos>=0&&edgex+xpos<hogcols&&edgey+ypos-1>=0&&edgey+ypos<hogrows &&
                  inferred[edgey+ypos-1+hogrows*(edgex+xpos)]&&inferred[edgey+ypos+hogrows*(edgex+xpos)]){
              mxArray* othcorresp=mxGetCell(corresp,edgey+ypos+(edgex+xpos)*hogrows);
              mxArray* othcorresp2=mxGetCell(corresp,edgey+ypos-1+(edgex+xpos)*hogrows);
              Map<Vector2d> othmu(mxGetPr(mxGetFieldByNumber(othcorresp, 0, 0))+2*pyridx);
              Map<Matrix2d> othcovar(mxGetPr(mxGetFieldByNumber(othcorresp, 0, 1))+4*pyridx);
              Map<Vector2d> othmu2(mxGetPr(mxGetFieldByNumber(othcorresp2, 0, 0))+2*pyridx);
              Map<Matrix2d> othcovar2(mxGetPr(mxGetFieldByNumber(othcorresp2, 0, 1))+4*pyridx);
              ysigma=ysigma+(othcovar.inverse()+othcovar2.inverse());
              yb=yb+(othcovar.inverse()+othcovar2.inverse())*(othmu-othmu2);
            }
          }
        }
        transfs[ypos+xpos*hogrows].leftCols(1)=(xsigma).colPivHouseholderQr().solve(xb);
        transfs[ypos+xpos*hogrows].rightCols(1)=(ysigma).colPivHouseholderQr().solve(yb);
        if(DEBUG){cout<<"xsig\n"<<xsigma<<"\nxb\n"<<xb<<"\ntransfs\n"<<transfs[ypos+xpos*hogrows]<<"\n";}
      }
      mustcomptransf.clear();

      if(DEBUG){cout<<"start computing sigma's\n";}
      mxArray *mypyr=getfield(mxGetCell(prhs[2],pyridx),"features");
      // Get the linear index of the next mu/sigma pair to optimize, and store the estimated
      // change.
      double heapdist=cellheap.minKey();
      int infidx=cellheap.deleteMin();
      mxArray* mycorresp=mxGetCell(corresp,infidx);
      int edgex=1;
      int edgey=0;
      int xpos=infidx/hogrows;
      int ypos=infidx-xpos*hogrows;
      if(DEBUG){mexPrintf("get predictor feats\n");mexEvalString("drawnow;");}
      if(DEBUG){mexPrintf("%f\n",mxGetPr(mxGetFieldByNumber(mycorresp,0,2))[pyridx]);}
      mxArray* othhogarr=mxGetCell(mypyr,mxGetPr(mxGetFieldByNumber(mycorresp,0,2))[pyridx]-1);
      double* othhog=mxGetPr(othhogarr);
      const int* othhogdims=mxGetDimensions(othhogarr);
      int othrows=othhogdims[1];
      int othcols=othhogdims[2];
      // relsigma aggregates both Sigma in equation 21 and the outer product in equation 23.
      Matrix2d relsigma = Matrix2d::Zero();
      // relsigmainv aggregates inv(Sigma) in equation 21.
      Matrix2d relsigmainv = Matrix2d::Zero();
      if(DEBUG){mexPrintf("iterate over neighbors\n");mexEvalString("drawnow;");}
      Map<Vector2d> mymu(mxGetPr(mxGetFieldByNumber(mycorresp, 0, 0))+2*pyridx);
      Map<Matrix2d> mycovar(mxGetPr(mxGetFieldByNumber(mycorresp, 0, 1))+4*pyridx);
      // We first iterate over edges
      for(int edgeidx=0; edgeidx<4; ++edgeidx){
        if(DEBUG){mexPrintf("%d %d %d %d %d %d\n",edgex, edgey, xpos, ypos, hogrows, hogcols);}
        if(edgex+xpos>=0&&edgex+xpos<hogcols&&edgey+ypos>=0&&edgey+ypos<hogrows &&
          inferred[edgey+ypos+hogrows*(edgex+xpos)]){
          mxArray* othcorresp=mxGetCell(corresp,edgey+ypos+(edgex+xpos)*hogrows);
          Map<Vector2d> othmu(mxGetPr(mxGetFieldByNumber(othcorresp, 0, 0))+2*pyridx);
          Map<Matrix2d> othcovar(mxGetPr(mxGetFieldByNumber(othcorresp, 0, 1))+4*pyridx);
          Vector2d edge;
          edge << edgex, edgey;
          Matrix2d relsigmainv2=Matrix2d::Zero();
          Matrix2d relsigma2=Matrix2d::Zero();
          double n=0.0;
          // Then we iterate over the \alpha's that give us constraints over these edges
          for(int transfx=max(0,max(xpos,xpos+edgex)-2);transfx<=min(hogcols-1,min(xpos,xpos+edgex)+2); ++transfx){
            for(int transfy=max(0,max(ypos,ypos+edgey)-2);transfy<=min(hogrows-1,min(ypos,ypos+edgey)+2); ++transfy){
                if(DEBUG){cout << "tx " << transfx << " ty " <<transfy << "\n";}
              if(inferred[transfy+hogrows*(transfx)]){
                Vector2d mudiff=othmu-mymu-transfs[transfy+(transfx)*hogrows]*edge;
                if(DEBUG){cout << "othmu\n" << othmu << "\n" << mymu << "\n" << edge << "\ntransf\n"<<transfs[transfy+(transfx)*hogrows]<<"\n";}
                relsigma2=relsigma2+(othcovar+mudiff*mudiff.transpose())*lambdaprime;
                relsigmainv2=relsigmainv2+othcovar.inverse()*lambdaprime;
                n+=1.0;
              }
            }
          }
          // n is very nearly a constant and so in the paper I just rolled into lambdaprime.  However,
          // near the edge n will get smaller, and so dividing by n will give the warping
          // a little boost near the edge where there are fewer alpha's affecting each
          // edge.
          relsigma=relsigma+relsigma2/n;
          relsigmainv=relsigmainv+relsigmainv2/n;
          Vector2d mudiff=othmu-mymu-edge;
          relsigma=relsigma+(othcovar+mudiff*mudiff.transpose())*lambda;
          relsigmainv=relsigmainv+othcovar.inverse()*lambda;
        }
        // We rotate the edge to get the next edge in the sequence.
        int tmp=edgey;
        edgey=edgex;
        edgex=-tmp;
      }

      // Next we need to aggregate the unary terms.  We actually do an e-step (equation 13) as we
      // go along; the values computed on the e-step are not stored.
      //
      // First we need to compute a small region over which we should aggregate our statistics.
      // we use the same major/minor axis trick that we used in contextpredict to get
      // the window.
      SelfAdjointEigenSolver<Matrix2d> eigensolver(mycovar);
      Matrix2d V=eigensolver.eigenvectors().cwiseAbs();
      Vector2d D=eigensolver.eigenvalues();
      double mycovardet=mycovar.determinant();
      if(DEBUG){mexPrintf("computing eigenvalues 1\n");mexEvalString("drawnow;");}
      Matrix2d mycovarinv=mycovar.inverse();
      double dist=sqrt(fabs(1/(((1/(log(.0001)+log(mycovardet)/2))*(V.block(0,0,2,1).transpose()*mycovarinv*V.block(0,0,2,1)/2))(0))));
      double dist2=sqrt(fabs(1/(((1/(log(.0001)+log(mycovardet)/2))*(V.block(0,1,2,1).transpose()*mycovarinv*V.block(0,1,2,1)/2))(0))));
      if(DEBUG){mexPrintf("%f %f\n",dist,dist2);}
      Vector2d distvec = dist*V.block(0,0,2,1).cwiseMax(dist2*V.block(0,1,2,1));
      Map<VectorXd> mydata(hogdata+ndims*ypos+ndims*hogrows*xpos,ndims-1);
      double myconst = *(hogdata+ndims*ypos+ndims*hogrows*xpos+ndims-1);
      VectorXd mydatatransf=mydata;
      double sumprob=0;
      Matrix2d sumwt = Matrix2d::Zero();
      if(DEBUG){mexPrintf("5\n");mexEvalString("drawnow;");}
      //NOTE: window is in 1-indexed coordinates
      if(DEBUG){cout << distvec << "\n";}
      int xmin=floor(mymu(0)-min(distvec(0),30.0));
      int xmax=ceil(mymu(0)+min(distvec(0),30.0));
      int ymin=floor(mymu(1)-min(distvec(1),30.0));
      int ymax=ceil(mymu(1)+min(distvec(1),30.0));
      if(DEBUG){mexPrintf("5.1\n");mexEvalString("drawnow;");}

      // Now that we have the window, check to see if we have all the dot
      // products in that window cached.  We make sure that no more than 1/5
      // of the region that we need to compute is hanging outside of the
      // region stored in the cache on each side.
      CacheCell& mycache=cache[pyridx][hogrows*xpos+ypos];
      bool createcache=false;
      double* cachedata;
      if(DEBUG){mexPrintf("5.2\n");mexEvalString("drawnow;");}
      if(mycache.data!=NULL&&mycache.xmin<xmin+(xmax-xmin)/5&&mycache.ymin<ymin+(ymax-ymin)/5&&
         mycache.xmax>xmin+((xmax-xmin)*4+4)/5&&mycache.ymax>ymin+((ymax-ymin)*4+4)/5){
         ++cachehit;
      }else{
        ++cachemiss;
        createcache=true;
        mycache.xmin=xmin;
        mycache.ymin=ymin;
        mycache.xmax=xmax;
        mycache.ymax=ymax;
        if(mycache.data!=NULL){
          delete mycache.data;
        }
        mycache.data=new double[(ymax-ymin+1)*(xmax-xmin+1)];
      }
      if(DEBUG){mexPrintf("5.3\n");mexEvalString("drawnow;");}
      cachedata=mycache.data;
      if(DEBUG){mexPrintf("5.4\n");mexEvalString("drawnow;");}

      // Loop over the cache or the window, computing the inner products
      // and summing the total probability as required in equation 19.
      for(int windowx = max(xmin,mycache.xmin); windowx<=min(xmax,mycache.xmax); ++windowx){
        for(int windowy = max(ymin,mycache.ymin); windowy<=min(ymax,mycache.ymax); ++windowy){
          Vector2d window;
          window << windowx,windowy;
          Vector2d pt=window-mymu;
          double prob=1/(2*M_PI*sqrt(mycovardet))*exp((-pt.transpose()*mycovarinv*pt)(0)/2);
          if(prob!=prob){
            if(DEBUG){mexPrintf("%d\n",mycovardet);}
            errmsg="nan prob";
            if(DEBUG){mexErrMsgTxt(errmsg.c_str());}
            error=1;
            continue;
          }
          double tmpprob;
          if(createcache){
            int idx=ndims*(min(othrows-1,max(0,windowy-1))+min(othcols-1,max(0,windowx-1))*othrows);
            Map<VectorXd> othdata(othhog+idx,ndims-1);
            tmpprob = exp(-(mydatatransf.dot(othdata)+myconst)/2);
            tmpprob=tmpprob/(tmpprob+.01*(*(othhog+idx+ndims-1)));
            cachedata[(windowx-xmin)*(ymax-ymin+1)+windowy-ymin]=tmpprob;
            if(prob!=prob){
              if(DEBUG){mexPrintf("%f %f %f %f\n",tmpprob,myconst,mydatatransf.dot(othdata),*(othhog+idx+ndims-1));}
              if(DEBUG){cout<<"othdata\n"<<othdata<<"\nmydata\n"<<mydata<<"\n";}
              errmsg="nan prob v2";
              if(DEBUG){mexErrMsgTxt(errmsg.c_str());}
              error=1;
              continue;
            }
          }else{
            tmpprob=cachedata[(windowx-mycache.xmin)*(mycache.ymax-mycache.ymin+1)+windowy-mycache.ymin];
          }

          prob=prob*tmpprob;
          sumprob+=prob;
          sumwt=sumwt+prob*pt*(pt.transpose());
        }
      }
      if(DEBUG){mexPrintf("6\n");mexEvalString("drawnow;");}
      if(DEBUG){mexPrintf("sumwt\n");mexEvalString("drawnow;");}
      if(DEBUG){cout << sumwt << "\n" << sumprob << "\n";}
      if(DEBUG){mexPrintf("relsigma\n");mexEvalString("drawnow;");}
      if(DEBUG){cout << relsigma << "\n";}
      if(DEBUG){mexPrintf("sumwt\n");mexEvalString("drawnow;");}
      if(DEBUG){cout << sumwt << "\n";}
      if(DEBUG){mexPrintf("relsigmainv\n");mexEvalString("drawnow;");}
      if(DEBUG){cout << relsigmainv << "\n";}
      if(DEBUG){mexPrintf("mycovarinv\n");mexEvalString("drawnow;");}
      if(DEBUG){cout << mycovarinv << "\n";}
      double myconfidence = confidences[ypos+hogrows*(xpos)];
      if(DEBUG){mexPrintf("myconfidence\n");mexEvalString("drawnow;");}
      if(DEBUG){cout << myconfidence << "\n";}

      // Finally, put it all together
      Matrix2d nablasigma=(mycovarinv*(relsigma+sumwt*myconfidence)*mycovarinv-relsigmainv-mycovarinv*sumprob*myconfidence)/2;
      double covarstepsz=.1;
      if(nablasigma(0)!=nablasigma(0)){
        errmsg="nans in nablasigma";
        if(DEBUG){mexErrMsgTxt(errmsg.c_str());}
        error=1;
        continue;
      }
      // As is often the case with EM algorithms, we can end up with a covariance matrix shrinking to zero and
      // causing problems.  The right thing to do is to put a constraint on the algorithm that the determinant
      // of the inverse covariance must be less than 1, and correctly project onto the constraint set.  The wrong
      // thing to do is to not take gradient steps when the determinant of the covariance matrix ends up less
      // than 1.  However, out of sheer laziness I've done the latter.  In practice this happens very rarely.
      while(-nablasigma(0)*covarstepsz>mycovar(0)||-nablasigma(3)*covarstepsz>mycovar(3)||(mycovar+nablasigma*covarstepsz).determinant()<1){
        covarstepsz=covarstepsz*.5;
      }
      if(DEBUG){mexPrintf("nablasigma\n");mexEvalString("drawnow;");}
      if(DEBUG){cout << nablasigma << "\n";}
      Matrix2d newcovar=mycovar+nablasigma*covarstepsz;
      if(newcovar(0)>1e4||newcovar(4)>1e4){
        errmsg="exploding sigma";
        if(DEBUG){mexErrMsgTxt(errmsg.c_str());}
        error=1;
        continue;
      }
      Matrix2d newcovarinv=newcovar.inverse();
      Vector2d constr=Vector2d::Zero();
      Matrix2d syst=Matrix2d::Zero();

      if(DEBUG){mexPrintf("starting mu update\n");mexEvalString("drawnow;");}
      // Just like before, we start off by iterating over edges, this time
      // aggregating system of linear equations that's defined by taking the
      // derivative of equation 17 (note that this system also needs to have
      // some terms from  equation 16 added in, but
      // we'll add those terms to the system later) and setting it to 0.
      //
      // We ultimately end up with an equation of the form syst*mu=constr,
      // where each term in the sums of equation 17 adds to both syst and constr.
      edgex=1;
      edgey=0;
      for(int edgeidx=0; edgeidx<4; ++edgeidx){
        if(DEBUG){mexPrintf("%d %d %d %d %d %d\n",edgex, edgey, xpos, ypos, hogrows, hogcols);mexEvalString("drawnow;");}
        if(edgex+xpos>=0&&edgex+xpos<hogcols&&edgey+ypos>=0&&edgey+ypos<hogrows &&
           inferred[edgey+ypos+hogrows*(edgex+xpos)]){
          mxArray* othcorresp=mxGetCell(corresp,edgey+ypos+(edgex+xpos)*hogrows);
          Map<Vector2d> othmu(mxGetPr(mxGetFieldByNumber(othcorresp, 0, 0))+2*pyridx);
          Map<Matrix2d> othcovar(mxGetPr(mxGetFieldByNumber(othcorresp, 0, 1))+4*pyridx);
          Vector2d edge;
          edge << edgex, edgey;
          if(DEBUG){mexPrintf("17\n");mexEvalString("drawnow;");}
          if(DEBUG){cout << othcovar << "\n";}
          if(DEBUG){cout << othmu << "\n";}
          Matrix2d othinvcovar=othcovar.inverse();
          Matrix2d syst2=Matrix2d::Zero();
          Vector2d constr2=Vector2d::Zero();
          double n=0.0;
          for(int transfx=max(0,max(xpos,xpos+edgex)-2);transfx<=min(hogcols-1,min(xpos,xpos+edgex)+2); ++transfx){
            for(int transfy=max(0,max(ypos,ypos+edgey)-2);transfy<=min(hogrows-1,min(ypos,ypos+edgey)+2); ++transfy){
              if(inferred[transfy+hogrows*(transfx)]){
                Vector2d mudiff=othmu-transfs[transfy+(transfx)*hogrows]*edge;
                constr2=constr2+(othinvcovar+newcovarinv)*(mudiff)*lambdaprime;
                syst2=syst2+(othinvcovar+newcovarinv)*lambdaprime;
                n+=1.0;
              }
            }
          }
          // n is very nearly a constant and so in the paper I just rolled it into lambdaprime.  However,
          // near the edge n will get smaller, and so dividing by n will give the warping
          // a little boost near the edge where there are fewer alpha's affecting each
          // edge.
          constr=constr+constr2/n;
          syst=syst+syst2/n;

          Vector2d mudiff=othmu-edge;
          constr=constr+(othinvcovar+newcovarinv)*(mudiff)*lambda;
          syst=syst+(othinvcovar+newcovarinv)*lambda;
        }
        int tmp=edgey;
        edgey=edgex;
        edgex=-tmp;
      }
      if(DEBUG){mexPrintf("11\n");mexEvalString("drawnow;");}
      if(DEBUG){cout << constr << "\n";}
      if(DEBUG){cout << syst << "\n";}

      double newcovardet=newcovar.determinant();
      Vector2d mn = Vector2d::Zero();
      sumprob=0;

      // Now aggregate the contribution from the system from equation 16.  Computing
      // all the stuff for the e-step (including the window, the cache, etc.) is the
      // same as above.
      xmin=floor(mymu(0)-min(distvec(0),30.0));
      xmax=ceil(mymu(0)+min(distvec(0),30.0));
      ymin=floor(mymu(1)-min(distvec(1),30.0));
      ymax=ceil(mymu(1)+min(distvec(1),30.0));
      mycache=cache[pyridx][hogrows*xpos+ypos];
      createcache=false;
      if(mycache.xmin<xmin+(xmax-xmin)/5&&mycache.ymin<ymin+(ymax-ymin)/5&&
         mycache.xmax>xmin+((xmax-xmin)*4+4)/5&&mycache.ymax>ymin+((ymax-ymin)*4+4)/5){
         ++cachehit;
      }else{
        ++cachemiss;
        createcache=true;
        mycache.xmin=xmin;
        mycache.ymin=ymin;
        mycache.xmax=xmax;
        mycache.ymax=ymax;
        if(mycache.data!=NULL){
          delete mycache.data;
        }
        mycache.data=new double[(ymax-ymin+1)*(xmax-xmin+1)];
      }
      cachedata=mycache.data;

      for(int windowx = max(xmin,mycache.xmin); windowx<=min(xmax,mycache.xmax); ++windowx){
        for(int windowy = max(ymin,mycache.ymin); windowy<=min(ymax,mycache.ymax); ++windowy){
          Vector2d window;
          window << windowx,windowy;
          Vector2d pt=window-mymu;
          double prob=1/(2*M_PI*sqrt(newcovardet))*exp((-pt.transpose()*newcovarinv*pt)(0)/2);
          double tmpprob;
          if(createcache){
            int idx=ndims*(min(othrows-1,max(0,windowy-1))+min(othcols-1,max(0,windowx-1))*othrows);
            Map<VectorXd> othdata(othhog+idx,ndims-1);
            tmpprob = exp(-(mydatatransf.dot(othdata)+myconst)/2);
            tmpprob=tmpprob/(tmpprob+.01*(*(othhog+idx+ndims-1)));
            cachedata[(windowx-xmin)*(ymax-ymin+1)+windowy-ymin]=tmpprob;
            if(prob!=prob){
              if(DEBUG){mexPrintf("%f %f %f %f\n",tmpprob,myconst,mydatatransf.dot(othdata),*(othhog+idx+ndims-1));}
              if(DEBUG){cout<<"othdata\n"<<othdata<<"\nmydata\n"<<mydata<<"\n";}
              errmsg="nan prob v3";
              if(DEBUG){mexErrMsgTxt(errmsg.c_str());}
              error=1;
              continue;
            }
          }else{
            tmpprob=cachedata[(windowx-mycache.xmin)*(mycache.ymax-mycache.ymin+1)+windowy-mycache.ymin];
          }
          prob=prob*tmpprob;
          sumprob+=prob;
          mn=mn+prob*pt;
        }
      }
      if(DEBUG){cout << mn << "\n" << mymu << "\n";}
      if(DEBUG){cout << sumprob << "\n";}
      constr=constr+newcovarinv*(mn+mymu*sumprob)*myconfidence;
      syst=syst+newcovarinv*sumprob*myconfidence;
      // Now put it all together.
      Vector2d newmu=syst.inverse()*constr;
      if(DEBUG){mexPrintf("Computed newmu, saving\n");mexEvalString("drawnow;");}
      if(DEBUG){mexPrintf("%f\n",infidx);mexEvalString("drawnow;");}

      mxArray* newcorresp=mxGetCell(corresp,infidx);
      Map<Vector2d> infmu(mxGetPr(mxGetFieldByNumber(newcorresp, 0, 0))+2*pyridx);
      Map<Matrix2d> infcovar(mxGetPr(mxGetFieldByNumber(newcorresp, 0, 1))+4*pyridx);
      if(DEBUG){cout << newmu << "\n";}
      if(DEBUG){cout << newcovar << "\n";}
      if(DEBUG){mexPrintf("%d\n",ntransfs);mexEvalString("drawnow;");}
      if(newmu(0)!=newmu(0)){
        errmsg="nans in newmu";
        if(DEBUG){mexErrMsgTxt(errmsg.c_str());}
        error=1;
        continue;
      }

      // At this point, actually updating the mu and sigma is finished. All that's left is to
      // put this cell back in the queue, and update all the other nodes that it's connected
      // to by an edge (since when this node moves, it will probably make its neighbors move).
      // In general, we aim to overestimate. so we don't have to be afraid of thinking we're
      // converged when we're not.
      double diff=(infmu-newmu).norm()*myconfidence;
      infmu=newmu;
      infcovar=newcovar;
      // We guess that, next time, the current node will move as much as it did this time.
      cellheap.insert(infidx,-diff);
      // Now, for each neighbor, we increment our estimate of how much it will move by the
      // amount that this node moved, divided by the number of neighbors that node has.
      // This is in general a pretty bad estimate, but it's far better than just sweeping
      // over everything.
      for(int edge=0; edge<4; ++edge){
        if(DEBUG){mexPrintf("%d %d %d %d %d %d\n",edgex, edgey, xpos, ypos, hogrows, hogcols);mexEvalString("drawnow;");}
        if(edgex+xpos>=0&&edgex+xpos<hogcols&&edgey+ypos>=0&&edgey+ypos<hogrows &&
           inferred[edgey+ypos+hogrows*(edgex+xpos)]){
             cellheap.decreaseKey(edgey+ypos+hogrows*(edgex+xpos),cellheap.keyOf(edgey+ypos+hogrows*(edgex+xpos))-diff/numneighbors[edgey+ypos+hogrows*(edgex+xpos)]);
        }
        int tmp=edgey;
        edgey=edgex;
        edgex=-tmp;
      }
      // On the next round, re-infer some alphas.  Specifically, only re-infer the alpha
      // associated with the current node.  There's probably a better way to do this, but
      // in practice this seems to work well enough.
      mustcomptransf.push_back(infidx);

      // If we've run out of computation budget or we don't seem to be
      // moving very much anymore, return.  Copy the set of computed alpha's
      // from the stack to the heap so we can return them.
      if(diff>=.1){
        nsincebigmove=0;
      }
      if((nupdates>=maxupdates||nsincebigmove==100)&&heapdist>-100000){
        if(nlhs>=1){
          transf_out[pyridx]=new double[hogrows*hogcols*4];
          ++ntransfs;
          for(int i=0;i<hogrows*hogcols; ++i){
            Map<Matrix2d> transf_out_i(transf_out[pyridx]+4*i);
            transf_out_i=transfs[i];
          }
        }
        break;
      }
    }
  }
  if(error){
    mexErrMsgTxt(errmsg.c_str());
  }

  // Copy the alpha's from the heap to Matlab.
  plhs[0]=corresp;
  if(nlhs>=1){
    mxArray* tout=mxCreateCellMatrix(npyrs,1);
    for(int i=0; i<npyrs; ++i){
      int dims[4];
      dims[0]=2;
      dims[1]=2;
      dims[2]=hogrows;
      dims[3]=hogcols;
      mxArray* data=mxCreateNumericArray(4, dims, mxDOUBLE_CLASS, mxREAL);
      double* transf_out2=mxGetPr(data);
      for(int j=0; j<dims[2]*dims[3]*4; ++j){
        transf_out2[j]=transf_out[i][j];
      }
      mxSetCell(tout,i,data);
      delete transf_out[i];
    }
    plhs[1]=tout;
  }
  // Useful for debugging/performance analysis.
  mexPrintf("cache hit %d miss %d\n",cachehit,cachemiss);
  delete transf_out;
}


/* CirKit: A circuit toolkit
 * Copyright (C) 2009-2015  University of Bremen
 * Copyright (C) 2015-2016  EPFL
 *
 * Permission is hereby granted, free of charge, to any person
 * obtaining a copy of this software and associated documentation
 * files (the "Software"), to deal in the Software without
 * restriction, including without limitation the rights to use,
 * copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following
 * conditions:
 *
 * The above copyright notice and this permission notice shall be
 * included in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
 * OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
 * HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
 * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
 * OTHER DEALINGS IN THE SOFTWARE.
 */

#include "enumerate_linear_transformations.hpp"

#include <fstream>
#include <iostream>
#include <unordered_map>

#include <boost/dynamic_bitset.hpp>
#include <boost/format.hpp>
#include <boost/graph/graphviz.hpp>
#include <boost/graph/isomorphism.hpp>
#include <boost/graph/vf2_sub_graph_iso.hpp>

#include <igraph/igraph.h>

#include <core/utils/bitset_utils.hpp>
#include <core/utils/graph_utils.hpp>
#include <core/utils/range_utils.hpp>
#include <core/utils/string_utils.hpp>
#include <classical/utils/truth_table_utils.hpp>

#define L(x) { if ( verbose ) { std::cout << x << std::endl; } }

namespace cirkit
{

/******************************************************************************
 * Types                                                                      *
 ******************************************************************************/

struct boolean_square_matrix
{
public:
  boolean_square_matrix( const boost::dynamic_bitset<>& matrix, unsigned dimension )
    : matrix( matrix ),
      dimension( dimension )
  {
  }

  boolean_square_matrix( const std::vector<boost::dynamic_bitset<>>& rows )
  {
    assert( !rows.empty() );

    dimension = rows.front().size();
    matrix.resize( dimension * dimension );

    unsigned index = 0u;

    for ( const auto& r : rows )
    {
      for ( auto c = 0u; c < dimension; ++c )
      {
        matrix[index++] = r[c];
      }
    }
  }

  static boolean_square_matrix identity( unsigned dimension )
  {
    std::vector<boost::dynamic_bitset<>> rows( dimension, boost::dynamic_bitset<>( dimension ) );
    for ( auto c = 0u; c < dimension; ++c )
    {
      rows[c].set( c );
    }
    return boolean_square_matrix( rows );
  }

  inline bool at( unsigned row, unsigned column ) const
  {
    return matrix.test( row * dimension + column );
  }

  inline void set( unsigned row, unsigned column, bool value )
  {
    matrix.set( row * dimension + column, value );
  }

  std::vector<boost::dynamic_bitset<>> get_row_vectors() const
  {
    std::vector<boost::dynamic_bitset<>> rows( dimension, boost::dynamic_bitset<>( dimension ) );

    for ( auto r = 0u; r < dimension; ++r )
    {
      for ( auto c = 0u; c < dimension; ++c )
      {
        rows[r][c] = at( r, c );
      }
    }

    return rows;
  }

  boolean_square_matrix get_submatrix_by_removing( unsigned row, unsigned column ) const
  {
    boost::dynamic_bitset<> subm( ( dimension - 1u ) * ( dimension - 1u ) );

    unsigned index = 0;
    for ( auto r = 0u; r < dimension; ++r )
    {
      if ( r == row ) continue;
      for ( auto c = 0u; c < dimension; ++c )
      {
        if ( c == column ) continue;
        subm[index++] = at( r, c );
      }
    }

    return boolean_square_matrix( subm, dimension - 1u );
  }

  bool determinant()
  {
    if ( dimension == 2u )
    {
      return ( at( 0, 0 ) && at( 1, 1 ) ) != ( at( 0, 1 ) && at( 1, 0 ) );
    }
    else
    {
      bool result = false;
      for ( auto c = 0u; c < dimension; ++c )
      {
        result = result != ( at( 0u, c ) * get_submatrix_by_removing( 0u, c ).determinant() );
      }
      return result;
    }
  }

  bool is_sorted()
  {
    const auto rows = get_row_vectors();

    auto v = 0u;

    for ( const auto& r : rows )
    {
      const auto rv = r.to_ulong();
      if ( rv < v ) return false;
      v = rv;
    }

    return true;
  }

  unsigned row_sum() const
  {
    boost::dynamic_bitset<> sum( dimension );
    for ( const auto& r : get_row_vectors() )
    {
      for ( auto c = 0u; c < dimension; ++c )
      {
        sum[c] = sum[c] != r[c];
      }
    }
    return sum.to_ulong();
  }

  boolean_square_matrix swap( unsigned row1, unsigned row2 ) const
  {
    boolean_square_matrix other = *this;

    for ( auto c = 0u; c < dimension; ++c )
    {
      other.set( row1, c, at( row2, c ) );
      other.set( row2, c, at( row1, c ) );
    }

    return other;
  }

  boolean_square_matrix add_to( unsigned from_row, unsigned to_row ) const
  {
    boolean_square_matrix other = *this;

    for ( auto c = 0u; c < dimension; ++c )
    {
      other.set( to_row, c, at( from_row, c ) != at( to_row, c ) );
    }

    return other;
  }

  boolean_square_matrix sort() const
  {
    auto rows = get_row_vectors();

    std::sort( rows.begin(), rows.end(), []( const boost::dynamic_bitset<>& bs1, const boost::dynamic_bitset<>& bs2 ) { return bs1.to_ulong() < bs2.to_ulong(); } );

    return boolean_square_matrix( rows );
  }

  /* multiply a column vector */
  boost::dynamic_bitset<> multiply( const boost::dynamic_bitset<>& v ) const
  {
    boost::dynamic_bitset<> y( dimension );

    for ( auto r = 0u; r < dimension; ++r )
    {
      auto value = false;
      for ( auto c = 0u; c < dimension; ++c )
      {
        value = value != ( v[c] && at( r, c ) );
      }
      y[r] = value;
    }

    return y;
  }

  /* other can be sorted, but this must be as we want it
     we assume it's possible */
  std::pair<unsigned, unsigned> get_add_to_rows( const boolean_square_matrix& other ) const
  {
    const auto vectors = get_row_vectors();
    const auto ovectors = other.get_row_vectors();

    auto to = 0u, from = 0u;
    boost::dynamic_bitset<> to_row;
    boost::dynamic_bitset<> from_row;

    for ( const auto& r : index( vectors ) )
    {
      const auto it = std::find( ovectors.begin(), ovectors.end(), r.value );
      if ( it == ovectors.end() )
      {
        to = r.index;
        to_row = r.value;
      }
    }

    for ( const auto& r : ovectors )
    {
      const auto it = std::find( vectors.begin(), vectors.end(), r );
      if ( it == vectors.end() )
      {
        from_row = r;
      }
    }

    from_row ^= to_row;
    from = std::distance( vectors.begin(), std::find( vectors.begin(), vectors.end(), from_row ) );

    //std::cout << boost::format( "[i] added %s [%d] to %s [%d]" ) % to_string( from_row ) % from % to_string( to_row ) % to << std::endl;

    return {from, to};
  }

  std::pair<unsigned, unsigned> compute_swap_permutation( const boolean_square_matrix& other ) const
  {
    std::vector<unsigned> perm( 1u << dimension ); /* permutation of this matrix */
    std::vector<unsigned> perm2( 1u << dimension );

    boost::dynamic_bitset<> bs( dimension );

    do
    {
      perm[bs.to_ulong()] = multiply( bs ).to_ulong();

      inc( bs );
    } while ( bs.any() );

    do
    {
      perm2[other.multiply( bs ).to_ulong()] = bs.to_ulong();

      inc( bs );
    } while ( bs.any() );

    unsigned delta = 0u;
    unsigned theta = 0u;

    for ( auto i = 0u; i < perm.size(); ++i )
    {
      const auto j = perm2[perm[i]];
      if ( i < j )
      {
        delta = j - i;
        theta |= ( 1 << i );
      }
    }

    return {delta, theta};
  }

  inline unsigned long value() const
  {
    return matrix.to_ulong();
  }

  friend std::ostream& operator<<( std::ostream& os, const boolean_square_matrix& m )
  {
    for ( auto r = 0u; r < m.dimension; ++r )
    {
      for ( auto c = 0u; c < m.dimension; ++c )
      {
        if ( c > 0u )
        {
          os << " ";
        }
        os << m.at( r, m.dimension - c - 1 );
      }
      os << std::endl;
    }

    return os;
  }

private:
  boost::dynamic_bitset<> matrix;
  unsigned                dimension;
};

/******************************************************************************
 * Private functions                                                          *
 ******************************************************************************/

unsigned factorial( unsigned n )
{
  return ( n == 1u || n == 0u ) ? 1u : factorial( n - 1u ) * n;
}

unsigned compute_order( unsigned n )
{
  unsigned order = 1u;

  for ( unsigned i = 0; i < n; ++i )
  {
    order *= ( 1 << n ) - ( 1 << i );
  }

  return order;
}

int igraph_is_strongly_regular( const igraph_t *graph, igraph_bool_t *res, igraph_integer_t *k, igraph_integer_t *lambda, igraph_integer_t *mu )
{
  int err;
  *lambda = -1;
  *mu = -1;

  /* check g is undirected */
  if ( igraph_is_directed( graph ) ) { return IGRAPH_EINVAL; }

  /* special case if empty */
  if ( igraph_vcount( graph ) == 0u )
  {
    *res = 1;
    return IGRAPH_SUCCESS;
  }

  /* check same degree property */
  igraph_vector_t degrees;
  igraph_vector_init( &degrees, 0 );
  igraph_vs_t vertices;
  igraph_vs_all( &vertices );
  if ( ( err = igraph_degree( graph, &degrees, vertices, IGRAPH_ALL, 0 ) ) != IGRAPH_SUCCESS )
  {
    //return err;
  }
  igraph_vs_destroy( &vertices );

  *k = VECTOR( degrees )[0];
  for ( auto i = 1u; i < igraph_vector_size( &degrees ); ++i )
  {
    if ( VECTOR( degrees )[i] != *k )
    {
      igraph_vector_destroy( &degrees );
      std::cout << "[w] not all vertices have same degree" << std::endl;
      *res = 0;
      return IGRAPH_SUCCESS;
    }
  }

  igraph_vector_destroy( &degrees );

  /* check that neighbors have the same number of neighbors */
  for ( auto node = 0; node < igraph_vcount( graph ); ++node )
  {
    igraph_vector_t neis;
    igraph_vector_init( &neis, 0 );
    if ( ( err = igraph_neighbors( graph, &neis, node, IGRAPH_ALL ) ) != IGRAPH_SUCCESS )
    {
      // return err;
    }

    /* point to the first neighbor */
    auto cur = 0;
    auto conn = 0;
    while ( cur < igraph_vector_size( &neis ) && VECTOR( neis )[cur] < node + 1 ) ++cur;

    for ( auto onode = node + 1; onode < igraph_vcount( graph ); ++onode )
    {
      if ( onode == node ) continue;

      if ( cur < igraph_vector_size( &neis ) && VECTOR( neis )[cur] == onode )
      {
        /* connected */
        conn = 1;
        ++cur;
      }
      else
      {
        /* not connected */
        conn = 0;
      }

      igraph_vector_t oneis;
      igraph_vector_init( &oneis, 0 );
      igraph_neighbors( graph, &oneis, onode, IGRAPH_ALL );

      igraph_vector_t inter;
      igraph_vector_init( &inter, 0 );
      igraph_vector_intersect_sorted( &neis, &oneis, &inter );

      if ( conn )
      {
        if ( *lambda == -1 )
        {
          *lambda = igraph_vector_size( &inter );
        }
        else if ( *lambda != igraph_vector_size( &inter ) )
        {
          igraph_vector_destroy( &inter );
          igraph_vector_destroy( &oneis );
          igraph_vector_destroy( &neis );

          std::cout << "[w] not same lambda" << std::endl;
          *res = 0;
          return IGRAPH_SUCCESS;
        }
      }
      else
      {
        if ( *mu == -1 )
        {
          *mu = igraph_vector_size( &inter );
        }
        else if ( *mu != igraph_vector_size( &inter ) )
        {
          igraph_vector_destroy( &inter );
          igraph_vector_destroy( &oneis );
          igraph_vector_destroy( &neis );

          std::cout << "[w] not same mu" << std::endl;
          *res = 0;
          return IGRAPH_SUCCESS;
        }
      }

      igraph_vector_destroy( &inter );

      igraph_vector_destroy( &oneis );
    }

    igraph_vector_destroy( &neis );
  }

  *res = 1;
  return IGRAPH_SUCCESS;
}

struct enumerate_hamiltonians_is_equivalent
{
  enumerate_hamiltonians_is_equivalent( const std::vector<boolean_square_matrix>& node_to_matrix )
    : node_to_matrix( node_to_matrix )
  {
  }

  bool operator()( const vertex_t<graph_t<>>& vs, const vertex_t<graph_t<>>& vl ) const
  {
    if ( vs == 0u )
    {
      return vl == 0u;
    }
    if ( vs < prefix.size() && prefix[vs] != 0 )
    {
      return node_to_matrix[vl].row_sum() == prefix[vs];
    }
    else
    {
      return true;
    }
  }

private:
  //const std::vector<unsigned> prefix = {1, 2, 4, 3, 6, 7, 5, 3, 0, 0, 0, 0, 0, 0, 5, 0, 0, 0, 0, 0, 0, 7};
  //const std::vector<unsigned> prefix = {1, 2, 4, 3, 6, 7, 5, 3, 2, 4, 7, 1, 5, 6, 5, 4, 6, 1, 3, 7, 2, 7, 4, 1, 3, 5, 2, 6};
  const std::vector<unsigned> prefix = {7, 6, 5, 2, 4, 1, 3, 5, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 3};
  const std::vector<boolean_square_matrix>& node_to_matrix;
};

struct enumerate_hamiltonians_printer
{
public:
  enumerate_hamiltonians_printer( const graph_t<>& sg, const std::vector<boolean_square_matrix>& node_to_matrix )
    : sg( sg ),
      node_to_matrix( node_to_matrix )
  {
  }

  template<typename CorrespondenceMap1To2, typename CorrespondenceMap2To1>
  bool operator()( const CorrespondenceMap1To2& f, const CorrespondenceMap2To1& g ) const
  {
    std::vector<std::vector<unsigned>> parts( 4u, std::vector<unsigned>( 7u ) );

    for ( auto v = 0u; v < boost::num_vertices( sg ); ++v )
    {
      parts[v / 7][v % 7] = node_to_matrix[boost::get( f, v )].row_sum();
    }

    for ( auto& p : parts )
    {
      std::sort( p.begin(), p.end() );
      for ( auto j = 0u; j < 7u; ++j )
      {
        if ( p[j] != j + 1 )
        {
          return true;
        }
      }
    }

    for ( auto v = 0u; v < boost::num_vertices( sg ); ++v )
    {
      std::cout << " " << node_to_matrix[boost::get( f, v )].row_sum();
    }
    std::cout << std::endl;

    for ( auto v = 0u; v < boost::num_vertices( sg ); ++v )
    {
      std::cout << " " << boost::get( f, v );
    }
    std::cout << std::endl;

    for ( auto i = 0u; i < 4u; ++i )
    {
      std::cout << " " << node_to_matrix[boost::get( f, 7 * i )].row_sum();
    }
    std::cout << std::endl << std::endl;

    return true;
  }

private:
  const graph_t<>& sg;
  const std::vector<boolean_square_matrix>& node_to_matrix;
};

void enumerate_hamiltonians( const igraph_t* graph, const std::vector<boolean_square_matrix>& node_to_matrix )
{
  graph_t<> g;
  for ( auto i = 0; i < igraph_vcount( graph ); ++i )
  {
    add_vertex( g );
  }

  igraph_es_t es;
  igraph_es_all( &es, IGRAPH_EDGEORDER_ID );

  igraph_eit_t eit;
  igraph_eit_create( graph, es, &eit );

  while ( !IGRAPH_EIT_END( eit ) )
  {
    const auto eid = IGRAPH_EIT_GET( eit );

    igraph_integer_t from, to;
    igraph_edge( graph, eid, &from, &to );

    add_edge( from, to, g );

    IGRAPH_EIT_NEXT( eit );
  }

  igraph_eit_destroy( &eit );
  igraph_es_destroy( &es );

  const auto rg = ring_graph( boost::num_vertices( g ) );
  boost::vf2_subgraph_mono( rg, g,
                            enumerate_hamiltonians_printer( rg, node_to_matrix ),
                            boost::vertex_order_by_mult( rg ),
                            boost::vertices_equivalent( enumerate_hamiltonians_is_equivalent( node_to_matrix ) ) );
}

/******************************************************************************
 * Public functions                                                           *
 ******************************************************************************/

void enumerate_linear_transformations( unsigned n, const properties::ptr& settings, const properties::ptr& statistics )
{
  /* settings */
  const auto abstract    = get( settings, "abstract",    true );
  const auto hamiltonian = get( settings, "hamiltonian", false );
  const auto hamilenum   = get( settings, "hamilenum",   false );
  const auto frompath    = get( settings, "frompath",    std::string() );
  const auto dotname     = get( settings, "dotname",     std::string() );
  const auto verbose     = get( settings, "verbose",     false );

  std::vector<boolean_square_matrix> node_to_matrix;
  std::unordered_map<unsigned long, unsigned> matrix_to_node;
  std::vector<unsigned> edge_to_operation;

  igraph_t graph, * p_g = &graph;

  auto order = compute_order( n );
  if ( abstract )
  {
    order /= factorial( n );
  }
  L( "[i] order = " << order );
  igraph_empty( p_g, order, IGRAPH_UNDIRECTED );

  boost::dynamic_bitset<> bs( n * n );

  /* create vertices */
  foreach_bitset( n * n, [&matrix_to_node, &node_to_matrix, n, abstract, verbose]( const boost::dynamic_bitset<>& bs ) {
      boolean_square_matrix m( bs, n );
      if ( m.determinant() != 0u && ( !abstract || m.is_sorted() ) )
      {
        matrix_to_node[m.value()] = node_to_matrix.size();
        node_to_matrix.push_back( m );
      }
    } );

  /* create edges */
  for ( size_t v = 0; v < node_to_matrix.size(); ++v )
  {
    const auto& bm = node_to_matrix[v];

    /* adding */
    for ( auto from_row = 0u; from_row < n; ++from_row )
    {
      for ( auto to_row = 0u; to_row < n; ++to_row )
      {
        if ( from_row == to_row ) continue;

        auto other = bm.add_to( from_row, to_row );
        if ( abstract )
        {
          other = other.sort();
        }

        const auto w = matrix_to_node[other.value()];
        if ( w > v )
        {
          igraph_add_edge( p_g, v, w );

          igraph_integer_t eid;
          igraph_get_eid( p_g, &eid, v, w, 0, 0 );
          edge_to_operation.push_back( from_row * 3 + to_row );
        }
      }
    }

    /* swapping */
    if ( !abstract )
    {
      for ( unsigned row2 = 1u; row2 < n; ++row2 )
      {
        for ( unsigned row1 = 0u; row1 < row2; ++row1 )
        {
          const auto other = bm.swap( row1, row2 );
          const auto w = matrix_to_node[other.value()];
          if ( w > v )
          {
            igraph_add_edge( p_g, v, w );
          }
        }
      }
    }
  }

  /* check some properties */
  igraph_bool_t b_connected;
  igraph_is_connected( p_g, &b_connected, IGRAPH_STRONG );
  L( "[i] is strongly connected: " << b_connected );

  igraph_bool_t b_strongly_regular;
  igraph_integer_t k, lambda, mu;
  assert( !igraph_is_strongly_regular( p_g, &b_strongly_regular, &k, &lambda, &mu ) );
  L( "[i] is strongly regular: " << b_strongly_regular );

  /* Hamiltonian path */
  if ( hamiltonian )
  {
    igraph_t rgraph, * p_rg = &rgraph;
    igraph_ring( p_rg, order, 0, 0, 1 );

    igraph_bool_t b_iso;
    igraph_vector_t map;
    igraph_vector_init( &map, 0 );
    igraph_subisomorphic_lad( p_rg, p_g, nullptr, &b_iso, &map, nullptr, 0, 0 );
    L( "[i] found isomorphism: " << b_iso );

    igraph_vector_print( &map );
    igraph_vector_destroy( &map );
    igraph_destroy( p_rg );
  }

  if ( hamilenum )
  {
    enumerate_hamiltonians( p_g, node_to_matrix );
  }

  /* derive delta-swap sequence from path */
  if ( !frompath.empty() )
  {
    std::cout << "[i] derive delta-swap sequence from path: " << frompath << std::endl;
    std::vector<unsigned> path;
    parse_string_list( path, frompath, " " );

    unsigned delta, theta;
    const auto& swaps = tt_store::i().swaps( n );

    auto curm = node_to_matrix[path[0u]];

    std::cout << curm;
    for ( auto i = 0u; i < swaps.size(); ++i )
    {
      auto nextm = curm.swap( swaps[i], swaps[i] + 1 );
      std::tie( delta, theta ) = nextm.compute_swap_permutation( curm );
      std::cout << boost::format( "[i] swap %d with %d (%d, %d)" ) % swaps[i] % ( swaps[i] + 1 ) % delta % theta << std::endl;
      std::cout << nextm;
      curm = nextm;
    }

    auto next = 1u;
    while ( next < path.size() )
    {
      // const auto v = std::min( path[next - 1], path[next] );
      // const auto w = std::max( path[next - 1], path[next] );
      // igraph_integer_t eid;
      // igraph_get_eid( p_g, &eid, v, w, 0, 0 );
      // const auto op = edge_to_operation[eid];
      // const auto to = op % n;
      // const auto from = ( op - to ) / n;
      // std::cout << boost::format( "[i] add %d to %d" ) % from % to << std::endl;

      unsigned from_row, to_row;
      std::tie( from_row, to_row ) = curm.get_add_to_rows( node_to_matrix[path[next]] );

      auto nextm = curm.add_to( from_row, to_row );
      std::tie( delta, theta ) = nextm.compute_swap_permutation( curm );
      std::cout << boost::format( "[i] add %d to %d (%d, %d)" ) % from_row % to_row % delta % theta << std::endl;
      std::cout << nextm;
      curm = nextm;

      for ( auto i = 0u; i < swaps.size(); ++i )
      {
        auto nextm = curm.swap( swaps[i], swaps[i] + 1 );
        std::tie( delta, theta ) = nextm.compute_swap_permutation( curm );
        std::cout << boost::format( "[i] swap %d with %d (%d, %d)" ) % swaps[i] % ( swaps[i] + 1 ) % delta % theta << std::endl;
        std::cout << nextm;
        curm = nextm;
      }
      ++next;
    }
  }

  /* row value inspection */
  if ( false )
  {
    const auto from_sum = 2u;
    const auto to_sum = 3u;
    for ( const auto& bm : node_to_matrix )
    {
      if ( bm.row_sum() == from_sum )
      {
        std::cout << "[i] matrix with row sum " << from_sum << ": " << std::endl << bm;
        //std::cout << "[i] adjacent matrices have row sums:";

        igraph_vector_t neis;
        igraph_vector_init( &neis, 0 );
        igraph_neighbors( p_g, &neis, matrix_to_node[bm.value()], IGRAPH_ALL );

        for ( auto i = 0u; i < igraph_vector_size( &neis ); ++i )
        {
          const auto& om = node_to_matrix[VECTOR( neis )[i]];
          if ( om.row_sum() == to_sum )
          {
            std::cout << "[i] connects to: " << std::endl << om;
          }
        }
        //std::cout << std::endl;
        igraph_vector_destroy( &neis );
      }
    }
  }

  if ( !dotname.empty() )
  {
    auto * fp = fopen( dotname.c_str(), "w" );
    igraph_write_graph_dot( p_g, fp );
    fclose( fp );
  }

  igraph_destroy( p_g );
}

}

// Local Variables:
// c-basic-offset: 2
// eval: (c-set-offset 'substatement-open 0)
// eval: (c-set-offset 'innamespace 0)
// End: