// Copyright 2020 Speech Technology Center www.speechpro.com
//
// Licensed under the Apache License, Version 2.0 http://www.apache.org/licenses/LICENSE-2.0
//
// This code is provided *as is* basis, without warranties or conditions of any kind.
////////////////////////////////////////////////////////////////////////////////////////////
// If you find this code useful for your research or production, please cite our paper:
//
// @inproceedings{Medennikov_mixup2018,
//   author={Ivan Medennikov and Yuri Khokhlov and Aleksei Romanenko and Dmitry Popov and Natalia Tomashenko and Ivan Sorokin and Alexander Zatvornitskiy},
//   title={An Investigation of Mixup Training Strategies for Acoustic Models in ASR},
//   year=2018,
//   booktitle={Proc. Interspeech 2018},
//   pages={2903--2907},
//   doi={10.21437/Interspeech.2018-2191},
//   url={http://dx.doi.org/10.21437/Interspeech.2018-2191}
// }
//
// (corresponding author: medennikov@speechpro.com)
////////////////////////////////////////////////////////////////////////////////////////////

#include <limits>
#include <algorithm>
#include <fst/union.h>
#include <fst/topsort.h>
#include <boost/shared_ptr.hpp>
#include <boost/format.hpp>
#include <boost/lexical_cast.hpp>
#include <boost/random/mersenne_twister.hpp>
#include <boost/random/uniform_int_distribution.hpp>
#include <boost/random/uniform_real_distribution.hpp>
#include <boost/random/beta_distribution.hpp>
#include <boost/algorithm/string/split.hpp>
#include <boost/algorithm/string/case_conv.hpp>
#include <boost/algorithm/string/classification.hpp>
#include "base/kaldi-common.h"
#include "util/common-utils.h"
#include "nnet3/nnet-chain-example.h"

