/*
==== Author:

Relja Arandjelovic (relja@robots.ox.ac.uk)
Visual Geometry Group,
Department of Engineering Science
University of Oxford

==== Copyright:

The library belongs to Relja Arandjelovic and the University of Oxford.
No usage or redistribution is allowed without explicit permission.
*/

#include "det_ransac.h"

#include <Eigen/Dense>

#include "putative.h"

inline double mysqr( double x ){ return x*x; }

double
detRansac::matchWords_Hard(
    sameRandomUint32 const &sameRandomObj,

    uint32_t &nInliers,

    std::vector<quantDesc> const &ids1,
    std::vector<ellipse> const &ellipses1,
    std::vector<uint32_t> const *presortedInds1,
    std::vector<quantDesc> const &ids2,
    std::vector<ellipse> const &ellipses2,

    double errorThr,
    double lowAreaChange, double highAreaChange,
    uint32_t nReest,

    homography *H,
    matchesType *inlierInds

    ){

    std::vector<uint32_t> ids1_flat, ids2_flat;
    quantDesc::flattenHard( ids1, ids1_flat );
    quantDesc::flattenHard( ids2, ids2_flat );

    std::vector< std::pair<uint32_t, uint32_t> > putativeMatches;
    putative_quantized::getPutativeMatches_Hard( ids1_flat, ids2_flat, putativeMatches, presortedInds1 );

    return detRansac::match( sameRandomObj, nInliers, ellipses1, ellipses2, putativeMatches, NULL, errorThr, lowAreaChange, highAreaChange, nReest, H, inlierInds );

}

double
detRansac::matchWords_Soft(
    sameRandomUint32 const &sameRandomObj,

    uint32_t &nInliers,

    std::vector<quantDesc> const &ids1,
    std::vector<ellipse> const &ellipses1,
    std::vector<uint32_t> const *presortedInds1,
    std::vector<quantDesc> const &ids2,
    std::vector<ellipse> const &ellipses2,

    double errorThr,
    double lowAreaChange, double highAreaChange,
    uint32_t nReest,

    homography *H,
    matchesType *inlierInds

    ){

    std::vector<wordWeightPair> ids1_flat, ids2_flat;
    std::vector<uint32_t> qDInds1, qDInds2;
    quantDesc::flatten( ids1, ids1_flat, &qDInds1, true );
    quantDesc::flatten( ids2, ids2_flat, &qDInds2, true );

    std::vector< std::pair<uint32_t, uint32_t> > putativeMatches;
    std::vector<double> PMweights;
    putative_quantized::getPutativeMatches_Soft( ids1_flat, qDInds1, ids2_flat, qDInds2, putativeMatches, PMweights, presortedInds1 );

//     return detRansac::match( sameRandomObj, nInliers, ellipses1, ellipses2, putativeMatches, NULL, errorThr, lowAreaChange, highAreaChange, nReest, H, inlierInds );
    return detRansac::match( sameRandomObj, nInliers, ellipses1, ellipses2, putativeMatches, &PMweights, errorThr, lowAreaChange, highAreaChange, nReest, H, inlierInds );

}

double
detRansac::matchDesc(
    sameRandomUint32 const &sameRandomObj,

    uint32_t &nInliers,

    float const* desc1,
    std::vector<ellipse> const &ellipses1,
    float const* desc2,
    std::vector<ellipse> const &ellipses2,
    uint32_t nDims,

    double errorThr,
    double lowAreaChange, double highAreaChange,
    uint32_t nReest,

    bool useLowe,
    float deltaSq,
    float epsilon,

    homography *H,
    matchesType *inlierInds

    ){

    std::vector< std::pair<uint32_t, uint32_t> > putativeMatches;

    putative_desc<float>::getPutativeMatches( desc1, ellipses1.size(), desc2, ellipses2.size(), nDims, putativeMatches, useLowe, deltaSq, epsilon );

    return detRansac::match( sameRandomObj, nInliers, ellipses1, ellipses2, putativeMatches, NULL, errorThr, lowAreaChange, highAreaChange, nReest, H, inlierInds);

}

double
detRansac::match(
    sameRandomUint32 const &sameRandomObj,

    uint32_t &bestNInliers,

    std::vector<ellipse> const &ellipses1,
    std::vector<ellipse> const &ellipses2,
    matchesType const &putativeMatches,
    std::vector<double> const *PMweights,

    double errorThr,
    double lowAreaChange, double highAreaChange,
    uint32_t nReest,

    homography *H,
    matchesType *inlierInds

    ){

    bestNInliers= 0;

    uint32_t nPutativeMatches= putativeMatches.size();

    if (nPutativeMatches<3)
        return 0;

    //------- prepare

    bool delWeights= false;
    if (PMweights==NULL) {
        PMweights= new std::vector<double>( nPutativeMatches, 1.0 );
        delWeights= true;
    }

    //------- generate Hs

    detRansac::inlierFinder inlierFinder_obj( ellipses1, ellipses2, putativeMatches, *PMweights, errorThr, lowAreaChange, highAreaChange );

    // initialize with identity
    homography Hident; Hident.setIdentity();
    homography bestH= Hident;

    double bestScore= 0.0;
    bestNInliers= 0;
    bestScore= inlierFinder_obj.getScore( bestH, bestNInliers, NULL );

    std::vector<homography> Hs;
    Hs.reserve( nPutativeMatches );

    for (std::vector< std::pair<uint32_t, uint32_t> >::const_iterator itPM= putativeMatches.begin();
         itPM!=putativeMatches.end();
         ++itPM){

        Hs.push_back( homography( ellipses1[itPM->first], ellipses2[itPM->second] ) );

    }

    uint32_t globalNIter= 0;

    sameRandomObj.shuffle<homography>( Hs.begin(), Hs.end() );

    {
        //------- RANSAC core

        static const double pFail= 0.001;

        uint32_t nInliers, iH= 0;
        double score;

        for (std::vector<homography>::const_iterator itH= Hs.begin();
             itH!=Hs.end() && globalNIter < detRansac::getNStopping(pFail, nPutativeMatches, bestNInliers);
             ++itH, ++globalNIter, ++iH){

            score= inlierFinder_obj.getScore( *itH, nInliers, NULL );

            if (nInliers>3 && score>bestScore) {
                bestNInliers= nInliers;
                bestScore= score;
                bestH= *itH;
            }

        }

    }

    if (bestNInliers > 3){

        matchesType bestInliers;
        bestInliers.reserve(bestNInliers);
        uint32_t nInliers_new;
        inlierFinder_obj.getScore( bestH, nInliers_new, &bestInliers );

        //------- reestimate

        homography H_new= bestH;
        double score_new;
        matchesType inliers_new= bestInliers;

        uint32_t iReest;
        for (iReest= 0; nInliers_new>3 && iReest<nReest; ++iReest){
            detRansac::getH( ellipses1, ellipses2, inliers_new, H_new );
            score_new= inlierFinder_obj.getScore( H_new, nInliers_new, &inliers_new );
            if (nInliers_new>3 && score_new>bestScore){
                bestScore= score_new;
                bestNInliers= nInliers_new;
                bestH= H_new;
                bestInliers= inliers_new; // don't swap as getH in next iteration needs it!
            }
        }

        if (inlierInds)
            inlierInds->swap( bestInliers );

        if (H)
            *H= bestH;

    }

    //------- cleanup

    if (delWeights)
        delete PMweights;

    return bestScore;

}

