//// DO NOT USE. UNDER CONSTRUCTION.

#include <cstdio>
#include <string>
#include <limits>
#include <set>

#include <boost/algorithm/cxx11/all_of.hpp>
#include <boost/algorithm/cxx11/any_of.hpp>
#include <cctbx/error.h>
// iostream is included by this one
#include <cctbx/sgtbx/direct_space_asu/proto/asymmetric_unit.h>

#include "cctbx/maptbx/skeletons.h"

namespace cctbx { namespace maptbx
{

inline bool any_lt(const int3_t &a, const int3_t &b)
{
  for(short j=0; j<3; ++j)
    if( a[j]<b[j] )
      return true;
  return false;
}

inline bool any_lt(const int3_t &a, int b)
{
  for(short j=0; j<3; ++j)
    if( a[j]<b )
      return true;
  return false;
}

inline bool any_ge(const int3_t &a, const int3_t &b)
{
  for(short j=0; j<3; ++j)
    if( a[j]>=b[j] )
      return true;
  return false;
}

// faster version of small size std::set ?
template<typename T, unsigned short N=1000> class array_as_set :
  private scitbx::af::small<T,N>
{
public:
  typedef scitbx::af::small<T,N> buf_t;
  using buf_t::begin;
  using buf_t::end;
  using buf_t::empty;
  using buf_t::size;
  using typename buf_t::iterator;
  using typename buf_t::const_iterator;
  void insert(const T &v)
  {
    iterator l=std::lower_bound(this->begin(), this->end(),v);
    // *l >= v;
    if( l!=this->end() )
    {
      if( *l != v )
        this->buf_t::insert(l,v);
    }
    else
        this->buf_t::insert(l,v);
    //CCTBX_ASSERT( !this->empty() );
    //CCTBX_ASSERT( std::is_sorted(this->begin(), this->end()) );
    //CCTBX_ASSERT(std::adjacent_find(this->begin(),this->end())==this->end());
  }
};

skeleton swanson(const cctbx::maptbx::asymmetric_map &amap, double sigma)
{
  double mean=0., esd=0., mx=-9.E200;
  const auto &map_data = amap.data().const_ref();
  for(std::size_t ii=0; ii<map_data.size(); ++ii)
  {
    //! @todo discard points outside asu
    double m = map_data[ii];
    mean += m;
    esd += m*m;
    if( m > mx )
      mx = m;
  }
  mean /= map_data.size();
  esd = esd/map_data.size() - mean*mean;
  esd = std::sqrt(esd);
  double mapcutoff = mean + esd * sigma;

  std::vector<xyzm_t> xyzm;
  xyzm.reserve(10000);
  //! @todo code duplication: mmtbx::masks::atom_mask::mask_asu
  std::size_t inside=0, tot=0;
  const auto &opt_asu = amap.optimized_asu();
  const double *md = map_data.begin();
  for(auto i3=amap.grid_begin(); !i3.over(); i3.incr(), ++md)
  {
    int3_t p=i3();
    double m = *md;
    if( m>mapcutoff )
    {
      ++tot;
      if( opt_asu.where_is(p) != 0 ) // inside or on the face
      {
        ++inside;
        xyzm_t x(p,m);
        xyzm.push_back( x );
      }
    }
  }
  std::sort(xyzm.begin(), xyzm.end(),
    [](const xyzm_t &a, const xyzm_t &b) {return (get<1>(a)) > (get<1>(b));}
  );

  CCTBX_ASSERT( get<1>(xyzm.front()) == mx );
  marks_t marks(map_data.accessor(), 0UL); // 0 means no mark
  std::size_t nmarks=0, nbonds=0, min_count = 0, grows_count=0, join_count=0;
  const unsigned short cube_size = 26; // 3*3 + (3*3-1) + (3*3) == 3^3 - 1
  typedef int3_t i3t;
  int3_t cube[cube_size] = {
    i3t(-1,-1,-1),  i3t(0,-1,-1),  i3t(1,-1,-1),  i3t(-1,0,-1),  i3t(0,0,-1),
    i3t(1,0,-1),    i3t(-1,1,-1),  i3t(0,1,-1),   i3t(1,1,-1),   i3t(-1,-1,0),
    i3t(0,-1,0),    i3t(1,-1,0),   i3t(-1,0,0),   i3t(1,0,0),    i3t(-1,1,0),
    i3t(0,1,0),     i3t(1,1,0),    i3t(-1,-1,1),  i3t(0,-1,1),   i3t(1,-1,1),
    i3t(-1,0,1),    i3t(0,0,1),    i3t(1,0,1),    i3t(-1,1,1),   i3t(0,1,1),
    i3t(1,1,1)
  };

  auto ibox_min = amap.box_begin();
  auto ibox_max = amap.box_end();
  skeleton result;
  result.maximums.reserve(1000);
  for(std::size_t jj=0; jj<xyzm.size(); ++jj)
  {
    // typedef std::set<std::size_t> feature_set_t;
    typedef array_as_set<std::size_t> feature_set_t;
    feature_set_t featureset;
    int3_t x = get<0>(xyzm[jj]);
    for(unsigned short ic=0; ic<cube_size; ++ic)
    {
      int3_t neighbor = x + cube[ic]; //! @todo check if in cell ?
      if( any_lt(neighbor,ibox_min) || any_ge(neighbor,ibox_max) )
        continue;
      unsigned mark = marks(neighbor);
      if( mark != 0U ) // 0 is not a mark
        featureset.insert(mark);
    }

    const unsigned fssize = featureset.size();
    if( fssize==0U ) // all_eq(featureset, 0) )
    {
      // 2a: new maximum
      ++nmarks;
      marks(x) = nmarks;
      result.maximums.push_back(x);
    }
    else if( fssize == 1 )
    {
      // 2b: part of the growing nodule
      std::size_t mark = *featureset.begin();
      CCTBX_ASSERT( mark != 0 );
      marks(x) = mark;
      ++grows_count;
    }
    else if( fssize == cube_size )
    {
      // 2d: local minimum will not be in the neighborhood of any point
      ++min_count;
      marks(x) = *featureset.begin(); // highest nodule
    }
    else if( fssize>1 )
    {
      // 2c: two or more nodules merging
      ++join_count;
      marks(x) = *featureset.begin(); // highest nodule
      for(feature_set_t::const_iterator i=featureset.begin();
          i!=featureset.end(); ++i)
      {
        std::size_t fi = *i;
        CCTBX_ASSERT( fi>0U && fi<=nmarks );
        feature_set_t::const_iterator j=i;
        ++j;
        for( ; j!=featureset.end(); ++j)
        {
          std::size_t fj = *j;
          CCTBX_ASSERT( fj>0U && fj<=nmarks );
          CCTBX_ASSERT( fi<fj );
          join_t join;
          join.ilt = fi;
          join.igt = fj;
          result.joins.insert(join); // inserts only new joins
        }
      }
    }
  }
  CCTBX_ASSERT( nmarks + grows_count + join_count + min_count == xyzm.size() );
  result.min_count = min_count;
  result.grows_count = grows_count;
  result.join_count = join_count;
  result.marks = std::move(marks);
  result.mapcutoff = mapcutoff;
  return result;
}

std::vector<std::size_t> find_clusters(const skeleton &skelet)
{
  std::size_t nmaxs = skelet.maximums.size(), mol_id=0;
  std::vector<std::size_t> atoms(nmaxs,0);
  for(joins_t::const_iterator i=skelet.joins.begin(); i!=skelet.joins.end();
    ++i)
  {
    const join_t &join = *i;
    long i_cat = join.ilt;
    long i_an = join.igt;
    CCTBX_ASSERT( i_cat < i_an && i_cat>0 && i_an>0 && i_cat<=nmaxs
      && i_an<=nmaxs );
    if( atoms[i_cat-1U]==0 && atoms[i_an-1U]==0 )
    {
      ++mol_id;
      atoms[i_cat] = mol_id;
      atoms[i_an] = mol_id;
    }
    else
    {
      if(atoms[i_cat]!=0 && atoms[i_an]!=0 && atoms[i_cat]!=atoms[i_an])
      {
        std::size_t imn = std::min(atoms[i_cat],atoms[i_an]);
        atoms[i_cat] = imn;
        atoms[i_an] = imn;
      }
      else
      {
        if( atoms[i_cat]!=0 )
          atoms[i_an] = atoms[i_cat];
        if( atoms[i_an]!=0 )
          atoms[i_cat] = atoms[i_an];
      }
    }
  }
  std::size_t imn = 0, imx=0;
  do
  {
    imx = *std::max_element(atoms.begin(), atoms.end());
    for(joins_t::const_iterator i=skelet.joins.begin(); i!=skelet.joins.end();
      ++i)
    {
      const join_t &join = *i;
      long i_cat = join.ilt;
      long i_an = join.igt;
      if( atoms[i_cat]==0 || atoms[i_an]==0 )
        throw std::logic_error("disaster");
      if( atoms[i_cat] != atoms[i_an] )
      {
        imn = std::min(atoms[i_cat],atoms[i_an]);
        atoms[i_cat] = imn;
        atoms[i_an] = imn;
      }
    }
  }while( imn!=0 );
  imx = *std::max_element(atoms.begin(), atoms.end());

  std::vector< array<std::size_t,2> > molecules;
  molecules.reserve(imx);
  std::size_t imxsz = 0;
  for(std::size_t i=0; i<imx; ++i)
  {
    imn = std::count(atoms.begin(), atoms.end(), i);
    array< std::size_t,2 > m = {imn,0};
    molecules.push_back(m);
    if( imn!=0 )
      ++mol_id;
    if( imn>imxsz )
    {
      imx = i;
      imxsz = imn;
    }
  }
  if( imx!=0 )
  {
    mol_id = 1;
    for(std::size_t i=0; i<molecules.size(); ++i)
    {
      get<1>(molecules[i]) = mol_id;
      ++mol_id;
    }
  }
  std::vector<std::size_t> molids;
  for(std::size_t i=0; i<atoms.size(); ++i)
  {
    if( atoms[i] == 0 )
      continue; // ?
    molids.push_back( get<1>(molecules[atoms[i]]) );
  }
  return molids;
}

}} // cctbx::maptbx


//#include <boost/program_options.hpp>
//namespace po = boost::program_options;

//#include <vector>
#include <iostream>
#include <complex>

//#include <stdio.h>

#include <fftw3.h>

//#include <tiffio.h>

#define cimg_use_tiff
#include <CImg.h>
using namespace cimg_library;

#ifdef _WIN32
#define _USE_MATH_DEFINES
#define rint(fp) (int)((fp) >= 0 ? (fp) + 0.5 : (fp) - 0.5)
#endif

#include <math.h>

extern int save_tiff(TIFF *tif, const unsigned int directory, int colind, const int nwaves, int width, int height, float * buffer);
extern int load_tiff(TIFF *const tif, const unsigned int directory, const unsigned colind, float *const buffer);

float fitparabola(float a1, float a2, float a3);
void determine_center_and_background(float *stack5phases, int nx, int ny, int nz, float *xc, float *yc, float *zc, float *background);
void shift_center(std::complex<float> *bands, int nx, int ny, int nz, float xc, float yc, float zc);
void cleanup(std::complex<float> *otfkxkz, int nx, int nz, float dkr, float dkz, float linespacing, int lambdanm, int twolens, float NA, float NIMM);
void radialft(std::complex<float> *bands, int nx, int ny, int nz, std::complex<float> *avg);