namespace kaldi { namespace nnet3 {

typedef Matrix<BaseFloat> KaldiMatrix;
typedef Vector<BaseFloat> KaldiVector;
typedef SparseMatrix<BaseFloat> KaldiSparMatrix;
typedef SparseVector<BaseFloat> KaldiSparVector;
typedef SubMatrix<BaseFloat> KaldiSubMatrix;
typedef SubVector<BaseFloat> KaldiSubVector;
typedef boost::shared_ptr<NnetChainExample> ExamplePtr;
typedef std::pair<std::string, ExamplePtr> ExamplePair;

class IRandomScale {
public:
    virtual float Value() = 0;
};

class RandomScaleUniform: public IRandomScale {
protected:
    typedef boost::random::uniform_real_distribution<float> real_distrib_t;
    typedef boost::shared_ptr<real_distrib_t> distrib_ptr_t;

protected:
    boost::random::mt19937& rand_gen;
    distrib_ptr_t distrib;

public:
    RandomScaleUniform(boost::random::mt19937& _rand_gen, const std::string& _params);
    float Value() override;
};

RandomScaleUniform::RandomScaleUniform(boost::random::mt19937& _rand_gen, const std::string& _params): rand_gen(_rand_gen), distrib() {
    std::vector<std::string> parts;
    boost::split(parts, _params, boost::is_any_of(":"));
    if (parts.size() != 2) {
        KALDI_ERR << "Wrong uniform distribution parameters string: \"" << _params << "\".";
    }
    if (parts.front() != "uniform") {
        KALDI_ERR << "Wrong uniform distribution parameters string: \"" << _params << "\".";
    }
    std::vector<std::string> min_max;
    boost::split(min_max, parts.back(), boost::is_any_of(","));
    if (parts.size() != 2) {
        KALDI_ERR << "Wrong uniform distribution parameters string: \"" << _params << "\".";
    }
    const auto min_value = boost::lexical_cast<float>(min_max.front());
    const auto max_value = boost::lexical_cast<float>(min_max.back());
    if (min_value > max_value) {
        KALDI_ERR << "min_value must be less or equal max_value (" << _params << ")";
    }
    if ((min_value < 0.0f) || (min_value > 1.0f)) {
        KALDI_ERR << "min_value must be in the range [0, 1] (" << _params << ")";
    }
    if ((max_value < 0.0f) || (max_value > 1.0f)) {
        KALDI_ERR << "max_value must be in the range [0, 1] (" << _params << ")";
    }
    distrib.reset(new real_distrib_t(min_value, max_value));
}

float RandomScaleUniform::Value() {
    return (*distrib)(rand_gen);
}

class RandomScaleBeta: public IRandomScale {
protected:
    typedef boost::random::beta_distribution<double> real_distrib_t;
    typedef boost::shared_ptr<real_distrib_t> distrib_ptr_t;

protected:
    boost::random::mt19937& rand_gen;
    distrib_ptr_t distrib;

public:
    RandomScaleBeta(boost::random::mt19937& _rand_gen, const std::string& _params);
    float Value() override;
};

RandomScaleBeta::RandomScaleBeta(boost::random::mt19937& _rand_gen, const std::string& _params): rand_gen(_rand_gen), distrib() {
    std::vector<std::string> parts;
    boost::split(parts, _params, boost::is_any_of(":"));
    if (parts.size() != 2) {
        KALDI_ERR << "Wrong beta distribution parameters string: \"" << _params << "\".";
    }
    if ((parts.front() != "beta") && (parts.front() != "beta2")) {
        KALDI_ERR << "Wrong beta distribution parameters string: \"" << _params << "\".";
    }
    const auto alpha = boost::lexical_cast<float>(parts.back());
    if (alpha <= 0.0f) {
        KALDI_ERR << "alpha must be a positive value (" << _params << ")";
    }
    distrib.reset(new real_distrib_t(alpha, alpha));
}

float RandomScaleBeta::Value() {
    return (float)(*distrib)(rand_gen);
}

class RandomScaleBeta2: public RandomScaleBeta {
public:
    RandomScaleBeta2(boost::random::mt19937& _rand_gen, const std::string& _params);
    float Value() override;
};

RandomScaleBeta2::RandomScaleBeta2(boost::random::mt19937& _rand_gen, const std::string& _params):
    RandomScaleBeta(_rand_gen, _params)
{}

float RandomScaleBeta2::Value() {
    const double value = (*distrib)(rand_gen);
    if (value <= 0.5) {
        return (float) value;
    } else {
        return (float)(1.0 - value);
    }
}

class RandomScale {
protected:
    typedef boost::shared_ptr<IRandomScale> rand_scale_t;

protected:
    rand_scale_t rand_scale;

public:
    RandomScale(boost::random::mt19937& _rand_gen, const std::string& _params);
    float operator()();
};

RandomScale::RandomScale(boost::random::mt19937& _rand_gen, const std::string& _params): rand_scale() {
    if (_params.find("uniform") == 0) {
        rand_scale.reset(new RandomScaleUniform(_rand_gen, _params));
    } else if (_params.find("beta2") == 0) {
        rand_scale.reset(new RandomScaleBeta2(_rand_gen, _params));
    } else if (_params.find("beta") == 0) {
        rand_scale.reset(new RandomScaleBeta(_rand_gen, _params));
    } else {
        KALDI_ERR << "Unknown random scale generator ID: \"" << _params << "\".";
    }
}

float RandomScale::operator()() {
    return rand_scale->Value();
}

template<class Type>
class TCounter {
protected:
    Type minimum;
    Type maximum;
    double summa;
    size_t count;

public:
    TCounter():
        minimum(std::numeric_limits<Type>::max()),
        maximum(std::numeric_limits<Type>::min()),
        summa(0.0), count(0)
    {}
    Type Minimum() const {return minimum;}
    Type Maximum() const {return maximum;}
    double Average() const {return summa / count;}
    size_t Count() const {return count;}
    bool Valid() const {return (count > 0);}
    void operator+=(Type _value) {
        minimum = std::min(_value, minimum);
        maximum = std::max(_value, maximum);
        summa += _value;
        ++count;
    }
};

typedef TCounter<int> IntCounter;
typedef TCounter<float> FloatCounter;

typedef fst::StdVectorFst fst_t;
typedef fst_t::StateId state_t;
typedef fst_t::Arc arc_t;
typedef arc_t::Label label_t;
typedef arc_t::Weight wght_t;
typedef fst::ArcIterator<fst_t> iter_t;

class ExampleMixer {
protected:
    typedef boost::random::mt19937 rand_gen_t;
    typedef boost::random::uniform_int_distribution<int32_t> int_distrib_t;
    typedef boost::random::uniform_int_distribution<size_t> uint_distrib_t;
    typedef boost::random::uniform_real_distribution<float> real_distrib_t;
    typedef std::vector<ExamplePair> egs_buffer_t;
    typedef unordered_map<NnetChainExample*, egs_buffer_t, NnetChainExampleStructureHasher, NnetChainExampleStructureCompare> eg_to_egs_t;

protected:
    struct MixupData {
        int32_t row_main;
        int32_t row_admx;
        int32_t label_main;
        int32_t label_admx;
        float scale_main;
        float scale_admx;
        MixupData(): row_main(0), row_admx(), label_main(-1), label_admx(), scale_main(0.0f), scale_admx() {}
        MixupData(int32_t _row_main,int32_t _label_main, float _scale_main):
            row_main(_row_main), row_admx(),
            label_main(_label_main), label_admx(),
            scale_main(_scale_main), scale_admx()
        {}
        MixupData(
            int32_t _row_main, int32_t _rows_admx,
            int32_t _label_main, int32_t _labels_admx,
            float _scale_main, float _scale_admx
        ):
            row_main(_row_main), row_admx(_rows_admx),
            label_main(_label_main), label_admx(_labels_admx),
            scale_main(_scale_main), scale_admx(_scale_admx)
        {}
    };

protected:
    const std::string mix_mode;
    NnetChainExampleWriter& example_writer;
    std::string scale_fst_algo;
    bool swap_scales;
    float scale_boost;
    float scale_eps;
    bool mix_ivect;
    bool mix_feats;
    bool max_super;
    int32_t min_shift;
    int32_t max_shift;
    float fixed;
    size_t buff_size;
    int32_t frame_shift;
    std::vector<std::string> exclude_names;
    bool compress;
    bool test_mode;
    rand_gen_t rand_gen;
    uint_distrib_t int_distrib;
    int_distrib_t shift_distrib;
    real_distrib_t float_distrib;
    RandomScale scale_distrib;
    eg_to_egs_t eg_to_egs;
    egs_buffer_t egs_buffer;
    FloatCounter scale_count;
    IntCounter shift_count;
    size_t num_mixed;
    size_t num_untouched;
    size_t num_accepted;
    size_t num_wrote;

public:
    ExampleMixer(
        std::string _mix_mode, const std::string& _distrib,
        NnetChainExampleWriter& _example_writer,
        std::string _scale_fst_algo, bool _swap_scales,
        bool _mix_ivect, bool _mix_feats,
        bool _max_super, int32_t _min_shift, int32_t _max_shift,
        float _fixed, size_t _buff_size, int32_t _frame_shift,
        bool _compress, bool _test_mode
    );

protected:
    const std::vector<Index>& FindIndexes(const std::string& _name, const std::vector<NnetIo>& _nnet_io) const;
    GeneralMatrix& FindFeatures(const std::string& _name, std::vector<NnetIo>& _nnet_io) const;
    const GeneralMatrix& FindFeatures(const std::string& _name, const std::vector<NnetIo>& _nnet_io) const;
    GeneralMatrix* FindIVector(std::vector<NnetIo>& _nnet_io) const;
    const GeneralMatrix* FindIVector(const std::vector<NnetIo>& _nnet_io) const;
    void CheckConsistence(const NnetChainExample& _example) const;
    void CheckConsistence(const NnetChainExample& _example1, const NnetChainExample& _example2) const;
    void ScaleGraph(fst_t& _target, float _scale) const;
    void UnionGraphs(const fst_t& _admixture, fst_t& _example) const;
    void FuseGraphs(const fst_t& _admixture, float _admx_scale, fst_t& _example) const;
    void AdmixGlobal(const NnetChainExample& _admixture, float _admx_scale, ExamplePair& _example);
    void FlushGlobal(egs_buffer_t& _buffer);
    void ShiftAndMixup(ExamplePair& _example);

public:
    void AcceptExample(ExamplePair& _example);
    void Finish();
    const FloatCounter& ScaleCount() const {return scale_count;}
    const IntCounter& ShiftCount() const {return shift_count;}
    size_t NumMixed() const {return num_mixed;}
    size_t NumUntouched() const {return num_untouched;}
    size_t NumAccepted() const {return num_accepted;}
    size_t NumWrote() const {return num_wrote;}
};

ExampleMixer::ExampleMixer(
    std::string _mix_mode, const std::string& _distrib,
    NnetChainExampleWriter& _example_writer,
    std::string _scale_fst_algo, bool _swap_scales,
    bool _mix_ivect, bool _mix_feats,
    bool _max_super, int32_t _min_shift, int32_t _max_shift,
    float _fixed, size_t _buff_size, int32_t _frame_shift,
    bool _compress, bool _test_mode
):
    mix_mode(std::move(_mix_mode)), example_writer(_example_writer),
    scale_fst_algo(std::move(_scale_fst_algo)), swap_scales(_swap_scales),
    scale_boost(1.0f), scale_eps(1e-3f),
    mix_ivect(_mix_ivect), mix_feats(_mix_feats),
    max_super(_max_super), min_shift(_min_shift), max_shift(_max_shift),
    fixed(_fixed), buff_size(_buff_size), frame_shift(_frame_shift),
    exclude_names(), compress(_compress && !_test_mode), test_mode(_test_mode),
    rand_gen(), int_distrib(0, 100000), shift_distrib(_min_shift, _max_shift),
    float_distrib(0.0f, 1.0f), scale_distrib(rand_gen, _distrib), eg_to_egs(), egs_buffer(),
    scale_count(), shift_count(), num_mixed(0), num_untouched(0), num_accepted(0), num_wrote(0)
{
    rand_gen.seed(static_cast<unsigned int>(std::time(0)));
    const size_t indx = scale_fst_algo.find(':');
    if (indx != std::string::npos) {
        std::string value = scale_fst_algo.substr(indx + 1);
        if (value.empty()) {
            KALDI_ERR << "Invalid FST scaling parameners string format: \"" << scale_fst_algo << "\".";
        }
        std::vector<std::string> parts;
        boost::split(parts, value, boost::is_any_of(","), boost::token_compress_on);
        scale_fst_algo = scale_fst_algo.substr(0, indx);
        if (scale_fst_algo.empty() || (scale_fst_algo == "noscale")) {
            scale_eps = boost::lexical_cast<float>(parts.at(0));
        } else {
            scale_boost = boost::lexical_cast<float>(parts.at(0));
            if (parts.size() > 1) {
                scale_eps = boost::lexical_cast<float>(parts.at(1));
            }
        }
    }
    exclude_names.emplace_back(std::string("ivector"));
    KALDI_LOG << "mix_mode: " << mix_mode;
    KALDI_LOG << "distrib: " << _distrib;
    KALDI_LOG << "scale_fst_algo: \"" << scale_fst_algo << "\"";
    KALDI_LOG << "swap_scales: " << (swap_scales? "yes": "no");
    KALDI_LOG << "scale_boost: " << scale_boost;
    KALDI_LOG << "scale_eps: " << scale_eps;
    KALDI_LOG << "mix_ivect: " << (mix_ivect? "yes": "no");
    KALDI_LOG << "mix_feats: " << (mix_feats? "yes": "no");
    KALDI_LOG << "max_super: " << (max_super? "yes": "no");
    if (mix_mode == "shift") {
        KALDI_LOG << "min_shift: " << min_shift;
        KALDI_LOG << "max_shift: " << max_shift;
    }
    KALDI_LOG << "fixed: " << fixed;
    KALDI_LOG << "buff_size: " << buff_size;
    KALDI_LOG << "frame_shift: " << frame_shift;
    KALDI_LOG << "compress: " << (compress? "yes": "no");
    KALDI_LOG << "test_mode: " << (test_mode? "yes": "no");
}

const std::vector<Index>& ExampleMixer::FindIndexes(const std::string& _name, const std::vector<NnetIo>& _nnet_io) const {
    for (const auto & nnet_io : _nnet_io) {
        if (nnet_io.name == _name) {
            if ((nnet_io.indexes.size() != nnet_io.features.NumRows())) {
                KALDI_ERR << "Data indexes have wrong dimension " << nnet_io.indexes.size() << " (must be " << nnet_io.features.NumRows() << ").";
            }
            return nnet_io.indexes;
        }
    }
    KALDI_ERR << "Failed to find example indexes with name \"" << _name << "\".";
}

GeneralMatrix& ExampleMixer::FindFeatures(const std::string& _name, std::vector<NnetIo>& _nnet_io) const {
    for (auto & nnet_io : _nnet_io) {
        if (nnet_io.name == _name) {
            if ((nnet_io.indexes.size() != nnet_io.features.NumRows())) {
                KALDI_ERR << "Data indexes have wrong dimension " << nnet_io.indexes.size() << " (must be " << nnet_io.features.NumRows() << ").";
            }
            return nnet_io.features;
        }
    }
    KALDI_ERR << "Failed to find example features with name \"" << _name << "\".";
}

const GeneralMatrix& ExampleMixer::FindFeatures(const std::string& _name, const std::vector<NnetIo>& _nnet_io) const {
    for (const auto & nnet_io : _nnet_io) {
        if (nnet_io.name == _name) {
            if ((nnet_io.indexes.size() != nnet_io.features.NumRows())) {
                KALDI_ERR << "Data indexes have wrong dimension " << nnet_io.indexes.size() << " (must be " << nnet_io.features.NumRows() << ").";
            }
            return nnet_io.features;
        }
    }
    KALDI_ERR << "Failed to find example features with name \"" << _name << "\".";
}

GeneralMatrix* ExampleMixer::FindIVector(std::vector<NnetIo>& _nnet_io) const {
    for (auto & nnet_io : _nnet_io) {
        if (nnet_io.name == "ivector") {
            if ((nnet_io.indexes.size() != nnet_io.features.NumRows())) {
                KALDI_ERR << "I-vector indexes have wrong dimension " << nnet_io.indexes.size() << " (must be " << nnet_io.features.NumRows() << ").";
            }
            return &nnet_io.features;
        }
    }
    return nullptr;
}

const GeneralMatrix* ExampleMixer::FindIVector(const std::vector<NnetIo>& _nnet_io) const {
    for (const auto & nnet_io : _nnet_io) {
        if (nnet_io.name == "ivector") {
            if ((nnet_io.indexes.size() != nnet_io.features.NumRows())) {
                KALDI_ERR << "I-vector indexes have wrong dimension " << nnet_io.indexes.size() << " (must be " << nnet_io.features.NumRows() << ").";
            }
            return &nnet_io.features;
        }
    }
    return nullptr;
}

void ExampleMixer::CheckConsistence(const NnetChainExample& _example) const {
    if (_example.outputs.size() > 1) {
        KALDI_ERR << "Examples with multiple outputs are not supported.";
    }
    if (_example.outputs.front().name != "output") {
        KALDI_ERR << "Unexpected example output name \"" << _example.outputs.front().name << " (must be \"output\").";
    }
    const chain::Supervision& superv = _example.outputs.front().supervision;
    if (superv.num_sequences != 1) {
        KALDI_ERR << "Supervision has unexpected number of sequences " << superv.num_sequences << " (must be 1).";
    }
    if (_example.outputs.front().indexes.size() != superv.frames_per_sequence) {
        KALDI_ERR << "Unexpected number of supervision indexes " << _example.outputs.front().indexes.size() << " (must be " << superv.frames_per_sequence << ").";
    }
    if (_example.outputs.front().deriv_weights.Dim() != superv.frames_per_sequence) {
        KALDI_ERR << "Example output deriv_weights has unexpected dimension " << _example.outputs.front().deriv_weights.Dim() << " (must be " << superv.frames_per_sequence << ").";
    }
}

void ExampleMixer::CheckConsistence(const NnetChainExample& _example1, const NnetChainExample& _example2) const {
    CheckConsistence(_example1);
    CheckConsistence(_example2);
    const NnetChainSupervision& exam_sup1 = _example1.outputs.front();
    const NnetChainSupervision& exam_sup2 = _example2.outputs.front();
    if (exam_sup1.indexes != exam_sup2.indexes) {
        KALDI_ERR << "Examples have different indexes.";
    }
    const int32_t fp_seq1 = exam_sup1.supervision.frames_per_sequence;
    const int32_t fp_seq2 = exam_sup2.supervision.frames_per_sequence;
    if (fp_seq1 != fp_seq2) {
        KALDI_ERR << "Examples have different sequences lengths: " << fp_seq1 << " and " << fp_seq2 << ".";
    }
    if ((exam_sup1.supervision.fst.NumStates() == 0) && (exam_sup1.supervision.e2e_fsts.size() != exam_sup2.supervision.e2e_fsts.size())) {
        KALDI_ERR << "Examples have different number of e2e_fsts: " << exam_sup1.supervision.e2e_fsts.size() << " and " << exam_sup2.supervision.e2e_fsts.size() << ".";
    }
}

void ExampleMixer::ScaleGraph(fst_t& _target, float _scale) const {
    typedef fst::MutableArcIterator<fst_t> iter_t;
    const float scale = -(std::log(_scale) * scale_boost);
    for (state_t state = 0; state < _target.NumStates(); ++state) {
        for (iter_t iarc(&_target, state); !iarc.Done(); iarc.Next()) {
            const arc_t& arc = iarc.Value();
            if (_target.Final(arc.nextstate) != wght_t::Zero()) {
                iarc.SetValue(arc_t(arc.ilabel, arc.olabel, wght_t(arc.weight.Value() + scale), arc.nextstate));
            }
        }
    }
}

void ExampleMixer::UnionGraphs(const fst_t& _admixture, fst_t& _example) const {
    fst_t result;
    result.ReserveStates(_example.NumStates() + _admixture.NumStates());
    const state_t start = result.AddState();
    result.SetStart(start);
    std::vector<state_t> stmap(_example.NumStates(), -1);
    for (state_t state = 0; state < _example.NumStates(); ++state) {
        stmap[state] = (state == _example.Start())? start: result.AddState();
        const wght_t& wght = _example.Final(state);
        if (wght != wght_t::Zero()) {
            if (state == _example.Start()) {
                KALDI_WARN << "Final and start states are equal!";
            }
            result.SetFinal(stmap[state], wght);
        }
    }
    for (state_t state = 0; state < _example.NumStates(); ++state) {
        const state_t begin = (state == _example.Start())? start: stmap[state];
        for (iter_t iarc(_example, state); !iarc.Done(); iarc.Next()) {
            const arc_t& arc = iarc.Value();
            const state_t end = (arc.nextstate == _example.Start())? start: stmap[arc.nextstate];
            result.AddArc(begin, arc_t(arc.ilabel, arc.olabel, arc.weight, end));
        }
    }
    stmap.assign(_admixture.NumStates(), -1);
    for (state_t state = 0; state < _admixture.NumStates(); ++state) {
        stmap[state] = (state == _admixture.Start())? start: result.AddState();
        const wght_t& wght = _admixture.Final(state);
        if (wght != wght_t::Zero()) {
            if (state == _admixture.Start()) {
                KALDI_WARN << "Final and start states are equal!";
            }
            result.SetFinal(stmap[state], wght);
        }
    }
    for (state_t state = 0; state < _admixture.NumStates(); ++state) {
        const state_t begin = (state == _admixture.Start())? start: stmap[state];
        for (iter_t iarc(_admixture, state); !iarc.Done(); iarc.Next()) {
            const arc_t& arc = iarc.Value();
            const state_t end = (arc.nextstate == _admixture.Start())? start: stmap[arc.nextstate];
            result.AddArc(begin, arc_t(arc.ilabel, arc.olabel, arc.weight, end));
        }
    }
    _example = result;
}

void ExampleMixer::FuseGraphs(const fst_t& _admixture, float _admx_scale, fst_t& _example) const {
    if (_admx_scale < scale_eps) {
        return;
    } else if ((1.0f - _admx_scale) < scale_eps) {
        _example = _admixture;
        return;
    }
    if (scale_fst_algo.empty() || (scale_fst_algo == "noscale")) {
        UnionGraphs(_admixture, _example);
    } else if (scale_fst_algo == "default") {
        fst_t admixture(_admixture);
        if (swap_scales) {
            ScaleGraph(admixture, 1.0f - _admx_scale);
            ScaleGraph(_example, _admx_scale);
        } else {
            ScaleGraph(admixture, _admx_scale);
            ScaleGraph(_example, 1.0f - _admx_scale);
        }
        UnionGraphs(admixture, _example);
    } else if (scale_fst_algo == "balanced") {
        const double scale_norm = std::sqrt(_admx_scale * (1.0 - _admx_scale));
        const auto main_scale = (float)((1.0 - _admx_scale) / scale_norm);
        const auto admx_scale = (float)(_admx_scale / scale_norm);
        fst_t admixture(_admixture);
        if (swap_scales) {
            ScaleGraph(admixture, main_scale);
            ScaleGraph(_example, admx_scale);
        } else {
            ScaleGraph(admixture, admx_scale);
            ScaleGraph(_example, main_scale);
        }
        UnionGraphs(admixture, _example);
    } else {
        KALDI_ERR << "Unknown FST scaling algorithm ID: \"" << scale_fst_algo << "\".";
    }
}

void ExampleMixer::AdmixGlobal(const NnetChainExample& _admixture, float _admx_scale, ExamplePair& _example) {
    CheckConsistence(*_example.second, _admixture);
    const float exam_scale = 1.0f - _admx_scale;
    GeneralMatrix* ivector = FindIVector(_example.second->inputs);
    if (mix_ivect && (ivector != nullptr)) {
        KaldiMatrix ivec_main;
        ivector->GetMatrix(&ivec_main);
        ivec_main.Scale(exam_scale);
        KaldiMatrix ivec_admx;
        FindIVector(_admixture.inputs)->GetMatrix(&ivec_admx);
        ivec_main.AddMat(_admx_scale, ivec_admx);
        *ivector = ivec_main;
        if (compress) {
            ivector->Compress();
        }
    }
    if (mix_feats) {
        GeneralMatrix &features = FindFeatures("input", _example.second->inputs);
        KaldiMatrix feat_main;
        features.GetMatrix(&feat_main);
        feat_main.Scale(exam_scale);
        KaldiMatrix feat_admx;
        FindFeatures("input", _admixture.inputs).GetMatrix(&feat_admx);
        feat_main.AddMat(_admx_scale, feat_admx);
        features = feat_main;
        if (compress) {
            features.Compress();
        }
    }
    NnetChainSupervision& exam_nn_sup = _example.second->outputs.front();
    const NnetChainSupervision& admx_nn_sup = _admixture.outputs.front();
    if (max_super) {
        if (_admx_scale > 0.5f) {
            exam_nn_sup = admx_nn_sup;
        }
    } else {
        exam_nn_sup.deriv_weights.Scale(exam_scale);
        exam_nn_sup.deriv_weights.AddVec(_admx_scale, admx_nn_sup.deriv_weights);
        chain::Supervision &exam_sup = exam_nn_sup.supervision;
        const chain::Supervision &admx_sup = admx_nn_sup.supervision;
        exam_sup.weight = exam_sup.weight * exam_scale + admx_sup.weight * _admx_scale;
        if (admx_sup.fst.NumStates() == 0) {
            for (size_t i = 0; i < exam_sup.e2e_fsts.size(); ++i) {
                FuseGraphs(admx_sup.e2e_fsts.at(i), _admx_scale, exam_sup.e2e_fsts.at(i));
            }
        } else {
            FuseGraphs(admx_sup.fst, _admx_scale, exam_sup.fst);
        }
    }
    scale_count += _admx_scale;
}

void ExampleMixer::FlushGlobal(egs_buffer_t& _buffer) {
    egs_buffer_t buffer(_buffer.size());
    for (size_t i = 0; i < _buffer.size(); ++i) {
        const ExamplePair& pair = _buffer[i];
        buffer[i] = std::make_pair(pair.first, ExamplePtr(new NnetChainExample(*pair.second)));
    }
    for (size_t i = 0; i < _buffer.size(); ++i) {
        ExamplePair& example = _buffer.at(i);
        const bool mixup = ((float_distrib(rand_gen) > fixed) && (_buffer.size() > 1));
        if (mixup) {
            size_t indx = int_distrib(rand_gen) % buffer.size();
            if (indx == i) {
                indx = (indx == (buffer.size() - 1)) ? indx - 1 : indx + 1;
            }
            const float scale = scale_distrib();
            AdmixGlobal(*buffer.at(indx).second, scale, example);
            ++num_mixed;
        } else {
            ++num_untouched;
        }
        if (frame_shift != 0) {
            ShiftChainExampleTimes(frame_shift, exclude_names, example.second.get());
        }
        example_writer.Write(example.first, *example.second);
        ++num_wrote;
    }
}

void ExampleMixer::ShiftAndMixup(ExamplePair& _example) {
    CheckConsistence(*_example.second);
    GeneralMatrix& features = FindFeatures("input", _example.second->inputs);
    KaldiMatrix feat_main;
    features.GetMatrix(&feat_main);
    KaldiMatrix feat_admx(feat_main.NumRows(), feat_main.NumCols());
    const int32_t shift = ((float_distrib(rand_gen) > 0.5)? 1: -1) * shift_distrib(rand_gen);
    for (int32_t i = 0; i < feat_main.NumRows(); ++i) {
        int32_t admx_indx = i + shift;
        if (admx_indx < 0) {
            admx_indx = 0;
        } else if (admx_indx >= feat_main.NumRows()) {
            admx_indx = feat_main.NumRows() - 1;
        }
        feat_admx.Row(i).CopyRowFromMat(feat_main, admx_indx);
    }
    const float admx_scale = scale_distrib();
    const float exam_scale = 1.0f - admx_scale;
    feat_main.Scale(exam_scale);
    feat_main.AddMat(admx_scale, feat_admx);
    features = feat_main;
    if (compress) {
        features.Compress();
    }
    scale_count += admx_scale;
    shift_count += shift;
}

void ExampleMixer::AcceptExample(ExamplePair& _example) {
    if (mix_mode == "global") {
        egs_buffer_t &buffer = eg_to_egs[_example.second.get()];
        if (buffer.empty()) {
            buffer.reserve(buff_size);
        }
        buffer.push_back(_example);
        if (buffer.size() == buff_size) {
            egs_buffer_t buff_copy = buffer;
            eg_to_egs.erase(_example.second.get());
            FlushGlobal(buff_copy);
        }
    } else if (mix_mode == "shift") {
        const bool mixup = (float_distrib(rand_gen) > fixed);
        if (mixup) {
            ShiftAndMixup(_example);
            ++num_mixed;
        } else {
            ++num_untouched;
        }
        example_writer.Write(_example.first, *_example.second);
        ++num_wrote;
    } else {
        KALDI_ERR << "Unknown mixup mode: \"" << mix_mode << "\"";
    }
    ++num_accepted;
}

void ExampleMixer::Finish() {
    if (mix_mode == "global") {
        while (!eg_to_egs.empty()) {
            egs_buffer_t buffer = eg_to_egs.begin()->second;
            eg_to_egs.erase(eg_to_egs.begin());
            if (!buffer.empty()) {
                FlushGlobal(buffer);
            }
        }
    }
}

} }