void
detRansac::getH( std::vector<ellipse> const &ellipses1,
                 std::vector<ellipse> const &ellipses2,
                 matchesType const &inliers, homography &H ){

    Eigen::Matrix<double, Eigen::Dynamic, 7> A( inliers.size()*2, 7 );

    double *x1, *y1, *x2, *y2, *_temp;
    ellipse::getCentres( ellipses1, ellipses2, inliers, x1, y1, x2, y2, _temp );
    uint32_t nInliers= inliers.size();

    // normalize points
    homography Hnorm1, Hnorm2;
    normPoints(x1, y1, nInliers, Hnorm1);
    normPoints(x2, y2, nInliers, Hnorm2);
    double Hnorm2inv[9];
    Hnorm2.getInverse(Hnorm2inv);

    // fit homography
    uint32_t i=0;

    for (matchesType::const_iterator itIn= inliers.begin();
         itIn!=inliers.end();
         ++itIn, ++i){

        A.coeffRef( i*2   , 0 )=    0.0;
        A.coeffRef( i*2   , 1 )=    0.0;
        A.coeffRef( i*2   , 2 )=    0.0;
        A.coeffRef( i*2   , 3 )=  x1[i];
        A.coeffRef( i*2   , 4 )=  y1[i];
        A.coeffRef( i*2   , 5 )=    1.0;
        A.coeffRef( i*2   , 6 )= -y2[i];

        A.coeffRef( i*2+1 , 0 )= -x1[i];
        A.coeffRef( i*2+1 , 1 )= -y1[i];
        A.coeffRef( i*2+1 , 2 )=   -1.0;
        A.coeffRef( i*2+1 , 3 )=    0.0;
        A.coeffRef( i*2+1 , 4 )=    0.0;
        A.coeffRef( i*2+1 , 5 )=    0.0;
        A.coeffRef( i*2+1 , 6 )=  x2[i];

    }

    delete []x1; delete []y1; delete []x2; delete []y2; delete []_temp;

    typedef Eigen::Matrix<double, 7, 7> matrix7x7;
    matrix7x7 AtA;
    AtA.noalias()= A.transpose() * A;
    Eigen::SelfAdjointEigenSolver<matrix7x7> sol( AtA );
    Eigen::Matrix<double,7,1> x= sol.eigenvectors().col(0);

    H.H[0]= x.coeff(0,0); H.H[1]= x.coeff(1,0); H.H[2]= x.coeff(2,0);
    H.H[3]= x.coeff(3,0); H.H[4]= x.coeff(4,0); H.H[5]= x.coeff(5,0);
    H.H[6]=          0.0; H.H[7]=          0.0; H.H[8]= x.coeff(6,0);
    H.normLast();

    // denormalize H: H= Hnorm2inv * Hnorm * Hnorm1

    Eigen::Matrix<double, 3, 3> Hnorm2inv_, Hnorm1_, Hnorm_;
    for (int i=0; i<3; ++i)
        for (int j=0; j<3; ++j){
            Hnorm_.coeffRef(i,j)= H.H[i*3+j];
            Hnorm1_.coeffRef(i,j)= Hnorm1.H[i*3+j];
            Hnorm2inv_.coeffRef(i,j)= Hnorm2inv[i*3+j];
        }
    Eigen::Matrix<double, 3, 3> H_;
    H_.noalias()= Hnorm2inv_ * Hnorm_ * Hnorm1_;

    for (int i=0; i<3; ++i)
        for (int j=0; j<3; ++j)
            H.H[i*3+j]= H_.coeff(i,j);

#if 0
    // this is how it used to be done in OpenCV, not updating though so it might be obsolete..

    cv::Mat A( inliers.size()*2, 7, cv::DataType<double>::type ), x;

    double *x1, *y1, *x2, *y2, *_temp;
    ellipse::getCentres( ellipses1, ellipses2, inliers, x1, y1, x2, y2, _temp );
    uint32_t nInliers= inliers.size();

    // normalize points
    homography Hnorm1, Hnorm2;
    normPoints(x1, y1, nInliers, Hnorm1);
    normPoints(x2, y2, nInliers, Hnorm2);
    double Hnorm2inv[9];
    Hnorm2.getInverse(Hnorm2inv);

    // fit homography
    uint32_t i=0;

    for (matchesType::const_iterator itIn= inliers.begin();
         itIn!=inliers.end();
         ++itIn, ++i){

        // normalize transform

        A.at<double>( i*2   , 0 )=    0.0;
        A.at<double>( i*2   , 1 )=    0.0;
        A.at<double>( i*2   , 2 )=    0.0;
        A.at<double>( i*2   , 3 )=  x1[i];
        A.at<double>( i*2   , 4 )=  y1[i];
        A.at<double>( i*2   , 5 )=    1.0;
        A.at<double>( i*2   , 6 )= -y2[i];

        A.at<double>( i*2+1 , 0 )= -x1[i];
        A.at<double>( i*2+1 , 1 )= -y1[i];
        A.at<double>( i*2+1 , 2 )=   -1.0;
        A.at<double>( i*2+1 , 3 )=    0.0;
        A.at<double>( i*2+1 , 4 )=    0.0;
        A.at<double>( i*2+1 , 5 )=    0.0;
        A.at<double>( i*2+1 , 6 )=  x2[i];

    }

    delete []x1; delete []y1; delete []x2; delete []y2; delete []_temp;

    cv::SVD::solveZ(A, x);

    H.H[0]= x.at<double>(0,0); H.H[1]= x.at<double>(1,0); H.H[2]= x.at<double>(2,0);
    H.H[3]= x.at<double>(3,0); H.H[4]= x.at<double>(4,0); H.H[5]= x.at<double>(5,0);
    H.H[6]=               0.0; H.H[7]=               0.0; H.H[8]= x.at<double>(6,0);
    H.normLast();

    // denormalize H: H= Hnorm2inv * Hnorm * Hnorm1

    cv::Mat Hnorm2inv_(3,3, cv::DataType<double>::type),
            Hnorm1_(3,3, cv::DataType<double>::type),
            Hnorm_(3,3, cv::DataType<double>::type);
    for (int i=0; i<3; ++i)
        for (int j=0; j<3; ++j){
            Hnorm_.at<double>(i,j)= H.H[i*3+j];
            Hnorm1_.at<double>(i,j)= Hnorm1.H[i*3+j];
            Hnorm2inv_.at<double>(i,j)= Hnorm2inv[i*3+j];
        }
    cv::Mat H_= Hnorm2inv_ * Hnorm_ * Hnorm1_;

    for (int i=0; i<3; ++i)
        for (int j=0; j<3; ++j)
            H.H[i*3+j]= H_.at<double>(i,j);
#endif

}

void
detRansac::normPoints( double *x, double *y, uint32_t n, homography &Hnorm ){

    double EX= 0, EY= 0, stdX, stdY, EX2= 0, EY2= 0;
    uint32_t i;

    // get mean/variance
    for (i= 0; i<n; ++i){
        EX+= x[i];
        EY+= y[i];
        EX2+= mysqr(x[i]);
        EY2+= mysqr(y[i]);
    }
    EX/=n; EY/=n;
    EX2/=n; EY2/=n;
    // max is there just in case std=0 (as later norm=NaN in this case), in this case the dimension gets mapped to 0 (as mean is 0) so scaling makes no difference.. This is not nice however, as it means the problem is illposed
    stdX= std::max( sqrt(EX2-mysqr(EX)), 1e-4 );
    stdY= std::max( sqrt(EY2-mysqr(EY)), 1e-4 );

    double normX= sqrt(2.0)/stdX, normY= sqrt(2.0)/stdY;

    Hnorm.setIdentity();
    Hnorm.H[0]= normX; Hnorm.H[2]= -EX*normX;
    Hnorm.H[4]= normY; Hnorm.H[5]= -EY*normY;

    // normalize points
    for (i= 0; i<n; ++i){
        x[i]= (x[i]-EX)*normX;
        y[i]= (y[i]-EY)*normY;
    }

}

detRansac::inlierFinder::inlierFinder(
    std::vector<ellipse> const &aEllipses1,
    std::vector<ellipse> const &aEllipses2,
    matchesType const &aPutativeMatches,
    std::vector<double> const &aPMweights,
    double aErrorThr,
    double aLowAreaChange, double aHighAreaChange) :

    nIter(0),
    putativeMatches(&aPutativeMatches), PMweights(&aPMweights),
    errorThrSq(mysqr(aErrorThr)),
    lowAreaChangeSq(mysqr(aLowAreaChange)),
    highAreaChangeSq(mysqr(aHighAreaChange)),
    nPutativeMatches(aPutativeMatches.size()) {

    //------- get largest pIDs

    uint32_t maxpID1= 0, maxpID2= 0;

    for (matchesType::const_iterator itPM=putativeMatches->begin();
         itPM != putativeMatches->end();
         ++itPM){
        maxpID1= std::max(maxpID1, itPM->first);
        maxpID2= std::max(maxpID2, itPM->second);
    }
    point1Used.clear(); point1Used.resize(maxpID1+1,0);
    point1Used.clear(); point2Used.resize(maxpID2+1,0);

    ellipse::getCentres( aEllipses1, aEllipses2, *putativeMatches, x1, y1, x2, y2, areaDiffSq );

}

detRansac::inlierFinder::~inlierFinder(){
    delete []x1; delete []y1; delete []x2; delete []y2; delete []areaDiffSq;
}

double
detRansac::inlierFinder::getScore( homography const &H, uint32_t &nInliers, matchesType *inliers ){

    double score, detASq;

    score= 0.0;
    nInliers= 0;
    if (inliers)
        inliers->clear();

    detASq= mysqr( H.getDetAffine() );
    if (detASq<1e-4) return 0.0;

    double error;
    uint32_t pID1, pID2, iPM;
    double x, y, xi, yi;
    double lowAreaChangeSqByD =  lowAreaChangeSq / detASq;
    double highAreaChangeSqByD= highAreaChangeSq / detASq;

    double Hinv[9];
    H.getInverse( Hinv );

    ++nIter;

    iPM= 0;
    for (matchesType::const_iterator itPM= putativeMatches->begin();
         itPM!=putativeMatches->end();
         ++itPM, ++iPM){

        pID1= itPM->first;
        pID2= itPM->second;
        if (point1Used[ pID1 ]==nIter || point2Used[ pID2 ]==nIter)
            continue;

        homography::affTransform( H.H , x1[iPM], y1[iPM], x , y );
        homography::affTransform( Hinv, x2[iPM], y2[iPM], xi, yi );

        error= mysqr( x1[iPM]-xi ) + mysqr( y1[iPM]-yi ) +
               mysqr( x2[iPM]-x  ) + mysqr( y2[iPM]-y  );

        if (error < errorThrSq){
            /*
            areaChangeSq= detASq * ellipses2[ pID2 ].getPropAreaSq() / ellipses1[ pID1 ].getPropAreaSq();
            if ( areaChangeSq > lowAreaChangeSq && areaChangeSq < highAreaChangeSq ){
            */
            if (areaDiffSq[iPM] > lowAreaChangeSqByD && areaDiffSq[iPM] < highAreaChangeSqByD) {
                score+= PMweights->at(iPM);
                ++nInliers;
                point1Used[ pID1 ]= nIter;
                point2Used[ pID2 ]= nIter;
                if (inliers)
                    inliers->push_back(std::make_pair(pID1,pID2));
            }
        }

    }

    return score;

}


#pragma once
#include <Eigen/Core>

#include "ear/ear.hpp"
#include "tensorfile.hpp"

namespace bear {

/// holds two arbitrary values relating to the direct and diffuse paths
template <typename T>
struct DirectDiffuse {
  T direct;
  T diffuse;
};

/// holds two arbitrary values relating to the left and right channels/ears
template <typename T>
struct LeftRight {
  T left;
  T right;
};

/// reference to a single view from a multi-view BRIR set
struct SelectedBRIR {
  size_t brir_index;
  // TODO: add orientation here
};

using namespace Eigen;

/// Holds all non-user-configurable renderer information (like virtual
/// loudspeaker layouts, BRIR sets, delay sets, decorrelation filters), and
/// provides methods intended to be used to drive the baseline DSP (like
/// calculating gains and delays, or accessing the BRIR sets).
///
/// For Objects, the data flow should look like this:
///
///   type_metadata -calc_objects_gains-> gains
///   gains, selected_brir -get_vs_gains-> vs_gains
///   gains, selected_brir -get_direct_delays-> direct_delays
///   selected_brir -get_diffuse_delays-> diffuse_delays
///
/// vs_gains, direct_delays and diffuse_delays are the gains and delays used by
/// the DSP components.
///
/// Note that the gains (length num_gains()) and virtual loudspeaker gains
/// (length num_virtual_loudspeakers()) are separate, so that extra
/// processing (e.g. gain normalisation, virtual loudspeaker remapping) can be
/// inserted without affecting the delay calculation.
class Panner {
 public:
  Panner(const std::string &brir_file_name);

  size_t num_gains() const { return n_gains_; }
  size_t num_virtual_loudspeakers() const { return n_virtual_loudspeakers_; }
  size_t num_views() const { return n_views_; }

  void calc_objects_gains(const ear::ObjectsTypeMetadata &type_metadata,
                          DirectDiffuse<Ref<VectorXd>> gains) const;

  void get_vs_gains(const DirectDiffuse<const Ref<const VectorXd> &> &gains,
                    DirectDiffuse<Ref<VectorXd>> vs_gains,
                    SelectedBRIR selected_brir = {}) const;