bool fixorigin(std::complex<float> *otfkxkz, int nx, int nz, int kx2);
void rescale(std::complex<float> *otfkxkz, int nx, int nz);

#ifdef _WIN32
#ifdef RADIALFT_IMPORT
  #define RADIALFT_API __declspec( dllimport )
#else
  #define RADIALFT_API __declspec( dllexport )
#endif
#else
  #define RADIALFT_API
#endif

extern "C" {
RADIALFT_API int makeOTF(const char *const ifiles, const char *const ofiles,
      int lambdanm = 520, float dz = 0.102, int interpkr = 10,
      bool bUserBackground = false, float background = 90,
      float NA = 1.25, float NIMM = 1.3, float dr = 0.102,
      int krmax = 0, bool bDoCleanup = false);
}

//   ("na", po::value<float>(&NA)->default_value(1.25), "NA of detection objective")
//   ("nimm", po::value<float>(&NIMM)->default_value(1.3), "refractive index of immersion medium")
//   ("xyres", po::value<float>(&dr)->default_value(.102), "x-y pixel size")
//   ("zres", po::value<float>(&dz)->default_value(.102), "z pixel size")
//   ("wavelength", po::value<int>(&lambdanm)->default_value(520), "emission wavelength in nm")
//   ("fixorigin", po::value<int>(&interpkr)->default_value(10),
//    "for all kz, extrapolate using pixels kr=1 to this pixel to get value for kr=0")
//   ("krmax", po::value<int>(&krmax)->default_value(0),
//    "pixels outside this limit will be zeroed (overwriting estimated value from NA and NIMM)")
//   ("nocleanup", po::bool_switch(&bDoCleanup)->implicit_value(false), "elect not to do clean-up outside OTF support")
//   ("background", po::value<float>(&background), "use user-supplied background instead of the estimated")
//   ("input-file", po::value<std::string>(&ifiles)->required(), "input file")
//   ("output-file", po::value<std::string>(&ofiles)->required(), "output file")

std::string ifiles, ofiles;
int nx, ny, nz, nxy;
int i, j, z;
float dkr, dkz, background, estBackground, xcofm, ycofm, zcofm;
float *floatimage;
std::complex<float> *bands, *avg_output;
fftwf_plan rfftplan3d;

// float dr = 0.102;
// float dz = 0.102;
// int interpkr = 10;
// float NA = 1.25;
// float NIMM = 1.3;
// int lambdanm = 520;
// int krmax = 0;
// bool bDoCleanup = true;
// bool bUserBackground = false;

int makeOTF(const char *const ifiles, const char *const ofiles, int lambdanm,
    float dz, int interpkr, bool bUserBackground, float background,
    float NA, float NIMM, float dr, int krmax, bool bDoCleanup)
{
  printf("called");

  TIFFSetWarningHandler(NULL);

  CImg<> rawtiff(ifiles);

  nz = rawtiff.depth();
  ny = rawtiff.height();
  nx = rawtiff.width();

  printf("nx=%d, ny=%d, nz=%d\n", nx, ny, nz);

  dkr = 1/(ny*dr);
  dkz = 1/(nz*dz);

  nxy=(nx+2)*ny;

  floatimage = (float *) malloc(nxy*nz*sizeof(float));
  bands = (std::complex<float> *) floatimage;

 // printf("Reading data...\n\n");

  for(z=0; z<nz; z++)
    for (i=0; i<ny; i++) {
      for (j=0; j<nx; j++)
        floatimage[z*nxy+i*(nx+2)+j] = rawtiff(j, i, z);
    }

  /* Before FFT, estimate bead center position */
  determine_center_and_background(floatimage, nx, ny, nz, &xcofm, &ycofm, &zcofm, &estBackground);

  printf("Center of mass is (%.3f, %.3f, %.3f)\n", xcofm, ycofm, zcofm);

  if (!bUserBackground)
    background = estBackground;

  printf("Background is %.3f\n", background);

  for(z=0; z<nz; z++) {
    for(i=0;i<ny;i++)
      for(j=0;j<nx;j++)
        floatimage[z*nxy + i*(nx+2) + j] -= background;
  }

  rfftplan3d = fftwf_plan_dft_r2c_3d(nz, ny, nx, floatimage,
                                     (fftwf_complex *) floatimage, FFTW_ESTIMATE);

  //printf("Before fft\n");

  fftwf_execute_dft_r2c(rfftplan3d, floatimage, (fftwf_complex *) floatimage);

  fftwf_destroy_plan(rfftplan3d);

  //printf("After fft\n\n");

  /* modify the phase of bands, so that it corresponds to FFT of a bead at origin */
  //printf("Shifting center...\n");

  shift_center(bands, nx, ny, nz, xcofm, ycofm, zcofm);

  CImg<> output_tiff(nz*2, nx/2+1, 1, 1, 0.f);
  avg_output = (std::complex<float> *) output_tiff.data();

  radialft(bands, nx, ny, nz, avg_output);

  if (bDoCleanup)
    cleanup(avg_output, nx, nz, dkr, dkz, 1.0, lambdanm, krmax, NA, NIMM);

  if (interpkr > 0)
    try {
    while (!fixorigin(avg_output, nx, nz, interpkr)) {
      interpkr --;
      if (interpkr < 4)
        throw std::runtime_error("#pixels < 4 used in kr=0 extrapolation");
    }}
  catch (std::exception &e) {
    std::cout << "\n!!Error occurred: " << e.what() << std::endl;
    return 1;
  }

//  printf("%d\n", interpkr);
  rescale(avg_output, nx, nz);

  /* For side bands, combine bandre's and bandim's into bandplus */
  /* Shouldn't this be done later on the averaged bands? */

  output_tiff.save_tiff(ofiles);

  return 0;
}

/*  locate peak pixel to subpixel accuracy by fitting parabolas  */
void determine_center_and_background(float *stack3d, int nx, int ny, int nz, float *xc, float *yc, float *zc, float *background)
{
  int i, j, k, maxi, maxj, maxk, ind, nxy2, infocus_sec;
  int iminus, iplus, jminus, jplus, kminus, kplus;
  float maxval, reval, valminus, valplus;
  double sum;

  //printf("In determine_center_and_background()\n");
  nxy2 = (nx+2)*ny;

  /* Search for the peak pixel */
  /* Be aware that stack3d is of dimension (nx+2)xnyxnz */
  maxval=0.0;
  for(k=0;k<nz;k++)
    for(i=0;i<ny;i++)
      for(j=0;j<nx;j++) {
    ind=k*nxy2+i*(nx+2)+j;
    reval=stack3d[ind];
    if( reval > maxval ) {
      maxval = reval;
      maxi=i; maxj=j;
      maxk=k;
    }
      }

  iminus = maxi-1; iplus = maxi+1;
  if( iminus<0 ) iminus+=ny;
  if( iplus>=ny ) iplus-=ny;
  jminus = maxj-1; jplus = maxj+1;
  if( jminus<0 ) jminus+=nx;
  if( jplus>=nx ) jplus-=nx;
  kminus = maxk-1; kplus = maxk+1;
  if( kminus<0 ) kminus+=nz;
  if( kplus>=nz ) kplus-=nz;

  valminus = stack3d[kminus*nxy2+maxi*(nx+2)+maxj];
  valplus  = stack3d[kplus *nxy2+maxi*(nx+2)+maxj];
  *zc = maxk + fitparabola(valminus, maxval, valplus);

  *zc += 0.6;

  valminus = stack3d[maxk*nxy2+iminus*(nx+2)+maxj];
  valplus  = stack3d[maxk*nxy2+iplus *(nx+2)+maxj];
  *yc = maxi + fitparabola(valminus, maxval, valplus);

  valminus = stack3d[maxk*nxy2+maxi*(nx+2)+jminus];
  valplus  = stack3d[maxk*nxy2+maxi*(nx+2)+jplus];
  *xc = maxj + fitparabola(valminus, maxval, valplus);

  sum = 0;
  infocus_sec = floor(*zc);
  for (i=0; i<*yc-20; i++)
    for (j=0; j<nx; j++)
    sum += stack3d[infocus_sec*nxy2 + i*(nx+2) + j];
  *background = sum / ((*yc-20)*nx);
}

/***************************** fitparabola **********************************/
/*     Fits a parabola to the three points (-1,a1), (0,a2), and (1,a3).     */
/*     Returns the x-value of the max (or min) of the parabola.             */
/****************************************************************************/

float fitparabola( float a1, float a2, float a3 )
{
 float slope,curve,peak;

 slope = 0.5* (a3-a1);         /* the slope at (x=0). */
 curve = (a3+a1) - 2*a2;       /* (a3-a2)-(a2-a1). The change in slope per unit of x. */
 if( curve == 0 )
 {
   printf("no peak: a1=%f, a2=%f, a3=%f, slope=%f, curvature=%f\n",a1,a2,a3,slope,curve);
   return( 0.0 );
 }
 peak = -slope/curve;          /* the x value where slope = 0  */
 if( peak>1.5 || peak<-1.5 )
 {
   printf("bad peak position: a1=%f, a2=%f, a3=%f, slope=%f, curvature=%f, peak=%f\n",a1,a2,a3,slope,curve,peak);
   return( 0.0 );
 }
 return( peak );
}

/* To get rid of checkerboard effect in the OTF bands */
/* (xc, yc, zc) is the estimated center of the point source, which in most cases is the bead */
/* Converted from Fortran code. kz is treated differently than kx and ky. don't know why */
void shift_center(std::complex<float> *bands, int nx, int ny, int nz, float xc, float yc, float zc)
{
  int kin, iin, jin, indin, nxy, kz, kx, ky, kycent, kxcent, kzcent;
  std::complex<float> exp_iphi;
  float phi1, phi2, phi, dphiz, dphiy, dphix;

  kycent = ny/2;
  kxcent = nx/2;
  kzcent = nz/2;
  nxy = (nx/2+1)*ny;

  dphiz = 2*M_PI*zc/nz;
  dphiy = 2*M_PI*yc/ny;
  dphix = 2*M_PI*xc/nx;

  for (kin=0; kin<nz; kin++) {    /* the origin of Fourier space is at (0,0) */
    kz = kin;
    if (kz>kzcent) kz -= nz;
    phi1 = dphiz*kz;      /* first part of phi */
    for (iin=0; iin<ny; iin++) {
      ky = iin;
      if (iin>kycent) ky -= ny;
      phi2 = dphiy*ky;   /* second part of phi */
      for (jin=0; jin<kxcent+1; jin++) {
        kx = jin;
        indin = kin*nxy+iin*(nx/2+1)+jin;
        phi = phi1+phi2+dphix*kx;  /* third part of phi */
        /* kz part of Phi has a minus sign, I don't know why. */
        exp_iphi = std::complex<float> (cos(phi), sin(phi));
        bands[indin] = bands[indin] * exp_iphi;
      }
    }
  }
}