bool AsBool(const char* _value) {
    if (_value == nullptr) {
        KALDI_ERR << "Pointer to bool string is nullptr.";
    }
    std::string value(_value);
    boost::to_lower(value);
    return (value == "true") || (value == "yes") || (value == "on") || (value == "1");
}

// --test-mode=true --distrib=uniform:0.1,0.4 --frame-shift=1 ark:/media/work/coding/data/mgb3/train_data/cegs.10.ark ark:/dev/null
// --test-mode=true --distrib=uniform:0.1,0.4 --frame-shift=1 ark:/mnt/TOSHIBA/khokhlov/coding/temp/i-vect/temp/mixup/cegs.10.ark ark:/dev/null
// --test-mode=true --distrib=uniform:0.0,0.4 --scale-fst-algo=noscale:0.2 --frame-shift=1 ark:/mnt/TOSHIBA/khokhlov/coding/temp/i-vect/temp/mixup/cegs.10.ark ark:/dev/null
// --test-mode=true --distrib=uniform:0.0,0.4 --scale-fst-algo=default --swap-scales=true --frame-shift=1 ark:/mnt/TOSHIBA/khokhlov/coding/temp/i-vect/temp/mixup/cegs.10.ark ark:/dev/null
// --test-mode=true --distrib=uniform:0.0,0.4 --scale-fst-algo=default:1.5 --frame-shift=1 ark:/mnt/TOSHIBA/khokhlov/coding/temp/i-vect/temp/mixup/cegs.10.ark ark:/dev/null
// --test-mode=true --distrib=uniform:0.0,0.1 --scale-fst-algo=default:1.5,0.1 --frame-shift=1 ark:/mnt/TOSHIBA/khokhlov/coding/temp/i-vect/temp/mixup/cegs.10.ark ark:/dev/null
// --test-mode=true --distrib=uniform:0.1,0.7 --scale-fst-algo=balanced --frame-shift=1 ark:/mnt/TOSHIBA/khokhlov/coding/temp/i-vect/temp/mixup/cegs.10.ark ark:/dev/null
// --test-mode=true --distrib=uniform:0.0,1.0 --scale-fst-algo=balanced:1.5 --frame-shift=1 ark:/mnt/TOSHIBA/khokhlov/coding/temp/i-vect/temp/mixup/cegs.10.ark ark:/dev/null
// --test-mode=true --distrib=uniform:0.9,1.0 --scale-fst-algo=balanced:1.5,0.2 --frame-shift=1 ark:/mnt/TOSHIBA/khokhlov/coding/temp/i-vect/temp/mixup/cegs.10.ark ark:/dev/null
// --test-mode=true --distrib=beta:0.5 --max-super=false --frame-shift=1 ark:/mnt/TOSHIBA/khokhlov/coding/temp/i-vect/temp/mixup/cegs.10.ark ark:/dev/null
// --test-mode=true --distrib=beta:0.5 --max-super=true --frame-shift=1 ark:/mnt/TOSHIBA/khokhlov/coding/temp/i-vect/temp/mixup/cegs.10.ark ark:/dev/null

