/* main.cpp - a "main" file, just a debugging tool
 *
 * Copyright (C) 2021, LWE-PVSS
 *
 * 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.
 **/

// This file is just a convenience, a handy tool that lets us run
// small porgrams without having to use the awkward ctest syntax.
#include <cassert>
#include <cmath>
#include <random>
#include <chrono>
#include <string>
#include <sys/time.h>
#include <sys/resource.h>
using namespace std;

#include <NTL/version.h>
#include "regevEnc.hpp"
#include "regevProofs.hpp"

using namespace ALGEBRA;
using namespace REGEVENC;

int main(int argc, char** argv) {
    // std::cout << "- Found GMP version "<<__GNU_MP__ <<std::endl;
    // std::cout << "- Found NTL version "<<NTL_VERSION <<std::endl;
    // std::cout << "- Found Sodium version "<<SODIUM_VERSION_STRING<<std::endl;

    int nParties = 512;
    if (argc > 1) {
        nParties = std::stoi(argv[1]);
    }
    if (nParties < 32 || nParties > 4096)
        nParties = 512;
    std::cout << "nParties="<<nParties << std::endl;

    // The dimensions of the the CRX is k-by-m, but note that this is a
    // matrix over GF(p^2) so the lattice dimensions we get is twice that
    KeyParams kp(nParties);
    // kp.k=64;
    GlobalKey gpk("testContext", kp);
    std::cout <<"{ kay:"<<gpk.kay<<", enn:"<<gpk.enn
      <<", sigmaEnc1:"<<gpk.sigmaEnc1<<", sigmaEnc2:"<<gpk.sigmaEnc2<<" }\n";

    TernaryEMatrix::init();
    MerlinRegev mer;
    PedersenContext ped;
    SharingParams ssp(interval(1,gpk.enn+1), gpk.tee);
    VerifierData vd(gpk, ped, mer, ssp);
    ProverData pd(vd);

    // Generate/verify the proofs by the second party (idx=1)
    int partyIdx = 1;

    // Key generation for all the parties
    std::vector<ALGEBRA::EVector> kgNoise(gpk.enn);
    std::vector<ALGEBRA::EVector> sk(gpk.enn);
    std::vector<ALGEBRA::EVector> pk(gpk.enn);
    auto start = chrono::steady_clock::now();
    crsTicks = 0;
    for (int i=0; i<gpk.enn; i++) {
        std::tie(sk[i],pk[i]) = gpk.genKeys(&kgNoise[i]);
        gpk.addPK(pk[i]);
    }
    gpk.setKeyHash();
    auto end = chrono::steady_clock::now();
    auto ticks = chrono::duration_cast<chrono::milliseconds>(end - start).count();
    std::cout <<gpk.enn<<" keyGens in "<<ticks<<" milliseconds, avg="<<(ticks/double(gpk.enn))
        << " ("<< (crsTicks/double(gpk.enn)) << " for s x A)\n";

    // encryption
    std::vector<ALGEBRA::SVector> ptxt1(gpk.enn);
    std::vector<GlobalKey::CtxtPair> ctxt1(gpk.enn);
    // secret sharing of a random value , the secret itself is sshr[0]
    ALGEBRA::SVector sshr;
    ssp.randomSharing(sshr);
    for (int i=0; i<gpk.enn; i++) {
        resize(ptxt1[i], gpk.enn);
        for (int j=0; j<gpk.enn; j++) ptxt1[i][j] = sshr[i+1];
    }
    start = chrono::steady_clock::now();
    crsTicks = 0;
    for (int i=0; i<gpk.enn; i++) {
        ctxt1[i] = gpk.encrypt(ptxt1[i]);
    }
    end = chrono::steady_clock::now();
    ticks = chrono::duration_cast<chrono::milliseconds>(end - start).count();
    std::cout <<gpk.enn<<" encryptions in "<<ticks<<" milliseconds, avg="<<(ticks/double(gpk.enn))
        << " ("<< (crsTicks/double(gpk.enn)) << " for A x r)\n";

    // decryption at party #1
    ALGEBRA::SVector ptxt2;    resize(ptxt2, gpk.tee);
    ALGEBRA::EVector decNoise; resize(decNoise, gpk.tee);
    start = chrono::steady_clock::now();
    for (int i=0; i<gpk.tee; i++) { // decrypt 2nd entry in i'th ctxt
        ptxt2[i] = gpk.decrypt(sk[partyIdx], partyIdx, ctxt1[i], &(decNoise[i]));
    }
    end = chrono::steady_clock::now();
    ticks = chrono::duration_cast<chrono::milliseconds>(end - start).count();
    std::cout << gpk.tee << " decryptions in "<<ticks<< " milliseconds, avg="
        << (ticks/double(gpk.tee)) << std::endl;

    for (int i=0; i<gpk.tee; i++) { // decrypt 2nd entry in i'th ctxt
        if (ptxt2[i] != ptxt1[i][partyIdx])
            std::cout << "decryption error in "<<i<<"th ciphertext\n";
    }

    // re-encryption at party #1
    ALGEBRA::SVector ptxt3;
    resize(ptxt3, gpk.enn);
    for (int j=0; j<gpk.enn; j++) ptxt3[j] = sshr[j+1];
    ALGEBRA::EVector encRnd;
    REGEVENC::GlobalKey::CtxtPair eNoise;
    auto ctxt2 = gpk.encrypt(ptxt3, encRnd, eNoise);

    // Copy the first t ciphertexts into a k x t matrix and another t-vector
    EMatrix ctxtMat;
    resize(ctxtMat, gpk.kay, gpk.tee);
    EVector ctxtVec;
    resize(ctxtVec, gpk.tee);
    for (int i=0; i<gpk.tee; i++) {
        for (int j=0; j<gpk.kay; j++)
            ctxtMat[j][i] = ctxt1[i].first[j];
        ctxtVec[i] = ctxt1[i].second[partyIdx];
    }

    // prepare for proof, commit to the secret key
    DLPROOFS::Point::counter = 0;
    DLPROOFS::Point::timer = 0;
    int origSize = sk[partyIdx].length();

    vd.sk1Com = commit(sk[partyIdx], vd.sk1Idx, vd.Gs, pd.sk1Rnd);

    start = chrono::steady_clock::now();
    SVector lagrange = vd.sp->lagrangeCoeffs(interval(1,gpk.tee+1));

    crsTicks = 0;
    proveDecryption(pd, ctxtMat, ctxtVec, ptxt2, sk[partyIdx], decNoise);
    proveEncryption(pd, ctxt2.first, ctxt2.second, ptxt3, encRnd, eNoise.first, eNoise.second);
    proveKeyGen(pd, partyIdx, sk[partyIdx], kgNoise[partyIdx]);
    proveReShare(pd, lagrange, ptxt2, ptxt3);
    proveSmallness(pd);

    end = chrono::steady_clock::now();
    ticks = chrono::duration_cast<chrono::milliseconds>(end - start).count();
    std::cout << "preparing to prove and committing in "<<ticks<< " milliseconds, "
        << DLPROOFS::Point::counter << " exponentiations in "
        << ((500+DLPROOFS::Point::timer)/1000) << " milliseconds"
        << " ("<< crsTicks << " for xR x A)\n";

    // aggregate the constraints and flatten everything before proving
    DLPROOFS::Point::counter = 0;
    DLPROOFS::Point::timer = 0;
    std::cout<<"aggregting constraints\n";
    start = chrono::steady_clock::now();
    ReadyToProve rtp;
    rtp.aggregateProver(pd);

    // Make copies of the Merlin transcripts and specialize them
    // for the final constraints before proving/verifying them
    auto merLin = *vd.mer;
    merLin.processConstraint("linear", rtp.linCnstr);

    auto merQuad = *vd.mer;
    merQuad.processConstraint("quadratic", rtp.quadCnstr);

    // Flatten the statements, this relases the memory of the constraints
    // (hence the Merlin processing above must be done before doing this).
    std::cout<<"flatenning constraints\n";
    rtp.flattenLinPrv(pd);
    rtp.flattenQuadPrv(pd);

    ReadyToVerify rtv = rtp; // a copy without the secret variables

    // prove and verify the linear statement
    auto merLinVer = merLin; // another copy for verification
    DLPROOFS::LinPfTranscript pfL("Linear");
    pfL.C = rtp.linCom;

    end = chrono::steady_clock::now();
    ticks = chrono::duration_cast<chrono::milliseconds>(end - start).count();
    std::cout << "aggregating constaints in "<<ticks<< " milliseconds, "
        << DLPROOFS::Point::counter << " exponentiations in "
        << ((500+DLPROOFS::Point::timer)/1000) << " milliseconds\n";

    // The actual proof
    DLPROOFS::Point::counter = 0;
    DLPROOFS::Point::timer = 0;
    start = chrono::steady_clock::now();
    DLPROOFS::proveLinear(pfL, rtp.lComRnd, merLin, rtp.linWtns.data(),
            rtp.linStmnt.data(), rtp.linGs.data(), rtp.linGs.size());
    end = chrono::steady_clock::now();
    ticks = chrono::duration_cast<chrono::milliseconds>(end - start).count();
    std::cout << "proving linear in "<<ticks<< " milliseconds, "
        << DLPROOFS::Point::counter << " exponentiations in "
        << ((500+DLPROOFS::Point::timer)/1000) << " milliseconds\n";

    DLPROOFS::Point::counter = 0;
    DLPROOFS::Point::timer = 0;
    start = chrono::steady_clock::now();
    if (!DLPROOFS::verifyLinear(pfL, rtv.linStmnt.data(), rtv.linGs.data(),
                      rtv.linGs.size(), rtv.linCnstr.equalsTo, merLinVer))
        std::cout << "failed linear verification\n";
    end = chrono::steady_clock::now();
    ticks = chrono::duration_cast<chrono::milliseconds>(end - start).count();
    std::cout << "verifying linear in "<<ticks<< " milliseconds, "
        << DLPROOFS::Point::counter << " exponentiations in "
        << ((500+DLPROOFS::Point::timer)/1000) << " milliseconds\n";

    // prove and verify the quadratic statement
    DLPROOFS::Point::counter = 0;
    DLPROOFS::Point::timer = 0;
    auto merQuadVer = merQuad; // another copy for verification
    DLPROOFS::QuadPfTranscript pfQ("Quadratic");
    pfQ.C = rtp.quadCom;
     // The actual proof
    start = chrono::steady_clock::now();
    DLPROOFS::proveQuadratic(pfQ, rtp.qComRnd, merQuad, rtp.quadGs.data(),
                rtp.quadWtnsG.data(), rtp.quadHs.data(), rtp.quadWtnsH.data(),
                rtp.quadGs.size());
    end = chrono::steady_clock::now();
    ticks = chrono::duration_cast<chrono::milliseconds>(end - start).count();
    std::cout << "proving quadratic in "<<ticks<< " milliseconds, "
        << DLPROOFS::Point::counter << " exponentiations in "
        << ((500+DLPROOFS::Point::timer)/1000) << " milliseconds\n";

    // The actual verification
    DLPROOFS::Point::counter = 0;
    DLPROOFS::Point::timer = 0;
    start = chrono::steady_clock::now();
    if (!DLPROOFS::verifyQuadratic(pfQ, rtv.quadGs.data(), rtv.quadHs.data(),
                        rtp.quadGs.size(), rtv.quadCnstr.equalsTo, merQuadVer,
                        rtv.offstG.data(), rtv.offstH.data()))
        std::cout << "failed quadratic verification\n";
    end = chrono::steady_clock::now();
    ticks = chrono::duration_cast<chrono::milliseconds>(end - start).count();
    std::cout << "verifying quadratic in "<<ticks<< " milliseconds, "
        << DLPROOFS::Point::counter << " exponentiations in "
        << ((500+DLPROOFS::Point::timer)/1000) << " milliseconds\n";

    struct rusage ru;
    getrusage(RUSAGE_SELF, &ru);
    std::cout << " max mem: " << ru.ru_maxrss << " kilobytes\n";

    return 0;
}