void radialft(std::complex<float> *band, int nx, int ny, int nz, std::complex<float> *avg_output)
{
  int kin, iin, jin, indin, indout, indout_conj, kz, kx, ky, kycent, kxcent, kzcent;
  int *count, nxz, nxy;
  float rdist;

  printf("In radialft()\n");
  kycent = ny/2;
  kxcent = nx/2;
  kzcent = nz/2;
  nxy = (nx/2+1)*ny;
  nxz = (nx/2+1)*nz;

  count = (int *) calloc(nxz, sizeof(int));

  if (!count) {
    printf("No memory availale in radialft()\n");
    exit(-1);
  }

  for (kin=0; kin<nz; kin++) {
    kz = kin;
    if (kin>kzcent) kz -= nz;
    for (iin=0; iin<ny; iin++) {
      ky = iin;
      if (iin>kycent) ky -= ny;
      for (jin=0; jin<kxcent+1; jin++) {
        kx = jin;
        rdist = sqrt(kx*kx+ky*ky);
        if (rdist < nx/2+1) {
          indin = kin*nxy+iin*(nx/2+1)+jin;
          indout = rint(rdist)*nz+kin;
          if (indout < nxz) {
            avg_output[indout] += band[indin];
            count[indout] ++;
          }
          // printf("kz=%d, ky=%d, kx=%d, indout=%d\n", kz, ky, kx, indout);
        }
      }
    }
  }

  for (indout=0; indout<nxz; indout++) {
    if (count[indout]>0) {
      avg_output[indout] /= count[indout];
    }
  }

  /* Then complete the rotational averaging and scaling*/
  for (kx=0; kx<nx/2+1; kx++) {
    indout = kx*nz+0;
    avg_output[indout] = std::complex<float>(avg_output[indout].real(), 0);
    for (kz=1; kz<=nz/2; kz++) {
      indout = kx*nz+kz;
      indout_conj = kx*nz + (nz-kz);
      avg_output[indout] = (avg_output[indout] + conj(avg_output[indout_conj])) / 2.f;
      avg_output[indout_conj] = conj(avg_output[indout]);
    }
  }
  free(count);
}

void cleanup(std::complex<float> *otfkxkz, int nx, int nz, float dkr, float dkz, float linespacing, int lamdanm, int krmax_user, float NA, float NIMM)
{
  int ix, iz, kzstart, kzend, icleanup=nx/2+1;
  float lamda, sinalpha, cosalpha, kr, krmax, beta, kzedge;

  lamda = lamdanm * 0.001;
  sinalpha = NA/NIMM;
  cosalpha = cos(asin(sinalpha));
  krmax = 2*NA/lamda;
  if (krmax_user*dkr<krmax && krmax_user!=0)
    krmax = krmax_user*dkr;

  printf("krmax=%f, lambda=%f\n", krmax, lamda);
  for (ix=0; ix<icleanup; ix++) {
    kr = ix * dkr;
    if ( kr <= krmax ) {
      beta = asin( ( NA - kr*lamda ) /NIMM );
      kzedge = (NIMM/lamda) * ( cos(beta) - cosalpha );
      /* kzstart = floor((kzedge/dkz) + 1.999); */ /* In fortran, it's 2.999 */
      kzstart = rint((kzedge/dkz) + 1);
      kzend = nz - kzstart;
      for (iz=kzstart; iz<=kzend; iz++)
        otfkxkz[ix*nz+iz] = 0;
    }
    else {   /* outside of lateral resolution limit */
      for (iz=0; iz<nz; iz++)
        otfkxkz[ix*nz+iz] = 0;
    }
  }
}

bool fixorigin(std::complex<float> *otfkxkz, int nx, int nz, int kx2)
{
  // linear fit the value at kx=0 using kx in [1, kx2]
  double mean_kx = (kx2+1)/2.; // the mean of [1, 2, ..., n] is (n+1)/2

  // printf("In fixorigin(), kx2=%d\n", nx, nz, kx2);

  for (int z=0; z<nz; z++) {
    std::complex<double> mean_val=0;
    std::complex<double> slope_numerator=0;
    std::complex<double> slope_denominat=0;

    for (int x=1; x<=kx2; x++)
      mean_val += otfkxkz[x*nz + z];

    mean_val /= kx2;
    for (int x=1; x<=kx2; x++) {
      std::complex<double> complexval = otfkxkz [x*nz+z];
      slope_numerator += (x - mean_kx) * (complexval - mean_val);
      slope_denominat += (x - mean_kx) * (x - mean_kx);
    }
    std::complex<double> slope = slope_numerator / slope_denominat;
    otfkxkz[z] = mean_val - slope * mean_kx;  // intercept at kx=0
    if (z==0 && std::abs(otfkxkz[z]) <= std::abs(otfkxkz[nz+z])) {
      return false; // indicating kx2 may be too large
    }
  }
  return true;
}

void rescale(std::complex<float> *otfkxkz, int nx, int nz)
{
  int nxz, ind;
  float valmax=0, mag, scalefactor;

  nxz = (nx/2+1)*nz;
  for (ind=0; ind<nxz; ind++) {
    mag = abs(otfkxkz[ind]);
    if (mag > valmax)
      valmax = mag;
  }
  scalefactor = 1/valmax;

  for (ind=0; ind<nxz; ind++) {
    otfkxkz[ind] *= scalefactor;
  }
}


/* Copyright (C) 2012-2020 IBM Corp.
 * This program is Licensed under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *   http://www.apache.org/licenses/LICENSE-2.0
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License. See accompanying LICENSE file.
 */
// debugging.cpp - debugging utilities
#include <NTL/xdouble.h>
#include <helib/debugging.h>
#include <helib/norms.h>
#include <helib/Context.h>
#include <helib/Ctxt.h>
#include <helib/EncryptedArray.h>
//#include <helib/powerful.h>

namespace helib {

SecKey* dbgKey = nullptr;
std::shared_ptr<const EncryptedArray> dbgEa = nullptr;
NTL::ZZX dbg_ptxt;

// return the ratio between the real noise <sk,ct> and the estimated one
double realToEstimatedNoise(const Ctxt& ctxt, const SecKey& sk)
{
  NTL::xdouble noiseEst = ctxt.getNoiseBound();
  if (ctxt.isCKKS())
    noiseEst += ctxt.getRatFactor() * ctxt.getPtxtMag();
  NTL::xdouble actualNoise = embeddingLargestCoeff(ctxt, sk);

  return NTL::conv<double>(actualNoise / noiseEst);
}

double log2_realToEstimatedNoise(const Ctxt& ctxt, const SecKey& sk)
{
  NTL::xdouble noiseEst = ctxt.getNoiseBound();
  if (ctxt.isCKKS())
    noiseEst += ctxt.getRatFactor() * ctxt.getPtxtMag();
  NTL::xdouble actualNoise = embeddingLargestCoeff(ctxt, sk);

  return log(actualNoise / noiseEst) / log(2.0);
}

// check that real-to-estimated ratio is not too large, print warning otherwise
void checkNoise(const Ctxt& ctxt,
                const SecKey& sk,
                const std::string& msg,
                double thresh)
{
  double ratio;
  if ((ratio = realToEstimatedNoise(ctxt, sk)) > thresh) {
    std::cerr << "\n*** too much noise: " << msg << ": " << ratio << "\n";
  }
}

// Decrypt and find the l-infinity norm of the result in canonical embedding
NTL::xdouble embeddingLargestCoeff(const Ctxt& ctxt, const SecKey& sk)
{
  const Context& context = ctxt.getContext();
  NTL::ZZX p, pp;
  sk.Decrypt(p, ctxt, pp);
  return embeddingLargestCoeff(pp, context.zMStar);
}

void decryptAndPrint(std::ostream& s,
                     const Ctxt& ctxt,
                     const SecKey& sk,
                     const EncryptedArray& ea,
                     long flags)
{
  const Context& context = ctxt.getContext();
  std::vector<NTL::ZZX> ptxt;
  NTL::ZZX p, pp;
  sk.Decrypt(p, ctxt, pp);

  NTL::xdouble modulus = NTL::xexp(context.logOfProduct(ctxt.getPrimeSet()));
  NTL::xdouble actualNoise =
      embeddingLargestCoeff(pp, ctxt.getContext().zMStar);
  NTL::xdouble noiseEst = ctxt.getNoiseBound();
  if (ctxt.isCKKS())
    noiseEst += ctxt.getRatFactor() * ctxt.getPtxtMag();

  s << "plaintext space mod " << ctxt.getPtxtSpace()
    << ", bitCapacity=" << ctxt.bitCapacity() << ", \n           |noise|=q*"
    << (actualNoise / modulus) << ", |noiseBound|=q*" << (noiseEst / modulus);
  if (ctxt.isCKKS()) {
    s << ", \n           ratFactor=" << ctxt.getRatFactor()
      << ", ptxtMag=" << ctxt.getPtxtMag()
      << ", realMag=" << (actualNoise / ctxt.getRatFactor());
  }
  s << std::endl;

  if (flags & FLAG_PRINT_ZZX) {
    s << "   before mod-p reduction=";
    printZZX(s, pp) << std::endl;
  }
  if (flags & FLAG_PRINT_POLY) {
    s << "   after mod-p reduction=";
    printZZX(s, p) << std::endl;
  }
  if (flags & FLAG_PRINT_VEC) { // decode to a vector of ZZX
    ea.decode(ptxt, p);
    if (ea.getAlMod().getTag() == PA_zz_p_tag &&
        ctxt.getPtxtSpace() != ea.getAlMod().getPPowR()) {
      long g = NTL::GCD(ctxt.getPtxtSpace(), ea.getAlMod().getPPowR());
      for (long i = 0; i < ea.size(); i++)
        PolyRed(ptxt[i], g, true);
    }
    s << "   decoded to ";
    if (deg(p) < 40) // just pring the whole thing
      s << ptxt << std::endl;
    else if (ptxt.size() == 1) // a single slot
      printZZX(s, ptxt[0]) << std::endl;
    else { // print first and last slots
      printZZX(s, ptxt[0], 20) << "--";
      printZZX(s, ptxt[ptxt.size() - 1], 20) << std::endl;
    }
  } else if (flags & FLAG_PRINT_DVEC) { // decode to a vector of doubles
    const EncryptedArrayCx& eacx = ea.getCx();
    std::vector<double> v;
    eacx.decrypt(ctxt, sk, v);
    printVec(s << "           ", v, 20) << std::endl;
  } else if (flags & FLAG_PRINT_XVEC) { // decode to a vector of complex
    const EncryptedArrayCx& eacx = ea.getCx();
    std::vector<cx_double> v;
    eacx.decrypt(ctxt, sk, v);
    printVec(s << "           ", v, 20) << std::endl;
  }
}

bool decryptAndCompare(const Ctxt& ctxt,
                       const SecKey& sk,
                       const EncryptedArray& ea,
                       const PlaintextArray& pa)
{
  PlaintextArray ppa(ea);
  ea.decrypt(ctxt, sk, ppa);

  return equals(ea, pa, ppa);
}

// Compute decryption with.without mod-q on a vector of ZZX'es,
// useful when debugging bootstrapping (after "raw mod-switch")
void rawDecrypt(NTL::ZZX& plaintxt,
                const std::vector<NTL::ZZX>& zzParts,
                const DoubleCRT& sKey,
                long q)
{
  const Context& context = sKey.getContext();

  // Set to zzParts[0] + sKey * zzParts[1] "over the integers"
  DoubleCRT ptxt = sKey;
  ptxt *= zzParts[1];
  ptxt += zzParts[0];

  // convert to coefficient representation
  ptxt.toPoly(plaintxt);

  if (q > 1)
    PolyRed(plaintxt, q, false /*reduce to [-q/2,1/2]*/);
}

void CheckCtxt(const Ctxt& c, const char* label)
{
  std::cerr << "  " << label
            << ", log2(modulus/noise)=" << (-c.log_of_ratio() / log(2.0))
            << ", p^r=" << c.getPtxtSpace();

  if (dbgKey) {
    double ratio = log2_realToEstimatedNoise(c, *dbgKey);
    std::cerr << ", log2(noise/bound)=" << ratio;
    if (ratio > 0)
      std::cerr << " BAD-BOUND";
  }

#if 0
  // This is not really a useful test

  if (dbgKey && c.getContext().isBootstrappable()) {
    Ctxt c1(c);
    //c1.dropSmallAndSpecialPrimes();

    const Context& context = c1.getContext();
    const RecryptData& rcData = context.rcData;
    const PAlgebra& palg = context.zMStar;

    NTL::ZZX p, pp;
    dbgKey->Decrypt(p, c1, pp);
    NTL::Vec<NTL::ZZ> powerful;
    rcData.p2dConv->ZZXtoPowerful(powerful, pp);

    NTL::ZZ q;
    q = context.productOfPrimes(c1.getPrimeSet());
    vecRed(powerful, powerful, q, false);

    NTL::ZZX pp_alt;
    rcData.p2dConv->powerfulToZZX(pp_alt, powerful);

    NTL::xdouble max_coeff = NTL::conv<NTL::xdouble>(largestCoeff(pp));
    NTL::xdouble max_pwrfl = NTL::conv<NTL::xdouble>(largestCoeff(powerful));
    NTL::xdouble max_canon = embeddingLargestCoeff(pp_alt, palg);
    double ratio = log(max_pwrfl/max_canon)/log(2.0);

    //cerr << ", max_coeff=" << max_coeff;
    //cerr << ", max_pwrfl=" << max_pwrfl;
    //cerr << ", max_canon=" << max_canon;
    std::cerr << ", log2(max_pwrfl/max_canon)=" << ratio;
    if (ratio > 0) std::cerr << " BAD-BOUND";
  }
#endif

  std::cerr << std::endl;
}

} // namespace helib