// --mix-mode=shift --test-mode=true --distrib=uniform:0.1,0.4 --min-shift=2 --max-shift=5 --frame-shift=1 ark:/mnt/TOSHIBA/khokhlov/coding/temp/i-vect/temp/mixup/cegs.10.ark ark,t:/mnt/TOSHIBA/khokhlov/coding/temp/i-vect/temp/mixup/cegs.10.mix.ark
// --mix-mode=global --test-mode=true --distrib=beta:0.01 --frame-shift=1 ark:/media/work/coding/data/mgb3/train_data/cegs.10.ark ark:/media/work/coding/data/mgb3/train_data/cegs.10.mix.ark
// --mix-mode=global --distrib=beta2:0.5 --buff_size=8000 ark:/media/work/coding/data/mgb3/train_data/cegs.1.ark ark:/dev/null
// --mix-mode=global --distrib=beta2:0.5 --buff_size=8000 ark:/media/work/coding/data/mgb3/train_data/cegs.1.ark ark,t:/media/work/coding/data/mgb3/train_data/cegs.1.mix.ark

// --frame-shift=1 ark:/mnt/TOSHIBA/khokhlov/coding/mixup_github/temp/cegs.1.ark ark:/dev/null

/*

work dir: /mnt/diskD/khokhlov/temp/mixup
ark:cegs.33.ark ark:/dev/null
--mix-ivect=false ark:cegs.33.ark ark:/dev/null
--mix-feats=false ark:cegs.33.ark ark:/dev/null

*/