  // frontal delay used to initialise direct path
  double get_default_direct_delay() const;
  // frontal delay used to initialise static path
  double get_default_static_delay() const;

  LeftRight<double> get_direct_delays(const DirectDiffuse<const Ref<const VectorXd> &> &gains,
                                      SelectedBRIR selected_brir = {}) const;

  LeftRight<double> get_direct_speakers_delays(const Ref<const VectorXd> &gains,
                                               SelectedBRIR selected_brir = {}) const;

  void get_diffuse_delays(LeftRight<Ref<VectorXd>> diffuse_delays, SelectedBRIR selected_brir = {}) const;

  void calc_direct_speakers_gains(const ear::DirectSpeakersTypeMetadata &type_metadata,
                                  Ref<VectorXd> gains) const;

  size_t decorrelator_length() const;
  const float *get_decorrelator(size_t i) const;

  size_t brir_length() const;
  const float *get_brir(size_t view, size_t virtual_loudspeaker, size_t ear) const;

  Eigen::MatrixX3d get_views() const;

  bool has_gain_compensation() const;
  // delays in seconds
  double compensation_gain(double *gains,
                           LeftRight<double> direct_delays,
                           SelectedBRIR selected_brir = {}) const;
  double compensation_gain(DirectDiffuse<Ref<VectorXd>> gains,
                           LeftRight<double> direct_delays,
                           SelectedBRIR selected_brir) const;

  double compensation_gain_direct(const Ref<const VectorXd> &gains, SelectedBRIR selected_brir) const;

  const double decorrelation_delay() const;

  size_t n_hoa_channels() const { return n_hoa_channels_; }
  size_t hoa_order() const { return hoa_order_; }
  size_t hoa_ir_length() const { return hoa_irs->shape(2); }
  const float *get_hoa_ir(size_t channel, size_t ear) const;
  double hoa_delay() const { return hoa_delay_ / fs; }

 private:
  double get_real_gain_quick(double *gains,
                             LeftRight<double> direct_delays,
                             SelectedBRIR selected_brir) const;
  double get_expected_gain_quick(double *gains,
                                 LeftRight<double> direct_delays,
                                 SelectedBRIR selected_brir) const;

  double get_real_gain_quick_direct(const Ref<const VectorXd> &gains, SelectedBRIR selected_brir) const;
  double get_expected_gain_quick_direct(const Ref<const VectorXd> &gains, SelectedBRIR selected_brir) const;

  template <typename Derived>
  LeftRight<double> get_delays(const Eigen::DenseBase<Derived> &gains, SelectedBRIR selected_brir = {}) const;

  size_t n_gains_;
  size_t n_views_;
  size_t n_virtual_loudspeakers_;
  size_t brir_length_;
  size_t decorrelator_length_;
  size_t decorrelation_delay_;
  size_t front_loudspeaker_;
  size_t n_hoa_channels_;
  size_t hoa_order_;
  double hoa_delay_ = 0.0;
  enum class GainCompType { NONE, QUICK } gain_comp_type = GainCompType::NONE;

  std::shared_ptr<tensorfile::NDArrayT<float>> views;
  std::shared_ptr<tensorfile::NDArrayT<float>> brirs;
  std::shared_ptr<tensorfile::NDArrayT<float>> delays;
  std::shared_ptr<tensorfile::NDArrayT<float>> decorrelation_filters;
  std::shared_ptr<tensorfile::NDArrayT<float>> gain_comp_factors;
  std::shared_ptr<tensorfile::NDArrayT<float>> hoa_irs;
  double fs;

  std::unique_ptr<ear::GainCalculatorObjects> gain_calc;
  std::unique_ptr<ear::GainCalculatorDirectSpeakers> direct_speakers_gain_calc;

  mutable std::vector<double> temp_direct;
  mutable std::vector<double> temp_diffuse;
  mutable std::vector<double> temp_direct_diffuse;

  mutable std::vector<double> temp_direct_speakers;
};

}  // namespace bear


// Copyright (c) 2019 Graphcore Ltd. All rights reserved.
#include "poplin/Norms.hpp"
#include "poplibs_support/Tracepoint.hpp"
#include "poplibs_support/logging.hpp"
#include "poplin/ConvUtil.hpp"
#include "poplin/Convolution.hpp"
#include "popops/ElementWise.hpp"
#include "popops/Rearrange.hpp"
#include "popops/Reduce.hpp"
#include "popops/ScaledAdd.hpp"
#include "poputil/DebugInfo.hpp"
#include "poputil/TileMapping.hpp"
#include "poputil/Util.hpp"
#include "poputil/VertexTemplates.hpp"
#include "poputil/exceptions.hpp"
#include <boost/icl/interval_map.hpp>
#include <cassert>
#include <cmath>
#include <set>

using namespace poplar;
using namespace poplar::program;
using namespace poputil;
using namespace popops;
namespace logging = poplibs_support::logging;