#if 0
#include <vector>
#include <iostream>
#include "bulletproof.hpp"

int main(int, char**) {
    constexpr size_t pfSize = 13;

    // build a constraint: sum_i ai*bi = b = \sum_bi^2
    DLPROOFS::LinConstraint cnstrL;
    for (size_t i=0; i<pfSize; i++) {
        CRV25519::Scalar& x = cnstrL.terms[i+1].setInteger(i+1);
        cnstrL.equalsTo += x * x;
    }
    DLPROOFS::PtxtVec& xes = cnstrL.terms;
    DLPROOFS::LinPfTranscript pfL = proveLinear("blah", cnstrL, xes);
    std::cout << "linear: "<<verifyLinear(cnstrL, pfL) <<std::endl;

    DLPROOFS::QuadConstraint cnstrQ;
    for (auto& x : xes) {
        cnstrQ.indexes.insert(cnstrQ.indexes.end(), x.first);
        cnstrQ.equalsTo += x.second * x.second;
    }
    DLPROOFS::PtxtVec& ys = cnstrL.terms;
    DLPROOFS::QuadPfTranscript pfQ = proveQuadratic("blah", cnstrQ, xes, ys);
    std::cout << "quadratic: "<<verifyQuadratic(cnstrQ, pfQ) <<std::endl;

    std::set<size_t> indexes;
    for (auto& elem: xes) // elem = {idx:scalar}
        indexes.insert(indexes.end(), elem.first);

    auto [normSq, prNS] = DLPROOFS::proveNormSquared("blah", xes);
    std::cout << "norm: "<<verifyNormSquared(indexes,normSq,prNS)<<std::endl;

    return 0;
}
#endif


//==================================================================================================
/*!
  @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_FUNCTION_REVERSEBITS_HPP_INCLUDED
#define BOOST_SIMD_FUNCTION_REVERSEBITS_HPP_INCLUDED

#if defined(DOXYGEN_ONLY)
namespace boost { namespace simd
{

  /*!
    @ingroup group-bitwise
    This function object returns the bits in reverse order for each scalar entry.

    @see reverse,  swapbytes

    @par Header <boost/simd/function/reversebits.hpp>

    @par Example:

      @snippet reversebits.cpp reversebits

    @par Possible output:

      @snippet reversebits.txt reversebits

  **/
  IntegerValue reversebits(IntegerValue const& x);
} }
#endif

#include <boost/simd/function/scalar/reversebits.hpp>
#include <boost/simd/function/simd/reversebits.hpp>

#endif


#define BOOST_TEST_MODULE "Test GroupableTimeSeriesSet class"

#include <vector>
#include <boost/test/unit_test.hpp>

#include "GroupableTimeSeriesSet.hpp"
#include "Exception.hpp"
#include "distance/Distance.hpp"
#include "Group.hpp"

using namespace konex;

struct MockDataset
{
  std::string test_10_20_space = "datasets/test/test_10_20_space.txt";
  std::string test_3_10_space = "datasets/test/test_3_10_space.txt";
} data;

BOOST_AUTO_TEST_CASE( groupable_time_series_set_load_space )
{
  GroupableTimeSeriesSet tsSet;
  tsSet.loadData(data.test_10_20_space, 20, 0, " ");

  BOOST_CHECK_EQUAL( tsSet.getItemLength(), 20 );
  BOOST_CHECK_EQUAL( tsSet.getItemCount(), 10 );
  BOOST_CHECK( tsSet.getFilePath() == data.test_10_20_space );
}

BOOST_AUTO_TEST_CASE( groupable_time_series_grouping_and_reset )
{
  GroupableTimeSeriesSet tsSet;
  tsSet.loadData(data.test_3_10_space, 20, 0, " ");
  int groupCnt = tsSet.groupAllLengths("euclidean", 0.5, 1);
  BOOST_CHECK( groupCnt > 2 );
}

BOOST_AUTO_TEST_CASE( groupable_time_series_not_grouped_exception )
{
  GroupableTimeSeriesSet tsSet;
  tsSet.loadData(data.test_3_10_space, 20, 0, " ");

  // dataset not grouped:
  BOOST_CHECK_THROW( tsSet.getBestMatch(tsSet.getTimeSeries(0)), KOnexException );
}

BOOST_AUTO_TEST_CASE( basic_get_best_match )
{
  GroupableTimeSeriesSet tsSet;
  tsSet.loadData(data.test_3_10_space, 20, 0, " ");
  tsSet.groupAllLengths("euclidean", 0.5, 1);
  candidate_time_series_t best = tsSet.getBestMatch(tsSet.getTimeSeries(0));
  BOOST_TEST( best.dist == 0.0 );
}


/*!
 * *****************************************************************************
 *   \file genIntegratedFlowMap.cpp
 *   \author moennen
 *   \brief
 *   \date 2018-10-01
 *   *****************************************************************************/

#include "utils/imgFileLst.h"
#include "utils/cv_utils.h"

#include <libopticalFlow/oclVarOpticalFlow.h>

#include <glm/glm.hpp>

#include <boost/filesystem.hpp>

#include <iostream>

using namespace std;
using namespace cv;
using namespace glm;
using namespace boost;

namespace
{

//------------------------------------------------------------------------------
//
void resizeToMax( Mat& img, const uvec2 sampleSz )
{
   uvec2 imgSz( img.cols, img.rows );
   // random rescale
   const float ds = std::max( (float)sampleSz.y / imgSz.y, (float)sampleSz.x / imgSz.x );
   resize( img, img, Size(), ds, ds, CV_INTER_AREA );
}

//------------------------------------------------------------------------------
//
Mat flowToWrite( const Mat& flow, const Mat&flowErr )
{
  vector<Mat> splitFlow(3); split(flow, &splitFlow[0]);
  splitFlow[2] = flowErr;
  Mat out(flow.rows, flow.cols, CV_32FC3);
  merge(&splitFlow[0],3,out);
  return out;
}

//------------------------------------------------------------------------------
//
Mat flowToImg( const Mat& flow, const bool leg = false )
{
   Mat flow_split[2];
   Mat t_split[3];
   Mat bgr;
   split( flow, flow_split );
   if ( leg )
   {
      t_split[0] = Mat::zeros( flow_split[0].size(), flow_split[0].type() );
      t_split[1] = flow_split[1];
      t_split[2] = flow_split[0];
      merge( t_split, 3, bgr );
   }
   else
   {
      Mat hsv;
      Mat magnitude, angle;
      cartToPolar( flow_split[0], flow_split[1], magnitude, angle, true );
      normalize( magnitude, magnitude, 0, 1, NORM_MINMAX );
      t_split[0] = angle;  // already in degrees - no normalization needed
      t_split[1] = Mat::ones( angle.size(), angle.type() );
      t_split[2] = magnitude;
      merge( t_split, 3, hsv );
      cvtColor( hsv, bgr, COLOR_HSV2BGR );
   }
   return bgr;
}

//------------------------------------------------------------------------------
//
Mat flowToErr( const Mat& flow, const Mat& imgFrom, const Mat& imgTo )
{
   Mat lk( flow.rows, flow.cols, CV_32FC1 );

#pragma omp parallel for
   for ( size_t y = 0; y < flow.rows; y++ )
   {
      float* lk_data = lk.ptr<float>( y );
      const vec2* flow_data = flow.ptr<vec2>( y );
      const vec3* from_data = imgFrom.ptr<vec3>( y );

      for ( size_t x = 0; x < flow.cols; x++ )
      {
         const vec2 f = flow_data[x];
         const vec3 from = from_data[x];
         const vec3 to = cv_utils::imsample32FC3<vec3>( imgTo, vec2( x, y ) + f );
         const vec3 diff = ( from - to );
         lk_data[x] = dot(diff,diff);
      }
   }

   return lk;
}

//------------------------------------------------------------------------------
//
Mat integrateFlow( const Mat& flow0To1, const Mat& flow1ToN )
{
   Mat flow0ToN(flow0To1.rows, flow0To1.cols, CV_32FC2);

#pragma omp parallel for
   for ( size_t y = 0; y < flow0To1.rows; y++ )
   {
      const vec2* flow0To1_data = flow0To1.ptr<vec2>( y );
      vec2* flow0ToN_data = flow0ToN.ptr<vec2>( y );
      for ( size_t x = 0; x < flow0To1.cols; x++ )
      {
         const vec2 uv0To1 = flow0To1_data[x];
         const vec2 uv1ToN = cv_utils::imsample32F<vec2>( flow1ToN, vec2(x,y) + uv0To1 );
         flow0ToN_data[x] = uv0To1 + uv1ToN;
      }
   }

   return flow0ToN;
}

}

//------------------------------------------------------------------------------
//
//------------------------------------------------------------------------------

const string keys =
    "{help h usage ? |         | print this message   }"
    "{@imgFileLst    |         | images list   }"
    "{@imgRootDir    |         | images root dir   }"
    "{@imgOutDir     |         | images output dir   }"
    "{seqLength      |7        | number of images to integrated per sequence}"
    "{outWidth       |640      | output width }"
    "{outHeight      |380      | output height }"
    "{maxFlowMSErr   |0.5      | maximum tolerance for the flow mean square error}"
    "{backward       |         |    }"
    "{show           |         |    }"
    "{nowrite        |         |    }";