int main(int argc, char *argv[]) {
    try {
        using namespace kaldi;
        using namespace kaldi::nnet3;

        const char *usage =
            "Usage:  nnet3-chain-mixup-egs [options] <egs-rspecifier> <egs-wspecifier>\n"
            "\n"
            "e.g.\n"
            "nnet3-chain-mixup-egs ark:train.egs ark:mixup.egs\n";
        ParseOptions po(usage);

        std::string mix_mode("global");
        const char* env_var = getenv("MIXUP_MIX_MODE");
        if (env_var != nullptr) {
            mix_mode = env_var;
        }
        po.Register("mix-mode", &mix_mode, R"(Mixup mode ("global", "shift") MIXUP_MIX_MODE)");

        std::string distrib("uniform:0.0,0.5");
        env_var = getenv("MIXUP_DISTRIB");
        if (env_var != nullptr) {
            distrib = env_var;
        }
        po.Register("distrib", &distrib, R"(Mixup scaling factors distribution ("uniform:min,max", "beta:alpha", "beta2:alpha") MIXUP_DISTRIB)");

        std::string scale_fst_algo;
        env_var = getenv("MIXUP_SCALE_FST_ALGO");
        if (env_var != nullptr) {
            scale_fst_algo = env_var;
        }
        po.Register("scale-fst-algo", &scale_fst_algo, R"(Scale supervision FSTs algorithm ("default[:scale[,eps]]", "balanced[:scale[,eps]]") MIXUP_SCALE_FST_ALGO)");

        bool swap_scales = false;
        env_var = getenv("MIXUP_SWAP_SCALES");
        if (env_var != nullptr) {
            swap_scales = AsBool(env_var);
        }
        po.Register("swap-scales", &swap_scales, "Swap supervision FST scales MIXUP_SWAP_SCALES");

        bool mix_ivect = true;
        env_var = getenv("MIXUP_MIX_IVECT");
        if (env_var != nullptr) {
            mix_ivect = AsBool(env_var);
        }
        po.Register("mix-ivect", &mix_ivect, "Make i-vectors mixtures (MIXUP_MIX_IVECT)");

        bool mix_feats = true;
        env_var = getenv("MIXUP_MIX_FEATS");
        if (env_var != nullptr) {
            mix_feats = AsBool(env_var);
        }
        po.Register("mix-feats", &mix_feats, "Make features mixtures (MIXUP_MIX_FEATS)");

        bool max_super = false;
        env_var = getenv("MIXUP_MAX_SUPER");
        if (env_var != nullptr) {
            max_super = AsBool(env_var);
        }
        po.Register("max-super", &max_super, "Get supervision from example with maximum scale MIXUP_MAX_SUPER");

        int32_t min_shift = 1;
        env_var = getenv("MIXUP_MIN_SHIFT");
        if (env_var != nullptr) {
            min_shift = boost::lexical_cast<int32_t>(env_var);
        }
        po.Register("min-shift", &min_shift, "Minimum sequence shift size (shift mode) MIXUP_MIN_SHIFT");

        int32_t max_shift = 3;
        env_var = getenv("MIXUP_MAX_SHIFT");
        if (env_var != nullptr) {
            max_shift = boost::lexical_cast<int32_t>(env_var);
        }
        po.Register("max-shift", &max_shift, "Maximum sequence shift size (shift mode) MIXUP_MAX_SHIFT");

        float fixed = 0.10;
        env_var = getenv("MIXUP_FIXED");
        if (env_var != nullptr) {
            fixed = boost::lexical_cast<float>(env_var);
        }
        po.Register("fixed", &fixed, "The portion of the data to leave untouched MIXUP_FIXED");

        int32_t buff_size = 500;
        env_var = getenv("MIXUP_BUFF_SIZE");
        if (env_var != nullptr) {
            buff_size = boost::lexical_cast<int32_t>(env_var);
        }
        po.Register("buff-size", &buff_size, "Buffer size for data shuffling (global mode) MIXUP_BUFF_SIZE");

        int32_t frame_shift = 0;
        env_var = getenv("MIXUP_FRAME_SHIFT");
        if (env_var != nullptr) {
            frame_shift = boost::lexical_cast<int32_t>(env_var);
        }
        po.Register("frame-shift", &frame_shift, "Allows you to shift time values in the supervision data (excluding iVector data) - useful in augmenting data.  Note, the outputs will remain at the closest exact multiples of the frame subsampling factor MIXUP_FRAME_SHIFT");

        int32_t compress = 0;
        env_var = getenv("MIXUP_COMPRESS");
        if (env_var != nullptr) {
            compress = boost::lexical_cast<int32_t>(env_var);
        }
        po.Register("compress", &compress, "Compress features and i-vectors MIXUP_COMPRESS");

        bool test_mode = false;
        po.Register("test-mode", &test_mode, "Self testing mode");

        po.Read(argc, argv);
        if (po.NumArgs() < 2) {
            po.PrintUsage();
            exit(1);
        }
        if (mix_mode == "shift") {
            if (min_shift < 1) {
                KALDI_ERR << "min_num must be greater or equal 1";
            }
        }
        std::string examples_rspecifier = po.GetArg(1);
        std::string examples_wspecifier = po.GetArg(2);

        SequentialNnetChainExampleReader example_reader(examples_rspecifier);
        NnetChainExampleWriter example_writer(examples_wspecifier);

        ExampleMixer mixer(
            mix_mode, distrib, example_writer, scale_fst_algo, swap_scales, mix_ivect, mix_feats, max_super,
            min_shift, max_shift, fixed, (size_t) buff_size, frame_shift, compress != 0, test_mode != 0
        );
        size_t num_read = 0;
        for (; !example_reader.Done(); example_reader.Next(), num_read++) {
            const NnetChainExample& example = example_reader.Value();
            ExamplePair ex_pair(example_reader.Key(), ExamplePtr(new NnetChainExample(example)));
            mixer.AcceptExample(ex_pair);
        }
        mixer.Finish();

        const FloatCounter& scale_count = mixer.ScaleCount();
        const IntCounter& shift_count = mixer.ShiftCount();
        const size_t num_mixed = mixer.NumMixed();
        const size_t num_untouched = mixer.NumUntouched();
        const size_t num_accepted = mixer.NumAccepted();
        const size_t num_wrote = mixer.NumWrote();

        KALDI_LOG << "Num read: " << num_read << " examples";
        KALDI_LOG << "Num accepted: " << num_accepted << " examples ( " << (100.0 * num_accepted / num_read) << " % of read )";
        KALDI_LOG << "Num wrote: " << num_wrote << " examples ( " << (100.0 * num_wrote / num_read) << " % of read )";
        KALDI_LOG << "Num mixed: " << num_mixed << " examples ( " << (100.0 * num_mixed / num_accepted) << " % of accepted )";
        KALDI_LOG << "Num untouched: " << num_untouched << " examples ( " << (100.0 * num_untouched / num_accepted) << " % of accepted )";
        KALDI_LOG << "Average scale: " << (boost::format("%.4f ( min: %.4f, max: %.4f )") % scale_count.Average() % scale_count.Minimum() % scale_count.Maximum()).str();
        if (mix_mode == "shift") {
            KALDI_LOG << "Average shift: " << (boost::format("%.4f ( min: %.4f, max: %.4f )") % shift_count.Average() % shift_count.Minimum() % shift_count.Maximum()).str();
        }

        return (num_wrote == 0 ? 1 : 0);
    } catch(const std::exception &e) {
        std::cerr << e.what() << '\n';
        return -1;
    }
}


/*
 Copyright (C) 2016 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 <ored/portfolio/builders/capfloorediborleg.hpp>
#include <ored/utilities/log.hpp>

#include <boost/make_shared.hpp>

namespace ore {
namespace data {

boost::shared_ptr<FloatingRateCouponPricer> CapFlooredIborLegEngineBuilder::engineImpl(const Currency& ccy) {

    // Check if we already have a pricer for this ccy
    const string& ccyCode = ccy.code();
    if (engines_.find(ccyCode) == engines_.end()) {
        Handle<YieldTermStructure> yts = market_->discountCurve(ccyCode, configuration(MarketContext::pricing));
        QL_REQUIRE(!yts.empty(), "engineFactory error: yield term structure not found for currency " << ccyCode);
        Handle<OptionletVolatilityStructure> ovs = market_->capFloorVol(ccyCode, configuration(MarketContext::pricing));
        BlackIborCouponPricer::TimingAdjustment timingAdjustment = BlackIborCouponPricer::Black76;
        boost::shared_ptr<SimpleQuote> correlation = boost::make_shared<SimpleQuote>(1.0);
        // for backwards compatibility we do not require the additional timing adjustment fields
        if (engineParameters_.find("TimingAdjustment") != engineParameters_.end()) {
            string adjStr = engineParameters_.at("TimingAdjustment");
            if (adjStr == "Black76")
                timingAdjustment = BlackIborCouponPricer::Black76;
            else if (adjStr == "BivariateLognormal")
                timingAdjustment = BlackIborCouponPricer::BivariateLognormal;
            else {
                QL_FAIL("timing adjustment parameter (" << adjStr << ") not recognised.");
            }
            correlation->setValue(parseReal(engineParameters_.at("Correlation")));
        }
        boost::shared_ptr<FloatingRateCouponPricer> pricer =
            boost::make_shared<BlackIborCouponPricer>(ovs, timingAdjustment, Handle<Quote>(correlation));
        engines_[ccyCode] = pricer;
    }

    // Return the cached pricer
    return engines_[ccyCode];
}
} // namespace data
} // namespace ore


// ====================================================================
// This file is part of FlexibleSUSY.
//
// FlexibleSUSY is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published
// by the Free Software Foundation, either version 3 of the License,
// or (at your option) any later version.
//
// FlexibleSUSY is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with FlexibleSUSY.  If not, see
// <http://www.gnu.org/licenses/>.
// ====================================================================

// File generated at Thu 10 May 2018 14:43:46

/**
 * @file HSSUSY_mass_eigenstates.hpp
 *
 * @brief contains class for model with routines needed to solve boundary
 *        value problem using the two_scale solver by solving EWSB
 *        and determine the pole masses and mixings
 *
 * This file was generated at Thu 10 May 2018 14:43:46 with FlexibleSUSY
 * 2.0.1 (git commit: unknown) and SARAH 4.12.2 .
 */

#ifndef HSSUSY_MASS_EIGENSTATES_H
#define HSSUSY_MASS_EIGENSTATES_H

#include "HSSUSY_info.hpp"
#include "HSSUSY_physical.hpp"
#include "HSSUSY_soft_parameters.hpp"
#include "loop_corrections.hpp"
#include "threshold_corrections.hpp"
#include "error.hpp"
#include "problems.hpp"
#include "config.h"

#include <iosfwd>
#include <memory>
#include <string>

#include <Eigen/Core>

namespace flexiblesusy {

class HSSUSY_ewsb_solver_interface;
/**
 * @class HSSUSY_mass_eigenstates
 * @brief model class with routines for determing masses and mixinga and EWSB
 */
class HSSUSY_mass_eigenstates : public HSSUSY_soft_parameters {
public:
   explicit HSSUSY_mass_eigenstates(const HSSUSY_input_parameters& input_ = HSSUSY_input_parameters());
   HSSUSY_mass_eigenstates(const HSSUSY_mass_eigenstates&) = default;
   HSSUSY_mass_eigenstates(HSSUSY_mass_eigenstates&&) = default;
   virtual ~HSSUSY_mass_eigenstates() = default;
   HSSUSY_mass_eigenstates& operator=(const HSSUSY_mass_eigenstates&) = default;
   HSSUSY_mass_eigenstates& operator=(HSSUSY_mass_eigenstates&&) = default;

   /// number of EWSB equations
   static const int number_of_ewsb_equations = 1;

   void calculate_DRbar_masses();
   void calculate_pole_masses();
   void check_pole_masses_for_tachyons();
   virtual void clear() override;
   void clear_DRbar_parameters();
   Eigen::ArrayXd get_DRbar_masses() const;
   Eigen::ArrayXd get_DRbar_masses_and_mixings() const;
   Eigen::ArrayXd get_extra_parameters() const;
   void do_calculate_sm_pole_masses(bool);
   bool do_calculate_sm_pole_masses() const;
   void do_calculate_bsm_pole_masses(bool);
   bool do_calculate_bsm_pole_masses() const;
   void do_force_output(bool);
   bool do_force_output() const;
   void reorder_DRbar_masses();
   void reorder_pole_masses();
   void set_ewsb_iteration_precision(double);
   void set_ewsb_loop_order(int);
   void set_loop_corrections(const Loop_corrections&);
   const Loop_corrections& get_loop_corrections() const;
   void set_threshold_corrections(const Threshold_corrections&);
   const Threshold_corrections& get_threshold_corrections() const;
   void set_DRbar_masses(const Eigen::ArrayXd&);
   void set_DRbar_masses_and_mixings(const Eigen::ArrayXd&);
   void set_extra_parameters(const Eigen::ArrayXd&);
   void set_pole_mass_loop_order(int);
   int get_pole_mass_loop_order() const;
   void set_physical(const HSSUSY_physical&);
   double get_ewsb_iteration_precision() const;
   double get_ewsb_loop_order() const;
   const HSSUSY_physical& get_physical() const;
   HSSUSY_physical& get_physical();
   const Problems& get_problems() const;
   Problems& get_problems();
   void set_ewsb_solver(const std::shared_ptr<HSSUSY_ewsb_solver_interface>&);
   int solve_ewsb_tree_level();
   int solve_ewsb_one_loop();
   int solve_ewsb();            ///< solve EWSB at ewsb_loop_order level

