/*===================================================================

The Medical Imaging Interaction Toolkit (MITK)

Copyright (c) German Cancer Research Center.

All rights reserved.

This software is distributed WITHOUT ANY WARRANTY; without
even the implied warranty of MERCHANTABILITY or FITNESS FOR
A PARTICULAR PURPOSE.

See LICENSE.txt or http://www.mitk.org for details.

===================================================================*/

#include "mitkDiffusionFunctionCollection.h"
#include "mitkNumericTypes.h"

#include <boost/math/special_functions/legendre.hpp>
#include <boost/math/special_functions/spherical_harmonic.hpp>
#include <boost/version.hpp>
#include <boost/algorithm/string.hpp>
#include <itkPointShell.h>
#include "itkVectorContainer.h"
#include "vnl/vnl_vector.h"
#include <vtkBox.h>
#include <vtkMath.h>
#include <itksys/SystemTools.hxx>
#include <itkShToOdfImageFilter.h>
#include <itkTensorImageToOdfImageFilter.h>

// Intersect a finite line (with end points p0 and p1) with all of the
// cells of a vtkImageData
std::vector< std::pair< itk::Index<3>, double > > mitk::imv::IntersectImage(const itk::Vector<double,3>& spacing, itk::Index<3>& si, itk::Index<3>& ei, itk::ContinuousIndex<float, 3>& sf, itk::ContinuousIndex<float, 3>& ef)
{
  std::vector< std::pair< itk::Index<3>, double > > out;
  if (si == ei)
  {
    double d[3];
    for (int i=0; i<3; ++i)
      d[i] = static_cast<double>(sf[i]-ef[i])*spacing[i];
    double len = std::sqrt( d[0]*d[0] + d[1]*d[1] + d[2]*d[2] );
    out.push_back(  std::pair< itk::Index<3>, double >(si, len) );
    return out;
  }

  double bounds[6];

  double entrancePoint[3];
  double exitPoint[3];

  double startPoint[3];
  double endPoint[3];

  double t0, t1;
  for (unsigned int i=0; i<3; ++i)
  {
    startPoint[i] = static_cast<double>(sf[i]);
    endPoint[i] = static_cast<double>(ef[i]);

    if (si[i]>ei[i])
    {
      auto t = si[i];
      si[i] = ei[i];
      ei[i] = t;
    }
  }

  for (auto x = si[0]; x<=ei[0]; ++x)
    for (auto y = si[1]; y<=ei[1]; ++y)
      for (auto z = si[2]; z<=ei[2]; ++z)
      {
        bounds[0] = static_cast<double>(x) - 0.5;
        bounds[1] = static_cast<double>(x) + 0.5;
        bounds[2] = static_cast<double>(y) - 0.5;
        bounds[3] = static_cast<double>(y) + 0.5;
        bounds[4] = static_cast<double>(z) - 0.5;
        bounds[5] = static_cast<double>(z) + 0.5;

        int entryPlane;
        int exitPlane;
        int hit = vtkBox::IntersectWithLine(bounds,
                                            startPoint,
                                            endPoint,
                                            t0,
                                            t1,
                                            entrancePoint,
                                            exitPoint,
                                            entryPlane,
                                            exitPlane);
        if (hit)
        {
          if (entryPlane>=0 && exitPlane>=0)
          {
            double d[3];
            for (int i=0; i<3; ++i)
              d[i] = (exitPoint[i] - entrancePoint[i])*spacing[i];
            double len = std::sqrt( d[0]*d[0] + d[1]*d[1] + d[2]*d[2] );

            itk::Index<3> idx; idx[0] = x; idx[1] = y; idx[2] = z;
            out.push_back(  std::pair< itk::Index<3>, double >(idx, len) );
          }
          else if (entryPlane>=0)
          {
            double d[3];
            for (int i=0; i<3; ++i)
              d[i] = (static_cast<double>(ef[i]) - entrancePoint[i])*spacing[i];
            double len = std::sqrt( d[0]*d[0] + d[1]*d[1] + d[2]*d[2] );

            itk::Index<3> idx; idx[0] = x; idx[1] = y; idx[2] = z;
            out.push_back(  std::pair< itk::Index<3>, double >(idx, len) );
          }
          else if (exitPlane>=0)
          {
            double d[3];
            for (int i=0; i<3; ++i)
              d[i] = (exitPoint[i]-static_cast<double>(sf[i]))*spacing[i];
            double len = std::sqrt( d[0]*d[0] + d[1]*d[1] + d[2]*d[2] );

            itk::Index<3> idx; idx[0] = x; idx[1] = y; idx[2] = z;
            out.push_back(  std::pair< itk::Index<3>, double >(idx, len) );
          }
        }
      }
  return out;
}

//------------------------- SH-function ------------------------------------

double mitk::sh::factorial(int number) {
  if(number <= 1) return 1;
  double result = 1.0;
  for(int i=1; i<=number; i++)
    result *= i;
  return result;
}

void mitk::sh::Cart2Sph(double x, double y, double z, double *spherical)
{
  double phi, th, rad;
  rad = sqrt(x*x+y*y+z*z);
  if( rad < mitk::eps )
  {
    th = itk::Math::pi/2;
    phi = itk::Math::pi/2;
  }
  else
  {
    th = acos(z/rad);
    phi = atan2(y, x);
  }
  spherical[0] = phi;
  spherical[1] = th;
  spherical[2] = rad;
}

vnl_vector_fixed<double, 3> mitk::sh::Sph2Cart(const double& theta, const double& phi, const double& rad)
{
  vnl_vector_fixed<double, 3> dir;
  dir[0] = rad * sin(theta) * cos(phi);
  dir[1] = rad * sin(theta) * sin(phi);
  dir[2] = rad * cos(theta);
  return dir;
}

double mitk::sh::legendre0(int l)
{
  if( l%2 != 0 )
  {
    return 0;
  }
  else
  {
    double prod1 = 1.0;
    for(int i=1;i<l;i+=2) prod1 *= i;
    double prod2 = 1.0;
    for(int i=2;i<=l;i+=2) prod2 *= i;
    return pow(-1.0,l/2.0)*(prod1/prod2);
  }
}

double mitk::sh::Yj(int m, int k, float theta, float phi, bool mrtrix)
{
  if (!mrtrix)
  {
    if (m<0)
      return sqrt(2.0)*static_cast<double>(::boost::math::spherical_harmonic_r(static_cast<unsigned int>(k), -m, theta, phi));
    else if (m==0)
      return static_cast<double>(::boost::math::spherical_harmonic_r(static_cast<unsigned int>(k), m, theta, phi));
    else
      return pow(-1.0,m)*sqrt(2.0)*static_cast<double>(::boost::math::spherical_harmonic_i(static_cast<unsigned int>(k), m, theta, phi));
  }
  else
  {
    double plm = static_cast<double>(::boost::math::legendre_p<float>(k,abs(m),-cos(theta)));
    double mag = sqrt((2.0*k+1.0)/(4.0*itk::Math::pi)*::boost::math::factorial<double>(k-abs(m))/::boost::math::factorial<double>(k+abs(m)))*plm;
    if (m>0)
      return mag*static_cast<double>(cos(m*phi));
    else if (m==0)
      return mag;
    else
      return mag*static_cast<double>(sin(-m*phi));
  }

  return 0;
}

mitk::OdfImage::ItkOdfImageType::Pointer mitk::convert::GetItkOdfFromShImage(mitk::Image::Pointer mitkImage)
{
  mitk::ShImage::Pointer mitkShImage = dynamic_cast<mitk::ShImage*>(mitkImage.GetPointer());
  if (mitkShImage.IsNull())
    mitkThrow() << "Input image is not a SH image!";
  mitk::OdfImage::ItkOdfImageType::Pointer output;
  switch (mitkShImage->ShOrder())
  {
  case 2:
  {
    typedef itk::ShToOdfImageFilter< float, 2 > ShConverterType;
    typename ShConverterType::InputImageType::Pointer itkvol = ShConverterType::InputImageType::New();
    mitk::CastToItkImage(mitkImage, itkvol);
    typename ShConverterType::Pointer converter = ShConverterType::New();
    converter->SetInput(itkvol);
    converter->SetToolkit(mitkShImage->GetShConvention());
    converter->Update();
    output = converter->GetOutput();
    break;
  }
  case 4:
  {
    typedef itk::ShToOdfImageFilter< float, 4 > ShConverterType;
    typename ShConverterType::InputImageType::Pointer itkvol = ShConverterType::InputImageType::New();
    mitk::CastToItkImage(mitkImage, itkvol);
    typename ShConverterType::Pointer converter = ShConverterType::New();
    converter->SetInput(itkvol);
    converter->SetToolkit(mitkShImage->GetShConvention());
    converter->Update();
    output = converter->GetOutput();
    break;
  }
  case 6:
  {
    typedef itk::ShToOdfImageFilter< float, 6 > ShConverterType;
    typename ShConverterType::InputImageType::Pointer itkvol = ShConverterType::InputImageType::New();
    mitk::CastToItkImage(mitkImage, itkvol);
    typename ShConverterType::Pointer converter = ShConverterType::New();
    converter->SetInput(itkvol);
    converter->SetToolkit(mitkShImage->GetShConvention());
    converter->Update();
    output = converter->GetOutput();
    break;
  }
  case 8:
  {
    typedef itk::ShToOdfImageFilter< float, 8 > ShConverterType;
    typename ShConverterType::InputImageType::Pointer itkvol = ShConverterType::InputImageType::New();
    mitk::CastToItkImage(mitkImage, itkvol);
    typename ShConverterType::Pointer converter = ShConverterType::New();
    converter->SetInput(itkvol);
    converter->SetToolkit(mitkShImage->GetShConvention());
    converter->Update();
    output = converter->GetOutput();
    break;
  }
  case 10:
  {
    typedef itk::ShToOdfImageFilter< float, 10 > ShConverterType;
    typename ShConverterType::InputImageType::Pointer itkvol = ShConverterType::InputImageType::New();
    mitk::CastToItkImage(mitkImage, itkvol);
    typename ShConverterType::Pointer converter = ShConverterType::New();
    converter->SetInput(itkvol);
    converter->SetToolkit(mitkShImage->GetShConvention());
    converter->Update();
    output = converter->GetOutput();
    break;
  }
  case 12:
  {
    typedef itk::ShToOdfImageFilter< float, 12 > ShConverterType;
    typename ShConverterType::InputImageType::Pointer itkvol = ShConverterType::InputImageType::New();
    mitk::CastToItkImage(mitkImage, itkvol);
    typename ShConverterType::Pointer converter = ShConverterType::New();
    converter->SetInput(itkvol);
    converter->SetToolkit(mitkShImage->GetShConvention());
    converter->Update();
    output = converter->GetOutput();
    break;
  }
  default:
    mitkThrow() << "SH orders higher than 12 are not supported!";
  }

  return output;
}

mitk::OdfImage::Pointer mitk::convert::GetOdfFromShImage(mitk::Image::Pointer mitkImage)
{
  mitk::OdfImage::Pointer image = mitk::OdfImage::New();
  auto img = GetItkOdfFromShImage(mitkImage);
  image->InitializeByItk( img.GetPointer() );
  image->SetVolume( img->GetBufferPointer() );
  return image;
}

mitk::OdfImage::ItkOdfImageType::Pointer mitk::convert::GetItkOdfFromTensorImage(mitk::Image::Pointer mitkImage)
{
  typedef itk::TensorImageToOdfImageFilter< float, float > FilterType;
  FilterType::Pointer filter = FilterType::New();
  filter->SetInput( GetItkTensorFromTensorImage(mitkImage) );
  filter->Update();
  return filter->GetOutput();
}