/*
    All rights reserved.

    Redistribution and use in source and binary forms, with or without
    modification, are permitted provided that the following conditions are met:
        * Redistributions of source code must retain the above copyright
        notice, this list of conditions and the following disclaimer.
        * Redistributions in binary form must reproduce the above copyright
        notice, this list of conditions and the following disclaimer in the
        documentation and/or other materials provided with the distribution.
        * Neither the name of the copyright holder nor the
        names of its contributors may be used to endorse or promote products
        derived from this software without specific prior written permission.

    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY
    EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
    WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURposE ARE
    DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY
    DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
    (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
    LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
    ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
    SOFTWARE, EVEN IF ADVISED OF THE posSIBILITY OF SUCH DAMAGE.
*/
#include "itkio.h"
#include "miaImageProcessing.h"

#include <cstdio>
#include <cmath>

#include <vector>
#include <string>
#include <sstream> //for ostringstream.
#include <fstream>
#include "densecrf.h"

#include "ioparsing.hpp"
#include "parameters.h"
#include <boost/program_options.hpp>
#include <boost/filesystem.hpp>
#include <chrono>

namespace po = boost::program_options;
namespace fs = boost::filesystem;

using namespace mia;

std::vector<Image> load_multi_channel_image(std::vector<std::string> image_filenames)
{

	std::vector<Image> images;
	images.push_back(itkio::load(image_filenames[0]));

	for (int c = 1; c < image_filenames.size(); c++)
	{
		Image img = itkio::load(image_filenames[c]);
		if (img.size() != images[0].size())
		{
			std::cout << "Image modalities size mismatch! Aborting..." << std::endl;
			exit(0);
		}
		else
			images.push_back(img);
	}
	return images;
}

float cap_prob(float p)
{

	return std::fmax(0.05, std::fmin(p, 0.95));
}

//Loads the nii with the unary potentials and makes a matrix that the rest of the original code needs.
void setupUnaryEnergyMatrix(MatrixXf &unaryCosts, std::vector<Image> probmaps, int numberOfForegroundClasses, bool computeBackgroundProbmap = true)
{

	int numberOfVoxelsInImage = probmaps[0].size();
	int start_class_index = 0;
	int numberOfClasses = numberOfForegroundClasses + 1;
	if (computeBackgroundProbmap)
		start_class_index = 1;

	float eps = 0.001;
	unaryCosts.setZero();
	for (int c = 0; c < numberOfForegroundClasses; ++c)
	{

		int index = 0;
		for (int z = 0; z < probmaps[c].sizeZ(); z++)
		{
			for (int y = 0; y < probmaps[c].sizeY(); y++)
			{
				for (int x = 0; x < probmaps[c].sizeX(); x++)
				{
					float val = probmaps[c](x, y, z);

					val = -log(cap_prob(probmaps[c](x, y, z)));
					unaryCosts(c + start_class_index, index) = val;
					if (computeBackgroundProbmap)
						unaryCosts(0, index) += std::fmax(0, probmaps[c](x, y, z));
					index++;
				}
			}
		}
	}
	if (computeBackgroundProbmap)
	{
		int index = 0;
		for (int z = 0; z < probmaps[0].sizeZ(); z++)
		{
			for (int y = 0; y < probmaps[0].sizeY(); y++)
			{
				for (int x = 0; x < probmaps[0].sizeX(); x++)
				{

					unaryCosts(0, index) = -log(cap_prob(1 - unaryCosts(0, index)));
					index++;
				}
			}
		}
	}
}

void addPairwiseBilateralMultiMod(DenseCRF3D &crf3d, std::vector<float> sXYZ, std::vector<float> sMod, std::vector<Image> &images, LabelCompatibility *function = NULL, KernelType kernel_type = DIAG_KERNEL, NormalizationType normalization_type = NORMALIZE_SYMMETRIC)
{
	int numImages = images.size();
	int sizeX = images[0].sizeX();
	int sizeY = images[0].sizeY();
	int sizeZ = images[0].sizeZ();

	MatrixXf feature(3 + numImages, images[0].size());
	feature.setZero();
	int index = 0;
	for (int z = 0; z < images[0].sizeZ(); z++)
	{
		for (int y = 0; y < images[0].sizeY(); y++)
		{
			for (int x = 0; x < images[0].sizeX(); x++)
			{
				feature(0, index) = x / sXYZ[0];
				feature(1, index) = y / sXYZ[1];
				feature(2, index) = z / sXYZ[2];

				for (int imgIdx = 0; imgIdx < numImages; imgIdx++)
					feature(3 + imgIdx, index) = images[imgIdx](x, y, z) / sMod[imgIdx];

				index++;
			}
		}
	}
	crf3d.addPairwiseEnergy(feature, function, kernel_type, normalization_type);
}

int main(int argc, char *argv[])
{

	Parameters params;

	std::string config_file;

	std::vector<std::string> simages;
	std::vector<std::string> simagelists;
	std::vector<std::string> sprobmaps;
	std::vector<std::string> sprobmaplists;

	std::string output_path = ".";

	int numberOfModalities = 0;
	int numberOfForegroundClasses = 0;
	bool computeBackgroundProbmap = true;
	bool overwrite_output = false;
	bool do_output_probmaps = false;
	bool do_output_unary = false;
	std::vector<std::string> sbilateralXYZStds;
	std::vector<std::string> sposXYZStds;
	std::vector<std::string> sbilateralModStds;

	float minIntensity = -3, maxIntensity = +3;

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

	try
	{
		// Declare the supported options.
		po::options_description generic("generic options");
		generic.add_options()("help", "produce help message")("config", po::value<std::string>(&config_file), "configuration file")("probs", po::bool_switch(&do_output_probmaps)->default_value(false), "")("unary", po::bool_switch(&do_output_unary)->default_value(false), "")("overwrite", po::bool_switch(&overwrite_output)->default_value(false), "");

		po::options_description config("specific options");
		config.add_options()("output", po::value<std::string>(&output_path), "output path")("image", po::value<std::vector<std::string>>(&simages)->multitoken(), "filename(s) of multi-modality images")("imagelist", po::value<std::vector<std::string>>(&simagelists)->multitoken(), "text file(s) listing images")("probmap", po::value<std::vector<std::string>>(&sprobmaps)->multitoken(), "filename(s) of multi-modality images")("probmaplist", po::value<std::vector<std::string>>(&sprobmaplists)->multitoken(), "text file(s) listing images")("posXYZStds", po::value<std::vector<std::string>>(&sposXYZStds), "posXYZStds")("biModStds", po::value<std::vector<std::string>>(&sbilateralModStds), "bilateralModStds")("biXYZStds", po::value<std::vector<std::string>>(&sbilateralXYZStds), "bilateralXYZStds")("numberOfForegroundClasses", po::value<int>(&numberOfForegroundClasses), "numberOfForegroundClasses")("posW", po::value<float>(&params.posW), "posW")("biW", po::value<float>(&params.bilateralW), "biW")("maxIterations", po::value<int>(&params.maxIterations), "maxIteration");

		po::options_description cmdline_options("options");
		cmdline_options.add(generic).add(config);

		po::options_description config_file_options;
		config_file_options.add(config);

		po::variables_map vm;

		po::store(po::parse_command_line(argc, argv, cmdline_options), vm);
		po::notify(vm);

		if (vm.count("config"))
		{
			std::ifstream ifs(config_file.c_str());
			if (!ifs)
			{
				std::cout << "cannot open config file: " << config_file << std::endl;
				return 0;
			}
			else
			{
				po::store(parse_config_file(ifs, config_file_options), vm);
				po::notify(vm);
			}
		}

		if (vm.count("help"))
		{
			std::cout << cmdline_options << std::endl;
			return 0;
		}
	}
	catch (std::exception &e)
	{
		std::cout << e.what() << std::endl;
		return 1;
	}

	if (!fs::exists(output_path)) fs::create_directories(output_path);

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

	std::vector<std::string> images;
	strings_to_values(simages, images);

	std::vector<std::string> imagelists;
	strings_to_values(simagelists, imagelists);

	strings_to_values(sbilateralXYZStds, params.bilateralXYZStds, true);

	strings_to_values(sposXYZStds, params.posXYZStds, true);

	std::vector<std::string> probmaps;
	strings_to_values(sprobmaps, probmaps);

	std::vector<std::string> probmaplists;
	strings_to_values(sprobmaplists, probmaplists);

	std::vector<std::vector<std::string>> image_filenames = parse_multichannel_filepaths(images, imagelists);
	std::vector<std::vector<std::string>> probmap_filenames = parse_multichannel_filepaths(probmaps, probmaplists);

	int numberOfImages = image_filenames.size();
	int numberOfProbChannels = probmap_filenames.size();

	if (numberOfImages == 0)
	{
		std::cout << "No images provided. Aborting..." << std::endl;
	}
	if (numberOfProbChannels == 0)
	{
		std::cout << "No probmaps provided. Aborting..." << std::endl;
		exit(EXIT_FAILURE);
	}

	if (numberOfProbChannels == numberOfForegroundClasses)
		computeBackgroundProbmap = true;
	else if (numberOfForegroundClasses > 0 && numberOfProbChannels == numberOfForegroundClasses + 1)
		computeBackgroundProbmap = false;
	else
	{
		std::cout << "number of probmaps mismatches with provided number of foreground classes" << std::endl;
		exit(EXIT_FAILURE);
	}

	numberOfModalities = image_filenames[0].size();
	int numClasses = numberOfForegroundClasses + 1;

	params.setNumberOfModalities(numberOfModalities);
	strings_to_values(sbilateralModStds, params.bilateralModsStds, true);

	std::vector<std::vector<Image>> intensity_images;

	params.print();

	for (int imgIdx = 0; imgIdx < numberOfImages; imgIdx++)
	{

		fs::path input_path(image_filenames[imgIdx][0]);

		std::string basename = fs::basename(input_path);
		if (fs::extension(basename) != "")
			basename = fs::basename(basename);
		std::stringstream base_output_sfilename;
		std::stringstream labelmap_output_path;

		base_output_sfilename << output_path << "/" << basename;
		labelmap_output_path << base_output_sfilename.str() << "_labelmap.nii.gz";

		if (fs::exists(labelmap_output_path.str()))
		{
			if (!overwrite_output)
			{
				std::cout << labelmap_output_path.str() << " already exists. skipping." << std::endl;
				continue;
			}
			else
				std::cout << labelmap_output_path.str() << " already exists. Overwriting." << std::endl;
		}
		std::cout << labelmap_output_path.str() << std::endl;

		namespace ch = std::chrono;
		auto start = ch::high_resolution_clock::now();

		std::cout << "loading data " << imgIdx + 1 << " of " << image_filenames.size() << " with " << numberOfModalities << " channel(s)" << std::endl;

		std::vector<Image> multiChannelImage = load_multi_channel_image(image_filenames[imgIdx]);
		std::vector<Image> imageProbmaps = load_multi_channel_image(probmap_filenames[imgIdx]);

		int sizeX = multiChannelImage[0].sizeX();
		int sizeY = multiChannelImage[0].sizeY();
		int sizeZ = multiChannelImage[0].sizeZ();

		DenseCRF3D crf3d(sizeX, sizeY, sizeZ, numberOfForegroundClasses + 1);

		std::cout << "********** Setting up Unary. **********" << std::endl;
		MatrixXf unaryCost = MatrixXf(numClasses, multiChannelImage[0].size());
		setupUnaryEnergyMatrix(unaryCost, imageProbmaps, numberOfForegroundClasses, computeBackgroundProbmap);
		crf3d.setUnaryEnergy(unaryCost);

		std::cout << "********** Setting up Pairwise potentials **********" << std::endl;
		crf3d.addPairwiseGaussian(params.posXYZStds[0],
								  params.posXYZStds[1],
								  params.posXYZStds[2],
								  new PottsCompatibility(params.posW));

		addPairwiseBilateralMultiMod(crf3d,
									 params.bilateralXYZStds,
									 params.bilateralModsStds,
									 multiChannelImage,
									 new PottsCompatibility(params.bilateralW));

		std::cout << "++++++++++++++++++++++++++ Performing Inference ++++++++++++++++++++++++++" << std::endl;

		MatrixXf probMapsMatrix = crf3d.inference(params.maxIterations);

		if (do_output_probmaps)
		{
			Image output_probmap = imageProbmaps[0].clone();
			for (int i = 0; i < numClasses; ++i)
			{
				std::stringstream output_filename;
				output_filename << base_output_sfilename.str() << "_probmap" << i << ".nii.gz";
				int index = 0;
				for (int z = 0; z < imageProbmaps[0].sizeZ(); z++)
				{
					for (int y = 0; y < imageProbmaps[0].sizeY(); y++)
					{
						for (int x = 0; x < imageProbmaps[0].sizeX(); x++)
						{
							output_probmap(x, y, z) = probMapsMatrix(i, index++);
						}
					}
				}
				itkio::save(output_probmap, output_filename.str());
			}
		}

		if (do_output_unary)
		{
			Image output_probmap = imageProbmaps[0].clone();
			for (int i = 0; i < numClasses; ++i)
			{
				std::stringstream output_filename;
				output_filename << base_output_sfilename.str() << "_unary" << i << ".nii.gz";
				int index = 0;
				for (int z = 0; z < imageProbmaps[0].sizeZ(); z++)
				{
					for (int y = 0; y < imageProbmaps[0].sizeY(); y++)
					{
						for (int x = 0; x < imageProbmaps[0].sizeX(); x++)
						{
							output_probmap(x, y, z) = unaryCost(i, index++);
						}
					}
				}
				itkio::save(output_probmap, output_filename.str());
			}
		}

		VectorXs segmentationVector = crf3d.currentMap(probMapsMatrix);

		Image output_labelmap = imageProbmaps[0].clone();
		output_labelmap.dataType(mia::USHORT);
		int index = 0;
		for (int z = 0; z < imageProbmaps[0].sizeZ(); z++)
		{
			for (int y = 0; y < imageProbmaps[0].sizeY(); y++)
			{
				for (int x = 0; x < imageProbmaps[0].sizeX(); x++)
				{
					output_labelmap(x, y, z) = segmentationVector[index++];
				}
			}
		}

		itkio::save(output_labelmap, labelmap_output_path.str());

		std::cout << "++++++++++++++++++++++++++ Done. ++++++++++++++++++++++++++" << std::endl;
		auto stop = ch::high_resolution_clock::now();
		std::cout << "done. took " << ch::duration_cast<ch::milliseconds>(stop - start).count() << " ms" << std::endl;
	}
}