namespace poplin {

static Tensor normReduce(Graph &graph, const Tensor &actsUngrouped,
                         const Tensor &scale, bool doSquare, Sequence &prog,
                         const Type &, // partialsType,
                         const Type &outputType,
                         const Tensor *outputToCloneFrom,
                         const DebugNameAndId &dnai) {
  std::string layer = "ReduceResult";
  Tensor t;

  // The output tensor mapping may be specified or created
  if (outputToCloneFrom) {
    t = graph.clone(outputType, *outputToCloneFrom, {dnai, layer});
  } else {
    t = createBroadcastOperand(graph, actsUngrouped, outputType, 1, true,
                               {dnai, layer});
  }

  if (actsUngrouped.rank() < 2)
    throw poplibs_error("NormReduce with rank " +
                        std::to_string(actsUngrouped.rank()) + " expected >=2");

  std::vector<std::size_t> reduceDims(actsUngrouped.rank() - 1);
  std::iota(reduceDims.begin() + 1, reduceDims.end(), 2);

  popops::reduceWithOutput(
      graph, actsUngrouped, t, reduceDims,
      {doSquare ? popops::Operation::SQUARE_ADD : popops::Operation::ADD, false,
       scale},
      prog, {dnai});
  return t;
}

static Tensor computeInvStdDev(Graph &graph, const Tensor &mean,
                               const Tensor &power, float eps, float scaleVar,
                               Sequence &prog, const Type &invStdDevType,
                               bool stableAlgo, const DebugNameAndId &dnai) {
  const auto meanType = mean.elementType();
  const auto powerType = power.elementType();
  auto iStdDev = graph.clone(invStdDevType, mean, {dnai, "iStdDev"});

  const auto meanFlat = mean.flatten();
  const auto powerFlat = power.flatten();
  const auto iStdDevFlat = iStdDev.flatten();

  const auto &target = graph.getTarget();
  const auto numTiles = target.getNumTiles();
  const auto cs = graph.addComputeSet({dnai, "iStdDev"});

  const auto mapping = graph.getTileMapping(iStdDev);
  const auto grainSize = target.getVectorWidth(invStdDevType);

  for (auto tile = 0U; tile != numTiles; ++tile) {
    const auto tileContiguousRegions =
        graph.getSortedContiguousRegions(iStdDevFlat, mapping[tile]);
    auto vertexRegions = splitRegionsBetweenWorkers(
        target, tileContiguousRegions, grainSize, 2 * grainSize);

    for (const auto &regions : vertexRegions) {
      auto v = graph.addVertex(cs,
                               templateVertex("poplin::InverseStdDeviation",
                                              meanType, powerType,
                                              invStdDevType, stableAlgo),
                               {{"mean", meanFlat.slices(regions)},
                                {"power", powerFlat.slices(regions)},
                                {"iStdDev", iStdDevFlat.slices(regions)}});
      graph.setInitialValue(v["eps"], eps);
      graph.setInitialValue(v["scaleVar"], scaleVar);
      graph.setTileMapping(v, tile);
    }
  }
  prog.add(Execute(cs, {dnai}));
  return iStdDev;
}

static Tensor broadcastChannelToMatch(const Tensor &ref, const Tensor &t) {
  return t.flatten().expand(std::vector<std::size_t>(ref.rank() - 2, 1));
}

static bool shouldExecuteCallback(unsigned replicaNormSize, unsigned normSize,
                                  DistributedNormReduceCallback callback) {
  // The callback should only be invoked when there is distributed reduction to
  // be performed.
  bool executeCallback =
      normSize > replicaNormSize && callback != nullptr && replicaNormSize > 0;
  if (executeCallback && normSize % replicaNormSize) {
    throw poplibs_error("Norm batch size must be an integer multiple of "
                        "replica batch size");
  }
  if (executeCallback) {
    logging::poplin::info("All-reduce callback called with group size of {}",
                          normSize / replicaNormSize);
  }
  return executeCallback;
}

static std::pair<Tensor, Tensor>
normStatisticsImpl(Graph &graph, const Tensor &acts, float eps, Sequence &prog,
                   bool unbiasedVarEstimate, bool stableAlgo,
                   const Type &partialsType,
                   DistributedNormReduceCallback reduceCallback,
                   unsigned normSize, const DebugNameAndId &dnai) {
  const std::string layer = "Norm/statistics";
  logging::poplin::info(
      "normStatistics acts={}, eps={}, unbiasedVarEstimate={}, type={}, "
      "normSize={} name={}",
      acts.shape(), eps, unbiasedVarEstimate, partialsType, normSize,
      dnai.getPathName() + "/" + layer);

  size_t numElements = acts.numElements();
  const auto replicaSize = acts.dim(0);
  bool executeCallback =
      shouldExecuteCallback(replicaSize, normSize, reduceCallback);

  // Ideally, we would like the scaling due to the distributed all-reduced to
  // be done by the all-reduce (see T36825).
  // But we know it doesn't do that for now and hence the scaling is folded in
  // here.
  if (executeCallback) {
    numElements *= normSize / replicaSize;
  }

  // Avoid a possible divide by zero FP exception.
  // Note that numElements will be 0 if acts.dim(1) is 0.
  if (acts.dim(1) > 0)
    numElements /= acts.dim(1);

  float scale = 1.0f;
  float scaleVar = 1.0f;
  if (numElements > 1) {
    scale /= numElements;
    if (unbiasedVarEstimate)
      scaleVar = static_cast<float>(numElements) / (numElements - 1);
  }

  auto scaleTensor =
      graph.addConstant(FLOAT, {}, scale, {dnai, layer + "/scaleTensor"});
  graph.setTileMapping(scaleTensor, 0);

  if (acts.rank() < 2)
    throw poplibs_error("NormReduce with rank " + std::to_string(acts.rank()) +
                        " expected >=2");
  std::vector<std::size_t> dims(acts.rank() - 1);
  std::iota(dims.begin() + 1, dims.end(), 2);

  const auto meanOutputType = acts.elementType();
  // The actual output type for squared sum may be different as the dynamic
  // range is higher. The selection should be based on actual statistics
  // gathered from training experiments. For now keep it at reduced precision
  // to save memory
  const auto powerOutputType = partialsType;

  constexpr bool update = false;
  popops::ReduceParams meanParams{popops::Operation::ADD, update, scaleTensor};
  popops::ReduceParams powerParams{popops::Operation::SQUARE_ADD, update,
                                   scaleTensor};
  poplar::Tensor mean =
      createBroadcastOperand(graph, acts, meanOutputType, 1, true,
                             {dnai, layer + "/mean/ReduceResult"});
  poplar::Tensor power =
      graph.clone(powerOutputType, mean, {dnai, layer + "/power/ReduceResult"});
  if (stableAlgo) {
    reduceWithOutput(graph, acts, mean, dims, std::move(meanParams), prog,
                     {dnai, layer + "/mean"});
    logging::poplin::info("Stable statistics estimator used");
    if (executeCallback) {
      auto allReduceResult =
          reduceCallback(graph, {mean}, prog, normSize / replicaSize,
                         {dnai, layer + "/mean"}, {});
      mean = allReduceResult.at(0);
    }
    using namespace popops::expr;
    auto zeroMeanActs = popops::map(graph, _1 - Cast(_2, acts.elementType()),
                                    {acts, broadcastChannelToMatch(acts, mean)},
                                    prog, {dnai, layer + "/removeMean"});
    reduceWithOutput(graph, std::move(zeroMeanActs), power, std::move(dims),
                     std::move(powerParams), prog, {dnai, layer + "/power"});
  } else {
    std::vector<poplar::Tensor> outputs = {std::move(mean), std::move(power)};
    std::vector<popops::SingleReduceOp> reductions = {
        popops::SingleReduceOp{/*in     = */ acts,
                               /*dims   = */ dims,
                               /*params = */ std::move(meanParams),
                               /*debugName = */ "mean"},
        popops::SingleReduceOp{/*in     = */ acts,
                               /*dims   = */ std::move(dims),
                               /*params = */ std::move(powerParams),
                               /*debugName = */ "power"}};
    popops::reduceMany(graph, reductions, outputs, prog,
                       {dnai, layer + "/reduceMany"});
    mean = std::move(outputs[0]);
    power = std::move(outputs[1]);
  }

  if (executeCallback) {
    std::vector<Tensor> inputsToCallback;
    std::string str;
    if (stableAlgo) {
      str = "/power";
    } else {
      inputsToCallback.push_back(mean);
      str = "/meanAndPower";
    }
    inputsToCallback.push_back(power);

    auto allReduceResult =
        reduceCallback(graph, inputsToCallback, prog, normSize / replicaSize,
                       {dnai, layer + str}, {});
    if (stableAlgo) {
      power = allReduceResult.at(0);
    } else {
      mean = allReduceResult.at(0);
      power = allReduceResult.at(1);
    }
  }

  auto iStdDev = computeInvStdDev(graph, mean, power, eps, scaleVar, prog,
                                  acts.elementType(), stableAlgo, {dnai});
  return std::make_pair(mean, iStdDev);
}

std::pair<Tensor, Tensor>
normStatistics(Graph &graph, const Tensor &acts, float eps, Sequence &prog,
               bool unbiasedVarEstimate, bool stableAlgo,
               const Type &partialsType,
               const poplar::DebugContext &debugContext) {
  POPLIN_TRACEPOINT();
  poputil::PoplibsOpDebugInfo di(
      debugContext,
      DI_ARGS(acts, eps, unbiasedVarEstimate, stableAlgo, partialsType));

  auto [mean, iStdDev] = normStatisticsImpl(
      graph, acts, eps, prog, unbiasedVarEstimate, stableAlgo, partialsType,
      nullptr, acts.dim(0), {di, "nonDistributed"});
  di.addOutputs(DI_ARGS(mean, iStdDev));
  return std::make_pair(mean, iStdDev);
}

std::pair<poplar::Tensor, poplar::Tensor> distributedNormStatistics(
    poplar::Graph &graph, const poplar::Tensor &acts, float eps,
    poplar::program::Sequence &prog, bool unbiasedVarEstimate,
    DistributedNormReduceCallback callback, unsigned normSize, bool stableAlgo,
    const poplar::Type &partialsType,
    const poplar::DebugContext &debugContext) {
  POPLIN_TRACEPOINT();
  poputil::PoplibsOpDebugInfo di(
      debugContext,
      DI_ARGS(acts, eps, unbiasedVarEstimate, stableAlgo, partialsType));
  auto [mean, iStdDev] = normStatisticsImpl(
      graph, acts, eps, prog, unbiasedVarEstimate, stableAlgo, partialsType,
      callback, normSize, {di, "Distributed"});
  di.addOutputs(DI_ARGS(mean, iStdDev));
  return std::make_pair(mean, iStdDev);
}

Tensor createNormGamma(Graph &graph, const Tensor &acts, const Type &type,
                       const poplar::DebugContext &debugContext) {
  POPLIN_TRACEPOINT();
  poputil::PoplibsOpDebugInfo di(debugContext, DI_ARGS(acts, type));
  auto output =
      createBroadcastOperand(graph, acts, type, 1, true, {di, "gamma"});
  di.addOutput(output);
  return output;
}

Tensor createNormGamma(Graph &graph, const Tensor &acts,
                       const poplar::DebugContext &debugContext) {
  POPLIN_TRACEPOINT();
  poputil::PoplibsOpDebugInfo di(debugContext, DI_ARGS(acts));
  auto output = createNormGamma(graph, acts, acts.elementType(), {di});
  di.addOutput(output);
  return output;
}

Tensor createNormBeta(Graph &graph, const Tensor &acts, const Type &type,
                      const poplar::DebugContext &debugContext) {
  POPLIN_TRACEPOINT();
  poputil::PoplibsOpDebugInfo di(debugContext, DI_ARGS(acts, type));
  auto output =
      createBroadcastOperand(graph, acts, type, 1, true, {di, "beta"});
  di.addOutput(output);
  return output;
}

Tensor createNormBeta(Graph &graph, const Tensor &acts,
                      const poplar::DebugContext &debugContext) {
  poputil::PoplibsOpDebugInfo di(debugContext, DI_ARGS(acts));
  auto output = createNormBeta(graph, acts, acts.elementType(), {di});
  di.addOutput(output);
  return output;
}

std::pair<Tensor, Tensor>
createNormParams(Graph &graph, const Tensor &acts,
                 const poplar::DebugContext &debugContext) {
  POPLIN_TRACEPOINT();
  poputil::PoplibsOpDebugInfo di(debugContext, DI_ARGS(acts));
  auto gamma = createNormGamma(graph, acts, {di});
  auto beta = createNormBeta(graph, acts, {di});
  di.addOutputs(DI_ARGS(gamma, beta));
  return std::make_pair(gamma, beta);
}

Tensor normWhiten(Graph &graph, const Tensor &acts, const Tensor &mean,
                  const Tensor &iStdDev, Sequence &prog,
                  const poplar::DebugContext &debugContext) {
  POPLIN_TRACEPOINT();
  poputil::PoplibsOpDebugInfo di(debugContext, DI_ARGS(acts, mean, iStdDev));

  const std::string layer = "Whiten";
  logging::poplin::info("normWhiten acts={}, mean={}, iStdDev={}, name={}",
                        acts.shape(), mean.shape(), iStdDev.shape(),
                        debugContext.getPathName() + "/" + layer);

  auto meanBroadcast = broadcastChannelToMatch(acts, mean);
  auto actsWhitened =
      sub(graph, acts, meanBroadcast, prog, {di, layer + "/mean"});
  auto iStdDevBroadcast = broadcastChannelToMatch(actsWhitened, iStdDev);
  mulInPlace(graph, actsWhitened, iStdDevBroadcast, prog,
             {di, layer + "/istdDev"});
  di.addOutput(actsWhitened);
  return actsWhitened;
}

Tensor normalise(Graph &graph, const Tensor &actsWhitened, const Tensor &gamma,
                 const Tensor &beta, Sequence &prog,
                 const poplar::DebugContext &debugContext) {
  POPLIN_TRACEPOINT();
  poputil::PoplibsOpDebugInfo di(debugContext,
                                 DI_ARGS(actsWhitened, gamma, beta));

  const std::string layer = "Norm/normalise";
  logging::poplin::info("normalise actsWhitened={}, gamma={}, beta={}, name={}",
                        actsWhitened.shape(), gamma.shape(), beta.shape(),
                        debugContext.getPathName() + "/" + layer);

  auto gammaBroadcast = broadcastChannelToMatch(actsWhitened, gamma);
  auto actsNormalised =
      mul(graph, actsWhitened, gammaBroadcast, prog, {di, layer + "/gamma"});
  auto betaBroadcast = broadcastChannelToMatch(actsNormalised, beta);
  addInPlace(graph, actsNormalised, betaBroadcast, prog, {di, layer + "/beta"});
  di.addOutput(actsNormalised);
  return actsNormalised;
}

static std::pair<Tensor, Tensor>
normParamGradients(Graph &graph, const Tensor &actsWhitened,
                   const Tensor &gradsIn, float scale, Sequence &prog,
                   const Type &partialsType, bool attemptRegroup,
                   const DebugNameAndId &dnai) {
  const std::string layer = "Norm/deltas";
  logging::poplin::info(
      "normParamGradients actsWhitened={}, gradsIn={}, scale={}, "
      "type={}, attemptRegroup={}, name={}",
      actsWhitened.shape(), gradsIn.shape(), scale, partialsType,
      attemptRegroup, dnai.getPathName() + "/" + layer);

  auto gradsInMaybeRegrouped =
      attemptRegroup ? popops::rearrange::regroupIfBeneficial(
                           graph, gradsIn, actsWhitened, prog, {dnai})
                     : gradsIn;
  const auto gradsInMultActs =
      mul(graph, actsWhitened, gradsInMaybeRegrouped, prog, {dnai, layer});

  auto numChannels = gradsInMultActs.dim(1);
  const auto concatInputs = concat({gradsInMultActs, gradsInMaybeRegrouped}, 1);

  // For beta = Re{gradsIn} where Re{x} reduces the tensor x along the
  //                              second dimension to produce a vector
  //                              of length x.dim(1)
  // For gamma = Re{actsWhitened .* gradsIn}
  //                              .* is element-wise multiplication operator
  //                              Reduction along second dimension

  auto scaleTensor = graph.addConstant(FLOAT, {}, scale, {dnai, "scaleTensor"});
  graph.setTileMapping(scaleTensor, 0);

  const auto concatDeltas =
      normReduce(graph, concatInputs, scaleTensor, false, prog, partialsType,
                 gradsInMaybeRegrouped.elementType(), nullptr,
                 {dnai, layer + "/JointGammaDelta"});

  return std::make_pair(concatDeltas.slice(0, numChannels),
                        concatDeltas.slice(numChannels, 2 * numChannels));
}

std::pair<Tensor, Tensor>
normParamGradients(Graph &graph, const Tensor &actsWhitened,
                   const Tensor &gradsIn, Sequence &prog,
                   const Type &partialsType,
                   const poplar::DebugContext &debugContext) {
  POPLIN_TRACEPOINT();
  poputil::PoplibsOpDebugInfo di(debugContext,
                                 DI_ARGS(actsWhitened, gradsIn, partialsType));

  auto outputs = normParamGradients(graph, actsWhitened, gradsIn, 1.0, prog,
                                    partialsType, true, {di});
  di.addOutputs({{"gammaGrad", toProfileValue(outputs.first)},
                 {"betaGrad", toProfileValue(outputs.second)}});
  return outputs;
}

Tensor normGradients(Graph &graph, const Tensor &gradsIn, const Tensor &gamma,
                     Sequence &prog, const poplar::DebugContext &debugContext) {
  poputil::PoplibsOpDebugInfo di(debugContext, DI_ARGS(gradsIn, gamma));

  const auto layer = "NormGrad";
  logging::poplin::info("normGradients gradsIn={}, gamma={}, name={}",
                        gradsIn.shape(), gamma.shape(),
                        debugContext.getPathName() + "/" + layer);
  auto gammaBroadcast = broadcastChannelToMatch(gradsIn, gamma);
  auto output = mul(graph, gradsIn, gammaBroadcast, prog, {di, layer});
  di.addOutput(output);
  return output;
}

Tensor normStatisticsGradientsImpl(Graph &graph, const Tensor &actsWhitened,
                                   const Tensor &gradsIn,
                                   const Tensor &invStdDev, Sequence &prog,
                                   const Type &partialsType, // currently unused
                                   DistributedNormReduceCallback reduceCallback,
                                   unsigned normSize,
                                   const DebugNameAndId &dnai) {
  logging::poplin::info("normStatisticsGradients actsWhitened={}, gradsIn={}, "
                        "invStdDev={}, name={}",
                        actsWhitened.shape(), gradsIn.shape(),
                        invStdDev.shape(), dnai.getPathName());

  const auto replicaNormSize = actsWhitened.dim(0);
  auto executeCallback =
      shouldExecuteCallback(replicaNormSize, normSize, reduceCallback);
  const auto actsShape = actsWhitened.shape();
  auto numElements = actsWhitened.numElements() / actsWhitened.dim(1);

  // Ideally, we would like the scaling due to the distributed all-reduced to
  // be done by the all-reduce. But we know it doesn't do that for now and hence
  // the scaling is folded in here.
  if (executeCallback) {
    numElements *= normSize / replicaNormSize;
  }
  const float rScale = 1.0f / numElements;

  auto gradsInMaybeRegrouped = popops::rearrange::regroupIfBeneficial(
      graph, gradsIn, actsWhitened, prog, {dnai});

  // split rScale = rScale1 * rScale2;
  // TODO: T12898 Research what the optimal split would be dependent on model
  // and field size.
  const auto scaleSplit = 3.0f / 4;
  float rScale1 = std::pow(rScale, scaleSplit);
  float rScale2 = rScale / rScale1;
  const auto dType = actsWhitened.elementType();

  // If type is half, ensure that rScale2 is exactly representable in device
  // HALF type so that the fastest codelet is picked up when rScale2 is used
  // in the scaledAddTo below.
  if (dType == HALF) {
    rScale2 = castToDeviceHalfValue(graph.getTarget(), rScale2);
    // re-evaluate to get better combined precision
    rScale1 = rScale / rScale2;
  }
  Tensor varDelta, meanDelta;
  // See Description of Re{} operator in normParamGradients
  // varDelta = Re{actsWhitened .* gradsIn} * -rScale
  //   Size of varDelta is the size of inverse standard deviation
  // meanDelta = Re{gradsIn} * -rScale
  std::tie(varDelta, meanDelta) =
      normParamGradients(graph, actsWhitened, gradsInMaybeRegrouped, -rScale1,
                         prog, partialsType, false, {dnai});

  if (executeCallback) {
    auto reducedGrads = reduceCallback(graph, {varDelta, meanDelta}, prog,
                                       normSize / replicaNormSize, {dnai}, {});
    varDelta = reducedGrads.at(0);
    meanDelta = reducedGrads.at(1);
  }

  auto gradient = graph.clone(actsWhitened, {dnai, "/gradsIn"});
  prog.add(Copy(gradsInMaybeRegrouped, gradient, false, {dnai}));

  // gradOut = gradsIn - rScale * actsWhitened .* Br{varDelta}
  // where Br{x} broadcast x along all dimensions other than dim(1) of
  // actsWhitened
  // gradsOut = gradsIn - rScale * actsWhitened .* Br{varDelta} + Br{meanDelta}

  auto varDeltaBroadcast = broadcastChannelToMatch(actsWhitened, varDelta);
  auto varGrads =
      mul(graph, actsWhitened, varDeltaBroadcast, prog, {dnai, "varGrads"});
  mulInPlace(graph, meanDelta, rScale2, prog, {dnai, "scaleMeanDelta"});
  auto meanDeltaBroadcast = broadcastChannelToMatch(gradient, meanDelta);
  addInPlace(graph, gradient, meanDeltaBroadcast, prog, {dnai, "meanGrads"});
  // TODO: T12899 Once scaledAddTo is targeted efficiently in element-wise ops,
  // this should become a mapInPlace() expression.
  scaledAddTo(graph, gradient, varGrads, rScale2, prog, {dnai, "addGrads"});

  // Br{invStdDev} .* (gradsIn - rScale * actsWhitened .* Br{varDelta}
  //                   + Br{meanDelta})
  auto invStdDevBroadcast = broadcastChannelToMatch(gradient, invStdDev);
  mulInPlace(graph, gradient, invStdDevBroadcast, prog, {dnai});
  return gradient;
}

Tensor normStatisticsGradients(Graph &graph, const Tensor &actsWhitened,
                               const Tensor &gradsIn, const Tensor &invStdDev,
                               Sequence &prog,
                               const Type &partialsType, // currently unused
                               const poplar::DebugContext &debugContext) {
  POPLIN_TRACEPOINT();
  poputil::PoplibsOpDebugInfo di(
      debugContext, DI_ARGS(actsWhitened, gradsIn, invStdDev, partialsType));
  const std::string layer = "NonDistributedNorm/gradients";

  auto gradient = normStatisticsGradientsImpl(
      graph, actsWhitened, gradsIn, invStdDev, prog, partialsType, nullptr,
      actsWhitened.dim(0), {di, layer});
  di.addOutput(gradient);
  return gradient;
}

Tensor distributedNormStatisticsGradients(
    Graph &graph, const Tensor &actsWhitened, const Tensor &gradsIn,
    const Tensor &invStdDev, Sequence &prog,
    DistributedNormReduceCallback normReduceCallback, unsigned normSize,
    const Type &partialsType, const poplar::DebugContext &debugContext) {
  POPLIN_TRACEPOINT();
  poputil::PoplibsOpDebugInfo di(
      debugContext, DI_ARGS(actsWhitened, gradsIn, invStdDev, partialsType));
  const std::string layer = "DistributedNorm/gradients";
  auto gradient = normStatisticsGradientsImpl(
      graph, actsWhitened, gradsIn, invStdDev, prog, partialsType,
      normReduceCallback, normSize, {di, layer});
  di.addOutput(gradient);
  return gradient;
}

} // namespace poplin