mitk::TensorImage::ItkTensorImageType::Pointer mitk::convert::GetItkTensorFromTensorImage(mitk::Image::Pointer mitkImage)
{
  typedef mitk::ImageToItk< mitk::TensorImage::ItkTensorImageType > CasterType;
  CasterType::Pointer caster = CasterType::New();
  caster->SetInput(mitkImage);
  caster->Update();
  return caster->GetOutput();
}

mitk::PeakImage::ItkPeakImageType::Pointer mitk::convert::GetItkPeakFromPeakImage(mitk::Image::Pointer mitkImage)
{
  typedef mitk::ImageToItk< mitk::PeakImage::ItkPeakImageType > CasterType;
  CasterType::Pointer caster = CasterType::New();
  caster->SetInput(mitkImage);
  caster->SetCopyMemFlag(true);
  caster->Update();
  return caster->GetOutput();
}

mitk::OdfImage::Pointer mitk::convert::GetOdfFromTensorImage(mitk::Image::Pointer mitkImage)
{
  mitk::OdfImage::Pointer image = mitk::OdfImage::New();
  auto img = GetItkOdfFromTensorImage(mitkImage);
  image->InitializeByItk( img.GetPointer() );
  image->SetVolume( img->GetBufferPointer() );
  return image;
}

mitk::OdfImage::ItkOdfImageType::Pointer mitk::convert::GetItkOdfFromOdfImage(mitk::Image::Pointer mitkImage)
{
  typedef mitk::ImageToItk< mitk::OdfImage::ItkOdfImageType > CasterType;
  CasterType::Pointer caster = CasterType::New();
  caster->SetInput(mitkImage);
  caster->Update();
  return caster->GetOutput();
}

vnl_matrix<float> mitk::sh::CalcShBasisForDirections(unsigned int sh_order, vnl_matrix<double> U, bool mrtrix)
{
  vnl_matrix<float> sh_basis  = vnl_matrix<float>(U.cols(), (sh_order*sh_order + sh_order + 2)/2 + sh_order );
  for(unsigned int i=0; i<U.cols(); i++)
  {
    double x = U(0,i);
    double y = U(1,i);
    double z = U(2,i);
    double spherical[3];
    mitk::sh::Cart2Sph(x,y,z,spherical);
    U(0,i) = spherical[0];
    U(1,i) = spherical[1];
    U(2,i) = spherical[2];
  }

  for(unsigned int i=0; i<U.cols(); i++)
  {
    for(int k=0; k<=static_cast<int>(sh_order); k+=2)
    {
      for(int m=-k; m<=k; m++)
      {
        int j = (k*k + k + 2)/2 + m - 1;
        double phi = U(0,i);
        double th = U(1,i);
        sh_basis(i,j) = mitk::sh::Yj(m,k,th,phi, mrtrix);
      }
    }
  }

  return sh_basis;
}

unsigned int mitk::sh::ShOrder(int num_coeffs)
{
  int c=3, d=2-2*num_coeffs;
  int D = c*c-4*d;
  if (D>0)
  {
    int s = (-c+static_cast<int>(sqrt(D)))/2;
    if (s<0)
      s = (-c-static_cast<int>(sqrt(D)))/2;
    return static_cast<unsigned int>(s);
  }
  else if (D==0)
    return static_cast<unsigned int>(-c/2);
  return 0;
}

float mitk::sh::GetValue(const vnl_vector<float> &coefficients, const int &sh_order, const double theta, const double phi, const bool mrtrix)
{
  float val = 0;
  for(int k=0; k<=sh_order; k+=2)
  {
    for(int m=-k; m<=k; m++)
    {
      unsigned int j = static_cast<unsigned int>((k*k + k + 2)/2 + m - 1);
      val += coefficients[j] * mitk::sh::Yj(m, k, theta, phi, mrtrix);
    }
  }

  return val;
}

float mitk::sh::GetValue(const vnl_vector<float> &coefficients, const int &sh_order, const vnl_vector_fixed<double, 3> &dir, const bool mrtrix)
{
  double spherical[3];
  mitk::sh::Cart2Sph(dir[0], dir[1], dir[2], spherical);

  float val = 0;
  for(int k=0; k<=sh_order; k+=2)
  {
    for(int m=-k; m<=k; m++)
    {
      int j = (k*k + k + 2)/2 + m - 1;
      val += coefficients[j] * mitk::sh::Yj(m, k, spherical[1], spherical[0], mrtrix);
    }
  }

  return val;
}

//------------------------- gradients-function ------------------------------------

mitk::gradients::GradientDirectionContainerType::ConstPointer mitk::gradients::ReadBvalsBvecs(std::string bvals_file, std::string bvecs_file, double& reference_bval)
{
  mitk::gradients::GradientDirectionContainerType::Pointer directioncontainer = mitk::gradients::GradientDirectionContainerType::New();

  std::vector<float> bvec_entries;
  if (!itksys::SystemTools::FileExists(bvecs_file))
    mitkThrow() << "bvecs file not existing: " << bvecs_file;
  else
  {
    std::string line;
    std::ifstream myfile (bvecs_file.c_str());
    if (myfile.is_open())
    {
      while (std::getline(myfile, line))
      {
        std::vector<std::string> strs;
        boost::split(strs,line,boost::is_any_of("\t \n"));
        for (auto token : strs)
        {
          if (!token.empty())
          {
            try
            {
              bvec_entries.push_back(boost::lexical_cast<float>(token));
            }
            catch(...)
            {
              mitkThrow() << "Encountered invalid bvecs file entry >" << token << "<";
            }
          }
        }
      }
      myfile.close();
    }
    else
    {
      mitkThrow() << "bvecs file could not be opened: " << bvals_file;
    }
  }

  reference_bval = -1;
  std::vector<float> bval_entries;
  if (!itksys::SystemTools::FileExists(bvals_file))
    mitkThrow() << "bvals file not existing: " << bvals_file;
  else
  {
    std::string line;
    std::ifstream myfile (bvals_file.c_str());
    if (myfile.is_open())
    {
      while (std::getline(myfile, line))
      {
        std::vector<std::string> strs;
        boost::split(strs,line,boost::is_any_of("\t \n"));
        for (auto token : strs)
        {
          if (!token.empty())
          {
            try {
              bval_entries.push_back(boost::lexical_cast<float>(token));
              if (bval_entries.back()>reference_bval)
                reference_bval = bval_entries.back();
            }
            catch(...)
            {
              mitkThrow() << "Encountered invalid bvals file entry >" << token << "<";
            }
          }
        }
      }
      myfile.close();
    }
    else
    {
      mitkThrow() << "bvals file could not be opened: " << bvals_file;
    }
  }

  for(unsigned int i=0; i<bval_entries.size(); i++)
  {
    double b_val = bval_entries.at(i);

    mitk::gradients::GradientDirectionType vec;
    vec[0] = bvec_entries.at(i);
    vec[1] = bvec_entries.at(i+bval_entries.size());
    vec[2] = bvec_entries.at(i+2*bval_entries.size());

    // Adjust the vector length to encode gradient strength
    if (reference_bval>0)
    {
      double factor = b_val/reference_bval;
      if(vec.magnitude() > 0)
      {
        vec.normalize();
        vec[0] = sqrt(factor)*vec[0];
        vec[1] = sqrt(factor)*vec[1];
        vec[2] = sqrt(factor)*vec[2];
      }
    }

    directioncontainer->InsertElement(i,vec);
  }

  return GradientDirectionContainerType::ConstPointer(directioncontainer);
}

void mitk::gradients::WriteBvalsBvecs(std::string bvals_file, std::string bvecs_file, GradientDirectionContainerType::ConstPointer gradients, double reference_bval)
{
  std::ofstream myfile;
  myfile.open (bvals_file.c_str());
  for(unsigned int i=0; i<gradients->Size(); i++)
  {
    double twonorm = gradients->ElementAt(i).two_norm();
    myfile << std::round(reference_bval*twonorm*twonorm) << " ";
  }
  myfile.close();

  std::ofstream myfile2;
  myfile2.open (bvecs_file.c_str());
  for(int j=0; j<3; j++)
  {
    for(unsigned int i=0; i<gradients->Size(); i++)
    {
      GradientDirectionType direction = gradients->ElementAt(i);
      direction.normalize();
      myfile2 << direction.get(j) << " ";
    }
    myfile2 << std::endl;
  }
}

std::vector<unsigned int> mitk::gradients::GetAllUniqueDirections(const BValueMap & refBValueMap, GradientDirectionContainerType::ConstPointer refGradientsContainer )
{

  IndiciesVector directioncontainer;
  auto mapIterator = refBValueMap.begin();

  if(refBValueMap.find(0) != refBValueMap.end() && refBValueMap.size() > 1)
    mapIterator++; //skip bzero Values

  for( ; mapIterator != refBValueMap.end(); mapIterator++){

    IndiciesVector currentShell = mapIterator->second;

    while(currentShell.size()>0)
    {
      unsigned int wntIndex = currentShell.back();
      currentShell.pop_back();

      auto containerIt = directioncontainer.begin();
      bool directionExist = false;
      while(containerIt != directioncontainer.end())
      {
        if (fabs(dot_product(refGradientsContainer->ElementAt(*containerIt), refGradientsContainer->ElementAt(wntIndex)))  > 0.9998)
        {
          directionExist = true;
          break;
        }
        containerIt++;
      }
      if(!directionExist)
      {
        directioncontainer.push_back(wntIndex);
      }
    }
  }

  return directioncontainer;
}

bool mitk::gradients::CheckForDifferingShellDirections(const BValueMap & refBValueMap, GradientDirectionContainerType::ConstPointer refGradientsContainer)
{
  auto mapIterator = refBValueMap.begin();

  if(refBValueMap.find(0) != refBValueMap.end() && refBValueMap.size() > 1)
    mapIterator++; //skip bzero Values

  for( ; mapIterator != refBValueMap.end(); mapIterator++){

    auto mapIterator_2 = refBValueMap.begin();
    if(refBValueMap.find(0) != refBValueMap.end() && refBValueMap.size() > 1)
      mapIterator_2++; //skip bzero Values

    for( ; mapIterator_2 != refBValueMap.end(); mapIterator_2++){

      if(mapIterator_2 == mapIterator) continue;

      IndiciesVector currentShell = mapIterator->second;
      IndiciesVector testShell = mapIterator_2->second;
      for (unsigned int i = 0; i< currentShell.size(); i++)
        if (fabs(dot_product(refGradientsContainer->ElementAt(currentShell[i]), refGradientsContainer->ElementAt(testShell[i])))  <= 0.9998) { return true; }

    }
  }
  return false;
}

vnl_matrix<double> mitk::gradients::ComputeSphericalFromCartesian(const IndiciesVector & refShell, const GradientDirectionContainerType * refGradientsContainer)
{

  vnl_matrix<double> Q(3, refShell.size());
  Q.fill(0.0);

  for(unsigned int i = 0; i < refShell.size(); i++)
  {
    GradientDirectionType dir = refGradientsContainer->ElementAt(refShell[i]);
    double x = dir.normalize().get(0);
    double y = dir.normalize().get(1);
    double z = dir.normalize().get(2);
    double cart[3];
    mitk::sh::Cart2Sph(x,y,z,cart);
    Q(0,i) = cart[0];
    Q(1,i) = cart[1];
    Q(2,i) = cart[2];
  }
  return Q;
}