/*
 * Copyright (c) 2019, The Regents of the University of California (Regents).
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met:
 *
 *    1. Redistributions of source code must retain the above copyright
 *       notice, this list of conditions and the following disclaimer.
 *
 *    2. Redistributions in binary form must reproduce the above
 *       copyright notice, this list of conditions and the following
 *       disclaimer in the documentation and/or other materials provided
 *       with the distribution.
 *
 *    3. Neither the name of the copyright holder nor the names of its
 *       contributors may be used to endorse or promote products derived
 *       from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS AS IS
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 *
 * Please contact the author(s) of this library if you have any questions.
 * Authors: David Fridovich-Keil   ( dfk@eecs.berkeley.edu )
 */

///////////////////////////////////////////////////////////////////////////////
//
// 12D quadrotor dynamics.
//
///////////////////////////////////////////////////////////////////////////////

#include <quads/quadrotor12d.h>

#include <math.h>
#include <ros/ros.h>
#include <Eigen/Dense>
#include <string>

namespace quads {

Vector12d Quadrotor12D::operator()(const Vector12d& x,
                                   const Vector3d& u) const {
  ROS_ASSERT(initialized_);

  const double gx = std::sin(x(kThetaIdx)) * std::cos(x(kPhiIdx)) / m_;
  const double gy = -std::sin(x(kPhiIdx)) / m_;
  const double gz = std::cos(x(kPhiIdx)) * std::cos(x(kThetaIdx)) / m_;

  Vector12d xdot;
  xdot << x(kDxIdx), x(kDyIdx), x(kDzIdx), x(kQIdx), x(kRIdx), gx * x(kZetaIdx),
      gy * x(kZetaIdx), gz * x(kZetaIdx) - 9.81, x(kXiIdx), u(0), u(1) / Ix_,
      u(2) / Iy_;
  return xdot;
}

Matrix12x12d Quadrotor12D::StateJacobian(const Vector12d& x,
                                         const Vector3d& u) const {
  const double theta = x(kThetaIdx);
  const double phi = x(kPhiIdx);
  const double zeta = x(kZetaIdx);

  Matrix12x12d F;
  F << 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0,
      0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0,
      0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0,
      zeta * std::cos(phi) * std::cos(theta),
      -zeta * std::sin(phi) * std::sin(theta), 0, 0, 0,
      std::cos(phi) * std::sin(theta), 0, 0, 0, 0, 0, 0, 0,
      -zeta * std::cos(phi), 0, 0, 0, -std::sin(phi), 0, 0, 0, 0, 0, 0,
      -zeta * std::cos(phi) * std::sin(theta),
      -zeta * std::cos(theta) * std::sin(phi), 0, 0, 0,
      std::cos(phi) * std::cos(theta), 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0,
      0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
      0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0;
  return F;
}

Matrix5x12d Quadrotor12D::OutputJacobian(const Vector12d& x) const {
  Matrix5x12d H(Matrix5x12d::Zero());
  H(0, 0) = 1.0;
  H(1, 1) = 1.0;
  H(2, 2) = 1.0;
  H(3, 3) = 1.0;
  H(4, 4) = 1.0;
  return H;
}

bool Quadrotor12D::Initialize(const ros::NodeHandle& n) {
  name_ = ros::names::append(n.getNamespace(), "quadrotor12d");

  if (!LoadParameters(n)) {
    ROS_ERROR("%s: Failed to load parameters.", name_.c_str());
    return false;
  }

  initialized_ = true;
  return true;
}

bool Quadrotor12D::LoadParameters(const ros::NodeHandle& n) {
  ros::NodeHandle nl(n);

  // Mass and inertia.
  if (!nl.getParam("dynamics/m", m_)) return false;
  if (!nl.getParam("dynamics/Ix", Ix_)) return false;
  if (!nl.getParam("dynamics/Iy", Iy_)) return false;

  return true;
}

}  // namespace quads


// This file is part of the reference implementation for the paper
//   Bayesian Collaborative Denoising for Monte-Carlo Rendering
//   Malik Boughida and Tamy Boubekeur.
//   Computer Graphics Forum (Proc. EGSR 2017), vol. 36, no. 4, p. 137-153, 2017
//
// All rights reserved. Use of this source code is governed by a
// BSD-style license that can be found in the LICENSE.txt file.

// BCD headers.
#include "bcd/DeepImage.h"
#include "bcd/Denoiser.h"
#include "bcd/IDenoiser.h"
#include "bcd/ImageIO.h"
#include "bcd/MultiscaleDenoiser.h"
#include "bcd/SpikeRemovalFilter.h"
#include "bcd/Utils.h"

// Eigen headers.
#include <Eigen/Dense>

// Standard headers.
#include <cstdlib>
#include <ctime>
#include <fstream>
#include <iostream>
#include <memory>
#include <string>

using namespace std;
using namespace bcd;

static const char* g_pProgramPath;

class ProgramArguments
{
  public:
    ProgramArguments() :
        m_denoisedOutputFilePath(""),
        m_colorImage(), m_nbOfSamplesImage(), m_histogramImage(), m_covarianceImage(),
        m_histogramPatchDistanceThreshold(1.f),
        m_patchRadius(1), m_searchWindowRadius(6),
        m_minEigenValue(1.e-8f),
        m_useRandomPixelOrder(false),
        m_prefilterSpikes(false),
        m_prefilterThresholdStDevFactor(2.f),
        m_markedPixelsSkippingProbability(1.f),
        m_nbOfScales(3),
        m_nbOfCores(0)
    {
    }

    string m_denoisedOutputFilePath; // File path to the denoised image output
    Deepimf m_colorImage; // Pixel color values
    Deepimf m_nbOfSamplesImage; // Pixel number of samples
    Deepimf m_histogramImage; // Pixel histograms
    Deepimf m_covarianceImage; // Pixel covariances
    float m_histogramPatchDistanceThreshold; // Histogram patch distance threshold
    int m_patchRadius; // Patch has (1 + 2 x m_patchRadius)^2 pixels
    int m_searchWindowRadius; // Search windows (for neighbors) spreads across (1 + 2 x m_patchRadius)^2 pixels
    float m_minEigenValue; // Minimum eigen value for matrix inversion
    bool m_useRandomPixelOrder; // True means the pixel will be processed in a random order ; could be useful to remove some "grid" artifacts
    bool m_prefilterSpikes; // True means a spike removal prefiltering will be applied
    float m_prefilterThresholdStDevFactor; // See SpikeRemovalFilter::filter argument
    float m_markedPixelsSkippingProbability; // 1 means the marked centers of the denoised patches will be skipped to accelerate a lot the computations
    int m_nbOfScales;
    int m_nbOfCores; // Number of cores used by OpenMP. O means using the value defined in environment variable OMP_NUM_THREADS
};