// Copyright (c) 2014-2018 The Bitcoin Core developers
// Distributed under the MIT software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.

#include <key_io.h>

#include <base58.h>
#include <bech32.h>
#include <script/script.h>
#include <util/strencodings.h>

#include <boost/variant/apply_visitor.hpp>
#include <boost/variant/static_visitor.hpp>

#include <assert.h>
#include <string.h>
#include <algorithm>

/** All alphanumeric characters except for "0", "I", "O", and "l" */
static const char* pszBase58 = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz";

bool N_DecodeBase58(const char* psz, std::vector<unsigned char>& vch)
{
    // Skip leading spaces.
    while (*psz && isspace(*psz))
        psz++;
    // Skip and count leading '1's.
    int zeroes = 0;
    while (*psz == '1') {
        zeroes++;
        psz++;
    }
    // Allocate enough space in big-endian base256 representation.
    std::vector<unsigned char> b256(strlen(psz) * 733 / 1000 + 1); // log(58) / log(256), rounded up.
    // Process the characters.
    while (*psz && !isspace(*psz)) {
        // Decode base58 character
        const char* ch = strchr(pszBase58, *psz);
        if (ch == NULL)
            return false;
        // Apply "b256 = b256 * 58 + ch".
        int carry = ch - pszBase58;
        for (std::vector<unsigned char>::reverse_iterator it = b256.rbegin(); it != b256.rend(); it++) {
            carry += 58 * (*it);
            *it = carry % 256;
            carry /= 256;
        }
        assert(carry == 0);
        psz++;
    }
    // Skip trailing spaces.
    while (isspace(*psz))
        psz++;
    if (*psz != 0)
        return false;
    // Skip leading zeroes in b256.
    std::vector<unsigned char>::iterator it = b256.begin();
    while (it != b256.end() && *it == 0)
        it++;
    // Copy result into output vector.
    vch.reserve(zeroes + (b256.end() - it));
    vch.assign(zeroes, 0x00);
    while (it != b256.end())
        vch.push_back(*(it++));
    return true;
}