int main( int argc, char* argv[] )
{
   CommandLineParser parser( argc, argv, keys );
   if ( parser.has( "help" ) )
   {
      parser.printMessage();
      return ( 0 );
   }

   const unsigned uSeqLenght = parser.get<unsigned>( "seqLength" );
   const uvec2 minSz =
       uvec2( parser.get<unsigned>( "outWidth" ), parser.get<unsigned>( "outHeight" ) );

   const bool toLinear = false;
   const bool doBackward = parser.get<bool>( "backward" );
   const bool doShow = parser.get<bool>( "show" );
   const bool doWrite = !parser.get<bool>( "nowrite" );

   const float fMaxFlowMSErr = parser.get<float>( "maxFlowMSErr" );

   const filesystem::path outRootPath( parser.get<string>( "@imgOutDir" ) );

   // Create the list of image triplets
   ImgNFileLst imgLst(
       uSeqLenght,
       parser.get<string>( "@imgFileLst" ).c_str(),
       parser.get<string>( "@imgRootDir" ).c_str() );
   if ( imgLst.size() == 0 )
   {
      cerr << "Invalid dataset : " << parser.get<string>( "@imgFileLst" ) << endl;
      return -1;
   }

   // Create the optical flow estimator
   OclVarOpticalFlow::params_t ofParams = OclVarOpticalFlow::getDefaultParams();
   ofParams.nonLinearIter = 9;
   ofParams.robustIter = 5;
   ofParams.solverIter = 5;
   ofParams.lambda = 0.17;
   //ofParams.gamma = 150;

   OclVarOpticalFlow ofEstimator( minSz.x, minSz.y, false, ofParams );

   for ( size_t i = 0; i < imgLst.size(); ++i )
   {
      const string outBasename = filesystem::path( imgLst( i, doBackward ? uSeqLenght-1 : 0 ) ).stem().string();

      Mat matRefImg = cv_utils::imread32FC3( imgLst.filePath( i, doBackward ? uSeqLenght-1 : 0 ), toLinear, true );
      if (matRefImg.empty())
      {
        cerr << "Cannot load image : " << imgLst.filePath( i, doBackward ? uSeqLenght-1 : 0 ) << endl;
        continue;
      }

      Mat matCurrImg = matRefImg.clone();
      GaussianBlur( matCurrImg, matCurrImg, Size( 3, 3 ), 0.113 );
      Mat matFlowCurrToRef = Mat(matCurrImg.rows, matCurrImg.cols, CV_32FC2, 0.0f);
      Mat matFlowCurrToPrev( matFlowCurrToRef.rows, matFlowCurrToRef.cols, CV_32FC2 );

      ofEstimator.setImgSize( matFlowCurrToRef.cols, matFlowCurrToRef.rows );

      bool success = true;
      for ( size_t j = 1; j < uSeqLenght; ++j )
      {
        Mat matPrevImg = matCurrImg;
        Mat matCurrImg = cv_utils::imread32FC3( imgLst.filePath( i, doBackward ? uSeqLenght-(j+1) : j ), toLinear, true );
        if (matCurrImg.empty())
        {
          success = false;
          cerr << "Cannot load image : " << imgLst.filePath( i, doBackward ? uSeqLenght-(j+1) : j ) << endl;
          break;
        }
        GaussianBlur( matCurrImg, matCurrImg, Size( 3, 3 ), 0.113 );

        ofEstimator.compute(
          reinterpret_cast<const float*>( matPrevImg.ptr() ),
          reinterpret_cast<const float*>( matCurrImg.ptr() ),
          matFlowCurrToPrev.cols,
          matFlowCurrToPrev.rows,
          reinterpret_cast<float*>( matFlowCurrToPrev.ptr() ) );

        matFlowCurrToRef = integrateFlow(matFlowCurrToPrev, matFlowCurrToRef);
      }

      if (!success) continue;

      // reread the current image to get it clean (with no preprocessing)
      matCurrImg = cv_utils::imread32FC3( imgLst.filePath( i, doBackward ? 0 : uSeqLenght-1 ), toLinear, true  );

      // check the error
      Mat matFlowErr = flowToErr( matFlowCurrToRef, matCurrImg, matRefImg);
      if (cv::mean(matFlowErr)[0] > fMaxFlowMSErr ) continue;

      cerr << "MSERR : " << cv::mean(matFlowErr)[0] << endl;

      const filesystem::path fImgRef( outBasename + "_ref.png" );
      const filesystem::path fImgCurr( outBasename + "_curr.png" );
      const filesystem::path fFlow( outBasename + string( "_flow" ) + ".exr" );

      cvtColor(matRefImg, matRefImg, COLOR_BGR2RGB );
      cvtColor(matCurrImg, matCurrImg, COLOR_BGR2RGB );

      if ( doWrite )
      {
         imwrite( filesystem::path( outRootPath / fImgRef ).string().c_str(), matRefImg * 255.0 );
         imwrite( filesystem::path( outRootPath / fImgCurr ).string().c_str(), matCurrImg * 255.0 );
         imwrite( filesystem::path( outRootPath / fFlow ).string().c_str(), flowToWrite(matFlowCurrToRef,matFlowErr) );
      }

      cout << fImgRef.string() << " " << fImgCurr.string() << " " << fFlow.string() << " " << uSeqLenght << endl;

      // display
      if ( doShow )
      {
         imshow( "Reference", matRefImg );
         imshow( "Current", matCurrImg );
         imshow( "Flow", flowToImg(matFlowCurrToRef) );
         waitKey( cv::mean(matFlowErr)[0] > 0.15 ?  0 : 10 );
      }
   }

   return ( 0 );
}


#include <iostream>
#include <map>
#include <cassert>

#include <boost/static_assert.hpp>

#include "prob.h"
#include "sparse_vector.h"
#include "fast_sparse_vector.h"
#include "stringlib.h"

using namespace std;

template <typename T>
void MPrint(const T& x) {
  typename T::const_iterator it = x.begin();
  for (; it != x.end(); ++it) {
    cerr << it->first << ":" << it->second << " ";
  }
  cerr << endl;
}

template <typename T>
void test_unique() {
  T x;
  x.set_value(100, 1.0);
  x.set_value(100, 2.0);
  MPrint(x);
  assert(x.size() == 1);
  assert(x.value(100) == 2.0);
}

void test_logv() {
  FastSparseVector<prob_t> x;
  cerr << "FSV<prob_t> = " << sizeof(FastSparseVector<prob_t>) << endl;
  x.set_value(999, prob_t(0.5));
  x.set_value(0, prob_t());
  x.set_value(1, prob_t(1));
  MPrint(x);
  x += x;
  MPrint(x);
  x -= x;
  MPrint(x);
  FastSparseVector<prob_t> y;
  y = x;
  for (int i = 1; i < 10; ++i) {
    x.set_value(i, prob_t(i*1.3));
    y.set_value(i*2, prob_t(i*1.4));
  }
  swap(x,y);
  MPrint(y);
  MPrint(x);
  x = y;
  y = y;
  x = x;
  MPrint(y);
}

int main() {
  cerr << sizeof(prob_t) << " " << sizeof(LogVal<float>) << endl;
  cerr << " sizeof(FSV<float>) = " << sizeof(FastSparseVector<float>) << endl;
  cerr << "sizeof(FSV<double>) = " << sizeof(FastSparseVector<double>) << endl;
  test_unique<FastSparseVector<float> >();
  test_logv();
//  sranddev();
  int c = 0;
  FastSparseVector<float> p;
  for (int i = 0; i < 1000000; ++i) {
    FastSparseVector<float> x;
    for (int j = 0; j < 10; ++j) {
      const int k = rand() % 200;
      const float v = 1000.0f / rand();
      x.set_value(k,v);
    }
    if (x.value(50)) { c++; }
    p = x;
    p += p;
    FastSparseVector<float> y = p + p;
    y *= 2;
    y -= y;
  }
  cerr << "Counted " << c << " times\n";

  cerr << md5("this is a test") << endl;
  cerr << md5("some other ||| string is") << endl;
  map<string,string> x; x["id"] = "12"; x["grammar"] = "/path/to/grammar.gz";
  cerr << SGMLOpenSegTag(x) << endl;
  return 0;
}


// Copyright András Vukics 2006–2020. Distributed under the Boost Software License, Version 1.0. (See accompanying file LICENSE.txt)
#include "ModeCorrelations.h"

#include "LazyDensityOperator.h"
#include "Structure.h"

#include <boost/range/algorithm/copy.hpp>

using namespace boost;

ModeCorrelations::ModeCorrelations()
  : EA_Base("ModeCorrelations",{"<a number operator>","VAR(a number operator)","real(<a>)","imag(\")","<X^2>-<X>^2","<Y^2>-<Y>^2","<(XY+YX)/2>-<X><Y>",
                                "<b number operator>","VAR(b number operator)","real(<b>)","imag(\")","<Q^2>-<Q>^2","<P^2>-<P>^2","<(QP+PQ)/2>-<Q><P>",
                                "<XQ>-<X><Q>","<XP>-<X><P>","<YQ>-<Y><Q>","<YP>-<Y><P>"}
           ),
    averagedMode_()
{
}

namespace {

#include "details_BinaryHelper.h"

}

const ModeCorrelations::Averages
ModeCorrelations::average_v(structure::NoTime, const LazyDensityOperator& matrix) const
{
  using quantumdata::partialTrace;

  auto averages(initializedAverages());

  {
    auto fun=[&](const auto& psiS){return averagedMode_.average(0,psiS);};

    const Averages
      a0{partialTrace<V0>(matrix,fun)},
      a1{partialTrace<V1>(matrix,fun)};

    copy(a1,copy(a0,averages.begin()));
  }

  for (int n=0; n<int(matrix.getDimension(0)); n++) for (int m=1; m<int(matrix.getDimension(1)); m++) {
      if(n<int(matrix.getDimension(0))-1) {
        dcomp temp=sqrt(m*(n+1))*matrix(n,m)(n+1,m-1);
        averages(14)+=real(temp);
        averages(15)+=imag(temp);
      }
      if(n>0) {
        dcomp temp=sqrt(m*n)*matrix(n,m)(n-1,m-1);
        averages(16)+=real(temp);
        averages(17)+=imag(temp);
      }
    }

  return averages;

}

void
ModeCorrelations::process_v(Averages& averages) const
{
  {
    Averages ranged(averages(blitz::Range(0, 6)));
    averagedMode_.process(ranged);
  }
  {
    Averages ranged(averages(blitz::Range(7,13)));
    averagedMode_.process(ranged);
  }

  double
    xAvr=sqrt(2)*averages( 2),
    yAvr=sqrt(2)*averages( 3),
    qAvr=sqrt(2)*averages( 9),
    pAvr=sqrt(2)*averages(10);

  double
    xq= averages(14)+averages(16)-xAvr*qAvr,
    xp= averages(15)+averages(17)-xAvr*pAvr,
    yq=-averages(15)+averages(17)-yAvr*qAvr,
    yp= averages(14)-averages(16)-yAvr*pAvr;

  averages(14)=xq;
  averages(15)=xp;
  averages(16)=yq;
  averages(17)=yp;

}


/*=============================================================================
Copyright (c) 2016 Paul W. Bible

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)
==============================================================================*/
#ifndef TERM_DEPTH_MAP
#define TERM_DEPTH_MAP

#include <ggtk/GoGraph.hpp>
#include <ggtk/AnnotationData.hpp>

#include <boost/graph/breadth_first_search.hpp>
#include <boost/graph/reverse_graph.hpp>

#include <boost/accumulators/accumulators.hpp>
#include <boost/accumulators/statistics.hpp>
#include <boost/accumulators/statistics/min.hpp>
#include <boost/accumulators/statistics/max.hpp>