class Callbacks
  : public ICallbacks
{
  public:
    Callbacks()
    {
    }

    void progress(const float i_progress) const override
    {
    }

    bool isAborted() const override
    {
        return false;
    }

  private:
    void logInfo(const char* msg) const override
    {
        cout << "Info: " << msg << std::endl;
    }

    void logWarning(const char* msg) const override
    {
        cout << "Warning: " << msg << std::endl;
    }

    void logError(const char* msg) const override
    {
        cout << "Error: " << msg << std::endl;
    }

    void logDebug(const char* msg) const override
    {
        cout << "Debug: " << msg << std::endl;
    }
};

Callbacks g_callbacks;

void initializeRandomSeed()
{
    srand(static_cast<unsigned int>(time(0)));
}

static void printUsage()
{
    ProgramArguments defaultProgramArgs;
    cout << "Bayesian Collaborative Denoising"<< endl << endl;
    cout << "Usage: " << g_pProgramPath << " <arguments list>" << endl;
    cout << "Only EXR images are supported." << endl << endl;
    cout << "Required arguments list:" << endl;
    cout << "    -o <output>          The file path to the output image" << endl;
    cout << "    -i <input>           The file path to the input image" << endl;
    cout << "    -h <hist>            The file path to the input histograms buffer" << endl;
    cout << "    -c <cov>             The file path to the input covariance matrices buffer" << endl;
    cout << "Optional arguments list:" << endl;
    cout << "    -d <float>           Histogram patch distance threshold (default: " << defaultProgramArgs.m_histogramPatchDistanceThreshold << ")" << endl;
    cout << "    -b <int>             Radius of search windows (default: " << defaultProgramArgs.m_searchWindowRadius << ")" << endl;
    cout << "    -w <int>             Radius of patches (default: " << defaultProgramArgs.m_patchRadius << ")" << endl;
    cout << "    -r <0/1>             1 for random pixel order (in case of grid artifacts) (default: " << (defaultProgramArgs.m_useRandomPixelOrder ? 1 : 0) << ")" << endl;
    cout << "    -p <0/1>             1 for a spike removal prefiltering (default: " << (defaultProgramArgs.m_prefilterSpikes ? 1 : 0) << ")" << endl;
    cout << "    --p-factor <float>   Factor that is multiplied by standard deviation to get the threshold for classifying spikes during prefiltering. Put lower value to remove more spikes (default: " << defaultProgramArgs.m_prefilterThresholdStDevFactor << ")" << endl;
    cout << "    -m <float in [0,1]>  Probability of skipping marked centers of denoised patches. 1 accelerates a lot the computations. 0 helps removing potential grid artifacts (default: " << defaultProgramArgs.m_markedPixelsSkippingProbability << ")" << endl;
    cout << "    -s <int>             Number of Scales for Multi-Scaling (default: " << defaultProgramArgs.m_nbOfScales << ")" << endl;
    cout << "    --ncores <nbOfCores> Number of cores used by OpenMP (default: environment variable OMP_NUM_THREADS)" << endl;
    cout << "    -e <float>           Minimum eigen value for matrix inversion (default: " << defaultProgramArgs.m_minEigenValue << ")" << endl;
}

bool parseProgramArguments(int argc, const char** argv, ProgramArguments& o_rProgramArguments)
{
    int argIndex = 0;
    bool missingColor = true, missingHist = true, missingCov = true, missingOutput = true;
    string inputColorFilePath;
    while (++argIndex < argc)
    {
        if (strcmp(argv[argIndex], "-o") == 0)
        {
            argIndex++;
            if (argIndex == argc)
            {
                cout << "Error in program arguments: expecting file path to the output image after '-o'" << endl;
                return false;
            }
            o_rProgramArguments.m_denoisedOutputFilePath = string(argv[argIndex]);
            ofstream outputFile(o_rProgramArguments.m_denoisedOutputFilePath, ofstream::out | ofstream::app);
            if(!outputFile)
            {
                cout << "Error in program arguments: cannot write output file '" << o_rProgramArguments.m_denoisedOutputFilePath << "'" << endl;
                return false;
            }
            missingOutput = false;
        }
        else if (strcmp(argv[argIndex], "-i") == 0)
        {
            argIndex++;
            if (argIndex == argc)
            {
                cout << "Error in program arguments: expecting file path to the input color image after '-i'" << endl;
                return false;
            }
            inputColorFilePath = argv[argIndex];
            if (!ImageIO::loadEXR(o_rProgramArguments.m_colorImage, argv[argIndex]))
            {
                cout << "Error in program arguments: couldn't load input color image file '" << argv[argIndex] << "'" << endl;
                return false;
            }
            missingColor = false;
        }
        else if (strcmp(argv[argIndex], "-h") == 0)
        {
            argIndex++;
            if (argIndex == argc)
            {
                cout << "Error in program arguments: expecting file path to the input histogram image after '-h'" << endl;
                return false;
            }
            Deepimf histAndNbOfSamplesImage;
            if (!ImageIO::loadMultiChannelsEXR(histAndNbOfSamplesImage, argv[argIndex]))
            {
                cout << "Error in program arguments: couldn't load input histogram image file '" << argv[argIndex] << "'" << endl;
                return false;
            }
            Utils::separateNbOfSamplesFromHistogram(o_rProgramArguments.m_histogramImage, o_rProgramArguments.m_nbOfSamplesImage, histAndNbOfSamplesImage);
            missingHist = false;
        }
        else if (strcmp(argv[argIndex], "-c") == 0)
        {
            argIndex++;
            if (argIndex == argc)
            {
                cout << "Error in program arguments: expecting file path to the input covariance matrix image after '-c'" << endl;
                return false;
            }
            if (!ImageIO::loadMultiChannelsEXR(o_rProgramArguments.m_covarianceImage, argv[argIndex]))
            {
                cout << "Error in program arguments: couldn't load input covariance matrix image file '" << argv[argIndex] << "'" << endl;
                return false;
            }
            missingCov = false;
        }
        else if (strcmp(argv[argIndex], "-d") == 0)
        {
            argIndex++;
            if (argIndex == argc)
            {
                cout << "Error in program arguments: expecting histogram patch distance threshold after '-d'" << endl;
                return false;
            }
            istringstream iss(argv[argIndex]);
            iss >> o_rProgramArguments.m_histogramPatchDistanceThreshold;
        }
        else if (strcmp(argv[argIndex], "-b") == 0)
        {
            argIndex++;
            if (argIndex == argc)
            {
                cout << "Error in program arguments: expecting radius of search window after '-b'" << endl;
                return false;
            }
            istringstream iss(argv[argIndex]);
            iss >> o_rProgramArguments.m_searchWindowRadius;
        }
        else if (strcmp(argv[argIndex], "-w") == 0)
        {
            argIndex++;
            if (argIndex == argc)
            {
                cout << "Error in program arguments: expecting radius of patch after '-w'" << endl;
                return false;
            }
            istringstream iss(argv[argIndex]);
            iss >> o_rProgramArguments.m_patchRadius;
        }
        else if (strcmp(argv[argIndex], "-e") == 0)
        {
            argIndex++;
            if (argIndex == argc)
            {
                cout << "Error in program arguments: expecting minimum eigen value after '-e'" << endl;
                return false;
            }
            istringstream iss(argv[argIndex]);
            iss >> o_rProgramArguments.m_minEigenValue;
        }
        else if (strcmp(argv[argIndex], "-r") == 0)
        {
            argIndex++;
            if (argIndex == argc)
            {
                cout << "Error in program arguments: expecting 0 or 1 after '-r'" << endl;
                return false;
            }
            istringstream iss(argv[argIndex]);
            int useRandomPixelOrder;
            iss >> useRandomPixelOrder;
            if(useRandomPixelOrder != 0 && useRandomPixelOrder != 1)
            {
                cout << "Error in program arguments: expecting 0 or 1 after '-r'" << endl;
                return false;
            }
            o_rProgramArguments.m_useRandomPixelOrder = (useRandomPixelOrder==1);
        }
        else if (strcmp(argv[argIndex], "-p") == 0)
        {
            argIndex++;
            if (argIndex == argc)
            {
                cout << "Error in program arguments: expecting 0 or 1 after '-p'" << endl;
                return false;
            }
            istringstream iss(argv[argIndex]);
            int prefilterSpikes;
            iss >> prefilterSpikes;
            if(prefilterSpikes != 0 && prefilterSpikes != 1)
            {
                cout << "Error in program arguments: expecting 0 or 1 after '-p'" << endl;
                return false;
            }
            o_rProgramArguments.m_prefilterSpikes = (prefilterSpikes==1);
        }
        else if (strcmp(argv[argIndex], "--p-factor") == 0)
        {
            argIndex++;
            if (argIndex == argc)
            {
                cout << "Error in program arguments: expecting standard deviation factor for spike prefiltering threshold after '--p-factor'" << endl;
                return false;
            }
            istringstream iss(argv[argIndex]);
            iss >> o_rProgramArguments.m_prefilterThresholdStDevFactor;
        }
        else if(strcmp(argv[argIndex], "-m") == 0)
        {
            argIndex++;
            if(argIndex == argc)
            {
                cout << "Error in program arguments: expecting float in [0,1] after '-m'" << endl;
                return false;
            }
            istringstream iss(argv[argIndex]);
            float markedPixelsSkippingProbability;
            iss >> markedPixelsSkippingProbability;
            if(markedPixelsSkippingProbability < 0 || markedPixelsSkippingProbability > 1)
            {
                cout << "Error in program arguments: expecting float in [0,1] after '-m'" << endl;
                return false;
            }
            o_rProgramArguments.m_markedPixelsSkippingProbability = markedPixelsSkippingProbability;
        }
        else if(strcmp(argv[argIndex], "-s") == 0)
        {
            argIndex++;
            if (argIndex == argc)
            {
                cout << "Error in program arguments: expecting number of scales after '-s'" << endl;
                return false;
            }
            istringstream iss(argv[argIndex]);
            iss >> o_rProgramArguments.m_nbOfScales;
        }
        else if (strcmp(argv[argIndex], "--ncores") == 0)
        {
            argIndex++;
            if (argIndex == argc)
            {
                cout << "Error in program arguments: expecting number of cores for OpenMP after '--ncores'" << endl;
                return false;
            }
            istringstream iss(argv[argIndex]);
            iss >> o_rProgramArguments.m_nbOfCores;
        }
    }
    if(!missingColor)
    {
        if(missingHist)
        {
            string inputHistFilePath = inputColorFilePath.substr(0, inputColorFilePath.length() - 4) + "_hist.exr"; // "-4" for removing extension .exr
            cout << "Warning: input histogram file not provided by -h argument: assuming '" + inputHistFilePath + "'" << endl;
            Deepimf histAndNbOfSamplesImage;
            if (!ImageIO::loadMultiChannelsEXR(histAndNbOfSamplesImage, inputHistFilePath.c_str()))
            {
                cout << "Error in program arguments: couldn't load input histogram image file '" << inputHistFilePath << "'" << endl;
                return false;
            }
            Utils::separateNbOfSamplesFromHistogram(o_rProgramArguments.m_histogramImage, o_rProgramArguments.m_nbOfSamplesImage, histAndNbOfSamplesImage);
            missingHist = false;
        }
        if(missingCov)
        {
            string inputCovFilePath = inputColorFilePath.substr(0, inputColorFilePath.length() - 4) + "_cov.exr"; // "-4" for removing extension .exr
            cout << "Warning: input covariance file not provided by -c argument: assuming '" + inputCovFilePath + "'" << endl;
            if (!ImageIO::loadMultiChannelsEXR(o_rProgramArguments.m_covarianceImage, inputCovFilePath.c_str()))
            {
                cout << "Error in program arguments: couldn't load input covariance matrix image file '" << inputCovFilePath << "'" << endl;
                return false;
            }
            missingCov = false;
        }
    }
    if (missingColor || missingHist || missingCov || missingOutput)
    {
        cout << "Error: Missing required program argument(s):";
        if (missingColor)
            cout << " -i";
        if (missingHist)
            cout << " -h";
        if (missingCov)
            cout << " -c";
        if (missingOutput)
            cout << " -o";
        cout << endl << endl;
        printUsage();
        return false;
    }
    return true;
}