   void calculate_spectrum();
   void clear_problems();
   std::string name() const;
   void run_to(double scale, double eps = -1.0) override;
   void print(std::ostream& out = std::cerr) const override;
   void set_precision(double);
   double get_precision() const;

   double get_MVG() const { return MVG; }
   double get_MHp() const { return MHp; }
   const Eigen::Array<double,3,1>& get_MFv() const { return MFv; }
   double get_MFv(int i) const { return MFv(i); }
   double get_MAh() const { return MAh; }
   double get_Mhh() const { return Mhh; }
   const Eigen::Array<double,3,1>& get_MFd() const { return MFd; }
   double get_MFd(int i) const { return MFd(i); }
   const Eigen::Array<double,3,1>& get_MFu() const { return MFu; }
   double get_MFu(int i) const { return MFu(i); }
   const Eigen::Array<double,3,1>& get_MFe() const { return MFe; }
   double get_MFe(int i) const { return MFe(i); }
   double get_MVWp() const { return MVWp; }
   double get_MVP() const { return MVP; }
   double get_MVZ() const { return MVZ; }

   const Eigen::Matrix<std::complex<double>,3,3>& get_Vd() const { return Vd; }
   std::complex<double> get_Vd(int i, int k) const { return Vd(i,k); }
   const Eigen::Matrix<std::complex<double>,3,3>& get_Ud() const { return Ud; }
   std::complex<double> get_Ud(int i, int k) const { return Ud(i,k); }
   const Eigen::Matrix<std::complex<double>,3,3>& get_Vu() const { return Vu; }
   std::complex<double> get_Vu(int i, int k) const { return Vu(i,k); }
   const Eigen::Matrix<std::complex<double>,3,3>& get_Uu() const { return Uu; }
   std::complex<double> get_Uu(int i, int k) const { return Uu(i,k); }
   const Eigen::Matrix<std::complex<double>,3,3>& get_Ve() const { return Ve; }
   std::complex<double> get_Ve(int i, int k) const { return Ve(i,k); }
   const Eigen::Matrix<std::complex<double>,3,3>& get_Ue() const { return Ue; }
   std::complex<double> get_Ue(int i, int k) const { return Ue(i,k); }
   const Eigen::Matrix<double,2,2>& get_ZZ() const { return ZZ; }
   double get_ZZ(int i, int k) const { return ZZ(i,k); }

   double get_mass_matrix_VG() const;
   void calculate_MVG();
   double get_mass_matrix_Hp() const;
   void calculate_MHp();
   Eigen::Matrix<double,3,3> get_mass_matrix_Fv() const;
   void calculate_MFv();
   double get_mass_matrix_Ah() const;
   void calculate_MAh();
   double get_mass_matrix_hh() const;
   void calculate_Mhh();
   Eigen::Matrix<double,3,3> get_mass_matrix_Fd() const;
   void calculate_MFd();
   Eigen::Matrix<double,3,3> get_mass_matrix_Fu() const;
   void calculate_MFu();
   Eigen::Matrix<double,3,3> get_mass_matrix_Fe() const;
   void calculate_MFe();
   double get_mass_matrix_VWp() const;
   void calculate_MVWp();
   Eigen::Matrix<double,2,2> get_mass_matrix_VPVZ() const;
   void calculate_MVPVZ();

   double get_ewsb_eq_hh_1() const;