/*! \class TermDepthMap
	\brief A class to calculate the depth of a GO term in the ontology.

	This class provides a map that returns the depth of a GO term. This method
	 is used in graph and edge based similarity methods to calculate a node's depth
*/
class TermDepthMap{
public:
	//! A parameterized constructor
	/*!
		This constructor takes pointers to GoGraph and AnnotationData objects.
		  Only the parameterized construtor is allowed to ensure these objects are
		  created with valid parameters.
	*/
	inline TermDepthMap(GoGraph* graph){

		//Initialize an annotation list the size of verticies in go, each value is 0
		//_depths = std::vector<std::size_t>(graph->getNumVertices(),0);
		_depths = std::vector<double>(graph->getNumVertices(), 0);

		//get the (first) root of the ontology.
		//GoGraph::GoVertex root = graph->getRoot();
		//TESTING
		//std::cout << root << std::endl;
		//std::cout << graph->getTermStringIdByIndex(root) << std::endl;

		//get the boost graph from the GoGraph object. Must be done to utilize boost algorithms
		GoGraph::Graph* g = graph->getGraph();

		//wrap _depth with a vertex map
		GoGraph::VertexIndexMap vMap = boost::get(boost::vertex_index,*g);

		/* // Temparary fix until I can get SWIG to recognize std::size_t
		boost::iterator_property_map< std::vector<std::size_t>::iterator,
			                          GoGraph::VertexIndexMap >
									  d_map(_depths.begin(), vMap);
		*/
		boost::iterator_property_map< std::vector<double>::iterator,
			GoGraph::VertexIndexMap >
			d_map(_depths.begin(), vMap);

		//call the boost depth first search using our custom visitor
		// revering the graph is necessary otherwise the root vertex would have no edges.
		//boost::depth_first_search(boost::make_reverse_graph(*g),boost::visitor(vis).root_vertex(root));
		GoGraph::GoVertex bpRoot = graph->getVertexByName(GO::getRootTermBP());
		GoGraph::GoVertex mfRoot = graph->getVertexByName(GO::getRootTermMF());
		GoGraph::GoVertex ccRoot = graph->getVertexByName(GO::getRootTermCC());

		//Start at bproot, record depths
		// must reverse graph due to edge relationship direction
		boost::breadth_first_search(boost::make_reverse_graph(*g), bpRoot,
			boost::visitor(boost::make_bfs_visitor(
			  boost::record_distances(d_map, boost::on_tree_edge())
			  )));

		//Start at bproot, record depths
		// must reverse graph due to edge relationship direction
		boost::breadth_first_search(boost::make_reverse_graph(*g), mfRoot,
			boost::visitor(boost::make_bfs_visitor(
			  boost::record_distances(d_map, boost::on_tree_edge())
			  )));

		//Start at bproot, record depths
		// must reverse graph due to edge relationship direction
		boost::breadth_first_search(boost::make_reverse_graph(*g), ccRoot,
			boost::visitor(boost::make_bfs_visitor(
			  boost::record_distances(d_map, boost::on_tree_edge())
			  )));

		//initialize the term to index map
		_nameToIndex = boost::unordered_map<std::string,std::size_t>(boost::num_vertices(*g));

		// Vertex Iterators
		GoGraph::GoVertexIterator vi,vend;
		for(boost::tie(vi,vend) = boost::vertices(*g); vi != vend; ++vi){
			GoGraph::GoVertex v = *vi;
			_nameToIndex[graph->getTermStringIdByIndex(vMap[v])] = vMap[v];
			//std::cout << vMap[v] << " " << _depths.at(vMap[v]) << " " << graph->getTermNameByIndex(vMap[v]) << std::endl;
			//std::cin.get();
		}
	}//end constructor

	//! A default constructor
	/*!
		This constructor initialized the storage structures. Should not be used.
	*/
	inline TermDepthMap(){
		//creat empty containers
		_nameToIndex = boost::unordered_map<std::string,std::size_t>();
		//_depths = std::vector<std::size_t>();
		_depths = std::vector<double>();
	}

	//! A default desctructor
	/*!
		This desctructor clears the containters
	*/
	inline ~TermDepthMap(){
		//empty containers
		_nameToIndex.clear();
		_depths.clear();
	}

	//! Accessor for probablities vector
	/*!
		Get the vector of values
	*/
	inline std::vector<double> getValues(){
	//inline std::vector<std::size_t> getValues(){
		return _depths;
	}

	//! Function to return all the keys in the map
	/*!
		Returns all valid keys in the map.
	*/
	inline std::vector<std::string> getKeys(){
		std::vector<std::string> keys;
		boost::unordered_map<std::string, std::size_t>::iterator it;
		for (it = _nameToIndex.begin(); it != _nameToIndex.end(); ++it){
			keys.push_back(it->first);
		}
		return keys;
	}

	//! Method to test if the id exists in the map
	/*!
		Return true the id is found, false if not
	*/
	inline bool hasTerm(std::string testTerm){
		return _nameToIndex.find(testTerm) != _nameToIndex.end();
	}

	//! Overloaded [] bracket operator to mimic Map
	/*!
		This defines a bracket operator to access the data inside of the map.
		  This is done to mimic the behavior of the map class
	*/
	inline double operator[](std::string termId){
	//inline size_t& operator[](std::string termId){
		//get index
		std::size_t index = _nameToIndex[termId];
		//return the depth
		return _depths.at(index);
	}

	//! Mapping function to return the value mapped by key
	/*!
		Get the value mapped by the given key. A specified function for the [] operator
	*/
	inline double getValue(std::string termId){
		//get index
		std::size_t index = _nameToIndex[termId];
		//return the depth
		return _depths.at(index);
	}

protected:
	//! A private map that returns the index of a term.
	/*!
		This map takes string term ids and returns the index for annotation count access.
	*/
	boost::unordered_map<std::string,std::size_t> _nameToIndex;

	//! A private list of term depths
	/*!
		This vector of doubles holds the depth for each term
	*/
	std::vector<double> _depths;		//temparary fix until I can get SWIG to recognize std::size_t
	//std::vector<std::size_t> _depths;
};
#endif


#include <boost/uuid/uuid.hpp>
#include <boost/uuid/uuid_generators.hpp>
#include <boost/uuid/uuid_io.hpp>

#include <boost/iostreams/copy.hpp>
#include <boost/iostreams/filter/gzip.hpp>
#include <boost/iostreams/filtering_streambuf.hpp>

#include <fstream>
#include <iostream>
#include <stdio.h>
#include <sys/stat.h>

#include "SeedHist2D.hpp"

#include <TCanvas.h>
#include <TFile.h>
#include <TGaxis.h>
#include <TThread.h>

#include <boost/lexical_cast.hpp>
#include <boost/uuid/uuid.hpp>            // uuid class
#include <boost/uuid/uuid_generators.hpp> // generators
#include <boost/uuid/uuid_io.hpp>         // streaming operators etc.

SeedH2::SeedH2(bool _invert) {
  title = "_x <=> _y";
  _hist = NULL;
  init_done = false;
  has_data = false;
  invert = _invert;
  data.reserve(BURN_IN);
  x_min = -FLT_MAX;
  x_max = FLT_MAX;
  y_min = -FLT_MAX;
  y_max = FLT_MAX;
  printf("created a SeedH2::%s\n", title.c_str());
}
SeedH2::~SeedH2() {
  printf("%s x_min: %f x_max: %f; y_min: %f, y_max: %f\n", title.c_str(), x_min,
         x_max, y_min, y_max);

  if (_hist != NULL)
    delete _hist;
};

void SeedH2::update(float _x, float _y) {
  has_data = true;

  if (!init_done)
    data.push_back(std::make_tuple(_x, _y));
  else {
    // update a histogram
    _hist->Fill(_x, _y);
  }

  if (data.size() == BURN_IN)
    this->flush();
}

void SeedH2::flush() {
  if (data.size() == 0)
    return;

  double meanX = 0.0;
  double meanY = 0.0;

  // gather the statistics
  for (auto &x : data) {
    meanX += std::get<0>(x);
    meanY += std::get<1>(x);
  }

  meanX /= double(data.size());
  meanY /= double(data.size());

  double stdX = 0.0;
  double stdY = 0.0;

  for (auto &x : data) {
    double _x, _y;
    std::tie(_x, _y) = x;

    stdX += (_x - meanX) * (_x - meanX);
    stdY += (_y - meanY) * (_y - meanY);
  }

  stdX = sqrt(stdX / double(data.size()));
  stdY = sqrt(stdY / double(data.size()));

  x_min = meanX - 3.0 * stdX;
  x_max = meanX + 3.0 * stdX;

  y_min = meanY - 3.0 * stdY;
  y_max = meanY + 3.0 * stdY;

  printf("[%s] x_min: %f x_max: %f y_min: %f y_max: %f\n", title.c_str(), x_min,
         x_max, y_min, y_max);

  /*boost::uuids::uuid uuid = boost::uuids::random_generator()();
  std::string name = boost::lexical_cast<std::string>(uuid);*/

  _hist = new TH2D(name.c_str(), title.c_str(), NBINS2, x_min, x_max, NBINS2,
                   y_min, y_max);
  _hist->SetCanExtend(TH1::kAllAxes);
  _hist->SetStats(false);
  _hist->GetXaxis()->SetTitle(x_title.c_str());
  _hist->GetYaxis()->SetTitle(y_title.c_str());

  for (auto &x : data) {
    double _x, _y;
    std::tie(_x, _y) = x;

    _hist->Fill(_x, _y);
  }

  data.clear();
  init_done = true;
}

void SeedH2::ReverseYAxis(TH1 *h) {
  // Remove the current axis
  h->GetYaxis()->SetLabelOffset(999);
  h->GetYaxis()->SetTickLength(0);
  // Redraw the new axis
  gPad->Update();
  TGaxis *newaxis =
      new TGaxis(gPad->GetUxmin(), gPad->GetUymax(), gPad->GetUxmin() - 0.001,
                 gPad->GetUymin(), h->GetYaxis()->GetXmin(),
                 h->GetYaxis()->GetXmax(), 510, "+");
  newaxis->SetLabelOffset(-0.03);
  newaxis->Draw();
}

void SeedH2::ReverseYData(TH2 *h) {
  Int_t nx = h->GetNbinsX();
  Int_t ny = h->GetNbinsY();

  for (Int_t i = 0; i < nx; i++) {
    for (Int_t j = 0; j < ny / 2; j++) {
      Int_t a = h->GetBin(i, j);
      Int_t b = h->GetBin(i, ny - 1 - j);

      auto tmp = h->GetBinContent(a);
      auto tmp2 = h->GetBinContent(b);

      h->SetBinContent(a, tmp2);
      h->SetBinContent(b, tmp);
    }
  }

  h->ComputeIntegral();
}

void SeedH2::export_root(std::string uuid, std::string docs_root,
                         std::string type) {
  if (!has_data)
    return;

  std::string filename =
      docs_root + "/gaiawebql/DATA/" + uuid + "/" + type + ".root";

  std::cout << "saving " << title << " into " << filename << std::endl;

  // mkdir DATA/<uuid>.tmp
  std::string tmp = docs_root + "/gaiawebql/DATA/" + uuid + ".tmp";

  if (mkdir(tmp.c_str(), 0777) != 0) {
    // return only in case of errors other than "directory already exists"
    if (errno != EEXIST) {
      perror(title.c_str());
      return;
    }
  }

  filename = tmp + "/" + type + ".root";

  TThread::Lock();

  TFile outputFile(filename.c_str(), "RECREATE");
  outputFile.SetCompressionLevel(
      ROOT::RCompressionSetting::ELevel::kDefaultZLIB);
  _hist->Write();
  outputFile.Close();

  TThread::UnLock();
}

void SeedH2::save(std::string uuid, std::string docs_root, std::string type) {
  std::string filename =
      docs_root + "/gaiawebql/DATA/" + uuid + "/" + type + ".json";

  std::cout << "saving " << title << " into " << filename << std::endl;

  // mkdir DATA/<uuid>.tmp
  std::string tmp = docs_root + "/gaiawebql/DATA/" + uuid + ".tmp";

  if (mkdir(tmp.c_str(), 0777) != 0) {
    // return only in case of errors other than "directory already exists"
    if (errno != EEXIST) {
      perror(title.c_str());
      return;
    }
  }

  TAxis *axis;

  axis = _hist->GetXaxis();
  x_min = axis->GetBinCenter(0);
  x_max = axis->GetBinCenter(NBINS2 - 1);

  axis = _hist->GetYaxis();
  y_min = axis->GetBinCenter(0);
  y_max = axis->GetBinCenter(NBINS2 - 1);

  // write the bin data
  {
    filename = tmp + "/" + type + ".json";
    std::ofstream json(filename);

    json << "{\n";
    json << "\t\"NBINS2\" : " << NBINS2 << ",\n";
    json << "\t\"xmin\" : " << x_min << ",\n";
    json << "\t\"xmax\" : " << x_max << ",\n";
    json << "\t\"ymin\" : " << y_min << ",\n";
    json << "\t\"ymax\" : " << y_max << ",\n";

    json << "\t\"bins\" : [";

    for (int j = 0; j < NBINS2; j++) {
      json << "[";
      for (int i = 0; i < NBINS2; i++) {
        Int_t a = _hist->GetBin(i, j);
        auto value = _hist->GetBinContent(a);
        json << ((value >= 1.0) ? std::to_string(log10(value))
                                : "null"); // log-z
        // json << value;

        if (i != NBINS2 - 1)
          json << ",";
      }
      json << "]";

      if (j != NBINS2 - 1)
        json << ",";
    }

    json << "]}\n";
    json.close();

    // gz-compress with Boost
    std::ifstream inStream(filename, std::ios_base::in);
    std::ofstream outStream(filename + ".gz", std::ios_base::binary);
    boost::iostreams::filtering_streambuf<boost::iostreams::input> in;
    in.push(boost::iostreams::gzip_compressor());
    in.push(inStream);
    boost::iostreams::copy(in, outStream);
  }
};