int launchBayesianCollaborativeDenoising(int argc, const char** argv)
{
    ProgramArguments programArgs;
    if(!parseProgramArguments(argc, argv, programArgs))
        return 1;

    if(programArgs.m_prefilterSpikes)
        SpikeRemovalFilter::filter(
                programArgs.m_colorImage,
                programArgs.m_nbOfSamplesImage,
                programArgs.m_histogramImage,
                programArgs.m_covarianceImage,
                programArgs.m_prefilterThresholdStDevFactor);

    DenoiserInputs inputs;
    DenoiserOutputs outputs;
    DenoiserParameters parameters;

    inputs.m_pColors = &(programArgs.m_colorImage);
    inputs.m_pNbOfSamples = &(programArgs.m_nbOfSamplesImage);
    inputs.m_pHistograms = &(programArgs.m_histogramImage);
    inputs.m_pSampleCovariances = &(programArgs.m_covarianceImage);

    Deepimf outputDenoisedColorImage(programArgs.m_colorImage);
    outputs.m_pDenoisedColors = &outputDenoisedColorImage;

    parameters.m_histogramDistanceThreshold = programArgs.m_histogramPatchDistanceThreshold;
    parameters.m_patchRadius = programArgs.m_patchRadius;
    parameters.m_searchWindowRadius = programArgs.m_searchWindowRadius;
    parameters.m_minEigenValue = programArgs.m_minEigenValue;
    parameters.m_useRandomPixelOrder = programArgs.m_useRandomPixelOrder;
    parameters.m_markedPixelsSkippingProbability = programArgs.m_markedPixelsSkippingProbability;
    parameters.m_nbOfCores = programArgs.m_nbOfCores;

    unique_ptr<IDenoiser> uDenoiser = nullptr;

    if(programArgs.m_nbOfScales > 1)
        uDenoiser.reset(new MultiscaleDenoiser(programArgs.m_nbOfScales));
    else
        uDenoiser.reset(new Denoiser());

    uDenoiser->setCallbacks(&g_callbacks);

    uDenoiser->setInputs(inputs);
    uDenoiser->setOutputs(outputs);
    uDenoiser->setParameters(parameters);

    if (!uDenoiser->inputsOutputsAreOk())
        return 1;

    uDenoiser->denoise();

    ImageIO::writeEXR(outputDenoisedColorImage, programArgs.m_denoisedOutputFilePath.c_str());

    cout << "Written denoised output in file " << programArgs.m_denoisedOutputFilePath.c_str() << std::endl;

    return 0;
}

int main(int argc, const char** argv)
{
    g_pProgramPath = argv[0];
    return launchBayesianCollaborativeDenoising(argc, argv);
}


/**
 * Copyright (c) 2017 Melown Technologies SE
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * *  Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 *
 * *  Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */
/**
 *  @file triclip.hpp
 *  @author Tomas Drinovsky <tomas.drinovsky@citationtech.net>
 *
 *  Triangle clipping stuff based on Jakub Cerveny's code
 */

#ifndef geometry_triclip_hpp_included_
#define geometry_triclip_hpp_included_

#include <boost/numeric/ublas/vector.hpp>
#include "math/geometry_core.hpp"
#include "dbglog/dbglog.hpp"

namespace ublas = boost::numeric::ublas;

namespace geometry {

/** Textured 3D triangle representation suitable for clipping algorithm.
 *
 * Triangle holds positions and tex. coordinates of its vertices.
 */
struct ClipTriangle
{
        typedef std::vector<ClipTriangle> list;
        math::Point3 pos[3]; // vertices
        math::Point2 uv[3];  // texture coordinates

        bool texCoordsAvailable;

        ClipTriangle(math::Point3 a, math::Point3 b, math::Point3 c)
            : texCoordsAvailable(false)
        {
            pos[0] = a, pos[1] = b, pos[2] = c;
        }

        ClipTriangle(math::Point3 a, math::Point3 b, math::Point3 c,
                     math::Point2 ta, math::Point2 tb, math::Point2 tc)
            : texCoordsAvailable(true)
        {
            pos[0] = a, pos[1] = b, pos[2] = c;
            uv[0] = ta, uv[1] = tb, uv[2] = tc;
        }
};

/** Clipping plane
 *
 * Plane is defined by it's normal and the shift in the direction of the normal
 */
struct ClipPlane
{
        math::Point3 normal;
        double d;

        ClipPlane(double a, double b, double c, double d)
            : normal(a, b, c), d(d)
        {}

        ClipPlane()
            : normal(0, 0, 0), d(0)
        {}
};

namespace detail {
    inline double signedDistance(const math::Point3 &point, const ClipPlane &plane)
    {
        return ublas::inner_prod(point,plane.normal) - plane.d;
    }

    inline math::Point3 intersection(const math::Point3 &p1, const math::Point3 &p2,
                       const ClipPlane &plane, double& t)
    {
        // sort points to prevent numerical inaccuracies
        math::Point3 sp1(std::max(p1,p2));
        math::Point3 sp2(std::min(p1,p2));

        double dot1 = ublas::inner_prod(sp1,plane.normal);
        double dot2 = ublas::inner_prod(sp2,plane.normal);
        double den = dot1 - dot2;

        // line parallel with plane, return the midpoint
        if (std::abs(den) < 1e-10) {
            t = 0.5;
            return (sp1 + sp2) * t;
        }

        t = (dot1 - plane.d) / den;
        return (1.0 - t)*sp1 + t*sp2;
    }
}

/** Clip triangles with given plane
 *
 * \param triangles list of triangles to be clipped
 * \param plane plane used to clip triangles
 * \param triangleInfos additional triangle informations,
                        each newly created triangle will have same triangleInfo
                        as it's origin
 * \return list of the clipped triangles
 */
template<typename TriangleInfo>
ClipTriangle::list clipTriangles( const ClipTriangle::list &triangles
                                , const ClipPlane &plane
                                , std::vector<TriangleInfo> &triangleInfos)
{
    ClipTriangle::list result;
    std::vector<TriangleInfo> resultInfo;

    if (triangleInfos.size() && triangleInfos.size()!=triangles.size()) {
        LOGTHROW(err3, std::runtime_error)
            << "Triangle count and triangle informations count mismatch.";
    }

    for ( std::size_t tid=0; tid<triangles.size(); ++tid)
    {
        auto &tri(triangles[tid]);

        bool positive[3] = {
            detail::signedDistance(tri.pos[0], plane) >= 0,
            detail::signedDistance(tri.pos[1], plane) >= 0,
            detail::signedDistance(tri.pos[2], plane) >= 0
        };

        int count = 0;
        for (int i = 0; i < 3; i++) {
            if (positive[i]) count++;
        }

        // triangle completely on negative side - do nothing
        if (count == 0) continue;

        // trinagle completely on positive side - copy to result
        if (count == 3) {
            result.push_back(tri);
            if(triangleInfos.size()){
                resultInfo.emplace_back(triangleInfos[tid]);
            }
            continue;
        }

        int a = 0, b = 0, c = 0;
        double t = 0.0;

        // case 1: one vertex on positive side, just adjust the other two
        if (count == 1)
        {
            if (positive[0]) a = 0, b = 1, c = 2;
            else if (positive[1]) a = 1, b = 2, c = 0;
            else a = 2, b = 0, c = 1;

            math::Point3 x1pos(detail::intersection(tri.pos[a], tri.pos[b], plane, t));
            math::Point2 x1uv((1.0 - t)*tri.uv[a] + t*tri.uv[b]);

            math::Point3 x2pos(detail::intersection(tri.pos[c], tri.pos[a], plane, t));
            math::Point2 x2uv((1.0 - t)*tri.uv[c] + t*tri.uv[a]);

            result.emplace_back(tri.pos[a], x1pos, x2pos,
                                tri.uv[a],  x1uv,  x2uv);
            if(triangleInfos.size()){
                resultInfo.emplace_back(triangleInfos[tid]);
            }
        }
        // case 2: two vertices on positive side, adjust triangle and add one more
        else
        {
            if (!positive[0]) a = 0, b = 1, c = 2;
            else if (!positive[1]) a = 1, b = 2, c = 0;
            else a = 2, b = 0, c = 1;

            auto tmp(1.0 - t);
            math::Point3 x1pos(detail::intersection(tri.pos[a], tri.pos[b], plane, t));
            math::Point2 x1uv(tmp*tri.uv[a] + t*tri.uv[b]);

            math::Point3 x2pos(detail::intersection(tri.pos[c], tri.pos[a], plane, t));
            math::Point2 x2uv(tmp*tri.uv[c] + t*tri.uv[a]);

            result.emplace_back(x1pos, tri.pos[b], tri.pos[c],
                                x1uv,  tri.uv[b],  tri.uv[c]);
            if(triangleInfos.size()){
                resultInfo.emplace_back(triangleInfos[tid]);
            }

            result.emplace_back(x1pos, tri.pos[c], x2pos,
                                x1uv,  tri.uv[c],  x2uv);
            if(triangleInfos.size()){
                resultInfo.emplace_back(triangleInfos[tid]);
            }
        }
    }

    triangleInfos.swap(resultInfo);
    return result;
}

} // namespace geometry
#endif // geometry_triclip_hpp_included_


/**
 * @file DynamicsState.hpp
 * @author Brahayam Ponton (brahayam.ponton@tuebingen.mpg.de)
 * @license License BSD-3-Clause
 * @copyright Copyright (c) 2019, New York University and Max Planck Gesellschaft.
 * @date 2019-10-08
 */

#pragma once

#include <array>
#include <vector>
#include <string>
#include <iostream>
#include <Eigen/Eigen>

#include <momentumopt/setting/Definitions.hpp>
#include <momentumopt/cntopt/ContactState.hpp>

namespace momentumopt {

  /**
   * This class is a container for all variables required to define a
   * dynamic state: center of mass position, linear and angular momenta
   * and its rates; forces, torques and center of pressure of the end-
   * effectors; positions, orientations, activations and contact types
   * of each end-effector.
   */
  struct DynamicsState
  {
    public:
	  EIGEN_MAKE_ALIGNED_OPERATOR_NEW

    public:
	  DynamicsState();
	  ~DynamicsState(){}

	  // Center of mass, linear and angular momenta
	  double& time() { return dtime_; }
	  Eigen::Vector3d& centerOfMass() { return com_; }
	  Eigen::Vector3d& linearMomentum() { return lmom_; }
	  Eigen::Vector3d& angularMomentum() { return amom_; }
	  Eigen::Vector3d& linearMomentumRate() { return lmomd_; }
	  Eigen::Vector3d& angularMomentumRate() { return amomd_; }