   double CphhHpconjHp() const;
   double CpbargWpgZHp() const;
   double CpbargZgWpconjHp() const;
   double CpbargWpCgZconjHp() const;
   double CpbargZgWpCHp() const;
   double CpconjHpVPVWp() const;
   double CpconjHpVWpVZ() const;
   double CpAhAhHpconjHp() const;
   double CphhhhHpconjHp() const;
   double CpHpHpconjHpconjHp() const;
   std::complex<double> CpAhconjHpVWp() const;
   double CphhconjHpVWp() const;
   double CpHpconjHpVP() const;
   double CpHpconjHpVZ() const;
   double CpHpconjHpconjVWpVWp() const;
   std::complex<double> CpHpconjHpVZVZ() const;
   std::complex<double> CpbarFdFuconjHpPR(int gI1, int gI2) const;
   std::complex<double> CpbarFdFuconjHpPL(int gI1, int gI2) const;
   double CpbarFeFvconjHpPR(int , int ) const;
   std::complex<double> CpbarFeFvconjHpPL(int gI1, int gI2) const;
   double CpAhAhhh() const;
   std::complex<double> CpbargWpgWpAh() const;
   std::complex<double> CpbargWpCgWpCAh() const;
   double CpAhAhAhAh() const;
   double CpAhAhhhhh() const;
   std::complex<double> CpAhhhVZ() const;
   std::complex<double> CpAhHpconjVWp() const;
   double CpAhAhconjVWpVWp() const;
   std::complex<double> CpAhAhVZVZ() const;
   std::complex<double> CpbarFdFdAhPR(int gI1, int gI2) const;
   std::complex<double> CpbarFdFdAhPL(int gI1, int gI2) const;
   std::complex<double> CpbarFeFeAhPR(int gI1, int gI2) const;
   std::complex<double> CpbarFeFeAhPL(int gI1, int gI2) const;
   std::complex<double> CpbarFuFuAhPR(int gI1, int gI2) const;
   std::complex<double> CpbarFuFuAhPL(int gI1, int gI2) const;
   double Cphhhhhh() const;
   double CphhVZVZ() const;
   double CphhconjVWpVWp() const;
   double CpbargWpgWphh() const;
   double CpbargWpCgWpChh() const;
   double CpbargZgZhh() const;
   double Cphhhhhhhh() const;
   double CphhHpconjVWp() const;
   double CphhhhconjVWpVWp() const;
   std::complex<double> CphhhhVZVZ() const;
   std::complex<double> CpbarFdFdhhPR(int gI1, int gI2) const;
   std::complex<double> CpbarFdFdhhPL(int gI1, int gI2) const;
   std::complex<double> CpbarFeFehhPR(int gI1, int gI2) const;
   std::complex<double> CpbarFeFehhPL(int gI1, int gI2) const;
   std::complex<double> CpbarFuFuhhPR(int gI1, int gI2) const;
   std::complex<double> CpbarFuFuhhPL(int gI1, int gI2) const;
   std::complex<double> CpVGVGVG() const;
   std::complex<double> CpbargGgGVG() const;
   double CpbarFdFdVGPL(int gI1, int gI2) const;
   double CpbarFdFdVGPR(int gI1, int gI2) const;
   double CpbarFuFuVGPL(int gI1, int gI2) const;
   double CpbarFuFuVGPR(int gI1, int gI2) const;
   double CpVGVGVGVG1() const;
   double CpVGVGVGVG2() const;
   double CpVGVGVGVG3() const;
   double CpHpconjVWpVP() const;
   double CpbargWpgWpVP() const;
   double CpbargWpCgWpCVP() const;
   std::complex<double> CpHpconjHpVPVP() const;
   double CpconjVWpVPVWp() const;
   double CpbarFdFdVPPL(int gI1, int gI2) const;
   double CpbarFdFdVPPR(int gI1, int gI2) const;
   double CpbarFeFeVPPL(int gI1, int gI2) const;
   double CpbarFeFeVPPR(int gI1, int gI2) const;
   double CpbarFuFuVPPL(int gI1, int gI2) const;
   double CpbarFuFuVPPR(int gI1, int gI2) const;
   double CpconjVWpVPVPVWp3() const;
   double CpconjVWpVPVPVWp1() const;
   double CpconjVWpVPVPVWp2() const;
   double CpHpconjVWpVZ() const;
   double CpbargWpgWpVZ() const;
   double CpbargWpCgWpCVZ() const;
   double CpconjVWpVWpVZ() const;
   double CpbarFdFdVZPL(int gI1, int gI2) const;
   double CpbarFdFdVZPR(int gI1, int gI2) const;
   double CpbarFeFeVZPL(int gI1, int gI2) const;
   double CpbarFeFeVZPR(int gI1, int gI2) const;
   double CpbarFuFuVZPL(int gI1, int gI2) const;
   double CpbarFuFuVZPR(int gI1, int gI2) const;
   double CpbarFvFvVZPL(int gI1, int gI2) const;
   double CpbarFvFvVZPR(int , int ) const;
   double CpconjVWpVWpVZVZ1() const;
   double CpconjVWpVWpVZVZ2() const;
   double CpconjVWpVWpVZVZ3() const;
   double CpbargPgWpconjVWp() const;
   double CpbargWpCgPconjVWp() const;
   double CpbargWpCgZconjVWp() const;
   double CpbargZgWpconjVWp() const;
   std::complex<double> CpbarFdFuconjVWpPL(int gI1, int gI2) const;
   double CpbarFdFuconjVWpPR(int , int ) const;
   std::complex<double> CpbarFeFvconjVWpPL(int gI1, int gI2) const;
   double CpbarFeFvconjVWpPR(int , int ) const;
   double CpconjVWpconjVWpVWpVWp2() const;
   double CpconjVWpconjVWpVWpVWp1() const;
   double CpconjVWpconjVWpVWpVWp3() const;
   std::complex<double> CpbarUFdFdAhPL(int gO2, int gI1) const;
   std::complex<double> CpbarUFdFdAhPR(int gO1, int gI1) const;
   std::complex<double> CpbarUFdFdhhPL(int gO2, int gI2) const;
   std::complex<double> CpbarUFdFdhhPR(int gO1, int gI2) const;
   std::complex<double> CpbarUFdFdVGPR(int gO2, int gI2) const;
   std::complex<double> CpbarUFdFdVGPL(int gO1, int gI2) const;
   std::complex<double> CpbarUFdFdVPPR(int gO2, int gI2) const;
   std::complex<double> CpbarUFdFdVPPL(int gO1, int gI2) const;
   std::complex<double> CpbarUFdFdVZPR(int gO2, int gI2) const;
   std::complex<double> CpbarUFdFdVZPL(int gO1, int gI2) const;
   std::complex<double> CpbarUFdFuconjHpPL(int gO2, int gI2) const;
   std::complex<double> CpbarUFdFuconjHpPR(int gO1, int gI2) const;
   double CpbarUFdFuconjVWpPR(int , int ) const;
   std::complex<double> CpbarUFdFuconjVWpPL(int gO1, int gI2) const;
   std::complex<double> CpbarUFuFuAhPL(int gO2, int gI1) const;
   std::complex<double> CpbarUFuFuAhPR(int gO1, int gI1) const;
   std::complex<double> CpbarUFuFdHpPL(int gO2, int gI2) const;
   std::complex<double> CpbarUFuFdHpPR(int gO1, int gI2) const;
   double CpbarUFuFdVWpPR(int , int ) const;
   std::complex<double> CpbarUFuFdVWpPL(int gO1, int gI2) const;
   std::complex<double> CpbarUFuFuhhPL(int gO2, int gI2) const;
   std::complex<double> CpbarUFuFuhhPR(int gO1, int gI2) const;
   std::complex<double> CpbarUFuFuVGPR(int gO2, int gI2) const;
   std::complex<double> CpbarUFuFuVGPL(int gO1, int gI2) const;
   std::complex<double> CpbarUFuFuVPPR(int gO2, int gI2) const;
   std::complex<double> CpbarUFuFuVPPL(int gO1, int gI2) const;
   std::complex<double> CpbarUFuFuVZPR(int gO2, int gI2) const;
   std::complex<double> CpbarUFuFuVZPL(int gO1, int gI2) const;
   std::complex<double> CpbarUFeFeAhPL(int gO2, int gI1) const;
   std::complex<double> CpbarUFeFeAhPR(int gO1, int gI1) const;
   std::complex<double> CpbarUFeFehhPL(int gO2, int gI2) const;
   std::complex<double> CpbarUFeFehhPR(int gO1, int gI2) const;
   std::complex<double> CpbarUFeFeVPPR(int gO2, int gI2) const;
   std::complex<double> CpbarUFeFeVPPL(int gO1, int gI2) const;
   std::complex<double> CpbarUFeFeVZPR(int gO2, int gI2) const;
   std::complex<double> CpbarUFeFeVZPL(int gO1, int gI2) const;
   std::complex<double> CpbarUFeFvconjHpPL(int gO2, int gI2) const;
   double CpbarUFeFvconjHpPR(int , int ) const;
   double CpbarUFeFvconjVWpPR(int , int ) const;
   double CpbarUFeFvconjVWpPL(int gO1, int gI2) const;
   double CpbarFvFeHpPL(int , int ) const;
   std::complex<double> CpbarFvFeHpPR(int gO1, int gI2) const;
   double CpbarFvFeVWpPR(int , int ) const;
   std::complex<double> CpbarFvFeVWpPL(int gO1, int gI2) const;
   std::complex<double> CpbarFuFdHpPL(int gO2, int gI2) const;
   std::complex<double> CpbarFuFdHpPR(int gO1, int gI2) const;
   double CpbarFuFdVWpPR(int , int ) const;
   std::complex<double> CpbarFuFdVWpPL(int gO1, int gI2) const;
   std::complex<double> self_energy_Hp_1loop(double p ) const;
   std::complex<double> self_energy_Ah_1loop(double p ) const;
   std::complex<double> self_energy_hh_1loop(double p ) const;
   std::complex<double> self_energy_VG_1loop(double p ) const;
   std::complex<double> self_energy_VP_1loop(double p ) const;
   std::complex<double> self_energy_VZ_1loop(double p ) const;
   std::complex<double> self_energy_VWp_1loop(double p ) const;
   std::complex<double> self_energy_Fd_1loop_1(double p , int gO1, int gO2) const;
   Eigen::Matrix<std::complex<double>,3,3> self_energy_Fd_1loop_1(double p) const;
   std::complex<double> self_energy_Fd_1loop_PR(double p , int gO1, int gO2) const;
   Eigen::Matrix<std::complex<double>,3,3> self_energy_Fd_1loop_PR(double p) const;
   std::complex<double> self_energy_Fd_1loop_PL(double p , int gO1, int gO2) const;
   Eigen::Matrix<std::complex<double>,3,3> self_energy_Fd_1loop_PL(double p) const;
   std::complex<double> self_energy_Fu_1loop_1(double p , int gO1, int gO2) const;
   Eigen::Matrix<std::complex<double>,3,3> self_energy_Fu_1loop_1(double p) const;
   std::complex<double> self_energy_Fu_1loop_PR(double p , int gO1, int gO2) const;
   Eigen::Matrix<std::complex<double>,3,3> self_energy_Fu_1loop_PR(double p) const;
   std::complex<double> self_energy_Fu_1loop_PL(double p , int gO1, int gO2) const;
   Eigen::Matrix<std::complex<double>,3,3> self_energy_Fu_1loop_PL(double p) const;
   std::complex<double> self_energy_Fe_1loop_1(double p , int gO1, int gO2) const;
   Eigen::Matrix<std::complex<double>,3,3> self_energy_Fe_1loop_1(double p) const;
   std::complex<double> self_energy_Fe_1loop_PR(double p , int gO1, int gO2) const;
   Eigen::Matrix<std::complex<double>,3,3> self_energy_Fe_1loop_PR(double p) const;
   std::complex<double> self_energy_Fe_1loop_PL(double p , int gO1, int gO2) const;
   Eigen::Matrix<std::complex<double>,3,3> self_energy_Fe_1loop_PL(double p) const;
   std::complex<double> self_energy_Fv_1loop_1(double p , int gO1, int gO2) const;
   Eigen::Matrix<std::complex<double>,3,3> self_energy_Fv_1loop_1(double p) const;
   std::complex<double> self_energy_Fv_1loop_PR(double p , int gO1, int gO2) const;
   Eigen::Matrix<std::complex<double>,3,3> self_energy_Fv_1loop_PR(double p) const;
   std::complex<double> self_energy_Fv_1loop_PL(double p , int gO1, int gO2) const;
   Eigen::Matrix<std::complex<double>,3,3> self_energy_Fv_1loop_PL(double p) const;
   std::complex<double> self_energy_Fd_1loop_1_heavy_rotated(double p , int gO1, int gO2) const;
   Eigen::Matrix<std::complex<double>,3,3> self_energy_Fd_1loop_1_heavy_rotated(double p) const;
   std::complex<double> self_energy_Fd_1loop_PR_heavy_rotated(double p , int gO1, int gO2) const;
   Eigen::Matrix<std::complex<double>,3,3> self_energy_Fd_1loop_PR_heavy_rotated(double p) const;
   std::complex<double> self_energy_Fd_1loop_PL_heavy_rotated(double p , int gO1, int gO2) const;
   Eigen::Matrix<std::complex<double>,3,3> self_energy_Fd_1loop_PL_heavy_rotated(double p) const;
   std::complex<double> self_energy_Fe_1loop_1_heavy_rotated(double p , int gO1, int gO2) const;
   Eigen::Matrix<std::complex<double>,3,3> self_energy_Fe_1loop_1_heavy_rotated(double p) const;
   std::complex<double> self_energy_Fe_1loop_PR_heavy_rotated(double p , int gO1, int gO2) const;
   Eigen::Matrix<std::complex<double>,3,3> self_energy_Fe_1loop_PR_heavy_rotated(double p) const;
   std::complex<double> self_energy_Fe_1loop_PL_heavy_rotated(double p , int gO1, int gO2) const;
   Eigen::Matrix<std::complex<double>,3,3> self_energy_Fe_1loop_PL_heavy_rotated(double p) const;
   std::complex<double> self_energy_Fu_1loop_1_heavy_rotated(double p , int gO1, int gO2) const;
   Eigen::Matrix<std::complex<double>,3,3> self_energy_Fu_1loop_1_heavy_rotated(double p) const;
   std::complex<double> self_energy_Fu_1loop_PR_heavy_rotated(double p , int gO1, int gO2) const;
   Eigen::Matrix<std::complex<double>,3,3> self_energy_Fu_1loop_PR_heavy_rotated(double p) const;
   std::complex<double> self_energy_Fu_1loop_PL_heavy_rotated(double p , int gO1, int gO2) const;
   Eigen::Matrix<std::complex<double>,3,3> self_energy_Fu_1loop_PL_heavy_rotated(double p) const;
   std::complex<double> self_energy_Fu_1loop_1_heavy(double p , int gO1, int gO2) const;
   Eigen::Matrix<std::complex<double>,3,3> self_energy_Fu_1loop_1_heavy(double p) const;
   std::complex<double> self_energy_Fu_1loop_PR_heavy(double p , int gO1, int gO2) const;
   Eigen::Matrix<std::complex<double>,3,3> self_energy_Fu_1loop_PR_heavy(double p) const;
   std::complex<double> self_energy_Fu_1loop_PL_heavy(double p , int gO1, int gO2) const;
   Eigen::Matrix<std::complex<double>,3,3> self_energy_Fu_1loop_PL_heavy(double p) const;
   std::complex<double> tadpole_hh_1loop() const;

   /// calculates the tadpoles at current loop order
   void tadpole_equations(double[number_of_ewsb_equations]) const;
   /// calculates the tadpoles at current loop order
   Eigen::Matrix<double,number_of_ewsb_equations,1> tadpole_equations() const;
   /// calculates the tadpoles divided by VEVs at current loop order
   Eigen::Matrix<double,number_of_ewsb_equations,1> tadpole_equations_over_vevs() const;