/*
 Copyright (C) 2019 Quaternion Risk Management Ltd
 All rights reserved.

 This file is part of ORE, a free-software/open-source library
 for transparent pricing and risk analysis - http://opensourcerisk.org

 ORE is free software: you can redistribute it and/or modify it
 under the terms of the Modified BSD License.  You should have received a
 copy of the license along with this program.
 The license is also available online at <http://opensourcerisk.org>

 This program is distributed on the basis that it will form a useful
 contribution to risk analytics and model standardisation, but WITHOUT
 ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 FITNESS FOR A PARTICULAR PURPOSE. See the license for more details.
*/

#include <boost/assign.hpp>
// clang-format off
#include <boost/test/unit_test.hpp>
#include <boost/test/data/test_case.hpp>
// clang-format on
#include <boost/variant.hpp>

#include <test/capfloormarketdata.hpp>
#include <test/toplevelfixture.hpp>
#include <test/yieldcurvemarketdata.hpp>

#include <ql/instruments/makecapfloor.hpp>
#include <ql/pricingengines/capfloor/bacheliercapfloorengine.hpp>
#include <ql/pricingengines/capfloor/blackcapfloorengine.hpp>
#include <ql/time/date.hpp>
#include <qle/math/flatextrapolation.hpp>
#include <qle/termstructures/capfloorhelper.hpp>
#include <qle/termstructures/piecewiseoptionletcurve.hpp>

using namespace QuantExt;
using namespace QuantLib;
using namespace boost::unit_test_framework;

namespace bdata = boost::unit_test::data;
typedef QuantLib::BootstrapHelper<QuantLib::OptionletVolatilityStructure> helper;
using boost::assign::list_of;
using boost::assign::map_list_of;
using std::boolalpha;
using std::fixed;
using std::map;
using std::ostream;
using std::pair;
using std::setprecision;
using std::string;
using std::vector;

namespace {

// Variables to be used in the test
struct CommonVars : public qle::test::TopLevelFixture {

    // Constructor
    CommonVars()
        : referenceDate(5, Feb, 2016), settlementDays(0), calendar(TARGET()), bdc(Following),
          dayCounter(Actual365Fixed()), accuracy(1.0e-12), globalAccuracy(1.0e-10), tolerance(1.0e-10) {

        // Reference date
        Settings::instance().evaluationDate() = referenceDate;

        // Set cap floor ibor index to EUR-EURIBOR-6M and attach a forwarding curve
        iborIndex = boost::make_shared<Euribor6M>(testYieldCurves.forward6M);
    }

    // Valuation date for the test
    Date referenceDate;

    // Variables used in the optionlet volatility structure creation
    Natural settlementDays;
    Calendar calendar;
    BusinessDayConvention bdc;
    DayCounter dayCounter;

    // Accuracy for optionlet stripping
    Real accuracy;

    // Global accuracy for optionlet stripping
    Real globalAccuracy;

    // Test tolerance for comparing the NPVs
    Real tolerance;

    // Cap floor ibor index
    boost::shared_ptr<IborIndex> iborIndex;

    // EUR discount curve test data from file test/yieldcurvemarketdata.hpp
    YieldCurveEUR testYieldCurves;
};

// Struct to hold a cap floor volatility column and some associated meta data
struct VolatilityColumn {
    Real strike;
    vector<Period> tenors;
    vector<Real> volatilities;
    VolatilityType type;
    Real displacement;
};

// Needed for data driven test case below as it writes out the VolatilityColumn
ostream& operator<<(ostream& os, const VolatilityColumn& vc) {
    return os << "Column with strike: " << vc.strike << ", volatility type: " << vc.type
              << ", shift: " << vc.displacement;
}

// From the EUR cap floor test volatility data in file test/capfloormarketdata.hpp, create a vector of
// VolatilityColumns which will be the data in the data driven test below
vector<VolatilityColumn> generateVolatilityColumns() {

    // Will hold the generated volatility columns
    vector<VolatilityColumn> volatilityColumns;

    // EUR cap floor test volatility data from file test/capfloormarketdata.hpp
    CapFloorVolatilityEUR testVols;

    VolatilityColumn volatilityColumn;
    vector<Real> volatilities(testVols.tenors.size());
    volatilityColumn.tenors = testVols.tenors;

    // The normal volatilities
    volatilityColumn.type = Normal;
    volatilityColumn.displacement = 0.0;
    for (Size j = 0; j < testVols.strikes.size(); j++) {
        for (Size i = 0; i < testVols.tenors.size(); i++) {
            volatilities[i] = testVols.nVols[i][j];
        }
        volatilityColumn.strike = testVols.strikes[j];
        volatilityColumn.volatilities = volatilities;
        volatilityColumns.push_back(volatilityColumn);
    }

    // The shifted lognormal volatilities with shift 1
    volatilityColumn.type = ShiftedLognormal;
    volatilityColumn.displacement = testVols.shift_1;
    for (Size j = 0; j < testVols.strikes.size(); j++) {
        for (Size i = 0; i < testVols.tenors.size(); i++) {
            volatilities[i] = testVols.slnVols_1[i][j];
        }
        volatilityColumn.strike = testVols.strikes[j];
        volatilityColumn.volatilities = volatilities;
        volatilityColumns.push_back(volatilityColumn);
    }

    // The shifted lognormal volatilities with shift 2
    volatilityColumn.type = ShiftedLognormal;
    volatilityColumn.displacement = testVols.shift_2;
    for (Size j = 0; j < testVols.strikes.size(); j++) {
        for (Size i = 0; i < testVols.tenors.size(); i++) {
            volatilities[i] = testVols.slnVols_2[i][j];
        }
        volatilityColumn.strike = testVols.strikes[j];
        volatilityColumn.volatilities = volatilities;
        volatilityColumns.push_back(volatilityColumn);
    }

    return volatilityColumns;
}

// Cap floor helper types for the data driven test case
vector<CapFloorHelper::Type> helperTypes =
    list_of(CapFloorHelper::Cap)(CapFloorHelper::Floor)(CapFloorHelper::Automatic);

// Quote types for the data driven test case
vector<CapFloorHelper::QuoteType> quoteTypes = list_of(CapFloorHelper::Volatility)(CapFloorHelper::Premium);

// Interpolation types for the data driven test case
typedef boost::variant<Linear, BackwardFlat, QuantExt::LinearFlat, Cubic, QuantExt::CubicFlat> InterpolationType;

vector<InterpolationType> interpolationTypes = list_of(InterpolationType(Linear()))(InterpolationType(BackwardFlat()))(
    InterpolationType(QuantExt::LinearFlat()))(InterpolationType(Cubic()))(InterpolationType(QuantExt::CubicFlat()));

vector<InterpolationType> interpolationTypesCached =
    list_of(InterpolationType(Linear()))(InterpolationType(BackwardFlat()))(InterpolationType(QuantExt::LinearFlat()));

// If the built optionlet structure in the test has a floating or fixed reference date
vector<bool> isMovingValues = list_of(true)(false);

// If the optionlet structure has a flat first period or not
vector<bool> flatFirstPeriodValues = list_of(true)(false);

// So that I can reuse below
string to_string(const InterpolationType& interpolationType) {
    string result;
    switch (interpolationType.which()) {
    case 0:
        result = "Linear";
        break;
    case 1:
        result = "BackwardFlat";
        break;
    case 2:
        result = "LinearFlat";
        break;
    case 3:
        result = "Cubic";
        break;
    case 4:
        result = "CubicFlat";
        break;
    default:
        BOOST_FAIL("Unexpected interpolation type");
    }
    return result;
}

// Cached values for comparison below. We compare against cached values for a given interpolation type and setting of
// the bool flatFirstPeriod. The cached values are keyed on the Size: 2 * interpolationType + flatFirstPeriod.
// clang-format off
vector<Date> cachedDates = list_of(Date(5, Feb, 2016))(Date(7, Feb, 2017))(Date(6, Aug, 2020))(Date(5, Aug, 2022))(Date(7, Aug, 2025))(Date(7, Aug, 2035));
map<Size, vector<Real> > cachedValues = map_list_of
    // Linear, flat first = false
    (0, list_of(0.000000000000)(0.009939243164)(0.008398540935)(0.008216105988)(0.006859464219)(0.006598726907).convert_to_container<vector<Real> >())
    // Linear, flat first = true
    (1, list_of(0.009938000000)(0.009938000000)(0.008399019469)(0.008215852284)(0.006859635836)(0.006598586367).convert_to_container<vector<Real> >())
    // BackwardFlat, flat first = false
    (2, list_of(0.000000000000)(0.009938000000)(0.008799306892)(0.008279139515)(0.007401656494)(0.006715983817).convert_to_container<vector<Real> >())
    // BackwardFlat, flat first = true
    (3, list_of(0.009938000000)(0.009938000000)(0.008799306892)(0.008279139515)(0.007401656494)(0.006715983817).convert_to_container<vector<Real> >())
    // LinearFlat, flat first = false
    (4, list_of(0.000000000000)(0.009939243164)(0.008398540935)(0.008216105988)(0.006859464219)(0.006598726907).convert_to_container<vector<Real> >())
    // LinearFlat, flat first = true
    (5, list_of(0.009938000000)(0.009938000000)(0.008399019469)(0.008215852284)(0.006859635836)(0.006598586367).convert_to_container<vector<Real> >());
// clang-format on

// Cached values, on dates that are not curve nodes, for comparison below.
// We pick a value in the first curve period to check the flatFirstPeriod setting, in the middle of the curve to check
// the interpolation and after the last curve date to check the extrapolation.
// clang-format off
vector<Date> cachedNonNodeDates = list_of(Date(5, Aug, 2016))(Date(6, Aug, 2021))(Date(7, Aug, 2036));
map<Size, vector<Real> > cachedNonNodeValues = map_list_of
    // Linear, flat first = false
    (0, list_of(0.004915603956)(0.008307198335)(0.006572596059).convert_to_container<vector<Real> >())
    // Linear, flat first = true
    (1, list_of(0.009938000000)(0.008307310247)(0.006572424235).convert_to_container<vector<Real> >())
    // BackwardFlat, flat first = false
    (2, list_of(0.009938000000)(0.008279139515)(0.006715983817).convert_to_container<vector<Real> >())
    // BackwardFlat, flat first = true
    (3, list_of(0.009938000000)(0.008279139515)(0.006715983817).convert_to_container<vector<Real> >())
    // LinearFlat, flat first = false
    (4, list_of(0.004915603956)(0.008307198335)(0.006598726907).convert_to_container<vector<Real> >())
    // LinearFlat, flat first = true
    (5, list_of(0.009938000000)(0.008307310247)(0.006598586367).convert_to_container<vector<Real> >());
// clang-format on

} // namespace