vnl_matrix<double> mitk::gradients::ComputeSphericalHarmonicsBasis(const vnl_matrix<double> & QBallReference, const unsigned int & LOrder)
{
  vnl_matrix<double> SHBasisOutput(QBallReference.cols(), (LOrder+1)*(LOrder+2)*0.5);
  SHBasisOutput.fill(0.0);
  for(unsigned int i=0; i<SHBasisOutput.rows(); i++)
    for(int k = 0; k <= (int)LOrder; k += 2)
      for(int m =- k; m <= k; m++)
      {
        int j = ( k * k + k + 2 ) / 2.0 + m - 1;
        double phi = QBallReference(0,i);
        double th = QBallReference(1,i);
        double val = mitk::sh::Yj(m,k,th,phi);
        SHBasisOutput(i,j) = val;
      }
  return SHBasisOutput;
}

mitk::gradients::GradientDirectionContainerType::Pointer mitk::gradients::CreateNormalizedUniqueGradientDirectionContainer(const mitk::gradients::BValueMap & bValueMap,
                                                                                const GradientDirectionContainerType *origninalGradentcontainer)
{
  mitk::gradients::GradientDirectionContainerType::Pointer directioncontainer = mitk::gradients::GradientDirectionContainerType::New();
  auto mapIterator = bValueMap.begin();

  if(bValueMap.find(0) != bValueMap.end() && bValueMap.size() > 1){
    mapIterator++; //skip bzero Values
    vnl_vector_fixed<double, 3> vec;
    vec.fill(0.0);
    directioncontainer->push_back(vec);
  }

  for( ; mapIterator != bValueMap.end(); mapIterator++){

    IndiciesVector currentShell = mapIterator->second;

    while(currentShell.size()>0)
    {
      unsigned int wntIndex = currentShell.back();
      currentShell.pop_back();

      mitk::gradients::GradientDirectionContainerType::Iterator containerIt = directioncontainer->Begin();
      bool directionExist = false;
      while(containerIt != directioncontainer->End())
      {
        if (fabs(dot_product(containerIt.Value(), origninalGradentcontainer->ElementAt(wntIndex)))  > 0.9998)
        {
          directionExist = true;
          break;
        }
        containerIt++;
      }
      if(!directionExist)
      {
        GradientDirectionType dir(origninalGradentcontainer->ElementAt(wntIndex));
        directioncontainer->push_back(dir.normalize());
      }
    }
  }

  return directioncontainer;
}


//==================================================================================================
/*!
  @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_ARCH_COMMON_DETAIL_SIMD_TRIG_BASE_HPP_INCLUDED
#define BOOST_SIMD_ARCH_COMMON_DETAIL_SIMD_TRIG_BASE_HPP_INCLUDED

#include <boost/simd/arch/common/detail/simd/trig_reduction.hpp>
#include <boost/simd/function/any.hpp>
#include <boost/simd/function/abs.hpp>
#include <boost/simd/function/is_invalid.hpp>
#include <boost/simd/function/is_eqz.hpp>
#include <boost/simd/function/is_nez.hpp>
#include <boost/simd/function/is_flint.hpp>
#include <boost/simd/function/shift_left.hpp>
#include <boost/simd/function/shr.hpp>
#include <boost/simd/function/sqr.hpp>
#include <boost/simd/function/if_nan_else.hpp>
#include <boost/simd/function/if_allbits_else_zero.hpp>
#include <boost/simd/function/bitwise_xor.hpp>
#include <boost/simd/function/if_else.hpp>
#include <boost/simd/function/bitofsign.hpp>
#include <boost/simd/function/bitwise_and.hpp>
#include <boost/simd/function/is_greater.hpp>
#include <boost/simd/function/rec.hpp>
#include <boost/simd/function/oneminus.hpp>
#include <boost/simd/detail/constant/maxleftshift.hpp>
#include <boost/simd/constant/nan.hpp>
#include <boost/simd/constant/one.hpp>
#include <boost/simd/constant/two.hpp>
#include <boost/simd/constant/mtwo.hpp>
#include <boost/simd/constant/signmask.hpp>
#include <boost/simd/arch/common/detail/tags.hpp>
#include <boost/simd/meta/as_logical.hpp>
#include <boost/simd/detail/dispatch/meta/as_integer.hpp>
#include <boost/simd/detail/dispatch/meta/scalar_of.hpp>

namespace boost { namespace simd
{
  namespace detail
  {
    namespace bd =  boost::dispatch;
    namespace bs =  boost::simd;

    template < class A0, class unit_tag, class mode>
    struct trig_base < A0, unit_tag,  tag::simd_type, mode>
    {
      typedef typename bs::as_logical_t<A0>                        bA0; // logical type associated to A0
      typedef trig_reduction<A0,unit_tag, tag::simd_type, mode> redu_t;
      typedef trig_evaluation<A0,tag::simd_type>                eval_t;
      typedef typename bd::scalar_of<A0>::type                     sA0; // scalar version of A0
      typedef typename bd::as_integer<A0, signed>::type       int_type; // signed integer type associated to A0
      typedef typename bs::as_logical<int_type>::type        bint_type; // logical type associated to int_type
      typedef typename bd::scalar_of<int_type>::type         sint_type; // scalar version of the associated type
      typedef typename mode::type                                style;

      // for all functions the scalar algorithm is:
      // * range reduction
      // * computation of sign and evaluation selections flags
      // * evaluations of the two branches and selection using flags
      // * return with flag based corrections and inf and nan or specific invalid cases inputs considerations

      static BOOST_FORCEINLINE A0 cosa(const A0& a0){ return cosa(a0, style()); }
      static BOOST_FORCEINLINE A0 sina(const A0& a0){ return sina(a0, style()); }
      static BOOST_FORCEINLINE A0 tana(const A0& a0){ return tana(a0, style()); }
      static BOOST_FORCEINLINE A0 cota(const A0& a0){ return cota(a0, style()); }
      static BOOST_FORCEINLINE std::pair<A0,A0> sincosa(const A0& a0){ return sincosa(a0,style()); }

    private:
      static BOOST_FORCEINLINE A0 cosa(const A0& a0, const tag::restricted&)
      {
        const A0 x =  scale(a0);
        return  eval_t::cos_eval(sqr(x));
      }

      static BOOST_FORCEINLINE A0 cosa(const A0& a0, const tag::regular&)
      {
        const A0 x = bs::abs(a0);
        A0 xr = Nan<A0>();
        A0 n =  redu_t::reduce(x, xr);
        auto tmp = if_one_else_zero(n >= Two<A0>());
        auto swap_bit = (fma(Mtwo<A0>(), tmp, n));
        auto sign_bit = if_else_zero(bitwise_xor(swap_bit, tmp), Signmask<A0>());
        A0 z = sqr(xr);
        A0 se = eval_t::sin_eval(z, xr);
        A0 ce = eval_t::cos_eval(z);
        A0 z1 = if_else(is_nez(swap_bit), se, ce);
        return bitwise_xor(z1, sign_bit);
      }

      static BOOST_FORCEINLINE A0 sina(const A0& a0_n, const tag::restricted&)
      {
        A0 x =   scale(a0_n);
        A0 se = eval_t::sin_eval(sqr(x), x);
        return se;
      }

      static BOOST_FORCEINLINE A0 sina(const A0& a0, const tag::regular&)
      {
        const A0 x = bs::abs(a0);
        A0 xr = Nan<A0>();
        const A0 n = redu_t::reduce(x, xr);
        auto tmp = if_one_else_zero(n >= Two<A0>());
        auto swap_bit = (fma(Mtwo<A0>(), tmp, n));
        auto sign_bit = bitwise_xor(bitofsign(a0), if_else_zero(tmp, Signmask<A0>()));
        const A0 z = sqr(xr);
        const A0 se = eval_t::sin_eval(z, xr);
        const A0 ce = eval_t::cos_eval(z);
        const A0 z1 = if_else(is_eqz(swap_bit), se, ce);
        return bitwise_xor(z1, sign_bit);
      }

      static BOOST_FORCEINLINE A0 tana(const A0& a0, const tag::restricted&)
      {
        return eval_t::base_tancot_eval(scale(a0));
      }

      static BOOST_FORCEINLINE A0 tana(const A0& a0, const tag::regular&)
      {
        A0 x =  bs::abs(a0);
        A0 xr = Nan<A0>();
        A0 n = redu_t::reduce(x, xr);
        auto tmp = if_one_else_zero( n >= Two<A0>());
        auto swap_bit = (fma(Mtwo<A0>(), tmp, n));
        auto test = is_eqz(swap_bit);
        const A0 y = eval_t::tan_eval(xr, test);
        auto testnan = redu_t::tan_invalid(a0);
        return if_nan_else(testnan, bitwise_xor(y, bitofsign(a0)));
      }

      static BOOST_FORCEINLINE A0 cota(const A0& a0, const tag::restricted&)
      {
        return eval_t::base_tancot_eval(scale(a0));
      }

      static BOOST_FORCEINLINE A0 cota(const A0& a0, const tag::regular&)
      {
        const A0 x = bs::abs(a0);
        A0 xr = Nan<A0>();
        const A0 n = redu_t::reduce(x, xr);
        auto tmp = if_one_else_zero( n >= Two<A0>());
        auto swap_bit = (fma(Mtwo<A0>(), tmp, n));
        auto test = is_eqz(swap_bit);
        const A0 y = eval_t::cot_eval(xr, test);
        const bA0 testnan = redu_t::cot_invalid(a0);
        // this if_else is normally not needed but with clang the zero value if erroneous
        // if not there !
        return if_else(is_nez(a0), if_nan_else(testnan
                                              , bitwise_xor(y, bitofsign(a0))), rec(a0));
      }

      // simultaneous cosa and sina function
      static BOOST_FORCEINLINE std::pair<A0, A0> sincosa(const A0& a0, const tag::restricted&)
      {
        A0 x =  scale(a0);
        A0 z =  sqr(x);
        return {eval_t::sin_eval(z, x), eval_t::cos_eval(z)};
      }

      static BOOST_FORCEINLINE  std::pair<A0, A0> sincosa(const A0& a0, const tag::regular&)
      {
        A0 x =  bs::abs(a0);
        A0 xr = Nan<A0>();
        A0 n = redu_t::reduce(x, xr);
        auto tmp = if_one_else_zero(n >= Two<A0>());
        auto swap_bit = (fma(Mtwo<A0>(), tmp, n));
        auto cos_sign_bit = if_else_zero(bitwise_xor(swap_bit, tmp), Signmask<A0>());
        auto sin_sign_bit = bitwise_xor(bitofsign(a0),if_else_zero(tmp, Signmask<A0>()));

        A0 z = bs::sqr(xr);
        A0 t1 = eval_t::sin_eval(z, xr);
        A0 t2 = eval_t::cos_eval(z);
        auto test = is_nez(swap_bit);
        return { bitwise_xor(if_else(test, t2, t1),sin_sign_bit)
               , bitwise_xor(if_else(test, t1, t2),cos_sign_bit) };
      }

      static BOOST_FORCEINLINE A0 scale(const A0& a0)
      {
        return if_nan_else(is_greater(bs::abs(a0),
                              trig_ranges<A0,unit_tag>::max_range()), a0)
          *trig_ranges<A0,unit_tag>::scale();
      }

    };
  }
} }

#endif


#include <boost/python/module.hpp>
#include <boost/python/def.hpp>

#include <cctbx/covariance/covariance.h>

namespace cctbx { namespace covariance {
  namespace boost_python {

    void wrap_covariance_matrix() {
      using namespace boost::python;

      def("extract_covariance_matrix_for_sites", (
        af::versa<double, af::packed_u_accessor>(*)(
        af::const_ref<std::size_t> const &,
        af::const_ref<double, af::packed_u_accessor> const &,
        cctbx::xray::parameter_map<cctbx::xray::scatterer<double> > const &))
        extract_covariance_matrix_for_sites,
        (arg("i_seqs"), arg("matrix"), arg("parameter_map")));
      def("extract_covariance_matrix_for_u_aniso", (
        af::versa<double, af::packed_u_accessor>(*)(
        std::size_t,
        af::const_ref<double, af::packed_u_accessor> const &,
        cctbx::xray::parameter_map<cctbx::xray::scatterer<double> > const &))
        extract_covariance_matrix_for_u_aniso,
        (arg("i_seq"), arg("matrix"), arg("parameter_map")));
      def("variance_for_u_iso", (
        double(*)(
        std::size_t,
        af::const_ref<double, af::packed_u_accessor> const &,
        cctbx::xray::parameter_map<cctbx::xray::scatterer<double> > const &))
        variance_for_u_iso,
        (arg("i_seq"), arg("matrix"), arg("parameter_map")));
      def("orthogonalize_covariance_matrix", (
        af::versa<double, af::packed_u_accessor>(*)(
        af::const_ref<double, af::packed_u_accessor> const &,
        cctbx::uctbx::unit_cell const &,
        cctbx::xray::parameter_map<cctbx::xray::scatterer<double> > const &))
        orthogonalize_covariance_matrix,
        (arg("matrix"), arg("unit_cell"), arg("parameter_map")));
      def("covariance_orthogonalization_matrix", (
       scitbx::sparse::matrix<double>(*)(
        cctbx::uctbx::unit_cell const &,
        cctbx::xray::parameter_map<cctbx::xray::scatterer<double> > const &))
        covariance_orthogonalization_matrix,
        (arg("unit_cell"), arg("parameter_map")));
    }

    void init_module()
    {
      wrap_covariance_matrix();
    }

}}} // namespace cctbx::covariance::boost_python

BOOST_PYTHON_MODULE(cctbx_covariance_ext)
{
  cctbx::covariance::boost_python::init_module();
}


#include <stdlib.h>
#include <iostream>
#include <functional>
#include <Eigen/Eigenvalues>
#include <Eigen/Geometry>
#include <Eigen/Cholesky>
#include <type_traits>
#include <random>
#include <array>
#include <map>

#ifdef ENABLE_PARALLEL
#include <omp.h>
#endif

#ifndef ARC_HELPERS_HPP
#define ARC_HELPERS_HPP

// Branch prediction hints
// Figure out which compiler we have
#if defined(__clang__)
    /* Clang/LLVM */
    #define likely(x) __builtin_expect(!!(x), 1)
    #define unlikely(x) __builtin_expect(!!(x), 0)