std::string N_DecodeBase58(const char* psz)
{
    std::vector<unsigned char> vch;
    N_DecodeBase58(psz, vch);
    std::stringstream ss;
    ss << std::hex;

    for (unsigned int i = 0; i < vch.size(); i++) {
        unsigned char* c = &vch[i];
        ss << std::setw(2) << std::setfill('0') << (int)c[0];
    }

    return ss.str();
}

std::string N_EncodeBase58(const unsigned char* pbegin, const unsigned char* pend)
{
    // Skip & count leading zeroes.
    int zeroes = 0;
    while (pbegin != pend && *pbegin == 0) {
        pbegin++;
        zeroes++;
    }
    // Allocate enough space in big-endian base58 representation.
    std::vector<unsigned char> b58((pend - pbegin) * 138 / 100 + 1); // log(256) / log(58), rounded up.
    // Process the bytes.
    while (pbegin != pend) {
        int carry = *pbegin;
        // Apply "b58 = b58 * 256 + ch".
        for (std::vector<unsigned char>::reverse_iterator it = b58.rbegin(); it != b58.rend(); it++) {
            carry += 256 * (*it);
            *it = carry % 58;
            carry /= 58;
        }
        assert(carry == 0);
        pbegin++;
    }
    // Skip leading zeroes in base58 result.
    std::vector<unsigned char>::iterator it = b58.begin();
    while (it != b58.end() && *it == 0)
        it++;
    // Translate the result into a string.
    std::string str;
    str.reserve(zeroes + (b58.end() - it));
    str.assign(zeroes, '1');
    while (it != b58.end())
        str += pszBase58[*(it++)];
    return str;
}

std::string N_EncodeBase58(const std::vector<unsigned char>& vch)
{
    return N_EncodeBase58(&vch[0], &vch[0] + vch.size());
}

bool N_DecodeBase58(const std::string& str, std::vector<unsigned char>& vchRet)
{
    return N_DecodeBase58(str.c_str(), vchRet);
}

std::string N_EncodeBase58Check(const std::vector<unsigned char>& vchIn)
{
    // add 4-byte hash check to the end
    std::vector<unsigned char> vch(vchIn);
    uint256 hash = Hash(vch.begin(), vch.end());
    vch.insert(vch.end(), (unsigned char*)&hash, (unsigned char*)&hash + 4);
    return N_EncodeBase58(vch);
}

bool N_DecodeBase58Check(const char* psz, std::vector<unsigned char>& vchRet)
{
    if (!N_DecodeBase58(psz, vchRet) ||
        (vchRet.size() < 4)) {
        vchRet.clear();
        return false;
    }
    // re-calculate the checksum, insure it matches the included 4-byte checksum
    uint256 hash = Hash(vchRet.begin(), vchRet.end() - 4);
    if (memcmp(&hash, &vchRet.end()[-4], 4) != 0) {
        vchRet.clear();
        return false;
    }
    vchRet.resize(vchRet.size() - 4);
    return true;
}

bool N_DecodeBase58Check(const std::string& str, std::vector<unsigned char>& vchRet)
{
    return N_DecodeBase58Check(str.c_str(), vchRet);
}

//namespace
//{
class DestinationEncoder : public boost::static_visitor<std::string>
{
private:
    const CChainParams& m_params;

public:
    explicit DestinationEncoder(const CChainParams& params) : m_params(params) {}

    //TO_FIX: set body this operator
    std::string operator()(const CKeyID& id) const
    {
        std::vector<unsigned char> data = m_params.Base58Prefix(CChainParams::PUBKEY_ADDRESS);
        data.insert(data.end(), id.begin(), id.end());
        return EncodeBase58Check(data);
    }

    //TO_FIX: set body this operator
    std::string operator()(const CScriptID& id) const
    {
        std::vector<unsigned char> data = m_params.Base58Prefix(CChainParams::PUBKEY_ADDRESS);
        data.insert(data.end(), id.begin(), id.end());
        return EncodeBase58Check(data);
    }
    std::string operator()(const PKHash& id) const
    {
        std::vector<unsigned char> data = m_params.Base58Prefix(CChainParams::PUBKEY_ADDRESS);
        data.insert(data.end(), id.begin(), id.end());
        return EncodeBase58Check(data);
    }

    std::string operator()(const ScriptHash& id) const
    {
        std::vector<unsigned char> data = m_params.Base58Prefix(CChainParams::SCRIPT_ADDRESS);
        data.insert(data.end(), id.begin(), id.end());
        return N_EncodeBase58Check(data);
    }

    std::string operator()(const WitnessV0KeyHash& id) const
    {
        std::vector<unsigned char> data = {0};
        data.reserve(33);
        ConvertBits<8, 5, true>([&](unsigned char c) { data.push_back(c); }, id.begin(), id.end());
        //TO_FIX: Link with body
       // return bech32::Encode(m_params.Bech32HRP(), data);
    }

    std::string operator()(const WitnessV0ScriptHash& id) const
    {
        std::vector<unsigned char> data = {0};
        data.reserve(53);
        ConvertBits<8, 5, true>([&](unsigned char c) { data.push_back(c); }, id.begin(), id.end());
        //TO_FIX: Link with body
        //return bech32::Encode(m_params.Bech32HRP(), data);
    }

    std::string operator()(const WitnessUnknown& id) const
    {
        if (id.version < 1 || id.version > 16 || id.length < 2 || id.length > 40) {
            return {};
        }
        std::vector<unsigned char> data = {(unsigned char)id.version};
        data.reserve(1 + (id.length * 8 + 4) / 5);
        ConvertBits<8, 5, true>([&](unsigned char c) { data.push_back(c); }, id.program, id.program + id.length);
        //TO_FIX: Link with body
        //return bech32::Encode(m_params.Bech32HRP(), data);
    }

    std::string operator()(const CNoDestination& no) const { return {}; }
};

CTxDestination DecodeDestination(const std::string& str, const CChainParams& params)
{
    std::vector<unsigned char> data;
    uint160 hash;
    if (N_DecodeBase58Check(str, data)) {
        // base58-encoded Bitcoin addresses.
        // Public-key-hash-addresses have version 0 (or 111 testnet).
        // The data vector contains RIPEMD160(SHA256(pubkey)), where pubkey is the serialized public key.
        const std::vector<unsigned char>& pubkey_prefix = params.Base58Prefix(CChainParams::PUBKEY_ADDRESS);
        if (data.size() == hash.size() + pubkey_prefix.size() && std::equal(pubkey_prefix.begin(), pubkey_prefix.end(), data.begin())) {
            std::copy(data.begin() + pubkey_prefix.size(), data.end(), hash.begin());
            return PKHash(hash);
        }
        // Script-hash-addresses have version 5 (or 196 testnet).
        // The data vector contains RIPEMD160(SHA256(cscript)), where cscript is the serialized redemption script.
        const std::vector<unsigned char>& script_prefix = params.Base58Prefix(CChainParams::SCRIPT_ADDRESS);
        if (data.size() == hash.size() + script_prefix.size() && std::equal(script_prefix.begin(), script_prefix.end(), data.begin())) {
            std::copy(data.begin() + script_prefix.size(), data.end(), hash.begin());
            return ScriptHash(hash);
        }
    }
    data.clear();
    auto bech = bech32::Decode(str);
    if (bech.second.size() > 0 /*&& bech.first == params.Bech32HRP()*/) {
        // Bech32 decoding
        int version = bech.second[0]; // The first 5 bit symbol is the witness version (0-16)
        // The rest of the symbols are converted witness program bytes.
        data.reserve(((bech.second.size() - 1) * 5) / 8);
        if (ConvertBits<5, 8, false>([&](unsigned char c) { data.push_back(c); }, bech.second.begin() + 1, bech.second.end())) {
            if (version == 0) {
                {
                    WitnessV0KeyHash keyid;
                    if (data.size() == keyid.size()) {
                        std::copy(data.begin(), data.end(), keyid.begin());
                        return keyid;
                    }
                }
                {
                    /*WitnessV0ScriptHash scriptid;
                    if (data.size() == scriptid.size()) {
                        std::copy(data.begin(), data.end(), scriptid.begin());
                        return scriptid;
                    }*/
                }
                return CNoDestination();
            }
            if (version > 16 || data.size() < 2 || data.size() > 40) {
                return CNoDestination();
            }
            /*WitnessUnknown unk;
            unk.version = version;
            std::copy(data.begin(), data.end(), unk.program);
            unk.length = data.size();
            return unk;*/
        }
    }
    return CNoDestination();
}
//} // namespace

CKey DecodeSecret(const std::string& str)
{
    CKey key;
    std::vector<unsigned char> data;
    if (N_DecodeBase58Check(str, data)) {
        const std::vector<unsigned char>& privkey_prefix = Params().Base58Prefix(CChainParams::SECRET_KEY);
        if ((data.size() == 32 + privkey_prefix.size() || (data.size() == 33 + privkey_prefix.size() && data.back() == 1)) &&
            std::equal(privkey_prefix.begin(), privkey_prefix.end(), data.begin())) {
            bool compressed = data.size() == 33 + privkey_prefix.size();
            key.Set(data.begin() + privkey_prefix.size(), data.begin() + privkey_prefix.size() + 32, compressed);
        }
    }
    if (!data.empty()) {
        memory_cleanse(data.data(), data.size());
    }
    return key;
}

std::string EncodeSecret(const CKey& key)
{
    assert(key.IsValid());
    std::vector<unsigned char> data = Params().Base58Prefix(CChainParams::SECRET_KEY);
    data.insert(data.end(), key.begin(), key.end());
    if (key.IsCompressed()) {
        data.push_back(1);
    }
    std::string ret = N_EncodeBase58Check(data);
    memory_cleanse(data.data(), data.size());
    return ret;
}

CExtPubKey DecodeExtPubKey(const std::string& str)
{
    CExtPubKey key;
    std::vector<unsigned char> data;
    if (N_DecodeBase58Check(str, data)) {
        const std::vector<unsigned char>& prefix = Params().Base58Prefix(CChainParams::EXT_PUBLIC_KEY);
        if (data.size() == BIP32_EXTKEY_SIZE + prefix.size() && std::equal(prefix.begin(), prefix.end(), data.begin())) {
            key.Decode(data.data() + prefix.size());
        }
    }
    return key;
}

std::string EncodeExtPubKey(const CExtPubKey& key)
{
    std::vector<unsigned char> data = Params().Base58Prefix(CChainParams::EXT_PUBLIC_KEY);
    size_t size = data.size();
    data.resize(size + BIP32_EXTKEY_SIZE);
    key.Encode(data.data() + size);
    std::string ret = N_EncodeBase58Check(data);
    return ret;
}