// Needed for BOOST_DATA_TEST_CASE below
// https://stackoverflow.com/a/33965517/1771882
namespace boost {
namespace test_tools {
namespace tt_detail {
template <> struct print_log_value<InterpolationType> {
    void operator()(ostream& os, const InterpolationType& interpolationType) { os << to_string(interpolationType); }
};
} // namespace tt_detail
} // namespace test_tools
} // namespace boost

BOOST_FIXTURE_TEST_SUITE(QuantExtTestSuite, qle::test::TopLevelFixture)

BOOST_AUTO_TEST_SUITE(PiecewiseOptionletCurveTests)

BOOST_DATA_TEST_CASE_F(CommonVars, testPiecewiseOptionletStripping,
                       bdata::make(generateVolatilityColumns()) * bdata::make(helperTypes) * bdata::make(quoteTypes) *
                           bdata::make(interpolationTypes) * bdata::make(isMovingValues) *
                           bdata::make(flatFirstPeriodValues),
                       volatilityColumn, helperType, quoteType, interpolationType, isMoving, flatFirstPeriod) {

    BOOST_TEST_MESSAGE("Testing piecewise optionlet stripping of cap floor quotes along a strike column");

    BOOST_TEST_MESSAGE("Test inputs are:");
    BOOST_TEST_MESSAGE("  Cap floor helper type: " << helperType);
    BOOST_TEST_MESSAGE("  Cap floor strike: " << volatilityColumn.strike);
    BOOST_TEST_MESSAGE("  Quote type: " << quoteType);
    if (quoteType == CapFloorHelper::Volatility) {
        BOOST_TEST_MESSAGE("  Quote volatility type: " << volatilityColumn.type);
        BOOST_TEST_MESSAGE("  Quote displacement: " << volatilityColumn.displacement);
    }
    BOOST_TEST_MESSAGE("  Interpolation type: " << to_string(interpolationType));
    BOOST_TEST_MESSAGE("  Floating reference date: " << boolalpha << isMoving);
    BOOST_TEST_MESSAGE("  Flat first period: " << boolalpha << flatFirstPeriod);

    if (quoteType == CapFloorHelper::Premium && helperType == CapFloorHelper::Automatic) {

        // This is a combination that should throw an error when creating the helper
        // Don't care about the value of the premium quote
        Handle<Quote> quote(boost::make_shared<SimpleQuote>(0.01));
        BOOST_REQUIRE_THROW(
            boost::make_shared<CapFloorHelper>(helperType, volatilityColumn.tenors.front(), volatilityColumn.strike,
                                               quote, iborIndex, testYieldCurves.discountEonia, isMoving, Date(),
                                               quoteType, volatilityColumn.type, volatilityColumn.displacement),
            Error);

    } else {

        // Form the cap floor helper instrument for each tenor in the strike column
        vector<boost::shared_ptr<helper> > helpers(volatilityColumn.tenors.size());

        // Store each cap floor instrument in the strike column and its NPV using the flat cap floor volatilities
        vector<boost::shared_ptr<CapFloor> > instruments(volatilityColumn.tenors.size());
        vector<Real> flatNpvs(volatilityColumn.tenors.size());

        BOOST_TEST_MESSAGE("The input values at each tenor are:");
        for (Size i = 0; i < volatilityColumn.tenors.size(); i++) {

            // Get the relevant quote value
            Real volatility = volatilityColumn.volatilities[i];

            // Create the cap floor instrument and store its price using the quoted flat volatility
            CapFloor::Type capFloorType = CapFloor::Cap;
            if (helperType == CapFloorHelper::Floor) {
                capFloorType = CapFloor::Floor;
            }
            instruments[i] = MakeCapFloor(capFloorType, volatilityColumn.tenors[i], iborIndex, volatilityColumn.strike);
            if (volatilityColumn.type == ShiftedLognormal) {
                instruments[i]->setPricingEngine(boost::make_shared<BlackCapFloorEngine>(
                    testYieldCurves.discountEonia, volatility, dayCounter, volatilityColumn.displacement));
            } else {
                instruments[i]->setPricingEngine(
                    boost::make_shared<BachelierCapFloorEngine>(testYieldCurves.discountEonia, volatility, dayCounter));
            }
            flatNpvs[i] = instruments[i]->NPV();

            BOOST_TEST_MESSAGE("  (Cap/Floor, Tenor, Volatility, Flat NPV) = ("
                               << capFloorType << ", " << volatilityColumn.tenors[i] << ", " << fixed
                               << setprecision(13) << volatility << ", " << flatNpvs[i] << ")");

            // Create a volatility or premium quote
            RelinkableHandle<Quote> quote;
            if (quoteType == CapFloorHelper::Volatility) {
                quote.linkTo(boost::make_shared<SimpleQuote>(volatility));
            } else {
                quote.linkTo(boost::make_shared<SimpleQuote>(flatNpvs[i]));
            }

            // Create the helper instrument
            helpers[i] =
                boost::make_shared<CapFloorHelper>(helperType, volatilityColumn.tenors[i], volatilityColumn.strike,
                                                   quote, iborIndex, testYieldCurves.discountEonia, isMoving, Date(),
                                                   quoteType, volatilityColumn.type, volatilityColumn.displacement);
        }

        // Create the piecewise optionlet curve, with the given interpolation type, and fail if it is not created
        VolatilityType curveVolatilityType = Normal;
        Real curveDisplacement = 0.0;
        boost::shared_ptr<OptionletVolatilityStructure> ovCurve;
        switch (interpolationType.which()) {
        case 0:
            if (isMoving) {
                BOOST_TEST_MESSAGE("Using Linear interpolation with a moving reference date");
                BOOST_REQUIRE_NO_THROW(ovCurve = boost::make_shared<PiecewiseOptionletCurve<Linear> >(
                                           settlementDays, helpers, calendar, bdc, dayCounter, curveVolatilityType,
                                           curveDisplacement, flatFirstPeriod));
            } else {
                BOOST_TEST_MESSAGE("Using Linear interpolation with a fixed reference date");
                BOOST_REQUIRE_NO_THROW(ovCurve = boost::make_shared<PiecewiseOptionletCurve<Linear> >(
                                           referenceDate, helpers, calendar, bdc, dayCounter, curveVolatilityType,
                                           curveDisplacement, flatFirstPeriod));
            }
            break;
        case 1:
            if (isMoving) {
                BOOST_TEST_MESSAGE("Using BackwardFlat interpolation with a moving reference date");
                BOOST_REQUIRE_NO_THROW(ovCurve = boost::make_shared<PiecewiseOptionletCurve<BackwardFlat> >(
                                           settlementDays, helpers, calendar, bdc, dayCounter, curveVolatilityType,
                                           curveDisplacement, flatFirstPeriod));
            } else {
                BOOST_TEST_MESSAGE("Using BackwardFlat interpolation with a fixed reference date");
                BOOST_REQUIRE_NO_THROW(ovCurve = boost::make_shared<PiecewiseOptionletCurve<BackwardFlat> >(
                                           referenceDate, helpers, calendar, bdc, dayCounter, curveVolatilityType,
                                           curveDisplacement, flatFirstPeriod));
            }
            break;
        case 2:
            if (isMoving) {
                BOOST_TEST_MESSAGE("Using LinearFlat interpolation with a moving reference date");
                BOOST_REQUIRE_NO_THROW(ovCurve = boost::make_shared<PiecewiseOptionletCurve<LinearFlat> >(
                                           settlementDays, helpers, calendar, bdc, dayCounter, curveVolatilityType,
                                           curveDisplacement, flatFirstPeriod));
            } else {
                BOOST_TEST_MESSAGE("Using LinearFlat interpolation with a fixed reference date");
                BOOST_REQUIRE_NO_THROW(ovCurve = boost::make_shared<PiecewiseOptionletCurve<LinearFlat> >(
                                           referenceDate, helpers, calendar, bdc, dayCounter, curveVolatilityType,
                                           curveDisplacement, flatFirstPeriod));
            }
            break;
        case 3:
            if (isMoving) {
                BOOST_TEST_MESSAGE("Using Cubic interpolation with a moving reference date");
                BOOST_REQUIRE_NO_THROW(
                    ovCurve = boost::make_shared<PiecewiseOptionletCurve<Cubic> >(
                        settlementDays, helpers, calendar, bdc, dayCounter, curveVolatilityType, curveDisplacement,
                        flatFirstPeriod, Cubic(),
                        QuantExt::IterativeBootstrap<
                            PiecewiseOptionletCurve<Cubic, QuantExt::IterativeBootstrap>::this_curve>(accuracy,
                                                                                globalAccuracy)));
            } else {
                BOOST_TEST_MESSAGE("Using Cubic interpolation with a fixed reference date");
                BOOST_REQUIRE_NO_THROW(
                    ovCurve = boost::make_shared<PiecewiseOptionletCurve<Cubic> >(
                        referenceDate, helpers, calendar, bdc, dayCounter, curveVolatilityType, curveDisplacement,
                        flatFirstPeriod, Cubic(),
                        QuantExt::IterativeBootstrap<
                            PiecewiseOptionletCurve<Cubic, QuantExt::IterativeBootstrap>::this_curve>(accuracy,
                                                                                globalAccuracy)));
            }
            break;
        case 4:
            if (isMoving) {
                BOOST_TEST_MESSAGE("Using CubicFlat interpolation with a moving reference date");
                BOOST_REQUIRE_NO_THROW(
                    ovCurve = boost::make_shared<PiecewiseOptionletCurve<CubicFlat> >(
                        settlementDays, helpers, calendar, bdc, dayCounter, curveVolatilityType, curveDisplacement,
                        flatFirstPeriod, CubicFlat(),
                        QuantExt::IterativeBootstrap<
                            PiecewiseOptionletCurve<CubicFlat, QuantExt::IterativeBootstrap>::this_curve>(
                            accuracy, globalAccuracy)));
            } else {
                BOOST_TEST_MESSAGE("Using CubicFlat interpolation with a fixed reference date");
                BOOST_REQUIRE_NO_THROW(
                    ovCurve = boost::make_shared<PiecewiseOptionletCurve<CubicFlat> >(
                        referenceDate, helpers, calendar, bdc, dayCounter, curveVolatilityType, curveDisplacement,
                        flatFirstPeriod, CubicFlat(),
                        QuantExt::IterativeBootstrap<
                            PiecewiseOptionletCurve<CubicFlat, QuantExt::IterativeBootstrap>::this_curve>(
                            accuracy, globalAccuracy)));
            }
            break;
        default:
            BOOST_FAIL("Unexpected interpolation type");
        }
        Handle<OptionletVolatilityStructure> hovs(ovCurve);

        // Price each cap floor instrument using the piecewise optionlet curve and check it against the flat NPV
        BOOST_TEST_MESSAGE("The stripped values and differences at each tenor are:");
        Real strippedNpv;
        for (Size i = 0; i < volatilityColumn.tenors.size(); i++) {

            // May need to update instruments type if it is being chosen automatically in the bootstrap
            if (helperType == CapFloorHelper::Automatic && quoteType != CapFloorHelper::Premium) {
                Real volatility = volatilityColumn.volatilities[i];
                CapFloor::Type capFloorType =
                    boost::dynamic_pointer_cast<CapFloorHelper>(helpers[i])->capFloor()->type();
                if (capFloorType != instruments[i]->type()) {
                    // Need to update the instrument and the flat NPV for the test
                    instruments[i] =
                        MakeCapFloor(capFloorType, volatilityColumn.tenors[i], iborIndex, volatilityColumn.strike);
                    if (volatilityColumn.type == ShiftedLognormal) {
                        instruments[i]->setPricingEngine(boost::make_shared<BlackCapFloorEngine>(
                            testYieldCurves.discountEonia, volatility, dayCounter, volatilityColumn.displacement));
                    } else {
                        instruments[i]->setPricingEngine(boost::make_shared<BachelierCapFloorEngine>(
                            testYieldCurves.discountEonia, volatility, dayCounter));
                    }
                    flatNpvs[i] = instruments[i]->NPV();
                }
            }

            // Price the instrument using the stripped optionlet structure
            if (ovCurve->volatilityType() == ShiftedLognormal) {
                instruments[i]->setPricingEngine(
                    boost::make_shared<BlackCapFloorEngine>(testYieldCurves.discountEonia, hovs));
            } else {
                instruments[i]->setPricingEngine(
                    boost::make_shared<BachelierCapFloorEngine>(testYieldCurves.discountEonia, hovs));
            }
            strippedNpv = instruments[i]->NPV();

            BOOST_TEST_MESSAGE("  (Cap/Floor, Tenor, Volatility, Flat NPV, Stripped NPV, Flat - Stripped) = ("
                               << instruments[i]->type() << ", " << volatilityColumn.tenors[i] << ", " << fixed
                               << setprecision(13) << volatilityColumn.volatilities[i] << ", " << flatNpvs[i] << ", "
                               << strippedNpv << ", " << (flatNpvs[i] - strippedNpv) << ")");

            BOOST_CHECK_SMALL(fabs(flatNpvs[i] - strippedNpv), tolerance);
        }
    }
}