#elif defined(__ICC) || defined(__INTEL_COMPILER)
    /* Intel ICC/ICPC */
    #define likely(x) __builtin_expect(!!(x), 1)
    #define unlikely(x) __builtin_expect(!!(x), 0)
#elif defined(__GNUC__) || defined(__GNUG__)
    /* GNU GCC/G++ */
    #define likely(x) __builtin_expect(!!(x), 1)
    #define unlikely(x) __builtin_expect(!!(x), 0)
#elif defined(_MSC_VER)
    /* Microsoft Visual Studio */
    /* MSVC doesn't support branch prediction hints. Use PGO instead. */
    #define likely(x) (x)
    #define unlikely(x) (x)
#endif

// Macro to disable unused parameter compiler warnings
#define UNUSED(x) (void)(x)

namespace arc_helpers
{
    template <typename T>
    inline T SetBit(const T current, const uint32_t bit_position, const bool bit_value)
    {
        // Safety check on the type we've been called with
        static_assert((std::is_same<T, uint8_t>::value
                       || std::is_same<T, uint16_t>::value
                       || std::is_same<T, uint32_t>::value
                       || std::is_same<T, uint64_t>::value),
                      "Type must be a fixed-size unsigned integral type");
        // Do it
        T update_mask = 1;
        update_mask = update_mask << bit_position;
        if (bit_value)
        {
            return (current | update_mask);
        }
        else
        {
            update_mask = (~update_mask);
            return (current & update_mask);
        }
    }

    template <typename T>
    inline bool GetBit(const T current, const uint32_t bit_position)
    {
        const uint32_t mask = arc_helpers::SetBit((T)0, bit_position, true);
        if ((mask & current) > 0)
        {
            return true;
        }
        else
        {
            return false;
        }
    }

    template <class T>
    inline T ClampValue(const T& val, const T& min, const T& max)
    {
        return std::min(max, std::max(min, val));
    }

    template<typename Datatype, typename Allocator=std::allocator<Datatype>>
    static Eigen::MatrixXd BuildDistanceMatrix(const std::vector<Datatype, Allocator>& data, const std::function<double(const Datatype&, const Datatype&)>& distance_fn)
    {
        Eigen::MatrixXd distance_matrix(data.size(), data.size());
#ifdef ENABLE_PARALLEL
        #pragma omp parallel for schedule(guided)
#endif
        for (size_t idx = 0; idx < data.size(); idx++)
        {
            for (size_t jdx = idx; jdx < data.size(); jdx++)
            {
                const double distance = distance_fn(data[idx], data[jdx]);
                distance_matrix((ssize_t)idx, (ssize_t)jdx) = distance;
                distance_matrix((ssize_t)jdx, (ssize_t)idx) = distance;
            }
        }
        return distance_matrix;
    }

    class SplitMix64PRNG
    {
    private:

        uint64_t state_; /* The state can be seeded with any value. */

        inline uint64_t next(void)
        {
            uint64_t z = (state_ += UINT64_C(0x9E3779B97F4A7C15));
            z = (z ^ (z >> 30)) * UINT64_C(0xBF58476D1CE4E5B9);
            z = (z ^ (z >> 27)) * UINT64_C(0x94D049BB133111EB);
            return z ^ (z >> 31);
        }

    public:

        inline SplitMix64PRNG(const uint64_t seed_val)
        {
            seed(seed_val);
        }

        static constexpr uint64_t min(void)
        {
            return 0u;
        }

        static constexpr uint64_t max(void)
        {
            return std::numeric_limits<uint64_t>::max();
        }

        inline void seed(const uint64_t seed_val)
        {
            state_ = seed_val;
        }

        inline void discard(const unsigned long long z)
        {
            uint64_t temp __attribute__((unused)); // This suppresses "set but not used" warnings
            temp = 0u;
            for (unsigned long long i = 0; i < z; i++)
            {
                temp = next();
                __asm__ __volatile__(""); // This should prevent the compiler from optimizing out the loop
            }
        }

        inline uint64_t operator() (void)
        {
            return next();
        }
    };

    class XorShift128PlusPRNG
    {
    private:

        uint64_t state_1_;
        uint64_t state_2_;

        inline uint64_t next(void)
        {
            uint64_t s1 = state_1_;
            const uint64_t s0 = state_2_;
            state_1_ = s0;
            s1 ^= s1 << 23; // a
            state_2_ = s1 ^ s0 ^ (s1 >> 18) ^ (s0 >> 5); // b, c
            return state_2_ + s0;
        }

    public:

        inline XorShift128PlusPRNG(const uint64_t seed_val)
        {
            seed(seed_val);
        }

        static constexpr uint64_t min(void)
        {
            return 0u;
        }

        static constexpr uint64_t max(void)
        {
            return std::numeric_limits<uint64_t>::max();
        }

        inline void seed(const uint64_t seed_val)
        {
            SplitMix64PRNG temp_seed_gen(seed_val);
            state_1_ = temp_seed_gen();
            state_2_ = temp_seed_gen();
        }

        inline void discard(const unsigned long long z)
        {
            uint64_t temp __attribute__((unused)); // This suppresses "set but not used" warnings
            temp = 0u;
            for (unsigned long long i = 0; i < z; i++)
            {
                temp = next();
                __asm__ __volatile__(""); // This should prevent the compiler from optimizing out the loop
            }
        }

        inline uint64_t operator() (void)
        {
            return next();
        }
    };

    class XorShift1024StarPRNG
    {
    private:

        std::array<uint64_t, 16> state_;
        int32_t p;

        inline uint64_t next(void)
        {
            const uint64_t s0 = state_[(size_t)p];
            p = (p + 1) & 15;
            uint64_t s1 = state_[(size_t)p];
            s1 ^= s1 << 31; // a
            state_[(size_t)p] = s1 ^ s0 ^ (s1 >> 11) ^ (s0 >> 30); // b,c
            return state_[(size_t)p] * UINT64_C(1181783497276652981);
        }

    public:

        inline XorShift1024StarPRNG(const uint64_t seed_val)
        {
            seed(seed_val);
            p = 0;
        }

        static constexpr uint64_t min(void)
        {
            return 0u;
        }

        static constexpr uint64_t max(void)
        {
            return std::numeric_limits<uint64_t>::max();
        }

        inline void seed(const uint64_t seed_val)
        {
            SplitMix64PRNG temp_seed_gen(seed_val);
            for (size_t idx = 0u; idx < state_.size(); idx++)
            {
                state_[idx] = temp_seed_gen();
            }
        }

        inline void discard(const unsigned long long z)
        {
            uint64_t temp __attribute__((unused)); // This suppresses "set but not used" warnings
            temp = 0u;
            for (unsigned long long i = 0; i < z; i++)
            {
                temp = next();
                __asm__ __volatile__(""); // This should prevent the compiler from optimizing out the loop
            }
        }

        inline uint64_t operator() (void)
        {
            return next();
        }
    };

    class TruncatedNormalDistribution
    {
    protected:

        double mean_;
        double stddev_;
        double std_lower_bound_;
        double std_upper_bound_;

        enum CASES {TYPE_1, TYPE_2, TYPE_3, TYPE_4, NONE};
        CASES case_;
        std::uniform_real_distribution<double> uniform_unit_dist_;
        std::uniform_real_distribution<double> uniform_range_dist_;
        std::exponential_distribution<double> exponential_dist_;
        std::normal_distribution<double> normal_dist_;

        inline bool CheckSimple(const double lower_bound, const double upper_bound) const
        {
            // Init Values Used in Inequality of Interest
            const double val1 = (2 * sqrt(exp(1))) / (lower_bound + sqrt(pow(lower_bound, 2) + 4));
            const double val2 = exp((pow(lower_bound, 2) - lower_bound * sqrt(pow(lower_bound, 2) + 4)) / (4));
            if (upper_bound > lower_bound + val1 * val2)
            {
                return true;
            }
            else
            {
                return false;
            }
        }

        // Naive Accept-Reject algorithm
        template<typename Generator>
        inline double NaiveAcceptReject(const double lower_bound, const double upper_bound, Generator& prng)
        {
            while (true)
            {
                const double draw = normal_dist_(prng); // Rf_rnorm(0.0, 1.0) ; // Normal distribution (i.e. std::normal_distribution<double>)
                if ((draw <= upper_bound) && (draw >= lower_bound))
                {
                    return draw;
                }
            }
        }