CExtKey DecodeExtKey(const std::string& str)
{
    CExtKey key;
    std::vector<unsigned char> data;
    if (N_DecodeBase58Check(str, data)) {
        const std::vector<unsigned char>& prefix = Params().Base58Prefix(CChainParams::EXT_SECRET_KEY);
        if (data.size() == BIP32_EXTKEY_SIZE + prefix.size() && std::equal(prefix.begin(), prefix.end(), data.begin())) {
            key.Decode(data.data() + prefix.size());
        }
    }
    return key;
}

std::string EncodeExtKey(const CExtKey& key)
{
    std::vector<unsigned char> data = Params().Base58Prefix(CChainParams::EXT_SECRET_KEY);
    size_t size = data.size();
    data.resize(size + BIP32_EXTKEY_SIZE);
    key.Encode(data.data() + size);
    std::string ret = N_EncodeBase58Check(data);
    memory_cleanse(data.data(), data.size());
    return ret;
}

std::string EncodeDestination(const CTxDestination& dest)
{
    return boost::apply_visitor(DestinationEncoder(Params()), dest);
}

CTxDestination DecodeDestination(const std::string& str)
{
    return DecodeDestination(str, Params());
}

bool IsValidDestinationString(const std::string& str, const CChainParams& params)
{
    return IsValidDestination(DecodeDestination(str, params));
}

bool IsValidDestinationString(const std::string& str)
{
    return IsValidDestinationString(str, Params());
}

bool IsValidContractSenderAddressString(const std::string& str)
{
    return IsValidContractSenderAddress(DecodeDestination(str));
}

#ifdef ENABLE_BITCORE_RPC
bool DecodeIndexKey(const std::string &str, uint256 &hashBytes, int &type)
{
    CTxDestination dest = DecodeDestination(str);
    if (IsValidDestination(dest))
    {
        const PKHash *keyID = boost::get<PKHash>(&dest);
        if(keyID)
        {
            memcpy(&hashBytes, keyID, 20);
            type = 1;
            return true;
        }

        const ScriptHash *scriptID = boost::get<ScriptHash>(&dest);
        if(scriptID)
        {
            memcpy(&hashBytes, scriptID, 20);
            type = 2;
            return true;
        }

        const WitnessV0ScriptHash *witnessV0ScriptID = boost::get<WitnessV0ScriptHash>(&dest);
        if (witnessV0ScriptID) {
            memcpy(&hashBytes, witnessV0ScriptID, 32);
            type = 3;
            return true;
        }

        const WitnessV0KeyHash *witnessV0KeyID = boost::get<WitnessV0KeyHash>(&dest);
        if (witnessV0KeyID) {
            memcpy(&hashBytes, witnessV0KeyID, 20);
            type = 4;
            return true;
        }
    }

    return false;
}
#endif


//
// Copyright (c) 2002--2010
// Toon Knapen, Karl Meerbergen, Kresimir Fresl,
// Thomas Klimpel and Rutger ter Borg
//
// Distributed under the Boost Software License, Version 1.0.
// (See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)
//
// THIS FILE IS AUTOMATICALLY GENERATED
// PLEASE DO NOT EDIT!
//

#ifndef BOOST_NUMERIC_BINDINGS_LAPACK_AUXILIARY_LARFB_HPP
#define BOOST_NUMERIC_BINDINGS_LAPACK_AUXILIARY_LARFB_HPP

#include <boost/assert.hpp>
#include <boost/numeric/bindings/begin.hpp>
#include <boost/numeric/bindings/data_order.hpp>
#include <boost/numeric/bindings/detail/array.hpp>
#include <boost/numeric/bindings/detail/if_left.hpp>
#include <boost/numeric/bindings/is_column_major.hpp>
#include <boost/numeric/bindings/is_complex.hpp>
#include <boost/numeric/bindings/is_mutable.hpp>
#include <boost/numeric/bindings/is_real.hpp>
#include <boost/numeric/bindings/lapack/workspace.hpp>
#include <boost/numeric/bindings/remove_imaginary.hpp>
#include <boost/numeric/bindings/size.hpp>
#include <boost/numeric/bindings/stride.hpp>
#include <boost/numeric/bindings/trans_tag.hpp>
#include <boost/numeric/bindings/value_type.hpp>
#include <boost/static_assert.hpp>
#include <boost/type_traits/is_same.hpp>
#include <boost/type_traits/remove_const.hpp>
#include <boost/utility/enable_if.hpp>

//
// The LAPACK-backend for larfb is the netlib-compatible backend.
//
#include <boost/numeric/bindings/lapack/detail/lapack.h>
#include <boost/numeric/bindings/lapack/detail/lapack_option.hpp>

namespace boost {
namespace numeric {
namespace bindings {
namespace lapack {

//
// The detail namespace contains value-type-overloaded functions that
// dispatch to the appropriate back-end LAPACK-routine.
//
namespace detail {

//
// Overloaded function for dispatching to
// * netlib-compatible LAPACK backend (the default), and
// * float value-type.
//
template< typename Side, typename Trans >
inline std::ptrdiff_t larfb( const Side, const Trans, const char direct,
        const char storev, const fortran_int_t m, const fortran_int_t n,
        const fortran_int_t k, const float* v, const fortran_int_t ldv,
        const float* t, const fortran_int_t ldt, float* c,
        const fortran_int_t ldc, float* work, const fortran_int_t ldwork ) {
    fortran_int_t info(0);
    LAPACK_SLARFB( &lapack_option< Side >::value, &lapack_option<
            Trans >::value, &direct, &storev, &m, &n, &k, v, &ldv, t, &ldt, c,
            &ldc, work, &ldwork );
    return info;
}

//
// Overloaded function for dispatching to
// * netlib-compatible LAPACK backend (the default), and
// * double value-type.
//
template< typename Side, typename Trans >
inline std::ptrdiff_t larfb( const Side, const Trans, const char direct,
        const char storev, const fortran_int_t m, const fortran_int_t n,
        const fortran_int_t k, const double* v, const fortran_int_t ldv,
        const double* t, const fortran_int_t ldt, double* c,
        const fortran_int_t ldc, double* work, const fortran_int_t ldwork ) {
    fortran_int_t info(0);
    LAPACK_DLARFB( &lapack_option< Side >::value, &lapack_option<
            Trans >::value, &direct, &storev, &m, &n, &k, v, &ldv, t, &ldt, c,
            &ldc, work, &ldwork );
    return info;
}

//
// Overloaded function for dispatching to
// * netlib-compatible LAPACK backend (the default), and
// * complex<float> value-type.
//
template< typename Side, typename Trans >
inline std::ptrdiff_t larfb( const Side, const Trans, const char direct,
        const char storev, const fortran_int_t m, const fortran_int_t n,
        const fortran_int_t k, const std::complex<float>* v,
        const fortran_int_t ldv, const std::complex<float>* t,
        const fortran_int_t ldt, std::complex<float>* c,
        const fortran_int_t ldc, std::complex<float>* work,
        const fortran_int_t ldwork ) {
    fortran_int_t info(0);
    LAPACK_CLARFB( &lapack_option< Side >::value, &lapack_option<
            Trans >::value, &direct, &storev, &m, &n, &k, v, &ldv, t, &ldt, c,
            &ldc, work, &ldwork );
    return info;
}

//
// Overloaded function for dispatching to
// * netlib-compatible LAPACK backend (the default), and
// * complex<double> value-type.
//
template< typename Side, typename Trans >
inline std::ptrdiff_t larfb( const Side, const Trans, const char direct,
        const char storev, const fortran_int_t m, const fortran_int_t n,
        const fortran_int_t k, const std::complex<double>* v,
        const fortran_int_t ldv, const std::complex<double>* t,
        const fortran_int_t ldt, std::complex<double>* c,
        const fortran_int_t ldc, std::complex<double>* work,
        const fortran_int_t ldwork ) {
    fortran_int_t info(0);
    LAPACK_ZLARFB( &lapack_option< Side >::value, &lapack_option<
            Trans >::value, &direct, &storev, &m, &n, &k, v, &ldv, t, &ldt, c,
            &ldc, work, &ldwork );
    return info;
}

} // namespace detail

//
// Value-type based template class. Use this class if you need a type
// for dispatching to larfb.
//
template< typename Value, typename Enable = void >
struct larfb_impl {};

//
// This implementation is enabled if Value is a real type.
//
template< typename Value >
struct larfb_impl< Value, typename boost::enable_if< is_real< Value > >::type > {

    typedef Value value_type;
    typedef typename remove_imaginary< Value >::type real_type;

    //
    // Static member function for user-defined workspaces, that
    // * Deduces the required arguments for dispatching to LAPACK, and
    // * Asserts that most arguments make sense.
    //
    template< typename Side, typename MatrixV, typename MatrixT,
            typename MatrixC, typename WORK >
    static std::ptrdiff_t invoke( const Side side, const char direct,
            const char storev, const MatrixV& v, const MatrixT& t, MatrixC& c,
            const fortran_int_t ldwork, detail::workspace1< WORK > work ) {
        namespace bindings = ::boost::numeric::bindings;
        typedef typename result_of::data_order< MatrixT >::type order;
        typedef typename result_of::trans_tag< MatrixV, order >::type trans;
        BOOST_STATIC_ASSERT( (bindings::is_column_major< MatrixC >::value) );
        BOOST_STATIC_ASSERT( (boost::is_same< typename remove_const<
                typename bindings::value_type< MatrixV >::type >::type,
                typename remove_const< typename bindings::value_type<
                MatrixT >::type >::type >::value) );
        BOOST_STATIC_ASSERT( (boost::is_same< typename remove_const<
                typename bindings::value_type< MatrixV >::type >::type,
                typename remove_const< typename bindings::value_type<
                MatrixC >::type >::type >::value) );
        BOOST_STATIC_ASSERT( (bindings::is_mutable< MatrixC >::value) );
        BOOST_ASSERT( bindings::size(work.select(real_type())) >=
                min_size_work( ldwork, bindings::size_column(t) ));
        BOOST_ASSERT( bindings::size_minor(c) == 1 ||
                bindings::stride_minor(c) == 1 );
        BOOST_ASSERT( bindings::size_minor(t) == 1 ||
                bindings::stride_minor(t) == 1 );
        BOOST_ASSERT( bindings::size_minor(v) == 1 ||
                bindings::stride_minor(v) == 1 );
        BOOST_ASSERT( bindings::stride_major(c) >= std::max< std::ptrdiff_t >(1,
                bindings::size_row(c)) );
        BOOST_ASSERT( bindings::stride_major(t) >= bindings::size_column(t) );
        BOOST_ASSERT( direct == 'F' || direct == 'B' );
        BOOST_ASSERT( storev == 'C' || storev == 'R' );
        return detail::larfb( side, trans(), direct, storev,
                bindings::size_row(c), bindings::size_column(c),
                bindings::size_column(t), bindings::begin_value(v),
                bindings::stride_major(v), bindings::begin_value(t),
                bindings::stride_major(t), bindings::begin_value(c),
                bindings::stride_major(c),
                bindings::begin_value(work.select(real_type())), ldwork );
    }