BOOST_DATA_TEST_CASE_F(CommonVars, testCachedValues,
                       bdata::make(interpolationTypesCached) * bdata::make(flatFirstPeriodValues), interpolationType,
                       flatFirstPeriod) {

    BOOST_TEST_MESSAGE("Testing stripping of single strike column against cached values");

    CapFloorHelper::Type helperType = CapFloorHelper::Automatic;
    CapFloorHelper::QuoteType quoteType = CapFloorHelper::Volatility;

    // Use EUR cap floor test volatility data from file test/capfloormarketdata.hpp
    // Take the highest strike column of the normal volatility matrix
    CapFloorVolatilityEUR testVols;
    vector<Period> tenors = testVols.tenors;
    vector<Real> volatilities(tenors.size());
    Size strikeIdx = testVols.strikes.size() - 1;
    Rate strike = testVols.strikes[strikeIdx];
    for (Size i = 0; i < tenors.size(); i++) {
        volatilities[i] = testVols.nVols[i][strikeIdx];
    }
    VolatilityType volatilityType = Normal;
    Real displacement = 0.0;

    // Make first tenor 18M highlight differences introduced by flatFirstPeriod setting.
    // With first tenor = 1Y and index tenor = 6M, we were not seeing a difference.
    tenors[0] = 18 * Months;

    BOOST_TEST_MESSAGE("Test inputs are:");
    BOOST_TEST_MESSAGE("  Cap floor helper type: " << helperType);
    BOOST_TEST_MESSAGE("  Cap floor strike: " << strike);
    BOOST_TEST_MESSAGE("  Quote type: " << quoteType);
    BOOST_TEST_MESSAGE("  Quote volatility type: " << volatilityType);
    BOOST_TEST_MESSAGE("  Quote displacement: " << displacement);
    BOOST_TEST_MESSAGE("  Interpolation type: " << to_string(interpolationType));
    BOOST_TEST_MESSAGE("  Flat first period: " << boolalpha << flatFirstPeriod);

    // Form the cap floor helper instrument for each tenor in the strike column
    vector<boost::shared_ptr<helper> > helpers(tenors.size());
    for (Size i = 0; i < tenors.size(); i++) {
        Handle<Quote> quote(boost::make_shared<SimpleQuote>(volatilities[i]));
        helpers[i] = boost::make_shared<CapFloorHelper>(helperType, tenors[i], strike, quote, iborIndex,
                                                        testYieldCurves.discountEonia, true, Date(), quoteType,
                                                        volatilityType, displacement);
    }

    // Create the piecewise optionlet curve, with the given interpolation type.
    // Store the nodes of this optionlet curve to compare with cached values
    VolatilityType curveVolatilityType = Normal;
    Real curveDisplacement = 0.0;
    boost::shared_ptr<OptionletVolatilityStructure> ovs;
    vector<pair<Date, Real> > curveNodes;
    switch (interpolationType.which()) {
    case 0: {
        BOOST_TEST_MESSAGE("Using Linear interpolation");
        boost::shared_ptr<PiecewiseOptionletCurve<Linear> > ovCurve =
            boost::make_shared<PiecewiseOptionletCurve<Linear> >(referenceDate, helpers, calendar, bdc, dayCounter,
                                                                 curveVolatilityType, curveDisplacement,
                                                                 flatFirstPeriod);
        curveNodes = ovCurve->nodes();
        ovs = ovCurve;
    } break;
    case 1: {
        BOOST_TEST_MESSAGE("Using BackwardFlat interpolation");
        boost::shared_ptr<PiecewiseOptionletCurve<BackwardFlat> > ovCurve =
            boost::make_shared<PiecewiseOptionletCurve<BackwardFlat> >(referenceDate, helpers, calendar, bdc,
                                                                       dayCounter, curveVolatilityType,
                                                                       curveDisplacement, flatFirstPeriod);
        curveNodes = ovCurve->nodes();
        ovs = ovCurve;
    } break;
    case 2: {
        BOOST_TEST_MESSAGE("Using LinearFlat interpolation");
        boost::shared_ptr<PiecewiseOptionletCurve<LinearFlat> > ovCurve =
            boost::make_shared<PiecewiseOptionletCurve<LinearFlat> >(referenceDate, helpers, calendar, bdc, dayCounter,
                                                                     curveVolatilityType, curveDisplacement,
                                                                     flatFirstPeriod);
        curveNodes = ovCurve->nodes();
        ovs = ovCurve;
    } break;
    default:
        BOOST_FAIL("Unexpected interpolation type");
    }

    // Get the key for the cached results for the current test
    Size key = 2 * interpolationType.which() + static_cast<Size>(flatFirstPeriod);

    // Check stripped optionlet volatilities against cached values
    BOOST_REQUIRE_EQUAL(curveNodes.size(), cachedDates.size());
    BOOST_REQUIRE(cachedValues.count(key) == 1);
    BOOST_REQUIRE_EQUAL(curveNodes.size(), cachedValues.at(key).size());
    BOOST_TEST_MESSAGE("node_date,node_vol");
    for (Size i = 0; i < curveNodes.size(); i++) {
        // Check the date
        BOOST_CHECK_EQUAL(curveNodes[i].first, cachedDates[i]);
        // Check the value
        BOOST_CHECK_SMALL(fabs(curveNodes[i].second - cachedValues.at(key)[i]), accuracy);
        // Print out the curve
        BOOST_TEST_MESSAGE(io::iso_date(curveNodes[i].first)
                           << "," << fixed << setprecision(12) << curveNodes[i].second);
    }

    // Check stripped optionlet volatilities on non-node dates against cached values
    BOOST_REQUIRE(cachedNonNodeValues.count(key) == 1);
    BOOST_TEST_MESSAGE("date,vol,cached_vol,diff");
    for (Size i = 0; i < cachedNonNodeDates.size(); i++) {
        Date d = cachedNonNodeDates[i];

        // The last date has been picked past the maximum curve date to check extrapolation. Check that we get an
        // error and then turn on extrapolation.
        if (i == cachedNonNodeDates.size() - 1) {
            BOOST_CHECK_THROW(ovs->volatility(d, strike), Error);
            ovs->enableExtrapolation();
        }

        Volatility vol = ovs->volatility(d, strike);
        Volatility cachedVol = cachedNonNodeValues.at(key)[i];
        Volatility diff = fabs(vol - cachedVol);
        // Check the value
        BOOST_CHECK_SMALL(diff, accuracy);
        // Print out the curve
        BOOST_TEST_MESSAGE(io::iso_date(d)
                           << "," << fixed << setprecision(12) << vol << "," << cachedVol << "," << diff);

        // The strike should not matter so check that the same test passes above and below the strike
        Real shift = 0.0010;
        diff = fabs(ovs->volatility(d, strike - shift) - cachedVol);
        BOOST_CHECK_SMALL(diff, accuracy);
        diff = fabs(ovs->volatility(d, strike + shift) - cachedVol);
        BOOST_CHECK_SMALL(diff, accuracy);
    }
}

BOOST_AUTO_TEST_SUITE_END()

BOOST_AUTO_TEST_SUITE_END()


/* Copyright (c) 2021, the adamantine authors.
 *
 * This file is subject to the Modified BSD License and may not be distributed
 * without copyright and license information. Please refer to the file LICENSE
 * for the text and further information on this license.
 */

#include <DataAssimilator.hh>

#include <deal.II/dofs/dof_tools.h>
#include <deal.II/lac/linear_operator_tools.h>

#include <execution>