        // Accept-Reject Algorithm
        template<typename Generator>
        inline double SimpleAcceptReject(const double lower_bound, Generator& prng)
        {
            // Init Values
            const double alpha = (lower_bound + sqrt(pow(lower_bound, 2) + 4.0)) / (2.0) ;
            while (true)
            {
                const double e = exponential_dist_(prng); // Rf_rexp(1.0) ; // Exponential distribution (i.e. std::exponential_distribution<double>)
                const double z = lower_bound + e / alpha;
                const double rho = exp(-pow(alpha - z, 2) / 2);
                const double u = uniform_unit_dist_(prng); //  Rf_runif(0, 1) ; // Uniform distribution (i.e. std::uniform_real_distribution<double>)
                if (u <= rho)
                {
                    return z;
                }
            }
        }

        // Accept-Reject Algorithm
        template<typename Generator>
        inline double ComplexAcceptReject(const double lower_bound, const double upper_bound, Generator& prng)
        {
            while (true)
            {
                const double z = uniform_range_dist_(prng); // Rf_runif(lower_bound, upper_bound) ; // Uniform distribution (i.e. std::uniform_real_distribution<double>)
                double rho = 0.0;
                if (0 < lower_bound)
                {
                    rho = exp((pow(lower_bound, 2) - pow(z, 2)) / 2);
                }
                else if (upper_bound < 0)
                {
                    rho = exp((pow(upper_bound, 2) - pow(z, 2)) / 2);
                }
                else if (0 < upper_bound && lower_bound < 0)
                {
                    rho = exp(- pow(z, 2) / 2);
                }
                const double u = uniform_unit_dist_(prng); // Rf_runif(0, 1) ; // Uniform distribution (i.e. std::uniform_real_distribution<double>)
                if (u <= rho)
                {
                    return z;
                }
            }
        }

        template<typename Generator>
        inline double Sample(Generator& prng)
        {
            if (case_ == TYPE_1)
            {
                const double draw = NaiveAcceptReject(std_lower_bound_, std_upper_bound_, prng);
                return mean_ + stddev_ * draw;
            }
            else if (case_ == TYPE_2)
            {
                const double draw = SimpleAcceptReject(std_lower_bound_, prng);
                return mean_ + stddev_ * draw;
            }
            else if (case_ == TYPE_3)
            {
                while (true)
                {
                    const double draw = SimpleAcceptReject(std_lower_bound_, prng); // use the simple algorithm if it is more efficient
                    if (draw <= std_upper_bound_)
                    {
                        return mean_ + stddev_ * draw;
                    }
                }
            }
            else if (case_ == TYPE_4)
            {
                const double draw = ComplexAcceptReject(std_lower_bound_, std_upper_bound_, prng);
                return mean_ + stddev_ * draw;
            }
            else
            {
                assert(case_ == NONE);
                return mean_;
            }
        }

    public:

        inline TruncatedNormalDistribution(const double mean, const double stddev, const double lower_bound, const double upper_bound) : uniform_unit_dist_(0.0, 1.0), uniform_range_dist_(lower_bound, upper_bound), exponential_dist_(1.0), normal_dist_(0.0, 1.0)
        {
            // Set operating parameters
            mean_ = mean;
            stddev_ = stddev;
            if (fabs(stddev_) == 0.0)
            {
                case_ = NONE;
            }
            else
            {
                // Standardize the lower and upper bounds
                std_lower_bound_ = (lower_bound - mean_) / stddev_;
                std_upper_bound_ = (upper_bound - mean_) / stddev_;
                // Set the operating case - i.e. which sampling method we will use
                case_ = NONE;
                if (0.0 <= std_upper_bound_ && 0.0 >= std_lower_bound_)
                {
                    case_ = TYPE_1;
                }
                if (0.0 < std_lower_bound_ && std_upper_bound_ == INFINITY)
                {
                    case_ = TYPE_2;
                }
                if (0.0 > std_upper_bound_ && std_lower_bound_ == -INFINITY)
                {
                    std_lower_bound_ = -1 * std_upper_bound_;
                    std_upper_bound_ = INFINITY;
                    stddev_ = -1 * stddev_;
                    case_ = TYPE_2;
                }
                if ((0.0 > std_upper_bound_ || 0.0 < std_lower_bound_) && !(std_upper_bound_ == INFINITY || std_lower_bound_ == -INFINITY))
                {
                    if (CheckSimple(std_lower_bound_, std_upper_bound_))
                    {
                        case_ = TYPE_3;
                    }
                    else
                    {
                        case_ = TYPE_4;
                    }
                }
                assert((case_ == TYPE_1) || (case_ == TYPE_2) || (case_ == TYPE_3) || (case_ == TYPE_4));
            }
        }

        template<typename Generator>
        inline double operator()(Generator& prng)
        {
            return Sample(prng);
        }
    };

    class MultivariteGaussianDistribution
    {
    protected:
        const Eigen::VectorXd mean_;
        const Eigen::MatrixXd norm_transform_;

        std::normal_distribution<double> unit_gaussian_dist_;

        template<typename Generator>
        inline Eigen::VectorXd Sample(Generator& prng)
        {
            Eigen::VectorXd draw;
            draw.resize(mean_.rows());

            for (ssize_t idx = 0; idx < draw.rows(); idx++)
            {
                draw(idx) = unit_gaussian_dist_(prng);
            }

            return norm_transform_ * draw + mean_;
        }

        static Eigen::MatrixXd CalculateNormTransform(const Eigen::MatrixXd& covariance)
        {
            Eigen::MatrixXd norm_transform;

            Eigen::LLT<Eigen::MatrixXd> chol_solver(covariance);

            if (chol_solver.info() == Eigen::Success)
            {
                // Use cholesky solver
                norm_transform = chol_solver.matrixL();
            }
            else
            {
                // Use eigen solver
                Eigen::SelfAdjointEigenSolver<Eigen::MatrixXd> eigen_solver(covariance);
                norm_transform = eigen_solver.eigenvectors() * eigen_solver.eigenvalues().cwiseMax(0.0).cwiseSqrt().asDiagonal();
            }

            return norm_transform;
        }

    public:
        inline MultivariteGaussianDistribution(const Eigen::VectorXd& mean, const Eigen::MatrixXd& covariance) : mean_(mean), norm_transform_(CalculateNormTransform(covariance)), unit_gaussian_dist_(0.0, 1.0)
        {
            assert(mean.rows() == covariance.rows());
            assert(covariance.cols() == covariance.rows());

            assert(!(norm_transform_.unaryExpr([] (const double &val) { return std::isnan(val); })).any() && "NaN Found in norm_transform in MultivariateGaussianDistribution");
            assert(!(norm_transform_.unaryExpr([] (const double &val) { return std::isinf(val); })).any() && "Inf Found in norm_transform in MultivariateGaussianDistribution");
        }

        template<typename Generator>
        inline Eigen::VectorXd operator()(Generator& prng)
        {
            return Sample(prng);
        }
    };

    class RandomRotationGenerator
    {
    protected:

        std::uniform_real_distribution<double> uniform_unit_dist_;

        // From: "Uniform Random Rotations", Ken Shoemake, Graphics Gems III, pg. 124-132
        template<typename Generator>
        inline Eigen::Quaterniond GenerateUniformRandomQuaternion(Generator& prng)
        {
            const double x0 = uniform_unit_dist_(prng);
            const double r1 = sqrt(1.0 - x0);
            const double r2 = sqrt(x0);
            const double t1 = 2.0 * M_PI * uniform_unit_dist_(prng);
            const double t2 = 2.0 * M_PI * uniform_unit_dist_(prng);
            const double c1 = cos(t1);
            const double s1 = sin(t1);
            const double c2 = cos(t2);
            const double s2 = sin(t2);
            const double x = s1 * r1;
            const double y = c1 * r1;
            const double z = s2 * r2;
            const double w = c2 * r2;
            return Eigen::Quaterniond(w, x, y, z);
        }

        // From Effective Sampling and Distance Metrics for 3D Rigid Body Path Planning, by James Kuffner, ICRA 2004
        template<typename Generator>
        Eigen::Vector3d GenerateUniformRandomEulerAngles(Generator& prng)
        {
            const double roll = M_PI * (-2.0 * uniform_unit_dist_(prng) + 1.0);
            const double pitch = acos(1.0 - 2.0 * uniform_unit_dist_(prng)) - M_PI_2;
            const double yaw = M_PI * (-2.0 * uniform_unit_dist_(prng) + 1.0);
            return Eigen::Vector3d(roll, pitch, yaw);
        }

    public:

        inline RandomRotationGenerator() : uniform_unit_dist_(0.0, 1.0) {}

        template<typename Generator>
        inline Eigen::Quaterniond GetQuaternion(Generator& prng)
        {
            return GenerateUniformRandomQuaternion(prng);
        }

        template<typename Generator>
        inline std::vector<double> GetRawQuaternion(Generator& prng)
        {
            const Eigen::Quaterniond quat = GenerateUniformRandomQuaternion(prng);
            return std::vector<double>{quat.x(), quat.y(), quat.z(), quat.w()};
        }

        template<typename Generator>
        inline Eigen::Vector3d GetEulerAngles(Generator& prng)
        {
            return GenerateUniformRandomEulerAngles(prng);
        }

        template<typename Generator>
        inline std::vector<double> GetRawEulerAngles(Generator& prng)
        {
            const Eigen::Vector3d angles = GenerateUniformRandomEulerAngles(prng);
            return std::vector<double>{angles.x(), angles.y(), angles.z()};
        }
    };

    /////////////////////////////////////////////////////////////////////////////////////////////////////////
    /////                                       PROTOTYPES ONLY                                         /////
    ///// Specializations for specific types - if you want a specialization for a new type, add it here /////
    /////////////////////////////////////////////////////////////////////////////////////////////////////////

    template<typename T>
    inline uint64_t SerializeFixedSizePOD(const T& item_to_serialize, std::vector<uint8_t>& buffer);

    template<typename T>
    inline std::pair<T, uint64_t> DeserializeFixedSizePOD(const std::vector<uint8_t>& buffer, const uint64_t current);

    template<typename T, typename Allocator=std::allocator<T>>
    inline uint64_t SerializeVector(const std::vector<T, Allocator>& vec_to_serialize, std::vector<uint8_t>& buffer, const std::function<uint64_t(const T&, std::vector<uint8_t>&)>& item_serializer);

    template<typename T, typename Allocator=std::allocator<T>>
    inline std::pair<std::vector<T, Allocator>, uint64_t> DeserializeVector(const std::vector<uint8_t>& buffer, const uint64_t current, const std::function<std::pair<T, uint64_t>(const std::vector<uint8_t>&, const uint64_t)>& item_deserializer);

    template<typename Key, typename T, typename Compare = std::less<Key>, typename Allocator = std::allocator<std::pair<const Key, T>>>
    inline uint64_t SerializeMap(const std::map<Key, T, Compare, Allocator>& map_to_serialize, std::vector<uint8_t>& buffer, const std::function<uint64_t(const Key&, std::vector<uint8_t>&)>& key_serializer, const std::function<uint64_t(const T&, std::vector<uint8_t>&)>& value_serializer);

    template<typename Key, typename T, typename Compare = std::less<Key>, typename Allocator = std::allocator<std::pair<const Key, T>>>
    inline std::pair<std::map<Key, T, Compare, Allocator>, uint64_t> DeserializeMap(const std::vector<uint8_t>& buffer, const uint64_t current, const std::function<std::pair<Key, uint64_t>(const std::vector<uint8_t>&, const uint64_t)>& key_deserializer, const std::function<std::pair<T, uint64_t>(const std::vector<uint8_t>&, const uint64_t)>& value_deserializer);

    template<typename First, typename Second>
    inline uint64_t SerializePair(const std::pair<First, Second>& pair_to_serialize, std::vector<uint8_t>& buffer, const std::function<uint64_t(const First&, std::vector<uint8_t>&)>& first_serializer, const std::function<uint64_t(const Second&, std::vector<uint8_t>&)>& second_serializer);