	  const double& time() const { return dtime_; }
	  const Eigen::Vector3d& centerOfMass() const { return com_; }
	  const Eigen::Vector3d& linearMomentum() const { return lmom_; }
	  const Eigen::Vector3d& angularMomentum() const { return amom_; }
	  const Eigen::Vector3d& linearMomentumRate() const { return lmomd_; }
	  const Eigen::Vector3d& angularMomentumRate() const { return amomd_; }

	  void time(const double& dtime) { dtime_ = dtime; }
	  void centerOfMass(const Eigen::Vector3d& com) { com_ = com; }
	  void linearMomentum(const Eigen::Vector3d& lmom) { lmom_ = lmom; }
	  void angularMomentum(const Eigen::Vector3d& amom) { amom_ = amom; }
	  void linearMomentumRate(const Eigen::Vector3d& lmomd) { lmomd_ = lmomd; }
	  void angularMomentumRate(const Eigen::Vector3d& amomd) { amomd_ = amomd; }

	  void setCenterOfMass(const Eigen::Vector3d& com) {
	      std::cout << "setting CoM to: " << std::endl;
	      com_ = com;
	      std::cout << com_ << std::endl;
	  }

	  void setCenterOfMass2(double x, double y, double z) {
	      com_ << x, y, z;
	  }

      // Endeffector forces, torques and cops
      Eigen::Vector3d& endeffectorCoP(int eff_id) { return eff_cops_[eff_id]; }
      Eigen::Vector3d& endeffectorForce(int eff_id) { return eff_forces_[eff_id]; }
      Eigen::Vector3d& endeffectorTorque(int eff_id) { return eff_torques_[eff_id]; }
      Eigen::Vector3d& endeffectorTorqueAtContactPoint(int eff_id) { return eefs_trqs_contact_point_[eff_id]; }

      const long endeffectorNum() { return eff_cops_.size(); }
      const Eigen::Vector3d& endeffectorCoP(int eff_id) const { return eff_cops_[eff_id]; }
      const Eigen::Vector3d& endeffectorForce(int eff_id) const { return eff_forces_[eff_id]; }
      const Eigen::Vector3d& endeffectorTorque(int eff_id) const { return eff_torques_[eff_id]; }
      const Eigen::Vector3d& endeffectorTorqueAtContactPoint(int eff_id) const { return eefs_trqs_contact_point_[eff_id]; }

      const std::vector<Eigen::Vector3d>& pyEndeffectorCops() const { return eff_cops_; }
      const std::vector<Eigen::Vector3d>& pyEndeffectorForces() const { return eff_forces_; }
      const std::vector<Eigen::Vector3d>& pyEndeffectorTorques() const { return eff_torques_; }

      void pyEndeffectorCops(const std::vector<Eigen::Vector3d>& eff_cops) { eff_cops_ = eff_cops; }
      void pyEndeffectorForces(const std::vector<Eigen::Vector3d>& eff_forces) { eff_forces_ = eff_forces; }
      void pyEndeffectorTorques(const std::vector<Eigen::Vector3d>& eff_torques) { eff_torques_ = eff_torques; }

	  // Endeffector activations, activation and contact ids
      int& endeffectorContactId(int eff_id) { return cnt_ids_[eff_id]; }
      int& endeffectorActivationId(int eff_id) { return eff_ids_[eff_id]; }
      bool& endeffectorActivation(int eff_id) { return eff_activations_[eff_id]; }

      const int& endeffectorContactId(int eff_id) const { return cnt_ids_[eff_id]; }
      const int& endeffectorActivationId(int eff_id) const { return eff_ids_[eff_id]; }
      bool endeffectorActivation(int eff_id) const { return eff_activations_[eff_id]; }

      // endeffector poses
      Eigen::Vector3d& endeffectorPosition(int eff_id) { return eff_positions_[eff_id]; }
      Eigen::Vector3d& endeffectorVelocity(int eff_id) { return eff_velocities_[eff_id]; }
      Eigen::Vector3d& endeffectorAcceleration(int eff_id) { return eff_accelerations_[eff_id]; }
      Eigen::Quaternion<double>& endeffectorOrientation(int eff_id) { return eff_orientations_[eff_id]; }

      const Eigen::Vector3d& endeffectorPosition(int eff_id) const { return eff_positions_[eff_id]; }
      const Eigen::Vector3d& endeffectorVelocity(int eff_id) const { return eff_velocities_[eff_id]; }
      const Eigen::Vector3d& endeffectorAcceleration(int eff_id) const { return eff_accelerations_[eff_id]; }
      const Eigen::Quaternion<double>& endeffectorOrientation(int eff_id) const { return eff_orientations_[eff_id]; }

      const std::vector<Eigen::Vector3d>& pyEndeffectorPositions() const { return eff_positions_; }
      void pyEndeffectorPositions(const std::vector<Eigen::Vector3d>& eff_positions) { eff_positions_ = eff_positions; }

	  // Helper functions
	  std::string toString() const;
	  friend std::ostream& operator<<(std::ostream &os, const DynamicsState& obj) { return os << obj.toString(); }
	  void fillInitialRobotState(const std::string cfg_file, const std::string robot_state = "initial_robot_configuration");

    private:
	  double dtime_;
	  Eigen::Vector3d com_, amom_, lmom_, amomd_, lmomd_;
      std::array<bool, Problem::n_endeffs_> eff_activations_;
      std::array<int, Problem::n_endeffs_> eff_ids_, cnt_ids_;

      std::vector<Eigen::Quaternion<double>> eff_orientations_;
      std::vector<Eigen::Vector3d> eff_positions_, eff_velocities_, eff_accelerations_;
      std::vector<Eigen::Vector3d> eff_forces_, eff_torques_, eff_cops_, eefs_trqs_contact_point_;
  };

  /**
   * This class is a container for a sequence of dynamic states,
   * for all time steps in the optimization.
   */
  class DynamicsSequence
  {
    public:
      EIGEN_MAKE_ALIGNED_OPERATOR_NEW

    public:
	  DynamicsSequence(){}
	  ~DynamicsSequence(){}

	  void resize(int num_timesteps);
	  void clean() { dynamics_sequence_.clear(); }
	  int size() const { return dynamics_sequence_.size(); }

	  const std::vector<DynamicsState>& dynamicsSequence() const { return dynamics_sequence_; }
	  void dynamicsSequence(const std::vector<DynamicsState>& dynamics_sequence) { dynamics_sequence_ = dynamics_sequence; }

	  void setCoMAtTime(const Eigen::VectorXd& com, int time_id) { dynamics_sequence_[time_id].setCenterOfMass(com); }

      DynamicsState& dynamicsState(int time_id) { return dynamics_sequence_[time_id]; }
	  const DynamicsState& dynamicsState(int time_id) const { return dynamics_sequence_[time_id]; }

	  Eigen::Matrix<int, Problem::n_endeffs_, 1>& activeEndeffectorSteps() { return active_endeffector_steps_; }
	  const Eigen::Matrix<int, Problem::n_endeffs_, 1>& activeEndeffectorSteps() const { return active_endeffector_steps_; }

	  std::string toString() const;
  	  friend std::ostream& operator<<(std::ostream &os, const DynamicsSequence& obj) { return os << obj.toString(); }

    private:
      std::vector<DynamicsState> dynamics_sequence_;
      Eigen::Matrix<int, Problem::n_endeffs_, 1> active_endeffector_steps_;
  };

}


//==================================================================================================
/*!
  @file

  @copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
**/
//==================================================================================================
#ifndef BOOST_SIMD_ALGORITHM_REDUCE_HPP_INCLUDED
#define BOOST_SIMD_ALGORITHM_REDUCE_HPP_INCLUDED

#include <boost/simd/range/segmented_input_range.hpp>
#include <boost/simd/function/sum.hpp>
#include <boost/simd/pack.hpp>

namespace boost { namespace simd
{
  /*!
    @ingroup group-std

    Computes the sum over elements in the given Contiguous Range @range{first,last} and
    the initial value @c init.

    \notebox{The summation order can be different from the order of a sequential summation
            , thus leading to different results.
            }

    @par Example:
    @snippet reduce.simple.cpp reduce-simple
    Possible output:
    @code
    SIMD reduce     : 45
    STD  accumulate : 45
    @endcode

    @param first  Beginning of the range of elements to sum
    @param last   End of the range of elements to sum
    @param init   Initial value of the sum

    @return The sum of the given value and elements in the given range.
  **/
  template<typename T> T reduce(T const* first, T const* last, T init)
  {
    pack<T> acc(0);
    auto pr = segmented_input_range(first,last);

    for( auto const& e : std::get<0>(pr) ) init += e;
    for( auto const& e : std::get<1>(pr) ) acc  += e;
    for( auto const& e : std::get<2>(pr) ) init += e;

    return init + sum(acc);
  }

  /*!
    @ingroup group-std

    Computes the generalized sum of the elements in the given Contiguous Range @range{first,last}
    over the binary functions @c binop and @c reduce, using @c init as the initial value.

    While @c binop is applied over the result of dereferencing the input pointers, @c reduce is to
    be used in the final reduction of the SIMD part of the generalized sum.

    \notebox{The summation order can be different from the order of a sequential summation
            , thus leading to different results.
            }

    @par Example:
    @snippet reduce.phases.cpp reduce-phases
    Possible output:
    @code
    SIMD reduce     : 285
    @endcode

    @param first    Beginning of the range of elements to sum
    @param last     End of the range of elements to sum
    @param init     Initial value of the reduction
    @param binop    Binary function object that will be applied in unspecified order to the
                    result of dereferencing the input pointers, the results of other @c binop
                    and @ init.
    @param neutral  Value containing the neutral element of @c binop
    @param reduce   Binary function object that will be applied to complete the reduction

    @return The generalized sum of the given value and elements in the given range over @ binop.
  **/
  template<typename T, typename U, typename F, typename N, typename G>
  U reduce( T const* first, T const* last, U init, F binop, N neutral, G reduce )
  {
    pack<U> acc(neutral);
    auto pr = segmented_input_range(first,last);

    for( auto const& e : std::get<0>(pr) ) init = binop(init,e);
    for( auto const& e : std::get<1>(pr) ) acc  = binop(acc,e);
    for( auto const& e : std::get<2>(pr) ) init = binop(init,e);
    for( U           e : acc)              init = reduce(init,e);

    return init;
  }

  /*!
    @ingroup group-std

    Computes the generalized sum of the elements in the given Contiguous Range @range{first,last}
    over the binary function @c binop, using @c init as the initial value.

    @c binop is applied over the result of dereferencing the input pointers and in the final
    reduction of the SIMD part of the generalized sum.

    \notebox{The summation order can be different from the order of a sequential summation
            , thus leading to different results.
            }

    @par Example:
    @snippet reduce.phase.cpp reduce-phase
    Possible output:
    @code
    SIMD reduce     : 362880
    @endcode

    @param first    Beginning of the range of elements to sum
    @param last     End of the range of elements to sum
    @param init     Initial value of the reduction
    @param binop    Binary function object that will be applied in unspecified order to the
                    result of dereferencing the input pointers, the results of other @c binop
                    and @c init.
    @param neutral  Value containing the neutral element of @c binop

    @return The sum of the given value and elements in the given range.
  **/
  template<typename T, typename U, typename F, typename N>
  U reduce(T const* first, T const* last, U init, F binop, N neutral)
  {
    return reduce(first,last,init,binop,neutral,binop);
  }
} }

#endif