namespace adamantine
{

DataAssimilator::DataAssimilator(boost::property_tree::ptree const &database)
{
  // Set the solver parameters from the input database
  // PropertyTreeInput data_assimilation.solver.max_number_of_temp_vectors
  if (boost::optional<unsigned int> max_num_temp_vectors =
          database.get_optional<unsigned int>(
              "solver.max_number_of_temp_vectors"))
    _additional_data.max_n_tmp_vectors = *max_num_temp_vectors;

  // PropertyTreeInput data_assimilation.solver.max_iterations
  if (boost::optional<unsigned int> max_iterations =
          database.get_optional<unsigned int>("solver.max_iterations"))
    _solver_control.set_max_steps(*max_iterations);

  // PropertyTreeInput data_assimilation.solver.convergence_tolerance
  if (boost::optional<double> tolerance =
          database.get_optional<double>("solver.convergence_tolerance"))
    _solver_control.set_tolerance(*tolerance);
}

void DataAssimilator::update_ensemble(
    std::vector<dealii::LA::distributed::Vector<double>> &sim_data,
    std::vector<double> const &expt_data, dealii::SparseMatrix<double> &R)
{
  // Set some constants
  _num_ensemble_members = sim_data.size();
  if (sim_data.size() > 0)
  {
    _sim_size = sim_data[0].size();
  }
  else
  {
    _sim_size = 0;
  }
  _expt_size = expt_data.size();

  // Get the perturbed innovation, ( y+u - Hx )
  std::vector<dealii::Vector<double>> perturbed_innovation(
      _num_ensemble_members);
  for (unsigned int member = 0; member < _num_ensemble_members; ++member)
  {
    perturbed_innovation[member].reinit(_expt_size);
    fill_noise_vector(perturbed_innovation[member], R);
    dealii::Vector<double> temporary = calc_Hx(sim_data[member]);
    for (unsigned int i = 0; i < _expt_size; ++i)
    {
      perturbed_innovation[member][i] += expt_data[i] - temporary[i];
    }
  }

  // Apply the Kalman filter to the perturbed innovation, K ( y+u - Hx )
  std::vector<dealii::LA::distributed::Vector<double>> forecast_shift =
      apply_kalman_gain(sim_data, R, perturbed_innovation);

  // Update the ensemble, x = x + K ( y+u - Hx )
  for (unsigned int member = 0; member < _num_ensemble_members; ++member)
  {
    sim_data[member] += forecast_shift[member];
  }
}

std::vector<dealii::LA::distributed::Vector<double>>
DataAssimilator::apply_kalman_gain(
    std::vector<dealii::LA::distributed::Vector<double>> &vec_ensemble,
    dealii::SparseMatrix<double> &R,
    std::vector<dealii::Vector<double>> &perturbed_innovation)
{
  /*
   * Currently this function uses GMRES to apply the inverse of HPH^T+R in the
   * Kalman gain calculation for each ensemble member individually. Depending on
   * the size of the datasets, the number of ensembles, and other factors doing
   * a direct solve of (HPH^T+R)^-1 once and then applying to the perturbed
   * innovation from each ensemble member might be more efficient.
   */

  dealii::SparsityPattern pattern_H(_expt_size, _sim_size, _expt_size);
  dealii::SparseMatrix<double> H = calc_H(pattern_H);

  dealii::FullMatrix<double> P = calc_sample_covariance_dense(vec_ensemble);

  const auto op_H = dealii::linear_operator(H);
  const auto op_P = dealii::linear_operator(P);
  const auto op_R = dealii::linear_operator(R);

  const auto op_HPH_plus_R =
      op_H * op_P * dealii::transpose_operator(op_H) + op_R;

  std::vector<dealii::LA::distributed::Vector<double>> output(
      _num_ensemble_members,
      dealii::LA::distributed::Vector<double>(_sim_size));

  // Create non-member versions of these for use in the lambda function
  auto solver_control = _solver_control;
  auto additional_data = _additional_data;
  auto sim_size = _sim_size;

  // Apply the Kalman gain to the perturbed innovation for the ensemble members
  // in parallel
  std::transform(
      std::execution::par, perturbed_innovation.begin(),
      perturbed_innovation.end(), output.begin(),
      [&](dealii::Vector<double> entry) {
        dealii::SolverGMRES<dealii::Vector<double>> HPH_plus_R_inv_solver(
            solver_control, additional_data);

        auto op_HPH_plus_R_inv =
            dealii::inverse_operator(op_HPH_plus_R, HPH_plus_R_inv_solver);

        const auto op_K =
            op_P * dealii::transpose_operator(op_H) * op_HPH_plus_R_inv;

        // Apply the Kalman gain to each innovation vector
        dealii::Vector<double> temporary = op_K * entry;

        // Copy into a distributed vector, this is the only place where the
        // mismatch matters, using dealii::Vector for the experimental data
        // and dealii::LA::distributed::Vector for the simulation data.
        dealii::LA::distributed::Vector<double> output_member(sim_size);
        for (unsigned int i = 0; i < sim_size; ++i)
        {
          output_member(i) = temporary(i);
        }

        return output_member;
      });

  return output;
}

dealii::SparseMatrix<double>
DataAssimilator::calc_H(dealii::SparsityPattern &pattern) const
{
  int num_expt_dof_map_entries = _expt_to_dof_mapping.first.size();

  for (auto i = 0; i < num_expt_dof_map_entries; ++i)
  {
    auto sim_index = _expt_to_dof_mapping.second[i];
    auto expt_index = _expt_to_dof_mapping.first[i];
    pattern.add(expt_index, sim_index);
  }

  pattern.compress();

  dealii::SparseMatrix<double> H(pattern);

  for (auto i = 0; i < num_expt_dof_map_entries; ++i)
  {
    auto sim_index = _expt_to_dof_mapping.second[i];
    auto expt_index = _expt_to_dof_mapping.first[i];
    H.add(expt_index, sim_index, 1.0);
  }

  return H;
}

template <int dim>
void DataAssimilator::update_dof_mapping(
    dealii::DoFHandler<dim> const &dof_handler,
    std::pair<std::vector<int>, std::vector<int>> const &indices_and_offsets)
{
  std::map<dealii::types::global_dof_index, dealii::Point<dim>> indices_points;
  dealii::DoFTools::map_dofs_to_support_points(
      dealii::StaticMappingQ1<dim>::mapping, dof_handler, indices_points);
  // Change the format to the one used by ArborX
  std::vector<dealii::types::global_dof_index> dof_indices(
      indices_points.size());
  unsigned int pos = 0;
  for (auto map_it = indices_points.begin(); map_it != indices_points.end();
       ++map_it, ++pos)
  {
    dof_indices[pos] = map_it->first;
  }

  _expt_to_dof_mapping.first.resize(indices_and_offsets.first.size());
  _expt_to_dof_mapping.second.resize(indices_and_offsets.first.size());

  for (unsigned int i = 0; i < _expt_size; ++i)
  {
    for (int j = indices_and_offsets.second[i];
         j < indices_and_offsets.second[i + 1]; ++j)
    {
      _expt_to_dof_mapping.first[j] = i;
      _expt_to_dof_mapping.second[j] =
          dof_indices[indices_and_offsets.first[j]];
    }
  }
}

dealii::Vector<double> DataAssimilator::calc_Hx(
    const dealii::LA::distributed::Vector<double> &sim_ensemble_member) const
{
  int num_expt_dof_map_entries = _expt_to_dof_mapping.first.size();

  dealii::Vector<double> out_vec(_expt_size);

  // Loop through the observation map to get the observation indices
  for (auto i = 0; i < num_expt_dof_map_entries; ++i)
  {
    auto sim_index = _expt_to_dof_mapping.second[i];
    auto expt_index = _expt_to_dof_mapping.first[i];
    out_vec(expt_index) = sim_ensemble_member(sim_index);
  }

  return out_vec;
}

void DataAssimilator::fill_noise_vector(dealii::Vector<double> &vec,
                                        dealii::SparseMatrix<double> &R)
{
  auto vector_size = vec.size();

  // Do Cholesky decomposition
  dealii::FullMatrix<double> L(vector_size);
  dealii::FullMatrix<double> R_full(vector_size);
  R_full.copy_from(R);
  L.cholesky(R_full);

  // Get a vector of normally distributed values
  dealii::Vector<double> uncorrelated_noise_vector(vector_size);

  for (unsigned int i = 0; i < vector_size; ++i)
  {
    uncorrelated_noise_vector(i) = _normal_dist_generator(_prng);
  }

  L.vmult(vec, uncorrelated_noise_vector);
}

template <typename VectorType>
dealii::FullMatrix<double> DataAssimilator::calc_sample_covariance_dense(
    std::vector<VectorType> vec_ensemble) const
{
  unsigned int num_ensemble_members = vec_ensemble.size();
  unsigned int vec_size = 0;
  if (vec_ensemble.size() > 0)
  {
    vec_size = vec_ensemble[0].size();
  }

  // Calculate the mean
  dealii::Vector<double> mean(vec_size);
  for (unsigned int i = 0; i < vec_size; ++i)
  {
    double sum = 0.0;
    for (unsigned int sample = 0; sample < num_ensemble_members; ++sample)
    {
      sum += vec_ensemble[sample][i];
    }
    mean[i] = sum / num_ensemble_members;
  }

  // Calculate the anomaly
  dealii::FullMatrix<double> anomaly(vec_size, num_ensemble_members);
  for (unsigned int member = 0; member < num_ensemble_members; ++member)
  {
    for (unsigned int i = 0; i < vec_size; ++i)
    {
      anomaly(i, member) = (vec_ensemble[member][i] - mean[i]) /
                           std::sqrt(num_ensemble_members - 1.0);
    }
  }

  dealii::FullMatrix<double> cov(vec_size);
  anomaly.mTmult(cov, anomaly);

  return cov;
}

// Explicit instantiation
template void DataAssimilator::update_dof_mapping<2>(
    dealii::DoFHandler<2> const &dof_handler,
    std::pair<std::vector<int>, std::vector<int>> const &indices_and_offsets);
template void DataAssimilator::update_dof_mapping<3>(
    dealii::DoFHandler<3> const &dof_handler,
    std::pair<std::vector<int>, std::vector<int>> const &indices_and_offsets);
template dealii::FullMatrix<double>
DataAssimilator::calc_sample_covariance_dense<dealii::Vector<double>>(
    std::vector<dealii::Vector<double>> vec_ensemble) const;
template dealii::FullMatrix<double>
DataAssimilator::calc_sample_covariance_dense<
    dealii::LA::distributed::Vector<double>>(
    std::vector<dealii::LA::distributed::Vector<double>> vec_ensemble) const;

} // namespace adamantine


/*
 [auto_generated]
 boost/numeric/odeint/external/mtl4/mtl4_algebra_dispatcher.hpp

 [begin_description]
 specialization of the algebra dispatcher for mtl4
 [end_description]

 Copyright 2013 Karsten Ahnert
 Copyright 2013 Mario Mulansky

 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)
 */

#ifndef BOOST_NUMERIC_ODEINT_MTL4_MTL4_ALGEBRA_DISPATCHER_HPP_INCLUDED
#define BOOST_NUMERIC_ODEINT_MTL4_MTL4_ALGEBRA_DISPATCHER_HPP_INCLUDED

#include <boost/numeric/mtl/mtl.hpp>

#include <boost/numeric/odeint/algebra/vector_space_algebra.hpp>
#include <boost/numeric/odeint/algebra/algebra_dispatcher.hpp>

namespace boost {
namespace numeric {
namespace odeint {

template<typename Value, typename Parameters>
struct algebra_dispatcher< mtl::dense_vector< Value , Parameters > >
{
    typedef vector_space_algebra algebra_type;
};

template<typename Value, typename Parameters>
struct algebra_dispatcher< mtl::dense2D< Value , Parameters > >
{
    typedef vector_space_algebra algebra_type;
};

template<typename Value , size_t BitMask , typename Parameters>
struct algebra_dispatcher< mtl::morton_dense< Value , BitMask, Parameters > >
{
    typedef vector_space_algebra algebra_type;
};

template<typename Value, typename Parameters>
struct algebra_dispatcher< mtl::compressed2D< Value , Parameters > >
{
    typedef vector_space_algebra algebra_type;
};

// specialization of infinity norm calculation

template<typename Value, typename Parameters>
struct vector_space_norm_inf< mtl::dense_vector< Value , Parameters > >
{
    typedef Value result_type;
    Value operator()( const mtl::dense_vector< Value , Parameters > &x ) const
    {
        return mtl::infinity_norm(x);
    }
};

template<typename Value, typename Parameters>
struct vector_space_norm_inf< mtl::dense2D< Value , Parameters > >
{
    typedef Value result_type;
    Value operator()( const mtl::dense2D< Value , Parameters > &x ) const
    {
        return mtl::infinity_norm(x);
    }
};

template<typename Value , size_t BitMask , typename Parameters>
struct vector_space_norm_inf< mtl::morton_dense< Value , BitMask , Parameters > >
{
    typedef Value result_type;
    Value operator()( const mtl::morton_dense< Value , BitMask , Parameters > &x ) const
    {
        return mtl::infinity_norm(x);
    }
};

template<typename Value, typename Parameters>
struct vector_space_norm_inf< mtl::compressed2D< Value , Parameters > >
{
    typedef Value result_type;
    Value operator()( const mtl::compressed2D< Value , Parameters > &x ) const
    {
        return mtl::infinity_norm(x);
    }
};

}
}
}

#endif // BOOST_NUMERIC_ODEINT_MTL4_MTL4_ALGEBRA_DISPATCHER_INCLUDED