    template<typename First, typename Second>
    inline const std::pair<std::pair<First, Second>, uint64_t> DeserializePair(const std::vector<uint8_t>& buffer, const uint64_t current, const std::function<std::pair<First, uint64_t>(const std::vector<uint8_t>&, const uint64_t)>& first_deserializer, const std::function<std::pair<Second, uint64_t>(const std::vector<uint8_t>&, const uint64_t)>& second_deserializer);

    /////////////////////////////////////////////////////////////////////////////////////////////////////////
    /////                                   IMPLEMENTATIONS ONLY                                        /////
    ///// Specializations for specific types - if you want a specialization for a new type, add it here /////
    /////////////////////////////////////////////////////////////////////////////////////////////////////////

    template<typename T>
    inline uint64_t SerializeFixedSizePOD(const T& item_to_serialize, std::vector<uint8_t>& buffer)
    {
        const uint64_t start_buffer_size = buffer.size();
        // Fixed-size serialization via memcpy
        std::vector<uint8_t> temp_buffer(sizeof(item_to_serialize), 0x00);
        memcpy(&temp_buffer[0], &item_to_serialize, sizeof(item_to_serialize));
        // Move to buffer
        buffer.insert(buffer.end(), temp_buffer.begin(), temp_buffer.end());
        // Figure out how many bytes were written
        const uint64_t end_buffer_size = buffer.size();
        const uint64_t bytes_written = end_buffer_size - start_buffer_size;
        return bytes_written;
    }

    template<typename T>
    inline std::pair<T, uint64_t> DeserializeFixedSizePOD(const std::vector<uint8_t>& buffer, const uint64_t current)
    {
        T temp_item;
        assert(current <= buffer.size());
        assert((current + sizeof(temp_item)) <= buffer.size());
        memcpy(&temp_item, &buffer[current], sizeof(temp_item));
        return std::make_pair(temp_item, sizeof(temp_item));
    }

    template<typename T, typename Allocator>
    inline uint64_t SerializeVector(const std::vector<T, Allocator>& vec_to_serialize, std::vector<uint8_t>& buffer, const std::function<uint64_t(const T&, std::vector<uint8_t>&)>& item_serializer)
    {
        const uint64_t start_buffer_size = buffer.size();
        // First, write a uint64_t size header
        const uint64_t size = (uint64_t)vec_to_serialize.size();
        SerializeFixedSizePOD<uint64_t>(size, buffer);
        // Serialize the contained items
        for (size_t idx = 0; idx < size; idx++)
        {
            const T& current = vec_to_serialize[idx];
            item_serializer(current, buffer);
        }
        // Figure out how many bytes were written
        const uint64_t end_buffer_size = buffer.size();
        const uint64_t bytes_written = end_buffer_size - start_buffer_size;
        return bytes_written;
    }

    template<typename T, typename Allocator>
    inline std::pair<std::vector<T, Allocator>, uint64_t> DeserializeVector(const std::vector<uint8_t>& buffer, const uint64_t current, const std::function<std::pair<T, uint64_t>(const std::vector<uint8_t>&, const uint64_t)>& item_deserializer)
    {
        // First, try to load the header
        assert(current < buffer.size());
        uint64_t current_position = current;
        // Load the header
        const std::pair<uint64_t, uint64_t> deserialized_size = DeserializeFixedSizePOD<uint64_t>(buffer, current_position);
        const uint64_t size = deserialized_size.first;
        current_position += deserialized_size.second;
        // Deserialize the items
        std::vector<T, Allocator> deserialized;
        deserialized.reserve(size);
        for (uint64_t idx = 0; idx < size; idx++)
        {
            const std::pair<T, uint64_t> deserialized_item = item_deserializer(buffer, current_position);
            deserialized.push_back(deserialized_item.first);
            current_position += deserialized_item.second;
        }
        // Figure out how many bytes were read
        const uint64_t bytes_read = current_position - current;
        return std::make_pair(deserialized, bytes_read);
    }

    template<typename Key, typename T, typename Compare, typename Allocator>
    inline uint64_t SerializeMap(const std::map<Key, T, Compare, Allocator>& map_to_serialize, std::vector<uint8_t>& buffer, const std::function<uint64_t(const Key&, std::vector<uint8_t>&)>& key_serializer, const std::function<uint64_t(const T&, std::vector<uint8_t>&)>& value_serializer)
    {
        const uint64_t start_buffer_size = buffer.size();
        // First, write a uint64_t size header
        const uint64_t size = (uint64_t)map_to_serialize.size();
        SerializeFixedSizePOD<uint64_t>(size, buffer);
        // Serialize the contained items
        typename std::map<Key, T, Compare, Allocator>::const_iterator itr;
        for (itr = map_to_serialize.begin(); itr != map_to_serialize.end(); ++itr)
        {
            SerializePair<Key, T>(*itr, buffer, key_serializer, value_serializer);
        }
        // Figure out how many bytes were written
        const uint64_t end_buffer_size = buffer.size();
        const uint64_t bytes_written = end_buffer_size - start_buffer_size;
        return bytes_written;
    }

    template<typename Key, typename T, typename Compare, typename Allocator>
    inline std::pair<std::map<Key, T, Compare, Allocator>, uint64_t> DeserializeMap(const std::vector<uint8_t>& buffer, const uint64_t current, const std::function<std::pair<Key, uint64_t>(const std::vector<uint8_t>&, const uint64_t)>& key_deserializer, const std::function<std::pair<T, uint64_t>(const std::vector<uint8_t>&, const uint64_t)>& value_deserializer)
    {
        // First, try to load the header
        assert(current < buffer.size());
        uint64_t current_position = current;
        // Load the header
        const std::pair<uint64_t, uint64_t> deserialized_size = DeserializeFixedSizePOD<uint64_t>(buffer, current_position);
        const uint64_t size = deserialized_size.first;
        current_position += deserialized_size.second;
        // Deserialize the items
        std::map<Key, T, Compare, Allocator> deserialized;
        for (uint64_t idx = 0; idx < size; idx++)
        {
            std::pair<std::pair<Key, T>, uint64_t> deserialized_pair = DeserializePair(buffer, current_position, key_deserializer, value_deserializer);
            deserialized.insert(deserialized_pair.first);
            current_position += deserialized_pair.second;
        }
        // Figure out how many bytes were read
        const uint64_t bytes_read = current_position - current;
        return std::make_pair(deserialized, bytes_read);
    }

    template<typename First, typename Second>
    inline uint64_t SerializePair(const std::pair<First, Second>& pair_to_serialize, std::vector<uint8_t>& buffer, const std::function<uint64_t(const First&, std::vector<uint8_t>&)>& first_serializer, const std::function<uint64_t(const Second&, std::vector<uint8_t>&)>& second_serializer)
    {
        const uint64_t start_buffer_size = buffer.size();
        uint64_t running_total = 0u;
        // Write each element of the pair into the buffer
        running_total += first_serializer(pair_to_serialize.first, buffer);
        running_total += second_serializer(pair_to_serialize.second, buffer);
        // Figure out how many bytes were written
        const uint64_t end_buffer_size = buffer.size();
        const uint64_t bytes_written = end_buffer_size - start_buffer_size;
        assert(bytes_written == running_total);
        return bytes_written;
    }

    template<typename First, typename Second>
    inline const std::pair<std::pair<First, Second>, uint64_t> DeserializePair(const std::vector<uint8_t>& buffer, const uint64_t current, const std::function<std::pair<First, uint64_t>(const std::vector<uint8_t>&, const uint64_t)>& first_deserializer, const std::function<std::pair<Second, uint64_t>(const std::vector<uint8_t>&, const uint64_t)>& second_deserializer)
    {
        assert(current < buffer.size());
        // Deserialize each item in the pair individually
        uint64_t current_position = current;
        const std::pair<First, uint64_t> deserialized_first = first_deserializer(buffer, current_position);
        current_position += deserialized_first.second;
        const std::pair<Second, uint64_t> deserialized_second = second_deserializer(buffer, current_position);
        current_position += deserialized_second.second;
        // Build the resulting pair
        // TODO: Why can't I used make_pair here?
        const std::pair<First, Second> deserialized(deserialized_first.first, deserialized_second.first);
        // Figure out how many bytes were read
        const uint64_t bytes_read = current_position - current;
        return std::make_pair(deserialized, bytes_read);
    }

    inline void ConditionalPrint(const std::string& msg, const int32_t msg_level, const int32_t print_level)
    {
        if (unlikely(msg_level <= print_level))
        {
            std::cout << msg << std::endl;
        }
    }

    inline bool CheckAllStringsForSubstring(const std::vector<std::string>& strings, const std::string& substring)
    {
        for (size_t idx = 0; idx < strings.size(); idx++)
        {
            const std::string& candidate_string = strings[idx];
            const size_t found = candidate_string.find(substring);
            if (found == std::string::npos)
            {
                return false;
            }
        }
        return true;
    }
}

#endif // ARC_HELPERS_HPP


// smooth_feedback: Control theory on Lie groups
// https://github.com/pettni/smooth_feedback
//
// Licensed under the MIT License <http://opensource.org/licenses/MIT>.
//
// Copyright (c) 2021 Petter Nilsson
//
// 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.

#ifndef SMOOTH__FEEDBACK__NLP_HPP_
#define SMOOTH__FEEDBACK__NLP_HPP_

/**
 * @file
 * @brief Nonlinear program definition.
 */

#include <Eigen/Core>
#include <Eigen/Sparse>

#include <limits>
#include <optional>

namespace smooth::feedback {

/**
 * @brief Nonlinear Programming Problem
 * \f[
 *  \begin{cases}
 *   \min_{x}    & f(x)                    \\
 *   \text{s.t.} & x_l \leq x \leq x_u     \\
 *               & g_l \leq g(x) \leq g_u
 *  \end{cases}
 * \f]
 * for \f$ f : \mathbb{R}^n \rightarrow \mathbb{R} \f$ and
 * \f$ g : \mathbb{R}^n \rightarrow \mathbb{R}^m \f$.
 */
template<typename T>
concept NLP = requires(
  std::decay_t<T> & nlp,
  const Eigen::Ref<const Eigen::VectorXd> x,
  const Eigen::Ref<const Eigen::VectorXd> lambda)
{
  // clang-format off
  {nlp.n()} -> std::convertible_to<std::size_t>;
  {nlp.m()} -> std::convertible_to<std::size_t>;

  // variable bounds
  {nlp.xl()} -> std::convertible_to<Eigen::VectorXd>;
  {nlp.xu()} -> std::convertible_to<Eigen::VectorXd>;

  // objective and its derivatives
  {nlp.f(x)}     -> std::convertible_to<double>;
  {nlp.df_dx(x)} -> std::convertible_to<Eigen::SparseMatrix<double>>;

  // constraint, its bounds, and its derivatives
  {nlp.g(x)}     -> std::convertible_to<Eigen::VectorXd>;
  {nlp.gl()}     -> std::convertible_to<Eigen::VectorXd>;
  {nlp.gu()}     -> std::convertible_to<Eigen::VectorXd>;
  {nlp.dg_dx(x)} -> std::convertible_to<Eigen::SparseMatrix<double>>;
  // clang-format on
};

/**
 * @brief Nonlinear Programming Problem with Hessian information
 */
template<typename T>
concept HessianNLP =
  NLP<T> && requires(std::decay_t<T> & nlp, Eigen::VectorXd x, Eigen::VectorXd lambda)
{
  // clang-format off
  {nlp.d2f_dx2(x)}         -> std::convertible_to<Eigen::SparseMatrix<double>>;
  {nlp.d2g_dx2(x, lambda)} -> std::convertible_to<Eigen::SparseMatrix<double>>;
  // clang-format on
};

/**
 * @brief Solution to a Nonlinear Programming Problem
 */
struct NLPSolution
{
  /// @brief Solver status
  enum class Status {
    Optimal,
    PrimalInfeasible,
    DualInfeasible,
    MaxIterations,
    MaxTime,
    Unknown,
  };