   double self_energy_hh_2loop(double p) const;

   double self_energy_hh_3loop() const;

   void calculate_MVG_pole();
   void calculate_MFv_pole();
   void calculate_Mhh_pole();
   void calculate_MVP_pole();
   void calculate_MVZ_pole();
   void calculate_MFd_pole();
   void calculate_MFu_pole();
   void calculate_MFe_pole();
   void calculate_MVWp_pole();
   double calculate_MVWp_pole(double);
   double calculate_MVZ_pole(double);

   double calculate_MFv_DRbar(double, int) const;
   double calculate_MFe_DRbar(double, int) const;
   double calculate_MFu_DRbar(double, int) const;
   double calculate_MFd_DRbar(double, int) const;
   double calculate_MVP_DRbar(double);
   double calculate_MVZ_DRbar(double);
   double calculate_MVWp_DRbar(double);

   double ThetaW() const;

private:
   int ewsb_loop_order{2};           ///< loop order for EWSB
   int pole_mass_loop_order{2};      ///< loop order for pole masses
   bool calculate_sm_pole_masses{false};  ///< switch to calculate the pole masses of the Standard Model particles
   bool calculate_bsm_pole_masses{true};  ///< switch to calculate the pole masses of the BSM particles
   bool force_output{false};              ///< switch to force output of pole masses
   double precision{1.e-3};               ///< RG running precision
   double ewsb_iteration_precision{1.e-5};///< precision goal of EWSB solution
   HSSUSY_physical physical{}; ///< contains the pole masses and mixings
   Problems problems{HSSUSY_info::model_name,
                     &HSSUSY_info::particle_names_getter,
                     &HSSUSY_info::parameter_names_getter}; ///< problems
   Loop_corrections loop_corrections{}; ///< used pole mass corrections
   std::shared_ptr<HSSUSY_ewsb_solver_interface> ewsb_solver{};
   Threshold_corrections threshold_corrections{}; ///< used threshold corrections

   int get_number_of_ewsb_iterations() const;
   int get_number_of_mass_iterations() const;
   int solve_ewsb_tree_level_custom();
   void copy_DRbar_masses_to_pole_masses();

   // Passarino-Veltman loop functions
   double A0(double) const noexcept;
   double B0(double, double, double) const noexcept;
   double B1(double, double, double) const noexcept;
   double B00(double, double, double) const noexcept;
   double B22(double, double, double) const noexcept;
   double H0(double, double, double) const noexcept;
   double F0(double, double, double) const noexcept;
   double G0(double, double, double) const noexcept;

   // DR-bar masses
   double MVG{};
   double MHp{};
   Eigen::Array<double,3,1> MFv{Eigen::Array<double,3,1>::Zero()};
   double MAh{};
   double Mhh{};
   Eigen::Array<double,3,1> MFd{Eigen::Array<double,3,1>::Zero()};
   Eigen::Array<double,3,1> MFu{Eigen::Array<double,3,1>::Zero()};
   Eigen::Array<double,3,1> MFe{Eigen::Array<double,3,1>::Zero()};
   double MVWp{};
   double MVP{};
   double MVZ{};

   // DR-bar mixing matrices
   Eigen::Matrix<std::complex<double>,3,3> Vd{Eigen::Matrix<std::complex<double>,3,3>::Zero()};
   Eigen::Matrix<std::complex<double>,3,3> Ud{Eigen::Matrix<std::complex<double>,3,3>::Zero()};
   Eigen::Matrix<std::complex<double>,3,3> Vu{Eigen::Matrix<std::complex<double>,3,3>::Zero()};
   Eigen::Matrix<std::complex<double>,3,3> Uu{Eigen::Matrix<std::complex<double>,3,3>::Zero()};
   Eigen::Matrix<std::complex<double>,3,3> Ve{Eigen::Matrix<std::complex<double>,3,3>::Zero()};
   Eigen::Matrix<std::complex<double>,3,3> Ue{Eigen::Matrix<std::complex<double>,3,3>::Zero()};
   Eigen::Matrix<double,2,2> ZZ{Eigen::Matrix<double,2,2>::Zero()};

   // phases

   // extra parameters

};

std::ostream& operator<<(std::ostream&, const HSSUSY_mass_eigenstates&);

} // namespace flexiblesusy

#endif


#include <boost/polygon/polygon_90_set_data.hpp>


/****************************************************************************
 *  hipipe library
 *  Copyright (c) 2017, Cognexa Solutions s.r.o.
 *  Copyright (c) 2018, Iterait a.s.
 *  Author(s) Filip Matzner, Adam Blazek
 *
 *  This file is distributed under the MIT License.
 *  See the accompanying file LICENSE.txt for the complete license agreement.
 ****************************************************************************/

#include <hipipe/core/python/initialize.hpp>

#include <boost/python.hpp>
#include <opencv2/core/core.hpp>

#include <algorithm>
#include <vector>

// chessboard matrix //

// check c++ -> python
cv::Mat chessboard()
{
    std::vector<std::uint8_t> data = {
      1, 0, 1,
      0, 1, 0,
      1, 0, 1,
      0, 1, 0};
    cv::Mat mat(4, 3, CV_8UC1);
    memcpy(mat.data, data.data(), data.size());
    return mat;
}

// check python -> c++
bool check_chessboard(cv::Mat mat)
{
    return cv::countNonZero(mat != chessboard()) == 0;
}

// 3-channeled rgb matrix //

// check c++ -> python
cv::Mat rgb_sample()
{
    std::vector<float> data = {
      0.1, 0.2, 0.3,
      0.1, 0.2, 0.3,
      0.4, 0.5, 0.6,
      0.4, 0.5, 0.6,
      0.7, 0.8, 0.9,
      0.7, 0.8, 0.9};
    cv::Mat mat(2, 3, CV_32FC3);
    memcpy(mat.data, data.data(), data.size() * sizeof(float));
    return mat;
}

// check python -> c++
bool check_rgb_sample(cv::Mat mat)
{
    // since OpenCV 2 cannot compare matrices with more than one channel,
    // we will use std::equal instead
    cv::Mat gold = rgb_sample();
    return std::equal(mat.begin<float>(), mat.end<float>(), gold.begin<float>());
}

BOOST_PYTHON_MODULE(pyboost_cv_mat_converter_py_cpp)
{
    // initialize hipipe OpenCV converters, exceptions, etc.
    hipipe::python::initialize();

    // expose the functions
    namespace py = boost::python;
    py::def("chessboard", chessboard);
    py::def("check_chessboard", check_chessboard);
    py::def("rgb_sample", rgb_sample);
    py::def("check_rgb_sample", check_rgb_sample);
}


/*
 * Make some convenient Ecal function available in python
 *
 * \author Stefano Argiro
 * \version $Id: EcalPyUtils.cc,v 1.6 2012/07/17 09:17:11 davidlt Exp $
 */

#include "DataFormats/EcalDetId/interface/EBDetId.h"
#include "DataFormats/EcalDetId/interface/EEDetId.h"
#include "CondTools/Ecal/interface/EcalFloatCondObjectContainerXMLTranslator.h"
#include <boost/python.hpp>
#include <boost/python/suite/indexing/vector_indexing_suite.hpp>

using namespace boost::python;

namespace ecalpyutils{

  std::vector<int> hashedIndexToEtaPhi(int hashedindex){
    int ieta= EBDetId::unhashIndex(hashedindex).ieta();
    int iphi= EBDetId::unhashIndex(hashedindex).iphi();
    std::vector<int> ret;
    ret.push_back(ieta);
    ret.push_back(iphi);
    return ret;
  }

  std::vector<int> hashedIndexToXY(int hashedindex){
    int ix= EEDetId::unhashIndex(hashedindex).ix();
    int iy= EEDetId::unhashIndex(hashedindex).iy();
    int zside =  EEDetId::unhashIndex(hashedindex).zside();
    std::vector<int> ret;
    ret.push_back(ix);
    ret.push_back(iy);
    ret.push_back(zside);
    return ret;
  }

  int hashedIndex(int ieta, int iphi){
    EBDetId d(ieta,iphi);
    return d.hashedIndex();
  }

  int hashedIndexEE(int ix, int iy, int iz){

    if (EEDetId::validDetId(ix,iy,iz)) {
      EEDetId d(ix,iy,iz);
      return d.hashedIndex();
    }
    return 0;
  }

  int ism(int ieta, int iphi){
    EBDetId d(ieta,iphi);
    return d.ism();
  }

  std::string arraystoXML(const std::vector<float>& eb, const std::vector<float>& ee){
    EcalCondHeader h;
    return EcalFloatCondObjectContainerXMLTranslator::dumpXML(h,eb,ee);
  }
}

BOOST_PYTHON_MODULE(pluginEcalPyUtils) {

//   looks like these are already defined somewhere

//   python access to stl integer vectors
//   class_< std::vector<int> >("vectorInt")
//     .def(vector_indexing_suite<std::vector<int> >())
//   ;

//   class_< std::vector<float> >("vectorFloat")
//     .def(vector_indexing_suite<std::vector<float> >())
//   ;

  def("hashedIndexToEtaPhi",&ecalpyutils::hashedIndexToEtaPhi);
  def("hashedIndexToXY",&ecalpyutils::hashedIndexToXY);
  def("hashedIndex",&ecalpyutils::hashedIndex);
  def("hashedIndexEE",&ecalpyutils::hashedIndexEE);
  def("ism",&ecalpyutils::ism);
  def("barrelfromXML",&EcalFloatCondObjectContainerXMLTranslator::barrelfromXML);
  def("endcapfromXML",&EcalFloatCondObjectContainerXMLTranslator::endcapfromXML);
  def("arraystoXML",&ecalpyutils::arraystoXML);
}