    //
    // Static member function that
    // * Figures out the minimal workspace requirements, and passes
    //   the results to the user-defined workspace overload of the
    //   invoke static member function
    // * Enables the unblocked algorithm (BLAS level 2)
    //
    template< typename Side, typename MatrixV, typename MatrixT,
            typename MatrixC >
    static std::ptrdiff_t invoke( const Side side, const char direct,
            const char storev, const MatrixV& v, const MatrixT& t, MatrixC& c,
            const fortran_int_t ldwork, minimal_workspace ) {
        namespace bindings = ::boost::numeric::bindings;
        typedef typename result_of::data_order< MatrixT >::type order;
        typedef typename result_of::trans_tag< MatrixV, order >::type trans;
        bindings::detail::array< real_type > tmp_work( min_size_work( ldwork,
                bindings::size_column(t) ) );
        return invoke( side, direct, storev, v, t, c, ldwork,
                workspace( tmp_work ) );
    }

    //
    // Static member function that
    // * Figures out the optimal workspace requirements, and passes
    //   the results to the user-defined workspace overload of the
    //   invoke static member
    // * Enables the blocked algorithm (BLAS level 3)
    //
    template< typename Side, typename MatrixV, typename MatrixT,
            typename MatrixC >
    static std::ptrdiff_t invoke( const Side side, const char direct,
            const char storev, const MatrixV& v, const MatrixT& t, MatrixC& c,
            const fortran_int_t ldwork, optimal_workspace ) {
        namespace bindings = ::boost::numeric::bindings;
        typedef typename result_of::data_order< MatrixT >::type order;
        typedef typename result_of::trans_tag< MatrixV, order >::type trans;
        return invoke( side, direct, storev, v, t, c, ldwork,
                minimal_workspace() );
    }

    //
    // Static member function that returns the minimum size of
    // workspace-array work.
    //
    static std::ptrdiff_t min_size_work( const std::ptrdiff_t ldwork,
            const std::ptrdiff_t k ) {
        return ldwork * k;
    }
};

//
// This implementation is enabled if Value is a complex type.
//
template< typename Value >
struct larfb_impl< Value, typename boost::enable_if< is_complex< Value > >::type > {

    typedef Value value_type;
    typedef typename remove_imaginary< Value >::type real_type;

    //
    // Static member function for user-defined workspaces, that
    // * Deduces the required arguments for dispatching to LAPACK, and
    // * Asserts that most arguments make sense.
    //
    template< typename Side, typename MatrixV, typename MatrixT,
            typename MatrixC, typename WORK >
    static std::ptrdiff_t invoke( const Side side, const char direct,
            const char storev, const MatrixV& v, const MatrixT& t, MatrixC& c,
            const fortran_int_t ldwork, detail::workspace1< WORK > work ) {
        namespace bindings = ::boost::numeric::bindings;
        typedef typename result_of::data_order< MatrixT >::type order;
        typedef typename result_of::trans_tag< MatrixV, order >::type trans;
        BOOST_STATIC_ASSERT( (bindings::is_column_major< MatrixC >::value) );
        BOOST_STATIC_ASSERT( (boost::is_same< typename remove_const<
                typename bindings::value_type< MatrixV >::type >::type,
                typename remove_const< typename bindings::value_type<
                MatrixT >::type >::type >::value) );
        BOOST_STATIC_ASSERT( (boost::is_same< typename remove_const<
                typename bindings::value_type< MatrixV >::type >::type,
                typename remove_const< typename bindings::value_type<
                MatrixC >::type >::type >::value) );
        BOOST_STATIC_ASSERT( (bindings::is_mutable< MatrixC >::value) );
        BOOST_ASSERT( bindings::size(work.select(value_type())) >=
                min_size_work( ldwork, bindings::size_column(t) ));
        BOOST_ASSERT( bindings::size_minor(c) == 1 ||
                bindings::stride_minor(c) == 1 );
        BOOST_ASSERT( bindings::size_minor(t) == 1 ||
                bindings::stride_minor(t) == 1 );
        BOOST_ASSERT( bindings::size_minor(v) == 1 ||
                bindings::stride_minor(v) == 1 );
        BOOST_ASSERT( bindings::stride_major(c) >= std::max< std::ptrdiff_t >(1,
                bindings::size_row(c)) );
        BOOST_ASSERT( bindings::stride_major(t) >= bindings::size_column(t) );
        BOOST_ASSERT( direct == 'F' || direct == 'B' );
        BOOST_ASSERT( storev == 'C' || storev == 'R' );
        return detail::larfb( side, trans(), direct, storev,
                bindings::size_row(c), bindings::size_column(c),
                bindings::size_column(t), bindings::begin_value(v),
                bindings::stride_major(v), bindings::begin_value(t),
                bindings::stride_major(t), bindings::begin_value(c),
                bindings::stride_major(c),
                bindings::begin_value(work.select(value_type())), ldwork );
    }

    //
    // Static member function that
    // * Figures out the minimal workspace requirements, and passes
    //   the results to the user-defined workspace overload of the
    //   invoke static member function
    // * Enables the unblocked algorithm (BLAS level 2)
    //
    template< typename Side, typename MatrixV, typename MatrixT,
            typename MatrixC >
    static std::ptrdiff_t invoke( const Side side, const char direct,
            const char storev, const MatrixV& v, const MatrixT& t, MatrixC& c,
            const fortran_int_t ldwork, minimal_workspace ) {
        namespace bindings = ::boost::numeric::bindings;
        typedef typename result_of::data_order< MatrixT >::type order;
        typedef typename result_of::trans_tag< MatrixV, order >::type trans;
        bindings::detail::array< value_type > tmp_work( min_size_work( ldwork,
                bindings::size_column(t) ) );
        return invoke( side, direct, storev, v, t, c, ldwork,
                workspace( tmp_work ) );
    }

    //
    // Static member function that
    // * Figures out the optimal workspace requirements, and passes
    //   the results to the user-defined workspace overload of the
    //   invoke static member
    // * Enables the blocked algorithm (BLAS level 3)
    //
    template< typename Side, typename MatrixV, typename MatrixT,
            typename MatrixC >
    static std::ptrdiff_t invoke( const Side side, const char direct,
            const char storev, const MatrixV& v, const MatrixT& t, MatrixC& c,
            const fortran_int_t ldwork, optimal_workspace ) {
        namespace bindings = ::boost::numeric::bindings;
        typedef typename result_of::data_order< MatrixT >::type order;
        typedef typename result_of::trans_tag< MatrixV, order >::type trans;
        return invoke( side, direct, storev, v, t, c, ldwork,
                minimal_workspace() );
    }

    //
    // Static member function that returns the minimum size of
    // workspace-array work.
    //
    static std::ptrdiff_t min_size_work( const std::ptrdiff_t ldwork,
            const std::ptrdiff_t k ) {
        return ldwork * k;
    }
};

//
// Functions for direct use. These functions are overloaded for temporaries,
// so that wrapped types can still be passed and used for write-access. In
// addition, if applicable, they are overloaded for user-defined workspaces.
// Calls to these functions are passed to the larfb_impl classes. In the
// documentation, most overloads are collapsed to avoid a large number of
// prototypes which are very similar.
//

//
// Overloaded function for larfb. Its overload differs for
// * User-defined workspace
//
template< typename Side, typename MatrixV, typename MatrixT, typename MatrixC,
        typename Workspace >
inline typename boost::enable_if< detail::is_workspace< Workspace >,
        std::ptrdiff_t >::type
larfb( const Side side, const char direct, const char storev,
        const MatrixV& v, const MatrixT& t, MatrixC& c,
        const fortran_int_t ldwork, Workspace work ) {
    return larfb_impl< typename bindings::value_type<
            MatrixV >::type >::invoke( side, direct, storev, v, t, c, ldwork,
            work );
}

//
// Overloaded function for larfb. Its overload differs for
// * Default workspace-type (optimal)
//
template< typename Side, typename MatrixV, typename MatrixT, typename MatrixC >
inline typename boost::disable_if< detail::is_workspace< MatrixC >,
        std::ptrdiff_t >::type
larfb( const Side side, const char direct, const char storev,
        const MatrixV& v, const MatrixT& t, MatrixC& c,
        const fortran_int_t ldwork ) {
    return larfb_impl< typename bindings::value_type<
            MatrixV >::type >::invoke( side, direct, storev, v, t, c, ldwork,
            optimal_workspace() );
}

} // namespace lapack
} // namespace bindings
} // namespace numeric
} // namespace boost

#endif


//
// Created by eliane on 26/12/18.
//

#include "ReScorer.h"

#include <exception>

#include <Engines/Internals/InternalsUtilityFunctions.h>

#include <boost/log/core.hpp>
#include <boost/log/trivial.hpp>

namespace SmolDock {
    ReScorer::ReScorer(SmolDock::Protein &prot, SmolDock::Molecule &mol, Score::ScoringFunctionType scorFuncType) :
            protein(prot), molecule(mol),
            scoringFunctionType(scorFuncType) {

    }

    bool ReScorer::prepare() {

        this->iprotein = this->protein.getiProtein();
        this->iconformer = this->molecule.getInitialConformer();
        this->itransform = iTransformIdentityInit();
        this->scoringFunction = scoringFunctionFactory(this->scoringFunctionType, this->iconformer,
                                                       this->iprotein, this->itransform,
                                                       10e-3);

        return true;
    }

    double ReScorer::getScore() {
        if (this->prepared == false) {
            BOOST_LOG_TRIVIAL(error) << "Rescorer getScore() called before prepare() ";
            std::terminate();
        }
        arma::mat state = this->scoringFunction->getStartingConditions();
        double score = this->scoringFunction->Evaluate(state);
        return score;
    }

}