  /// @brief Solver status
  Status status;

  /// @brief Number of iterations
  std::size_t iter{0};

  /// @brief Variable values
  Eigen::VectorXd x;

  ///@{
  /// @brief Inequality multipliers
  Eigen::VectorXd zl, zu;
  ///@}

  /// @brief Constraint multipliers
  Eigen::VectorXd lambda;

  /// @brief Objective
  double objective{0};
};

}  // namespace smooth::feedback

#endif  // SMOOTH__FEEDBACK__NLP_HPP_


// -----------------------------------------------------------------------
// RTToolbox - DKFZ radiotherapy quantitative evaluation library
//
// Copyright (c) German Cancer Research Center (DKFZ),
// Software development for Integrated Diagnostics and Therapy (SIDT).
// ALL RIGHTS RESERVED.
// See rttbCopyright.txt or
// http://www.dkfz.de/en/sidt/projects/rttb/copyright.html
//
// This software is distributed WITHOUT ANY WARRANTY; without even
// the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
// PURPOSE.  See the above copyright notices for more information.
//
//------------------------------------------------------------------------

#include <boost/make_shared.hpp>

#include "rttbDVHCalculator.h"
#include "rttbNullPointerException.h"
#include "rttbInvalidParameterException.h"

namespace rttb
{
	namespace core
	{

		DVHCalculator::DVHCalculator(DoseIteratorPointer aDoseIterator, const IDType& aStructureID,
		                             const IDType& aDoseID,
		                             DoseTypeGy aDeltaD, const int aNumberOfBins)
		{
			if (aDoseIterator == nullptr)
			{
				throw NullPointerException("aDoseIterator must not be nullptr! ");
			}

			_doseIteratorPtr = aDoseIterator;
			_structureID = aStructureID;
			_doseID = aDoseID;

			if (aNumberOfBins <= 0 || aDeltaD < 0)
			{
				throw InvalidParameterException("aNumberOfBins/aDeltaD must be >0! ");
			}

			_numberOfBins = aNumberOfBins;
			_deltaD = aDeltaD;

			if (_deltaD == 0)
			{
				aDoseIterator->reset();
				DoseTypeGy max = 0;

				while (aDoseIterator->isPositionValid())
				{
					DoseTypeGy currentVal = 0;
					currentVal = aDoseIterator->getCurrentDoseValue();

					if (currentVal > max)
					{
						max = currentVal;
					}

					aDoseIterator->next();
				}

				_deltaD = (max * 1.5 / _numberOfBins);

				if (_deltaD == 0)
				{
					_deltaD = 0.1;
				}
			}

		}

		DVHCalculator::~DVHCalculator() = default;

		DVH::Pointer DVHCalculator::generateDVH()
		{

			std::deque<DoseCalcType> dataDifferential(_numberOfBins, 0);

			// calculate DVH
			_doseIteratorPtr->reset();

			while (_doseIteratorPtr->isPositionValid())
			{
				DoseTypeGy currentVal = 0;
				FractionType voxelProportion = _doseIteratorPtr->getCurrentRelevantVolumeFraction();
				currentVal = _doseIteratorPtr->getCurrentDoseValue();

				auto dose_bin = static_cast<int>(currentVal / _deltaD);

				if (dose_bin < _numberOfBins)
				{
					dataDifferential[dose_bin] += voxelProportion;
				}
				else
				{
					throw InvalidParameterException("_numberOfBins is too small: dose bin out of bounds! ");
				}

				_doseIteratorPtr->next();
			}

			if (boost::dynamic_pointer_cast<MaskedDoseIteratorPointer>(_doseIteratorPtr))
			{
				_dvh = boost::make_shared<DVH>(dataDifferential, _deltaD, _doseIteratorPtr->getCurrentVoxelVolume(),
				                               _structureID,
				                               _doseID, _doseIteratorPtr->getVoxelizationID());
			}
			else
			{
				_dvh = boost::make_shared<DVH>(dataDifferential, _deltaD, _doseIteratorPtr->getCurrentVoxelVolume(),
				                               _structureID,
				                               _doseID);
			}

			return _dvh;
		}

	}//end namespace core
}//end namespace rttb


/* -*- c++ -*- */
/*
 * Copyright 2009-2013 Free Software Foundation, Inc.
 *
 * This file is part of GNU Radio
 *
 * GNU Radio is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3, or (at your option)
 * any later version.
 *
 * GNU Radio is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with GNU Radio; see the file COPYING.  If not, write to
 * the Free Software Foundation, Inc., 51 Franklin Street,
 * Boston, MA 02110-1301, USA.
 */

#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include <cstdio>
#include <cmath>
#include <algorithm>

#include "pfb_clock_sync_ccf_impl.h"
#include <gnuradio/io_signature.h>
#include <gnuradio/math.h>
#include <boost/format.hpp>
#include <boost/math/special_functions/round.hpp>

namespace gr {
  namespace digital {

    pfb_clock_sync_ccf::sptr
    pfb_clock_sync_ccf::make(double sps, float loop_bw,
			     const std::vector<float> &taps,
			     unsigned int filter_size,
			     float init_phase,
			     float max_rate_deviation,
			     int osps)
    {
      return gnuradio::get_initial_sptr
	(new pfb_clock_sync_ccf_impl(sps, loop_bw, taps,
				     filter_size,
				     init_phase,
				     max_rate_deviation,
				     osps));
    }

    static int ios[] = {sizeof(gr_complex), sizeof(float), sizeof(float), sizeof(float)};
    static std::vector<int> iosig(ios, ios+sizeof(ios)/sizeof(int));
    pfb_clock_sync_ccf_impl::pfb_clock_sync_ccf_impl(double sps, float loop_bw,
						     const std::vector<float> &taps,
						     unsigned int filter_size,
						     float init_phase,
						     float max_rate_deviation,
						     int osps)
      : block("pfb_clock_sync_ccf",
		  io_signature::make(1, 1, sizeof(gr_complex)),
		  io_signature::makev(1, 4, iosig)),
	d_updated(false), d_nfilters(filter_size),
	d_max_dev(max_rate_deviation),
	d_osps(osps), d_error(0), d_out_idx(0)
    {
      if(taps.size() == 0)
        throw std::runtime_error("pfb_clock_sync_ccf: please specify a filter.\n");

      // Let scheduler adjust our relative_rate.
      //enable_update_rate(true);
      set_tag_propagation_policy(TPP_DONT);

      d_nfilters = filter_size;
      d_sps = floor(sps);

      // Set the damping factor for a critically damped system
      d_damping = 2*d_nfilters;

      // Set the bandwidth, which will then call update_gains()
      set_loop_bandwidth(loop_bw);

      // Store the last filter between calls to work
      // The accumulator keeps track of overflow to increment the stride correctly.
      // set it here to the fractional difference based on the initial phaes
      d_k = init_phase;
      d_rate = (sps-floor(sps))*(double)d_nfilters;
      d_rate_i = (int)floor(d_rate);
      d_rate_f = d_rate - (float)d_rate_i;
      d_filtnum = (int)floor(d_k);

      d_filters = std::vector<kernel::fir_filter_ccf*>(d_nfilters);
      d_diff_filters = std::vector<kernel::fir_filter_ccf*>(d_nfilters);

      // Create an FIR filter for each channel and zero out the taps
      std::vector<float> vtaps(1,0);
      for(int i = 0; i < d_nfilters; i++) {
	d_filters[i] = new kernel::fir_filter_ccf(1, vtaps);
	d_diff_filters[i] = new kernel::fir_filter_ccf(1, vtaps);
      }

      // Now, actually set the filters' taps
      std::vector<float> dtaps;
      create_diff_taps(taps, dtaps);
      set_taps(taps, d_taps, d_filters);
      set_taps(dtaps, d_dtaps, d_diff_filters);

      d_old_in = 0;
      d_new_in = 0;
      d_last_out = 0;

      set_relative_rate((float)d_osps/(float)d_sps);
    }

    pfb_clock_sync_ccf_impl::~pfb_clock_sync_ccf_impl()
    {
      for(int i = 0; i < d_nfilters; i++) {
	delete d_filters[i];
	delete d_diff_filters[i];
      }
    }

    bool
    pfb_clock_sync_ccf_impl::check_topology(int ninputs, int noutputs)
    {
      return noutputs == 1 || noutputs == 4;
    }

    void
    pfb_clock_sync_ccf_impl::forecast(int noutput_items,
                                      gr_vector_int &ninput_items_required)
    {
      unsigned ninputs = ninput_items_required.size ();
      for(unsigned i = 0; i < ninputs; i++)
        ninput_items_required[i] = (noutput_items + history()) * (d_sps/d_osps);
    }

    void
    pfb_clock_sync_ccf_impl::update_taps(const std::vector<float> &taps)
    {
      d_updated_taps = taps;
      d_updated = true;
    }

    /*******************************************************************
     SET FUNCTIONS
    *******************************************************************/

    void
    pfb_clock_sync_ccf_impl::set_loop_bandwidth(float bw)
    {
      if(bw < 0) {
	throw std::out_of_range("pfb_clock_sync_ccf: invalid bandwidth. Must be >= 0.");
      }

      d_loop_bw = bw;
      update_gains();
    }

    void
    pfb_clock_sync_ccf_impl::set_damping_factor(float df)
    {
      if(df < 0 || df > 1.0) {
	throw std::out_of_range("pfb_clock_sync_ccf: invalid damping factor. Must be in [0,1].");
      }

      d_damping = df;
      update_gains();
    }

    void
    pfb_clock_sync_ccf_impl::set_alpha(float alpha)
    {
      if(alpha < 0 || alpha > 1.0) {
	throw std::out_of_range("pfb_clock_sync_ccf: invalid alpha. Must be in [0,1].");
      }
      d_alpha = alpha;
    }

    void
    pfb_clock_sync_ccf_impl::set_beta(float beta)
    {
      if(beta < 0 || beta > 1.0) {
	throw std::out_of_range("pfb_clock_sync_ccf: invalid beta. Must be in [0,1].");
      }
      d_beta = beta;
    }

    /*******************************************************************
     GET FUNCTIONS
    *******************************************************************/

    float
    pfb_clock_sync_ccf_impl::loop_bandwidth() const
    {
      return d_loop_bw;
    }

    float
    pfb_clock_sync_ccf_impl::damping_factor() const
    {
      return d_damping;
    }

    float
    pfb_clock_sync_ccf_impl::alpha() const
    {
      return d_alpha;
    }

    float
    pfb_clock_sync_ccf_impl::beta() const
    {
      return d_beta;
    }

    float
    pfb_clock_sync_ccf_impl::clock_rate() const
    {
      return d_rate_f;
    }

    float
    pfb_clock_sync_ccf_impl::error() const
    {
      return d_error;
    }

    float
    pfb_clock_sync_ccf_impl::rate() const
    {
      return d_rate_f;
    }

    float
    pfb_clock_sync_ccf_impl::phase() const
    {
      return d_k;
    }

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

    void
    pfb_clock_sync_ccf_impl::update_gains()
    {
      float denom = (1.0 + 2.0*d_damping*d_loop_bw + d_loop_bw*d_loop_bw);
      d_alpha = (4*d_damping*d_loop_bw) / denom;
      d_beta = (4*d_loop_bw*d_loop_bw) / denom;
    }

    void
    pfb_clock_sync_ccf_impl::set_taps(const std::vector<float> &newtaps,
				      std::vector< std::vector<float> > &ourtaps,
				      std::vector<kernel::fir_filter_ccf*> &ourfilter)
    {
      int i,j;

      unsigned int ntaps = newtaps.size();
      d_taps_per_filter = (unsigned int)ceil((double)ntaps/(double)d_nfilters);

      // Create d_numchan vectors to store each channel's taps
      ourtaps.resize(d_nfilters);

      // Make a vector of the taps plus fill it out with 0's to fill
      // each polyphase filter with exactly d_taps_per_filter
      std::vector<float> tmp_taps;
      tmp_taps = newtaps;
      while((float)(tmp_taps.size()) < d_nfilters*d_taps_per_filter) {
	tmp_taps.push_back(0.0);
      }

      // Partition the filter
      for(i = 0; i < d_nfilters; i++) {
	// Each channel uses all d_taps_per_filter with 0's if not enough taps to fill out
	ourtaps[i] = std::vector<float>(d_taps_per_filter, 0);
	for(j = 0; j < d_taps_per_filter; j++) {
	  ourtaps[i][j] = tmp_taps[i + j*d_nfilters];
	}

	// Build a filter for each channel and add it's taps to it
	ourfilter[i]->set_taps(ourtaps[i]);
      }

      // Set the history to ensure enough input items for each filter
      set_history(d_taps_per_filter + d_sps + d_sps);

      // Make sure there is enough output space for d_osps outputs/input.
      set_output_multiple(d_osps);
    }

    void
    pfb_clock_sync_ccf_impl::create_diff_taps(const std::vector<float> &newtaps,
					      std::vector<float> &difftaps)
    {
      std::vector<float> diff_filter(3);
      diff_filter[0] = -1;
      diff_filter[1] = 0;
      diff_filter[2] = 1;

      float pwr = 0;
      difftaps.clear();
      difftaps.push_back(0);
      for(unsigned int i = 0; i < newtaps.size()-2; i++) {
	float tap = 0;
	for(unsigned int j = 0; j < diff_filter.size(); j++) {
	  tap += diff_filter[j]*newtaps[i+j];
	}
	difftaps.push_back(tap);
        pwr += fabsf(tap);
      }
      difftaps.push_back(0);

      // Normalize the taps
      for(unsigned int i = 0; i < difftaps.size(); i++) {
        difftaps[i] *= d_nfilters/pwr;
        if(difftaps[i] != difftaps[i]) {
          throw std::runtime_error("pfb_clock_sync_ccf::create_diff_taps produced NaN.");
        }
      }
    }

    std::string
    pfb_clock_sync_ccf_impl::taps_as_string() const
    {
      int i, j;
      std::stringstream str;
      str.precision(4);
      str.setf(std::ios::scientific);

      str << "[ ";
      for(i = 0; i < d_nfilters; i++) {
	str << "[" << d_taps[i][0] << ", ";
	for(j = 1; j < d_taps_per_filter-1; j++) {
	  str << d_taps[i][j] << ", ";
	}
	str << d_taps[i][j] << "],";
      }
      str << " ]" << std::endl;

      return str.str();
    }

    std::string
    pfb_clock_sync_ccf_impl::diff_taps_as_string() const
    {
      int i, j;
      std::stringstream str;
      str.precision(4);
      str.setf(std::ios::scientific);

      str << "[ ";
      for(i = 0; i < d_nfilters; i++) {
	str << "[" << d_dtaps[i][0] << ", ";
	for(j = 1; j < d_taps_per_filter-1; j++) {
	  str << d_dtaps[i][j] << ", ";
	}
	str << d_dtaps[i][j] << "],";
      }
      str << " ]" << std::endl;

      return str.str();
    }

    std::vector< std::vector<float> >
    pfb_clock_sync_ccf_impl::taps() const
    {
      return d_taps;
    }

    std::vector< std::vector<float> >
    pfb_clock_sync_ccf_impl::diff_taps() const
    {
      return d_dtaps;
    }

    std::vector<float>
    pfb_clock_sync_ccf_impl::channel_taps(int channel) const
    {
      std::vector<float> taps;
      for(int i = 0; i < d_taps_per_filter; i++) {
	taps.push_back(d_taps[channel][i]);
      }
      return taps;
    }

    std::vector<float>
    pfb_clock_sync_ccf_impl::diff_channel_taps(int channel) const
    {
      std::vector<float> taps;
      for(int i = 0; i < d_taps_per_filter; i++) {
	taps.push_back(d_dtaps[channel][i]);
      }
      return taps;
    }

    int
    pfb_clock_sync_ccf_impl::general_work(int noutput_items,
					  gr_vector_int &ninput_items,
					  gr_vector_const_void_star &input_items,
					  gr_vector_void_star &output_items)
    {
      gr_complex *in = (gr_complex *) input_items[0];
      gr_complex *out = (gr_complex *) output_items[0];

      if(d_updated) {
        std::vector<float> dtaps;
        create_diff_taps(d_updated_taps, dtaps);
        set_taps(d_updated_taps, d_taps, d_filters);
        set_taps(dtaps, d_dtaps, d_diff_filters);
	d_updated = false;
	return 0;		     // history requirements may have changed.
      }

      float *err = NULL, *outrate = NULL, *outk = NULL;
      if(output_items.size() == 4) {
	err = (float *) output_items[1];
	outrate = (float*)output_items[2];
	outk = (float*)output_items[3];
      }

      std::vector<tag_t> tags;
      get_tags_in_window(tags, 0, 0,
                        d_sps*noutput_items,
                        pmt::intern("time_est"));

      int i = 0, count = 0;
      float error_r, error_i;

      // produce output as long as we can and there are enough input samples
      while(i < noutput_items) {
        if(tags.size() > 0) {
          size_t offset = tags[0].offset-nitems_read(0);
          if((offset >= (size_t)count) && (offset < (size_t)(count + d_sps))) {
            float center = (float)pmt::to_double(tags[0].value);
            d_k = d_nfilters*(center + (offset - count));

            tags.erase(tags.begin());
          }
        }

	while(d_out_idx < d_osps) {

	  d_filtnum = (int)floor(d_k);

	  // Keep the current filter number in [0, d_nfilters]
	  // If we've run beyond the last filter, wrap around and go to next sample
	  // If we've gone below 0, wrap around and go to previous sample
	  while(d_filtnum >= d_nfilters) {
	    d_k -= d_nfilters;
	    d_filtnum -= d_nfilters;
	    count += 1;
	  }
	  while(d_filtnum < 0) {
	    d_k += d_nfilters;
	    d_filtnum += d_nfilters;
	    count -= 1;
	  }

	  out[i+d_out_idx] = d_filters[d_filtnum]->filter(&in[count+d_out_idx]);
	  d_k = d_k + d_rate_i + d_rate_f; // update phase

          // Manage Tags
          std::vector<tag_t> xtags;
          std::vector<tag_t>::iterator itags;
          d_new_in = nitems_read(0) + count + d_out_idx + d_sps;
          get_tags_in_range(xtags, 0, d_old_in, d_new_in);
          for(itags = xtags.begin(); itags != xtags.end(); itags++) {
            tag_t new_tag = *itags;
            //new_tag.offset = d_last_out + d_taps_per_filter/(2*d_sps) - 2;
            new_tag.offset = d_last_out + d_taps_per_filter/4 - 2;
            add_item_tag(0, new_tag);
          }
          d_old_in = d_new_in;
          d_last_out = nitems_written(0) + i + d_out_idx;

          d_out_idx++;

	  if(output_items.size() == 4) {
	    err[i] = d_error;
	    outrate[i] = d_rate_f;
	    outk[i] = d_k;
	  }

	  // We've run out of output items we can create; return now.
	  if(i+d_out_idx >= noutput_items) {
	    consume_each(count);
	    return i;
	  }
	}

	// reset here; if we didn't complete a full osps samples last time,
	// the early return would take care of it.
	d_out_idx = 0;

	// Update the phase and rate estimates for this symbol
	gr_complex diff = d_diff_filters[d_filtnum]->filter(&in[count]);
	error_r = out[i].real() * diff.real();
	error_i = out[i].imag() * diff.imag();
	d_error = (error_i + error_r) / 2.0;       // average error from I&Q channel

        // Run the control loop to update the current phase (k) and
        // tracking rate estimates based on the error value
        // Interpolating here to update rates for ever sps.
        for(int s = 0; s < d_sps; s++) {
          d_rate_f = d_rate_f + d_beta*d_error;
          d_k = d_k + d_rate_f + d_alpha*d_error;
        }

	// Keep our rate within a good range
	d_rate_f = gr::branchless_clip(d_rate_f, d_max_dev);

	i+=d_osps;
	count += (int)floor(d_sps);
      }

      consume_each(count);
      return i;
    }

    void
    pfb_clock_sync_ccf_impl::setup_rpc()
    {
#ifdef GR_CTRLPORT
      // Getters
      add_rpc_variable(
          rpcbasic_sptr(new rpcbasic_register_get<pfb_clock_sync_ccf, float>(
	      alias(), "error",
	      &pfb_clock_sync_ccf::error,
	      pmt::mp(-2.0f), pmt::mp(2.0f), pmt::mp(0.0f),
	      "", "Error signal of loop", RPC_PRIVLVL_MIN,
              DISPTIME | DISPOPTSTRIP)));

      add_rpc_variable(
          rpcbasic_sptr(new rpcbasic_register_get<pfb_clock_sync_ccf, float>(
	      alias(), "rate",
	      &pfb_clock_sync_ccf::rate,
	      pmt::mp(-2.0f), pmt::mp(2.0f), pmt::mp(0.0f),
	      "", "Rate change of phase", RPC_PRIVLVL_MIN,
              DISPTIME | DISPOPTSTRIP)));

      add_rpc_variable(
          rpcbasic_sptr(new rpcbasic_register_get<pfb_clock_sync_ccf, float>(
	      alias(), "phase",
	      &pfb_clock_sync_ccf::phase,
	      pmt::mp(0), pmt::mp((int)d_nfilters), pmt::mp(0),
	      "", "Current filter phase arm", RPC_PRIVLVL_MIN,
              DISPTIME | DISPOPTSTRIP)));

      add_rpc_variable(
          rpcbasic_sptr(new rpcbasic_register_get<pfb_clock_sync_ccf, float>(
	      alias(), "loop bw",
	      &pfb_clock_sync_ccf::loop_bandwidth,
	      pmt::mp(0.0f), pmt::mp(1.0f), pmt::mp(0.0f),
	      "", "Loop bandwidth",
	      RPC_PRIVLVL_MIN, DISPNULL)));

      // Setters
      add_rpc_variable(
          rpcbasic_sptr(new rpcbasic_register_set<pfb_clock_sync_ccf, float>(
	      alias(), "loop bw",
	      &pfb_clock_sync_ccf::set_loop_bandwidth,
	      pmt::mp(0.0f), pmt::mp(1.0f), pmt::mp(0.0f),
	      "", "Loop bandwidth",
	      RPC_PRIVLVL_MIN, DISPNULL)));
#endif /* GR_CTRLPORT */
    }

  } /* namespace digital */
} /* namespace gr */