#include <fstream>

#include <limits>

#include <boost/lexical_cast.hpp>

#include <boost/numeric/ublas/vector_proxy.hpp>
#include <boost/numeric/ublas/matrix_proxy.hpp>

#include "flame/constants.h"

#include "flame/moment.h"
#include "flame/moment_sup.h"
#include "flame/rf_cavity.h"
#include "flame/chg_stripper.h"

#define sqr(x)  ((x)*(x))
#define cube(x) ((x)*(x)*(x))

namespace ub = boost::numeric::ublas;

namespace {

// ARR should be an array-like object (std::vector or or ublas vector or matrix storage)
// fill 'to' using config.get<>(name)
// 'T' selects throw error (true) or return boolean (false)
template<typename ARR>
bool load_storage(ARR& to, const Config& conf, const std::string& name, bool T=true)
{
    try{
        const std::vector<double>& val(conf.get<std::vector<double> >(name));
        if(to.size()!=val.size()) {
            throw std::invalid_argument(SB()<<"Array "<<name<<" must have "<<to.size()<<" elements, not "<<val.size());
        }
        std::copy(val.begin(), val.end(), to.begin());

        return true;
    }catch(std::exception&){
        if(T)
            throw std::invalid_argument(name+" not defined or has wrong type (must be vector)");
        else
            return false;
        // default to identity
    }
}

} // namespace

std::ostream& operator<<(std::ostream& strm, const Particle& P)
{
    strm
      <<"IonZ="<<std::scientific << std::setprecision(10)<<P.IonZ
      <<" IonQ="<<P.IonQ
      <<" IonEs="<<P.IonEs
      <<" IonEk="<<P.IonEk
      <<" SampleIonK="<<P.SampleIonK
      <<" phis="<<P.phis
      <<" IonW="<<P.IonW
      <<" gamma="<<P.gamma
      <<" beta="<<P.beta
      <<" bg="<<P.bg
      ;
    return strm;
}

MomentState::MomentState(const Config& c)
    :StateBase(c)
    ,ref()
    ,real()
    ,moment0_env(maxsize, 0e0)
    ,moment0_rms(maxsize, 0e0)
    ,moment1_env(boost::numeric::ublas::identity_matrix<double>(maxsize))
{
    // hack.  getArray() promises that returned pointers will remain valid for our lifetime.
    // This may not be true if std::vectors are resized.
    // Reserve for up to 10 states and hope for the best...
    // either need to provide real limit to max. states, or change getArray() iface
    real.reserve(10);

    double icstate_f = 0.0;
    bool have_cstate = c.tryGet<double>("cstate", icstate_f);
    size_t icstate = (size_t)icstate_f;

    std::string vectorname(c.get<std::string>("vector_variable", "moment0"));
    std::string matrixname(c.get<std::string>("matrix_variable", "initial"));

    std::vector<double> ics, nchg;
    bool have_ics = c.tryGet<std::vector<double> >("IonChargeStates", ics);

    ref.IonEs      = c.get<double>("IonEs", 0e0),
    ref.IonEk      = c.get<double>("IonEk", 0e0),
    ref.SampleFreq = c.get<double>("SampleFreq", SampleFreqDefault);
    ref.recalc();

    if(!have_ics) {
        ref.IonZ = c.get<double>("IonZ", 0e0);
        ref.IonQ = c.get<double>("IonQ", 1e0);

        ics.push_back(ref.IonZ);
        nchg.push_back(ref.IonQ);

    } else {
        if(ics.empty())
            throw std::invalid_argument("IonChargeStates w/ length 0");
        if(icstate>=ics.size())
            throw std::invalid_argument("IonChargeStates[cstate] is out of bounds");

        nchg = c.get<std::vector<double> >("NCharge");
        if(nchg.size()!=ics.size())
            throw std::invalid_argument("NCharge[] and IonChargeStates[] must have equal length");

        bool have_ionz = c.tryGet<double>("IonZ", ref.IonZ);
        if(!have_ionz) {
            ref.IonZ = ics[0];
        }

        bool have_ionq = c.tryGet<double>("IonQ", ref.IonQ);
        if(!have_ionq) {
            ref.IonQ = nchg[0];
        }
    }

    /* Possible configurations
     * 1. Neither 'cstate' nor 'IonChargeStates' defined (empty Config).
     *    No charge states, must go through source element to be useful
     * 2. 'IonChargeStates' defined, but not 'cstate'.
     *    Load all charge states
     * 3. 'cstate' and 'IonChargeStates' defined.
     *    Load a single charge state
     */
    if(!have_cstate && !have_ics) {
        // no-op
    } else if(!have_cstate && have_ics) {
        // many charge states
        icstate = 0;

    } else if(have_cstate && have_ics) {
        // single charge state

        // drop other than selected state
        ics[0]  = ics[icstate];
        nchg[0] = nchg[icstate];
        ics.resize(1);
        nchg.resize(1);

    } else {
        throw std::invalid_argument("MomentState: must define IonChargeStates and NCharge when cstate is set");
    }

    if(have_ics) {
        real.resize(ics.size());
        moment0.resize(ics.size());
        moment1.resize(ics.size());
        transmat.resize(ics.size());

        for(size_t i=0; i<ics.size(); i++) {
            std::string num(boost::lexical_cast<std::string>(icstate+i));

            moment0[i].resize(maxsize);
            moment1[i].resize(maxsize, maxsize);
            moment1[i] = boost::numeric::ublas::identity_matrix<double>(maxsize);
            transmat[i].resize(maxsize, maxsize);
            transmat[i] = boost::numeric::ublas::identity_matrix<double>(maxsize);

            load_storage(moment0[i].data(), c, vectorname+num);
            load_storage(moment1[i].data(), c, matrixname+num);

            real[i] = ref;

            real[i].IonZ = ics[i];
            real[i].IonQ = nchg[i];

            real[i].phis      = moment0[i][PS_S];
            real[i].IonEk    += moment0[i][PS_PS]*MeVtoeV;

            real[i].recalc();
        }
    } else {
        real.resize(1); // hack, ensure at least one element so getArray() can return some pointer
        real[0] = ref;

        moment0.resize(1);
        moment1.resize(1);
        transmat.resize(1);
        moment0[0].resize(maxsize);
        moment1[0].resize(maxsize, maxsize);
        transmat[0].resize(maxsize, maxsize);

        load_storage(moment0[0].data(), c, vectorname, false);
        load_storage(moment1[0].data(), c, matrixname, false);
        transmat[0] = boost::numeric::ublas::identity_matrix<double>(maxsize);
    }

    last_caviphi0 = 0e0;
    calc_rms();
}

MomentState::~MomentState() {}

void MomentState::calc_rms()
{
    assert(real.size()>0);
    assert(moment0_env.size()==maxsize);
    assert(moment0_rms.size()==maxsize);
    assert(moment1_env.size1()==maxsize);
    assert(moment1_env.size2()==maxsize);

    double totQ = 0.0;
    for(size_t n=0; n<real.size(); n++) {
        totQ += real[n].IonQ;
        if(n==0)
            noalias(moment0_env)  = moment0[n]*real[n].IonQ;
        else
            noalias(moment0_env) += moment0[n]*real[n].IonQ;
    }
    moment0_env /= totQ;

    // Zero orbit terms.
    moment1_env = ub::zero_matrix<double>(state_t::maxsize);
    boost::numeric::ublas::slice S(0, 1, 6);
    for(size_t n=0; n<real.size(); n++) {
        const double Q = real[n].IonQ;
        state_t::vector_t m0diff(moment0[n]-moment0_env);

        ub::project(moment1_env, S, S) += ub::project(Q*(moment1[n]+ub::outer_prod(m0diff, m0diff)), S, S);
    }
    moment1_env /= totQ;

    for(size_t j=0; j<maxsize; j++) {
        moment0_rms[j] = sqrt(moment1_env(j,j));
    }
}

MomentState::MomentState(const MomentState& o, clone_tag t)
    :StateBase(o, t)
    ,ref(o.ref)
    ,real(o.real)
    ,moment0(o.moment0)
    ,moment1(o.moment1)
    ,transmat(o.transmat)
    ,moment0_env(o.moment0_env)
    ,moment0_rms(o.moment0_rms)
    ,moment1_env(o.moment1_env)
    ,last_caviphi0(o.last_caviphi0)
{}

void MomentState::assign(const StateBase& other)
{
    const MomentState *O = dynamic_cast<const MomentState*>(&other);
    if(!O)
        throw std::invalid_argument("Can't assign State: incompatible types");
    ref     = O->ref;
    real    = O->real;
    moment0 = O->moment0;
    moment1 = O->moment1;
    transmat = O->transmat;
    moment0_env = O->moment0_env;
    moment0_rms = O->moment0_rms;
    moment1_env = O->moment1_env;
    last_caviphi0 = O->last_caviphi0;
    StateBase::assign(other);
}

void MomentState::show(std::ostream& strm, int level) const
{
    if(real.empty()) {
        strm<<"State: empty";
        return;
    }

    if(level<=0) {
        strm<<"State: moment0 mean="<<moment0_env;
    }
    if(level>=1) {
        strm << std::scientific << std::setprecision(8)
             << "\nState:\n  energy [eV] =\n" << std::setw(20) << real[0].IonEk << "\n  moment0 mean =\n    ";
        for (size_t k = 0; k < MomentState::maxsize; k++)
            strm << std::scientific << std::setprecision(10) << std::setw(18) << moment0_env(k) << ",";
        strm << std::scientific << std::setprecision(10)
             << "\n  moment0 rms =\n    ";
        for (size_t k = 0; k < MomentState::maxsize; k++)
            strm << std::scientific << std::setprecision(10) << std::setw(18) << moment0_rms(k) << ",";
        strm << "\n  moment1 mean =\n";
        for (size_t j = 0; j < MomentState::maxsize; j++) {
            strm << "    ";
            for (size_t k = 0; k < MomentState::maxsize; k++) {
                strm << std::scientific << std::setprecision(10) << std::setw(18) << moment1_env(j, k) << ",";
            }
            if (j < MomentState::maxsize-1) strm << "\n";
        }
    }
    if(level>=2) {
        strm<< "\n  Reference state:\n    "<<ref<<"\n";
        for(size_t k=0; k<size(); k++) {
            strm<<"  Charge state "<<k<<"\n"
                  "    "<<real[k]<<"\n"
                  "    moment0 "<<moment0[k]<<"\n";
            if(level>=3)
                strm<<
                  "    moment1 "<<moment1[k]<<"\n";
        }
    }
}

bool MomentState::getArray(unsigned idx, ArrayInfo& Info) {
    unsigned I=0;
    if(idx==I++) {
        Info.name = "moment1_env";
        Info.ptr = &moment1_env(0,0);
        Info.type = ArrayInfo::Double;
        Info.ndim = 2;
        Info.dim[0] = moment1_env.size1();
        Info.dim[1] = moment1_env.size2();
        Info.stride[0] = sizeof(double)*moment1_env.size2();
        Info.stride[1] = sizeof(double);
        return true;
    } else if(idx==I++) {
        /* Slight evilness here
         * moment0 is vector of ublas::matrix
         * We assume ublax::matrix uses storage bounded_array<>, and that this storage
         * is really a C array, which means that everything is part of one big allocation.
         * Further we assume that all entries in the vector have the same shape.
         * If this isn't the case, then SIGSEGV here we come...
         */
        static_assert(sizeof(moment1[0])>=sizeof(double)*maxsize*maxsize,
                      "storage assumption violated");
        Info.name = "moment1";
        Info.ptr = &moment1[0](0,0);
        Info.type = ArrayInfo::Double;
        Info.ndim = 3;
        Info.dim[0] = moment1[0].size1();
        Info.dim[1] = moment1[0].size2();
        Info.dim[2] = moment1.size();
        Info.stride[0] = sizeof(double)*moment1_env.size2();
        Info.stride[1] = sizeof(double);
        Info.stride[2] = sizeof(moment1[0]);
        return true;
    } else if(idx==I++) {
        static_assert(sizeof(transmat[0])>=sizeof(double)*maxsize*maxsize,
                      "storage assumption violated");
        Info.name = "transmat";
        Info.ptr = &transmat[0](0,0);
        Info.type = ArrayInfo::Double;
        Info.ndim = 3;
        Info.dim[0] = transmat[0].size1();
        Info.dim[1] = transmat[0].size2();
        Info.dim[2] = transmat.size();
        Info.stride[0] = sizeof(double)*moment1_env.size2();
        Info.stride[1] = sizeof(double);
        Info.stride[2] = sizeof(transmat[0]);
        return true;
    } else if(idx==I++) {
        Info.name = "moment0_env";
        Info.ptr = &moment0_env(0);
        Info.type = ArrayInfo::Double;
        Info.ndim = 1;
        Info.dim[0] = moment0_env.size();
        Info.stride[0] = sizeof(double);
        return true;
    } else if(idx==I++) {
        Info.name = "moment0_rms";
        Info.ptr = &moment0_rms(0);
        Info.type = ArrayInfo::Double;
        Info.ndim = 1;
        Info.dim[0] = moment0_rms.size();
        Info.stride[0] = sizeof(double);
        return true;
    } else if(idx==I++) {
        // more evilness here, see above
        static_assert(sizeof(moment0[0])>=sizeof(double)*maxsize,
                "storage assumption violated");
        Info.name = "moment0";
        Info.ptr = &moment0[0][0];
        Info.type = ArrayInfo::Double;
        Info.ndim = 2;
        Info.dim[0] = moment0[0].size();
        Info.dim[1] = moment0.size();
        Info.stride[0] = sizeof(double);
        Info.stride[1] = sizeof(moment0[0]);
        return true;
    } else if(idx==I++) {
        Info.name = "ref_IonZ";
        Info.ptr = &ref.IonZ;
        Info.type = ArrayInfo::Double;
        Info.ndim = 0;
        return true;
    } else if(idx==I++) {
        Info.name = "ref_IonQ";
        Info.ptr = &ref.IonQ;
        Info.type = ArrayInfo::Double;
        Info.ndim = 0;
        return true;
    } else if(idx==I++) {
        Info.name = "ref_IonEs";
        Info.ptr = &ref.IonEs;
        Info.type = ArrayInfo::Double;
        Info.ndim = 0;
        return true;
    } else if(idx==I++) {
        Info.name = "ref_IonW";
        Info.ptr = &ref.IonW;
        Info.type = ArrayInfo::Double;
        Info.ndim = 0;
        return true;
    } else if(idx==I++) {
        Info.name = "ref_gamma";
        Info.ptr = &ref.gamma;
        Info.type = ArrayInfo::Double;
        Info.ndim = 0;
        return true;
    } else if(idx==I++) {
        Info.name = "ref_beta";
        Info.ptr = &ref.beta;
        Info.type = ArrayInfo::Double;
        Info.ndim = 0;
        return true;
    } else if(idx==I++) {
        Info.name = "ref_bg";
        Info.ptr = &ref.bg;
        Info.type = ArrayInfo::Double;
        Info.ndim = 0;
        return true;
    } else if(idx==I++) {
        Info.name = "ref_SampleFreq";
        Info.ptr = &ref.SampleFreq;
        Info.type = ArrayInfo::Double;
        Info.ndim = 0;
        return true;
    } else if(idx==I++) {
        Info.name = "ref_SampleIonK";
        Info.ptr = &ref.SampleIonK;
        Info.type = ArrayInfo::Double;
        Info.ndim = 0;
        return true;
    } else if(idx==I++) {
        Info.name = "ref_phis";
        Info.ptr = &ref.phis;
        Info.type = ArrayInfo::Double;
        Info.ndim = 0;
        return true;
    } else if(idx==I++) {
        Info.name = "ref_IonEk";
        Info.ptr = &ref.IonEk;
        Info.type = ArrayInfo::Double;
        Info.ndim = 0;
        return true;
    } else if(idx==I++) {
        Info.name = "IonZ";
        Info.ptr  = &real[0].IonZ;
        Info.type = ArrayInfo::Double;
        Info.ndim = 1;
        Info.dim   [0] = real.size();
        Info.stride[0] = sizeof(real[0]);
        // Note: this array is discontigious as we reference a single member from a Particle[]
        return true;
    } else if(idx==I++) {
        Info.name = "IonEs";
        Info.ptr  = &real[0].IonEs;
        Info.type = ArrayInfo::Double;
        Info.ndim = 1;
        Info.dim   [0] = real.size();
        Info.stride[0] = sizeof(real[0]);
        return true;
    } else if(idx==I++) {
        Info.name = "IonW";
        Info.ptr  = &real[0].IonW;
        Info.type = ArrayInfo::Double;
        Info.ndim = 1;
        Info.dim   [0] = real.size();
        Info.stride[0] = sizeof(real[0]);
        return true;
    } else if(idx==I++) {
        Info.name = "gamma";
        Info.ptr  = &real[0].gamma;
        Info.type = ArrayInfo::Double;
        Info.ndim = 1;
        Info.dim   [0] = real.size();
        Info.stride[0] = sizeof(real[0]);
        return true;
    } else if(idx==I++) {
        Info.name = "beta";
        Info.ptr  = &real[0].beta;
        Info.type = ArrayInfo::Double;
        Info.ndim = 1;
        Info.dim   [0] = real.size();
        Info.stride[0] = sizeof(real[0]);
        return true;
    } else if(idx==I++) {
        Info.name = "bg";
        Info.ptr  = &real[0].bg;
        Info.type = ArrayInfo::Double;
        Info.ndim = 1;
        Info.dim   [0] = real.size();
        Info.stride[0] = sizeof(real[0]);
        return true;
    } else if(idx==I++) {
        Info.name = "SampleFreq";
        Info.ptr  = &real[0].SampleFreq;
        Info.type = ArrayInfo::Double;
        Info.ndim = 1;
        Info.dim   [0] = real.size();
        Info.stride[0] = sizeof(real[0]);
        return true;
    } else if(idx==I++) {
        Info.name = "SampleIonK";
        Info.ptr  = &real[0].SampleIonK;
        Info.type = ArrayInfo::Double;
        Info.ndim = 1;
        Info.dim   [0] = real.size();
        Info.stride[0] = sizeof(real[0]);
        return true;
    } else if(idx==I++) {
        Info.name = "phis";
        Info.ptr  = &real[0].phis;
        Info.type = ArrayInfo::Double;
        Info.ndim = 1;
        Info.dim   [0] = real.size();
        Info.stride[0] = sizeof(real[0]);
        return true;
    } else if(idx==I++) {
        Info.name = "IonEk";
        Info.ptr  = &real[0].IonEk;
        Info.type = ArrayInfo::Double;
        Info.ndim = 1;
        Info.dim   [0] = real.size();
        Info.stride[0] = sizeof(real[0]);
        return true;
    } else if(idx==I++) {
        Info.name = "IonQ";
        Info.ptr  = &real[0].IonQ;
        Info.type = ArrayInfo::Double;
        Info.ndim = 1;
        Info.dim   [0] = real.size();
        Info.stride[0] = sizeof(real[0]);
        // Note: this array is discontigious as we reference a single member from a Particle[]
        return true;
    } else if(idx==I++) {
        Info.name = "last_caviphi0";
        Info.ptr = &last_caviphi0;
        Info.type = ArrayInfo::Double;
        Info.ndim = 0;
        // driven phase [degree]
        return true;
    }
    return StateBase::getArray(idx-I, Info);
}

MomentElementBase::MomentElementBase(const Config& c)
    :ElementVoid(c)
    ,dx   (c.get<double>("dx",    0e0)*MtoMM)
    ,dy   (c.get<double>("dy",    0e0)*MtoMM)
    ,pitch(c.get<double>("pitch", 0e0))
    ,yaw  (c.get<double>("yaw",   0e0))
    ,roll (c.get<double>("roll",  0e0))
    ,skipcache(c.get<double>("skipcache", 0.0)!=0.0)
    ,scratch(state_t::maxsize, state_t::maxsize)
{
}

MomentElementBase::~MomentElementBase() {}

void MomentElementBase::assign(const ElementVoid *other)
{
    const MomentElementBase *O = static_cast<const MomentElementBase*>(other);
    last_real_in = O->last_real_in;
    last_real_out= O->last_real_out;
    last_ref_in  = O->last_ref_in;
    last_ref_out = O->last_ref_out;
    transfer = O->transfer;
    misalign = O->misalign;
    misalign_inv = O->misalign_inv;
    dx = O->dx;
    dy = O->dy;
    pitch = O->pitch;
    yaw   = O->yaw;
    roll  = O->roll;
    skipcache = O->skipcache;
    ElementVoid::assign(other);
}

void MomentElementBase::show(std::ostream& strm, int level) const
{
    using namespace boost::numeric::ublas;
    ElementVoid::show(strm, level);
    /*
    strm<<"Length "<<length<<"\n"
          "Transfer: "<<transfer<<"\n"
          "Mis-align: "<<misalign<<"\n";
          */
}

void MomentElementBase::get_misalign(const state_t &ST, const Particle &real, value_t &M, value_t &IM) const
{
    state_t::matrix_t R,
              scl     = boost::numeric::ublas::identity_matrix<double>(state_t::maxsize),
              scl_inv = scl,
              T       = scl,
              T_inv   = scl,
              R_inv   = scl;

    scl(state_t::PS_S, state_t::PS_S)   /= -real.SampleIonK;
    scl(state_t::PS_PS, state_t::PS_PS) /= sqr(real.beta)*real.gamma*ST.ref.IonEs/MeVtoeV;

    inverse(scl_inv, scl);

    // Translate to center of element.
    T(state_t::PS_S,  6) = -length/2e0*MtoMM;
    T(state_t::PS_PS, 6) = 1e0;
    inverse(T_inv, T);

    RotMat(dx, dy, pitch, yaw, roll, R);

    M = prod(T, scl);
    M = prod(R, M);
    M = prod(T_inv, M);
    M = prod(scl_inv, M);

    inverse(R_inv, R);

    // Translate to center of element.
    T(state_t::PS_S,  6) = length/2e0*MtoMM;
    T(state_t::PS_PS, 6) = 1e0;
    inverse(T_inv, T);

    IM = prod(T, scl);
    IM = prod(R_inv, IM);
    IM = prod(T_inv, IM);
    IM = prod(scl_inv, IM);
}

unsigned MomentElementBase::get_flag(const Config& c, const std::string& name, const unsigned& def_value)
{
    unsigned read_value;
    double check_value;

    try {
        check_value = boost::lexical_cast<double>(c.get<std::string>(name));
    }catch (std::exception&){
        try {
            check_value = c.get<double>(name);
        }catch (std::exception&){
            check_value = boost::lexical_cast<double>(def_value);
        }
    }

    //Check the value is an unsigned integer
    try {
        read_value = boost::lexical_cast<unsigned>(check_value);
        if (boost::lexical_cast<double>(read_value) != check_value)
            throw std::runtime_error(SB()<< name << " must be an unsigned integer");
    }catch (std::exception&){
        throw  std::runtime_error(SB()<< name << " must be an unsigned integer");
    }
    return read_value;
}

void MomentElementBase::advance(StateBase& s)
{
    state_t&  ST = static_cast<state_t&>(s);
    using namespace boost::numeric::ublas;

    // IonEk is Es + E_state; the latter is set by user.
    ST.recalc();

    if(!check_cache(ST)){
        // need to re-calculate energy dependent terms
        last_ref_in = ST.ref;
        last_real_in = ST.real;
        resize_cache(ST);

        recompute_matrix(ST); // updates transfer and last_Kenergy_out

        ST.recalc();

        if(!ST.retreat){
            ST.ref.phis += ST.ref.SampleIonK*length*MtoMM;
            for(size_t k=0; k<last_real_in.size(); k++)
                ST.real[k].phis += ST.real[k].SampleIonK*length*MtoMM;
        } else {
            ST.ref.phis -= ST.ref.SampleIonK*length*MtoMM;
            for(size_t k=0; k<last_real_in.size(); k++)
                ST.real[k].phis -= ST.real[k].SampleIonK*length*MtoMM;
        }

        last_ref_out = ST.ref;
        last_real_out = ST.real;
    } else {
        ST.ref = last_ref_out;
        assert(last_real_out.size()==ST.real.size()); // should be true if check_cache() -> true
        std::copy(last_real_out.begin(),
                  last_real_out.end(),
                  ST.real.begin());
    }

    if(!ST.retreat){
        // Forward propagation
        ST.pos += length;

        for(size_t k=0; k<last_real_in.size(); k++) {
            ST.moment0[k] = prod(transfer[k], ST.moment0[k]);

            scratch = prod(transfer[k], ST.moment1[k]);
            ST.moment1[k] = prod(scratch, trans(transfer[k]));

            ST.transmat[k] = transfer[k];
        }
    } else {
        // Backward propagation
        ST.pos -= length;

        value_t invmat = boost::numeric::ublas::identity_matrix<double>(state_t::maxsize);
        for(size_t k=0; k<last_real_in.size(); k++) {
            inverse(invmat, transfer[k]);

            ST.moment0[k] = prod(invmat, ST.moment0[k]);

            scratch = prod(invmat, ST.moment1[k]);
            ST.moment1[k] = prod(scratch, trans(invmat));

            ST.transmat[k] = invmat;
        }
    }

    ST.calc_rms();
}

bool MomentElementBase::check_cache(const state_t& ST) const
{
    return !skipcache
            && last_real_in.size()==ST.size()
            && last_ref_in==ST.ref
            && std::equal(last_real_in.begin(),
                          last_real_in.end(),
                          ST.real.begin());
}

bool MomentElementBase::check_backward(const state_t& ST) const
{
    bool reals = true;
    if (last_real_out.size()==ST.size()) {
        reals = last_ref_out<=ST.ref;
        for(size_t k=0; k<last_real_out.size(); k++) {
            reals &= last_real_out[k]<=ST.real[k];
        }
    } else {
        reals = false;
    }

    return reals;
}

void MomentElementBase::resize_cache(const state_t& ST)
{
    transfer.resize(ST.real.size(), boost::numeric::ublas::identity_matrix<double>(state_t::maxsize));
    misalign.resize(ST.real.size(), boost::numeric::ublas::identity_matrix<double>(state_t::maxsize));
    misalign_inv.resize(ST.real.size(), boost::numeric::ublas::identity_matrix<double>(state_t::maxsize));
}

void MomentElementBase::recompute_matrix(state_t& ST)
{
    // Default, initialize as no-op

    for(size_t k=0; k<last_real_in.size(); k++) {
        transfer[k] = boost::numeric::ublas::identity_matrix<double>(state_t::maxsize);
    }
}

namespace {

struct ElementSource : public MomentElementBase
{
    typedef ElementSource            self_t;
    typedef MomentElementBase        base_t;
    typedef typename base_t::state_t state_t;

    ElementSource(const Config& c): base_t(c), istate(c) {}

    virtual void advance(StateBase& s)
    {
        state_t& ST = static_cast<state_t&>(s);
        if (!ST.retreat)
            // Replace state with our initial values
            ST.assign(istate);
    }

    virtual void show(std::ostream& strm, int level) const
    {
        ElementVoid::show(strm, level);
        strm<<"Initial: "<<istate.moment0_env<<"\n";
    }

    state_t istate;
    // note that 'transfer' is not used by this element type

    virtual ~ElementSource() {}

    virtual const char* type_name() const {return "source";}

    virtual void assign(const ElementVoid *other) {
        base_t::assign(other);
        const self_t* O=static_cast<const self_t*>(other);
        istate.assign(O->istate);
    }
};

struct ElementMark : public MomentElementBase
{
    // Transport (identity) matrix for Marker.
    typedef ElementMark            self_t;
    typedef MomentElementBase     base_t;
    typedef typename base_t::state_t state_t;

    ElementMark(const Config& c): base_t(c) {length = 0e0;}
    virtual ~ElementMark() {}
    virtual const char* type_name() const {return "marker";}

    virtual void assign(const ElementVoid *other) { base_t::assign(other); }
};

struct ElementBPM : public MomentElementBase
{
    // Transport (identity) matrix for BPM.
    typedef ElementBPM               self_t;
    typedef MomentElementBase       base_t;
    typedef typename base_t::state_t state_t;

    ElementBPM(const Config& c): base_t(c) {length = 0e0;}
    virtual ~ElementBPM() {}
    virtual const char* type_name() const {return "bpm";}

    virtual void assign(const ElementVoid *other) { base_t::assign(other); }
};

struct ElementDrift : public MomentElementBase
{
    // Transport matrix for Drift.
    typedef ElementDrift             self_t;
    typedef MomentElementBase       base_t;
    typedef typename base_t::state_t state_t;

    ElementDrift(const Config& c) : base_t(c) {}
    virtual ~ElementDrift() {}
    virtual const char* type_name() const {return "drift";}

    virtual void assign(const ElementVoid *other) { base_t::assign(other); }

    virtual void recompute_matrix(state_t& ST)
    {
        // Re-initialize transport matrix.

        const double L = length*MtoMM; // Convert from [m] to [mm].

        for(size_t i=0; i<last_real_in.size(); i++) {
            transfer[i] = boost::numeric::ublas::identity_matrix<double>(state_t::maxsize);
            transfer[i](state_t::PS_X, state_t::PS_PX) = L;
            transfer[i](state_t::PS_Y, state_t::PS_PY) = L;
            transfer[i](state_t::PS_S, state_t::PS_PS) =
                -2e0*M_PI/(ST.real[i].SampleLambda*ST.real[i].IonEs/MeVtoeV*cube(ST.real[i].bg))*L;
        }
    }
};

struct ElementOrbTrim : public MomentElementBase
{
    // Transport matrix for Orbit Trim.
    typedef ElementOrbTrim           self_t;
    typedef MomentElementBase       base_t;
    typedef typename base_t::state_t state_t;

    ElementOrbTrim(const Config& c) : base_t(c) {length = 0e0;}
    virtual ~ElementOrbTrim() {}
    virtual const char* type_name() const {return "orbtrim";}

    virtual void assign(const ElementVoid *other) { base_t::assign(other); }

    virtual void recompute_matrix(state_t& ST)
    {
        // Re-initialize transport matrix.
        double theta_x = conf().get<double>("theta_x", 0e0),
               theta_y = conf().get<double>("theta_y", 0e0);
        const double tm_xkick = conf().get<double>("tm_xkick", 0e0),
                     tm_ykick = conf().get<double>("tm_ykick", 0e0),
                     xyrotate = conf().get<double>("xyrotate", 0e0)*M_PI/180e0;
        const bool   realpara = conf().get<double>("realpara", 0e0) == 1e0;

        if (realpara) {
            double ecpi = ST.ref.IonZ*C0/sqrt(sqr(ST.ref.IonW) - sqr(ST.ref.IonEs));
            theta_x = tm_xkick*ecpi;
            theta_y = tm_ykick*ecpi;
        }

        for(size_t i=0; i<last_real_in.size(); i++) {
            transfer[i] = boost::numeric::ublas::identity_matrix<double>(state_t::maxsize);
            transfer[i](state_t::PS_PX, 6) = theta_x*ST.real[i].IonZ/ST.ref.IonZ;
            transfer[i](state_t::PS_PY, 6) = theta_y*ST.real[i].IonZ/ST.ref.IonZ;

            get_misalign(ST, ST.real[i], misalign[i], misalign_inv[i]);

            noalias(scratch)  = prod(transfer[i], misalign[i]);
            noalias(transfer[i]) = prod(misalign_inv[i], scratch);

            if (xyrotate != 0e0) {
                state_t::matrix_t R;
                RotMat(0e0, 0e0, 0e0, 0e0, xyrotate, R);
                noalias(scratch)  = transfer[i];
                noalias(transfer[i]) = prod(scratch, R);
            }

        }
    }
};

struct ElementSBend : public MomentElementBase
{
    // Transport matrix for Gradient Sector Bend; with edge focusing (cylindrical coordinates).
    // Note, TLM only includes energy offset for the orbit; not the transport matrix.

    typedef ElementSBend             self_t;
    typedef MomentElementBase       base_t;
    typedef typename base_t::state_t state_t;

    unsigned HdipoleFitMode;

    ElementSBend(const Config& c) : base_t(c), HdipoleFitMode(0) {

        HdipoleFitMode = get_flag(c, "HdipoleFitMode", 1);
        if (HdipoleFitMode != 0 && HdipoleFitMode != 1)
            throw std::runtime_error(SB()<< "Undefined HdipoleFitMode: " << HdipoleFitMode);
    }
    virtual ~ElementSBend() {}
    virtual const char* type_name() const {return "sbend";}

    virtual void assign(const ElementVoid *other) {
        base_t::assign(other);
        const self_t* O=static_cast<const self_t*>(other);
        HdipoleFitMode = O->HdipoleFitMode;
    }

    virtual void advance(StateBase& s)
    {
        state_t&  ST = static_cast<state_t&>(s);
        using namespace boost::numeric::ublas;

        // IonEk is Es + E_state; the latter is set by user.
        ST.recalc();

        if(!check_cache(ST)) {
            // need to re-calculate energy dependent terms
            last_ref_in = ST.ref;
            last_real_in = ST.real;
            resize_cache(ST);

            recompute_matrix(ST); // updates transfer and last_Kenergy_out

            ST.recalc();
            last_ref_out = ST.ref;
            last_real_out = ST.real;
        } else {
            ST.ref = last_ref_out;
            assert(last_real_out.size()==ST.real.size()); // should be true if check_cache() -> true
            std::copy(last_real_out.begin(),
                      last_real_out.end(),
                      ST.real.begin());
        }

        if(!ST.retreat){
            // Forward propagation
            ST.pos += length;
            ST.ref.phis += ST.ref.SampleIonK*length*MtoMM;

            for(size_t i=0; i<last_real_in.size(); i++) {
                double phis_temp = ST.moment0[i][state_t::PS_S];

                ST.moment0[i]          = prod(transfer[i], ST.moment0[i]);

                noalias(scratch)       = prod(transfer[i], ST.moment1[i]);
                noalias(ST.moment1[i]) = prod(scratch, trans(transfer[i]));

                double dphis_temp = ST.moment0[i][state_t::PS_S] - phis_temp;

                ST.real[i].phis  += ST.real[i].SampleIonK*length*MtoMM + dphis_temp;
                ST.transmat[i] = transfer[i];
            }
        } else {
            // Backward propagation
            ST.pos -= length;
            ST.ref.phis -= ST.ref.SampleIonK*length*MtoMM;

            value_t invmat = boost::numeric::ublas::identity_matrix<double>(state_t::maxsize);
            for(size_t i=0; i<last_real_in.size(); i++) {
                double phis_temp = ST.moment0[i][state_t::PS_S];

                inverse(invmat, transfer[i]);
                ST.moment0[i]          = prod(invmat, ST.moment0[i]);

                noalias(scratch)       = prod(invmat, ST.moment1[i]);
                noalias(ST.moment1[i]) = prod(scratch, trans(invmat));

                double dphis_temp = ST.moment0[i][state_t::PS_S] - phis_temp;

                ST.real[i].phis  -= ST.real[i].SampleIonK*length*MtoMM - dphis_temp;
                ST.transmat[i] = invmat;
            }
        }

        ST.calc_rms();
    }

    virtual void recompute_matrix(state_t& ST)
    {
        // Re-initialize transport matrix.

        double L     = conf().get<double>("L")*MtoMM,
               phi   = conf().get<double>("phi")*M_PI/180e0,
               phi1  = conf().get<double>("phi1")*M_PI/180e0,
               phi2  = conf().get<double>("phi2")*M_PI/180e0,
               K     = conf().get<double>("K", 0e0)/sqr(MtoMM);

        for(size_t i=0; i<last_real_in.size(); i++) {
            double qmrel = (ST.real[i].IonZ-ST.ref.IonZ)/ST.ref.IonZ;

            transfer[i] = boost::numeric::ublas::identity_matrix<double>(state_t::maxsize);

            if (L != 0.0) {
                if (!HdipoleFitMode) {
                    double dip_bg    = conf().get<double>("bg"),
                           // Dipole reference energy.
                           dip_Ek    = (sqrt(sqr(dip_bg)+1e0)-1e0)*ST.ref.IonEs,
                           dip_gamma = (dip_Ek+ST.ref.IonEs)/ST.ref.IonEs,
                           dip_beta  = sqrt(1e0-1e0/sqr(dip_gamma)),
                           d         = (ST.ref.gamma-dip_gamma)/(sqr(dip_beta)*dip_gamma) - qmrel,
                           dip_IonK  = 2e0*M_PI/(dip_beta*ST.ref.SampleLambda);

                    GetSBendMatrix(L, phi, phi1, phi2, K, ST.ref.IonEs, ST.ref.gamma, qmrel,
                                   dip_beta, dip_gamma, d, dip_IonK, transfer[i]);
                } else
                    GetSBendMatrix(L, phi, phi1, phi2, K, ST.ref.IonEs, ST.ref.gamma, qmrel,
                                   ST.ref.beta, ST.ref.gamma, - qmrel, ST.ref.SampleIonK, transfer[i]);

                get_misalign(ST, ST.real[i], misalign[i], misalign_inv[i]);

                noalias(scratch)     = prod(transfer[i], misalign[i]);
                noalias(transfer[i]) = prod(misalign_inv[i], scratch);
            }
        }
    }
};

struct ElementQuad : public MomentElementBase
{
    // Transport matrix for Quadrupole; K = B2/Brho.
    typedef ElementQuad              self_t;
    typedef MomentElementBase       base_t;
    typedef typename base_t::state_t state_t;

    ElementQuad(const Config& c) : base_t(c) {}
    virtual ~ElementQuad() {}
    virtual const char* type_name() const {return "quadrupole";}

    virtual void assign(const ElementVoid *other) { base_t::assign(other); }

    virtual void recompute_matrix(state_t& ST)
    {
        const double L = conf().get<double>("L")*MtoMM;
        const unsigned ncurve = get_flag(conf(), "ncurve", 0);

        if (ncurve != 0) {
            std::vector<std::vector<double> > Curves;
            std::vector<double> Scales;
            GetCurveData(conf(), ncurve, Scales, Curves);

            for(size_t i=0; i<last_real_in.size(); i++) {
                double K;
                double dL = L/double(Curves[0].size()),
                       Brho = ST.real[i].Brho();
                transfer[i] = boost::numeric::ublas::identity_matrix<double>(state_t::maxsize);
                for (size_t j=0; j<Curves[0].size(); j++){
                    K = 0.0;
                    for (size_t n=0; n<Curves.size(); n++) K += Scales[n]*Curves[n][j]/Brho/sqr(MtoMM);
                    value_t tmstep = boost::numeric::ublas::identity_matrix<double>(state_t::maxsize);
                    GetQuadMatrix(dL,  K, (unsigned)state_t::PS_X, tmstep);
                    GetQuadMatrix(dL, -K, (unsigned)state_t::PS_Y, tmstep);

                    tmstep(state_t::PS_S, state_t::PS_PS) =
                        -2e0*M_PI/(ST.real[i].SampleLambda*ST.real[i].IonEs/MeVtoeV*cube(ST.real[i].bg))*dL;

                    transfer[i] = prod(tmstep, transfer[i]);
                }
                get_misalign(ST, ST.real[i], misalign[i], misalign_inv[i]);
                noalias(scratch)     = prod(transfer[i], misalign[i]);
                noalias(transfer[i]) = prod(misalign_inv[i], scratch);
            }

        } else {
            const double B2= conf().get<double>("B2");
            for(size_t i=0; i<last_real_in.size(); i++) {
                // Re-initialize transport matrix.
                transfer[i] = boost::numeric::ublas::identity_matrix<double>(state_t::maxsize);

                double Brho = ST.real[i].Brho(),
                       K = B2/Brho/sqr(MtoMM);

                // Horizontal plane.
                GetQuadMatrix(L,  K, (unsigned)state_t::PS_X, transfer[i]);
                // Vertical plane.
                GetQuadMatrix(L, -K, (unsigned)state_t::PS_Y, transfer[i]);
                // Longitudinal plane.

                transfer[i](state_t::PS_S, state_t::PS_PS) =
                        -2e0*M_PI/(ST.real[i].SampleLambda*ST.real[i].IonEs/MeVtoeV*cube(ST.real[i].bg))*L;

                get_misalign(ST, ST.real[i], misalign[i], misalign_inv[i]);
                noalias(scratch)     = prod(transfer[i], misalign[i]);
                noalias(transfer[i]) = prod(misalign_inv[i], scratch);
            }
        }
    }
};

struct ElementSext : public MomentElementBase
{
    // Transport matrix for Sextupole; K = B3/Brho.
    typedef ElementSext              self_t;
    typedef MomentElementBase       base_t;
    typedef typename base_t::state_t state_t;

    ElementSext(const Config& c) : base_t(c) {}

    virtual ~ElementSext() {}
    virtual const char* type_name() const {return "sextupole";}

    virtual void assign(const ElementVoid *other) {base_t::assign(other); }

    virtual void advance(StateBase& s)
    {
        const double B3= conf().get<double>("B3"),
                     L = conf().get<double>("L")*MtoMM;
        const int    step = conf().get<double>("step", 1.0);
        const bool   thinlens = conf().get<double>("thinlens", 0.0) == 1.0,
                     dstkick = conf().get<double>("dstkick", 1.0) == 1.0;

        state_t&  ST = static_cast<state_t&>(s);
        using namespace boost::numeric::ublas;

        ST.recalc();

        last_ref_in = ST.ref;
        last_real_in = ST.real;
        resize_cache(ST);

        if(ST.retreat) throw std::runtime_error(SB()<<
            "Backward propagation error: Backward propagation does not support sextupole.");

        const double dL = L/step;

        for(size_t k=0; k<last_real_in.size(); k++) {

            transfer[k] = boost::numeric::ublas::identity_matrix<double>(state_t::maxsize);
            ST.transmat[k] = transfer[k];

            double Brho = ST.real[k].Brho(),
                   K = B3/Brho/cube(MtoMM);

            get_misalign(ST, ST.real[k], misalign[k], misalign_inv[k]);

            ST.moment0[k] = prod(misalign[k], ST.moment0[k]);
            scratch = prod(misalign[k], ST.moment1[k]);
            ST.moment1[k] = prod(scratch, trans(misalign[k]));

            for(int i=0; i<step; i++){
                double Dx = ST.moment0[k][state_t::PS_X],
                       Dy = ST.moment0[k][state_t::PS_Y],
                       D2x = ST.moment1[k](state_t::PS_X, state_t::PS_X),
                       D2y = ST.moment1[k](state_t::PS_Y, state_t::PS_Y),
                       D2xy = ST.moment1[k](state_t::PS_X, state_t::PS_Y);

                GetSextMatrix(dL, K, Dx, Dy, D2x, D2y, D2xy, thinlens, dstkick, transfer[k]);

                transfer[k](state_t::PS_S, state_t::PS_PS) =
                        -2e0*M_PI/(ST.real[k].SampleLambda*ST.real[k].IonEs/MeVtoeV*cube(ST.real[k].bg))*dL;

                ST.moment0[k] = prod(transfer[k], ST.moment0[k]);

                scratch = prod(transfer[k], ST.moment1[k]);
                ST.moment1[k] = prod(scratch, trans(transfer[k]));

                ST.transmat[k] = prod(transfer[k], ST.transmat[k]);
            }
            ST.moment0[k] = prod(misalign_inv[k], ST.moment0[k]);
            scratch = prod(misalign_inv[k], ST.moment1[k]);
            ST.moment1[k] = prod(scratch, trans(misalign_inv[k]));

            scratch = prod(ST.transmat[k], misalign[k]);
            ST.transmat[k] = prod(misalign_inv[k], scratch);
        }

        ST.recalc();

        for(size_t k=0; k<last_real_in.size(); k++)
            ST.real[k].phis  += ST.real[k].SampleIonK*length*MtoMM;
        ST.ref.phis   += ST.ref.SampleIonK*length*MtoMM;

        last_ref_out = ST.ref;
        last_real_out = ST.real;

        ST.pos += length;

        ST.calc_rms();
    }
};

struct ElementSolenoid : public MomentElementBase
{
    // Transport (identity) matrix for a Solenoid; K = B/(2 Brho).
    typedef ElementSolenoid          self_t;
    typedef MomentElementBase       base_t;
    typedef typename base_t::state_t state_t;

    ElementSolenoid(const Config& c) : base_t(c) {}
    virtual ~ElementSolenoid() {}
    virtual const char* type_name() const {return "solenoid";}

    virtual void assign(const ElementVoid *other) { base_t::assign(other); }

    virtual void recompute_matrix(state_t& ST)
    {
        const double L = conf().get<double>("L")*MtoMM;      // Convert from [m] to [mm].
        const unsigned ncurve = get_flag(conf(), "ncurve", 0);

        if (ncurve != 0) {
            std::vector<std::vector<double> > Curves;
            std::vector<double> Scales;
            GetCurveData(conf(), ncurve, Scales, Curves);

            for(size_t i=0; i<last_real_in.size(); i++) {
                double K;
                double dL = L/double(Curves[0].size()),
                       Brho = ST.real[i].Brho();
                transfer[i] = boost::numeric::ublas::identity_matrix<double>(state_t::maxsize);
                for (size_t j=0; j<Curves[0].size(); j++){
                    K = 0.0;
                    for (size_t n=0; n<Curves.size(); n++) K += Scales[n]*Curves[n][j]/(2e0*Brho)/MtoMM;
                    value_t tmstep = boost::numeric::ublas::identity_matrix<double>(state_t::maxsize);

                    GetSolMatrix(dL, K, tmstep);

                    tmstep(state_t::PS_S, state_t::PS_PS) =
                        -2e0*M_PI/(ST.real[i].SampleLambda*ST.real[i].IonEs/MeVtoeV*cube(ST.real[i].bg))*dL;

                    transfer[i] = prod(tmstep, transfer[i]);
                }
                get_misalign(ST, ST.real[i], misalign[i], misalign_inv[i]);
                noalias(scratch)     = prod(transfer[i], misalign[i]);
                noalias(transfer[i]) = prod(misalign_inv[i], scratch);
            }
        } else {
            const double B = conf().get<double>("B");
            for(size_t i=0; i<last_real_in.size(); i++) {
                // Re-initialize transport matrix.
                transfer[i] = boost::numeric::ublas::identity_matrix<double>(state_t::maxsize);

                double Brho = ST.real[i].Brho(),
                       K    = B/(2e0*Brho)/MtoMM;

                GetSolMatrix(L, K, transfer[i]);

                transfer[i](state_t::PS_S, state_t::PS_PS) =
                        -2e0*M_PI/(ST.real[i].SampleLambda*ST.real[i].IonEs/MeVtoeV*cube(ST.real[i].bg))*L;

                get_misalign(ST, ST.real[i], misalign[i], misalign_inv[i]);

                noalias(scratch)     = prod(transfer[i], misalign[i]);
                noalias(transfer[i]) = prod(misalign_inv[i], scratch);
            }
        }
    }
};

struct ElementEDipole : public MomentElementBase
{
    // Transport matrix for Electrostatic Dipole with edge focusing.
    typedef ElementEDipole           self_t;
    typedef MomentElementBase       base_t;
    typedef typename base_t::state_t state_t;

    ElementEDipole(const Config& c) : base_t(c) {}
    virtual ~ElementEDipole() {}
    virtual const char* type_name() const {return "edipole";}

    virtual void assign(const ElementVoid *other) { base_t::assign(other); }

    virtual void recompute_matrix(state_t& ST)
    {
        // Re-initialize transport matrix.

        //value_mat R;

        bool   ver         = conf().get<double>("ver") == 1.0;
        double L           = conf().get<double>("L")*MtoMM,
               phi         = conf().get<double>("phi")*M_PI/180e0,
               // fit to TLM unit.
               fringe_x    = conf().get<double>("fringe_x", 0e0)/MtoMM,
               fringe_y    = conf().get<double>("fringe_y", 0e0)/MtoMM,
               kappa       = conf().get<double>("asym_fac", 0e0),
               // spher: cylindrical - 0, spherical - 1.
               spher       = conf().get<double>("spher"),
               // magnetic - 0, electrostatic - 1.
               h           = 1e0,
               dip_beta    = conf().get<double>("beta", ST.ref.beta);

        unsigned HdipoleFitMode = get_flag(conf(), "HdipoleFitMode", 1);

        if (HdipoleFitMode != 0 && HdipoleFitMode != 1)
            throw std::runtime_error(SB()<< "Undefined HdipoleFitMode: " << HdipoleFitMode);

        if (HdipoleFitMode) dip_beta = ST.ref.beta;

        for(size_t i=0; i<last_real_in.size(); i++) {
            double eta0        = (1e0/sqrt(1e0 - sqr(dip_beta)) - 1e0)/2e0,
                   Erho        = sqr(ST.real[i].beta)/ST.real[i].IonZ,
                   Erho0       = sqr(dip_beta)/ST.ref.IonZ,
                   eL          = Erho/Erho0*L,
                   rho         = eL/phi,
                   Kx          = (1e0-spher+sqr(1e0+2e0*eta0))/sqr(rho),
                   Ky          = spher/sqr(rho),
                   delta_K     = (Erho/Erho0 - 1e0) - (ST.real[i].IonEk - ST.ref.IonEk)/ST.real[i].IonEk,
                   delta_KZ    = ST.ref.IonZ/ST.real[i].IonZ - 1e0,
                   SampleIonK  = 2e0*M_PI/(ST.real[i].beta*ST.real[i].SampleLambda);

            transfer[i] = boost::numeric::ublas::identity_matrix<double>(state_t::maxsize);

            if (L != 0e0) {
                GetEBendMatrix(eL, phi, fringe_x, fringe_y, kappa, Kx, Ky, ST.ref.IonEs, ST.real[i].gamma,
                               eta0, h, delta_K, delta_KZ, SampleIonK, transfer[i]);

                if (ver) {
                    // Rotate transport matrix by 90 degrees.
                    value_t
                    R = boost::numeric::ublas::identity_matrix<double>(state_t::maxsize);
                    R(state_t::PS_X,  state_t::PS_X)   =  0e0;
                    R(state_t::PS_PX, state_t::PS_PX)  =  0e0;
                    R(state_t::PS_Y,  state_t::PS_Y)   =  0e0;
                    R(state_t::PS_PY, state_t::PS_PY)  =  0e0;
                    R(state_t::PS_X,  state_t::PS_Y)   = -1e0;
                    R(state_t::PS_PX, state_t::PS_PY)  = -1e0;
                    R(state_t::PS_Y,  state_t::PS_X)   =  1e0;
                    R(state_t::PS_PY,  state_t::PS_PX) =  1e0;

                    noalias(scratch)     = prod(R, transfer[i]);
                    noalias(transfer[i]) = prod(scratch, trans(R));
                    //TODO: no-op code?  results are unconditionally overwritten
                }

                get_misalign(ST, ST.real[i], misalign[i], misalign_inv[i]);

                noalias(scratch)     = prod(transfer[i], misalign[i]);
                noalias(transfer[i]) = prod(misalign_inv[i], scratch);
            }
        }
    }
};

struct ElementEQuad : public MomentElementBase
{
    // Transport matrix for Electrostatic Quadrupole.
    typedef ElementEQuad             self_t;
    typedef MomentElementBase       base_t;
    typedef typename base_t::state_t state_t;

    ElementEQuad(const Config& c) : base_t(c) {}
    virtual ~ElementEQuad() {}
    virtual const char* type_name() const {return "equad";}

    virtual void assign(const ElementVoid *other) { base_t::assign(other); }

    virtual void recompute_matrix(state_t& ST)
    {
        const double   L      = conf().get<double>("L")*MtoMM;
        const unsigned ncurve = get_flag(conf(), "ncurve", 0);

        if (ncurve != 0) {
            std::vector<std::vector<double> > Curves;
            std::vector<double> Scales;
            GetCurveData(conf(), ncurve, Scales, Curves);

            for(size_t i=0; i<last_real_in.size(); i++) {
                double K;
                double dL = L/double(Curves[0].size()),
                       Brho = ST.real[i].Brho();
                transfer[i] = boost::numeric::ublas::identity_matrix<double>(state_t::maxsize);
                for (size_t j=0; j<Curves[0].size(); j++){
                    K = 0.0;
                    for (size_t n=0; n<Curves.size(); n++) K += 2e0*Scales[n]*Curves[n][j]/(C0*ST.real[i].beta)/Brho/sqr(MtoMM);
                    value_t tmstep = boost::numeric::ublas::identity_matrix<double>(state_t::maxsize);
                    GetQuadMatrix(dL,  K, (unsigned)state_t::PS_X, tmstep);
                    GetQuadMatrix(dL, -K, (unsigned)state_t::PS_Y, tmstep);

                    tmstep(state_t::PS_S, state_t::PS_PS) =
                        -2e0*M_PI/(ST.real[i].SampleLambda*ST.real[i].IonEs/MeVtoeV*cube(ST.real[i].bg))*dL;

                    transfer[i] = prod(tmstep, transfer[i]);
                }
                get_misalign(ST, ST.real[i], misalign[i], misalign_inv[i]);
                noalias(scratch)     = prod(transfer[i], misalign[i]);
                noalias(transfer[i]) = prod(misalign_inv[i], scratch);
            }

        } else {
            const double V0 = conf().get<double>("V"),
                         R  = conf().get<double>("radius");

            for(size_t i=0; i<last_real_in.size(); i++) {
                // Re-initialize transport matrix.
                // V0 [V] electrode voltage and R [m] electrode half-distance.
                transfer[i] = boost::numeric::ublas::identity_matrix<double>(state_t::maxsize);

                double Brho = ST.real[i].Brho(),
                       K    = 2e0*V0/(C0*ST.real[i].beta*sqr(R))/Brho/sqr(MtoMM);

                // Horizontal plane.
                GetQuadMatrix(L,  K, (unsigned)state_t::PS_X, transfer[i]);
                // Vertical plane.
                GetQuadMatrix(L, -K, (unsigned)state_t::PS_Y, transfer[i]);
                // Longitudinal plane.
                //        transfer(state_t::PS_S, state_t::PS_S) = L;

                transfer[i](state_t::PS_S, state_t::PS_PS) =
                        -2e0*M_PI/(ST.real[i].SampleLambda*ST.real[i].IonEs/MeVtoeV*cube(ST.real[i].bg))*L;

                get_misalign(ST, ST.real[i], misalign[i], misalign_inv[i]);

                noalias(scratch)     = prod(transfer[i], misalign[i]);
                noalias(transfer[i]) = prod(misalign_inv[i], scratch);
            }
        }
    }
};

struct ElementTMatrix : public MomentElementBase
{
    // Transport matrix by user input
    typedef ElementTMatrix          self_t;
    typedef MomentElementBase       base_t;
    typedef typename base_t::state_t state_t;

    ElementTMatrix(const Config& c) : base_t(c) {}
    virtual ~ElementTMatrix() {}
    virtual const char* type_name() const {return "tmatrix";}

    virtual void assign(const ElementVoid *other) { base_t::assign(other); }

    virtual void recompute_matrix(state_t& ST)
    {
        for(size_t i=0; i<last_real_in.size(); i++) {
            load_storage(transfer[i].data(), conf(), "matrix");
        }
    }
};

} // namespace

void registerMoment()
{
    Machine::registerState<MomentState>("MomentMatrix");

    Machine::registerElement<ElementSource                 >("MomentMatrix", "source");

    Machine::registerElement<ElementMark                   >("MomentMatrix", "marker");

    Machine::registerElement<ElementBPM                    >("MomentMatrix", "bpm");

    Machine::registerElement<ElementDrift                  >("MomentMatrix", "drift");

    Machine::registerElement<ElementOrbTrim                >("MomentMatrix", "orbtrim");

    Machine::registerElement<ElementSBend                  >("MomentMatrix", "sbend");

    Machine::registerElement<ElementQuad                   >("MomentMatrix", "quadrupole");

    Machine::registerElement<ElementSext                   >("MomentMatrix", "sextupole");

    Machine::registerElement<ElementSolenoid               >("MomentMatrix", "solenoid");

    Machine::registerElement<ElementRFCavity               >("MomentMatrix", "rfcavity");

    Machine::registerElement<ElementStripper               >("MomentMatrix", "stripper");

    Machine::registerElement<ElementEDipole                >("MomentMatrix", "edipole");

    Machine::registerElement<ElementEQuad                  >("MomentMatrix", "equad");

    Machine::registerElement<ElementTMatrix                >("MomentMatrix", "tmatrix");
}


#ifdef ADEPT_PKG
  #include <Aria.h>
#else
  #include <Aria/Aria.h>
#endif

#include "LaserPublisher.h"
#include "ArTimeToROSTime.h"

#include <math.h>

#include <boost/algorithm/string.hpp>

// TODO publish pointcloud of cumulative readings in separate topic?
// TODO generic pointcloud sensor publisher (seprate point cloud stuff there)
// TODO make  similar sonar publisher?

LaserPublisher::LaserPublisher(ArLaser *_l, ros::NodeHandle& _n, bool _broadcast_tf, const std::string& _tf_frame, const std::string& _parent_tf_frame, const std::string& _global_tf_frame) :
  laserReadingsCB(this, &LaserPublisher::readingsCB),
  node(_n),
  laser(_l),
  tfname(_tf_frame),
  parenttfname(_parent_tf_frame),
  globaltfname(_global_tf_frame),
  broadcast_tf(_broadcast_tf)
{
  assert(_l);
  laser->lockDevice();
  laser->addReadingCB(&laserReadingsCB);
  laser->unlockDevice();
  std::string laserscan_name(laser->getName());
  boost::erase_all(laserscan_name,".");
  laserscan_name += "_laserscan";
  std::string pointcloud_name(laser->getName());
  boost::erase_all(pointcloud_name,".");
  pointcloud_name += "_pointcloud";
  laserscan_pub = node.advertise<sensor_msgs::LaserScan>(laserscan_name, 20);
  pointcloud_pub = node.advertise<sensor_msgs::PointCloud>(pointcloud_name, 50);

  tf::Quaternion q;
  if(laser->hasSensorPosition())
  {
    lasertf.setOrigin(tf::Vector3(laser->getSensorPositionX()/1000.0, laser->getSensorPositionY()/1000.0, laser->getSensorPositionZ()/1000.0));
    q.setRPY(0, 0, ArMath::degToRad(laser->getSensorPositionTh()));
  }
  else
  {
    lasertf.setOrigin(tf::Vector3(0, 0, 0));
    q.setRPY(0, 0, 0);
  }
  lasertf.setRotation(q);

  laserscan.header.frame_id = "laser_frame";
  laserscan.angle_min = ArMath::degToRad(laser->getStartDegrees());
  laserscan.angle_max = ArMath::degToRad(laser->getEndDegrees());
  //laserscan.time_increment = ?
  laserscan.range_min = 0; //laser->getMinRange() / 1000.0;
  laserscan.range_max = laser->getMaxRange() / 1000.0;
  pointcloud.header.frame_id = globaltfname;

  // Get angle_increment of the laser
  laserscan.angle_increment = 0;
  if(laser->canSetIncrement()) {
    laserscan.angle_increment = laser->getIncrement();
  }
  else if(laser->getIncrementChoice() != NULL) {
    laserscan.angle_increment = laser->getIncrementChoiceDouble();
  }
  assert(laserscan.angle_increment > 0);
  laserscan.angle_increment *= M_PI/180.0;

  //readingsCallbackTime = new ArTime;
}

LaserPublisher::~LaserPublisher()
{
  laser->lockDevice();
  laser->remReadingCB(&laserReadingsCB);
  laser->unlockDevice();
  //delete readingsCallbackTime;
}

void LaserPublisher::readingsCB()
{
  //printf("readingsCB(): %lu ms since last readingsCB() call.\n", readingsCallbackTime->mSecSince());
  assert(laser);
  laser->lockDevice();
  publishLaserScan();
  publishPointCloud();
  laser->unlockDevice();
  if(broadcast_tf)
    transform_broadcaster.sendTransform(tf::StampedTransform(lasertf, convertArTimeToROS(laser->getLastReadingTime()), parenttfname, tfname));
  //readingsCallbackTime->setToNow();
}

void LaserPublisher::publishLaserScan()
{
  laserscan.header.stamp = convertArTimeToROS(laser->getLastReadingTime());
  const std::list<ArSensorReading*> *readings = laser->getRawReadings();
  assert(readings);
  //printf("laserscan: %lu readings\n", readings->size());
  laserscan.ranges.resize(readings->size());
  size_t n = 0;
  if (laser->getFlipped()) {
    // Reverse the data
    for(std::list<ArSensorReading*>::const_reverse_iterator r = readings->rbegin(); r != readings->rend(); ++r)
    {
      assert(*r);

      if ((*r)->getIgnoreThisReading()) {
	laserscan.ranges[n] = -1;
      }
      else {
	laserscan.ranges[n] = (*r)->getRange() / 1000.0;
      }

      ++n;
    }
  }
  else {
    for(std::list<ArSensorReading*>::const_iterator r = readings->begin(); r != readings->end(); ++r)
    {
      assert(*r);

      if ((*r)->getIgnoreThisReading()) {
	laserscan.ranges[n] = -1;
      }
      else {
	laserscan.ranges[n] = (*r)->getRange() / 1000.0;
      }

      ++n;
    }
  }

  laserscan_pub.publish(laserscan);
}

void LaserPublisher::publishPointCloud()
{
  assert(laser);
  pointcloud.header.stamp = convertArTimeToROS(laser->getLastReadingTime());
  assert(laser->getCurrentBuffer());
  const std::list<ArPoseWithTime*> *p = laser->getCurrentRangeBuffer()->getBuffer();
  assert(p);
  pointcloud.points.resize(p->size());
  size_t n = 0;
  for(std::list<ArPoseWithTime*>::const_iterator i = p->begin(); i != p->end(); ++i)
  {
    assert(*i);
    pointcloud.points[n].x = (*i)->getX() / 1000.0;
    pointcloud.points[n].y = (*i)->getY() / 1000.0;
    pointcloud.points[n].z = (laser->hasSensorPosition() ?  laser->getSensorPositionZ() / 1000.0 : 0.0);
    ++n;
  }
  pointcloud_pub.publish(pointcloud);
}


#include <dials/array_family/flex_table.h>
#include <dials/array_family/reflection_table.h>
#include <boost/python.hpp>
#include <boost/python/def.hpp>
#include <algorithm>
#include <numeric>
#include <scitbx/vec3.h>
#include <scitbx/mat3.h>
#include <iterator>

using namespace boost::python;
using namespace dials::af;

// Helper classes to speed up calculations in autoreduce.py

namespace dials { namespace refinement { namespace boost_python {
  /*
    Helper class used by the AutoReduce._panel_gp_surplus_reflections function in
    autoreduce.py
  */
  struct pg_surpl_iter {
    template <typename IntType>
    pg_surpl_iter(const scitbx::af::shared<IntType>& ref_ids,
                  const scitbx::af::shared<std::size_t>& panel_id,
                  const boost::python::list pnl_ids,
                  const scitbx::af::shared<std::size_t> exp_ids,
                  const int cutoff) {
      // Perform python list conversion once; improves read access
      scitbx::af::shared<std::size_t> pnl_ids_flex;
      for (int ii = 0; ii < len(pnl_ids); ++ii) {
        pnl_ids_flex.push_back(boost::python::extract<std::size_t>(pnl_ids[ii]));
      }

      std::size_t count = 0;

      // Work with raw pointers instead of overloadable objects
      const std::size_t* ptr_exp_ids = exp_ids.begin();
      const std::size_t* ptr_panel_id = panel_id.begin();

      // Predefine loop variables
      std::size_t exp_it, pnl;

      std::pair<const IntType*, const IntType*> refIDRange;
      std::pair<const std::size_t*, const std::size_t*> panelIDRange;
      unsigned int data_offset = 0;
      unsigned int data_range = 0;

      /* Search for the beginning and end indices for the index values and track their
         pointers. Use these pointers to narow the next search margin and count the
         range of elements that match. These values indicate the count and can be added
         to a running total.
      */
      for (exp_it = 0; exp_it < exp_ids.size();
           ++exp_it) {  // Iterate through experiment IDs
        // Get the pointer indices to the beginning and end of the values with the
        // requested experimental ID. Performs search in ~2logN
        refIDRange =
          std::equal_range(ref_ids.begin(), ref_ids.end(), ptr_exp_ids[exp_it]);

        // Calculate pointer offset for start and end of the data to examine.
        data_range = refIDRange.second - refIDRange.first;
        data_offset = refIDRange.first - ref_ids.begin();

        // Taking the previous pointer positions, and perfoming a search within this
        // range for the matching panel IDs, we then count the pointer range as all
        // valid entries and add to the accumlator.
        for (pnl = 0; pnl < pnl_ids_flex.size(); ++pnl) {  // Iterate through panel IDs
          panelIDRange = std::equal_range(ptr_panel_id + data_offset,
                                          ptr_panel_id + data_offset + data_range,
                                          pnl_ids_flex[pnl]);
          count += (unsigned int)(panelIDRange.second - panelIDRange.first);
        }
      }
      result = count - cutoff;
    }
    int result;
  };

  /*
    Helper class used by the AutoReduce._unit_cell_surplus_reflections function in
    autoreduce.py
  */
  struct uc_surpl_iter {
    template <typename IntType>
    uc_surpl_iter(const scitbx::af::shared<IntType>& ref_ids,
                  const scitbx::af::shared<cctbx::miller::index<> >& ref_miller,
                  const scitbx::af::shared<std::size_t>& exp_ids,
                  const scitbx::af::shared<mat3<double> > F_dbdp) {
      // Expose pointers for experiment flex arrays
      const std::size_t* ptr_exp_ids = exp_ids.begin();

      std::size_t exp_it;  // ref_id;
      std::vector<std::pair<const IntType*, const IntType*> > refIDRange(
        exp_ids.size());

      std::vector<unsigned int> data_range(exp_ids.size());
      std::vector<unsigned int> data_offset(exp_ids.size());

      // Perform subselection of data ranges
      for (exp_it = 0; exp_it < exp_ids.size(); ++exp_it) {
        refIDRange[exp_it] =
          std::equal_range(ref_ids.begin(), ref_ids.end(), ptr_exp_ids[exp_it]);
        data_range[exp_it] = refIDRange[exp_it].second - refIDRange[exp_it].first;
        data_offset[exp_it] = refIDRange[exp_it].first - ref_ids.begin();
      }

      // Pre-initialise loop variables and array for storing norms
      int ii, jj;
      scitbx::af::shared<double> mul_norm;

      // Calculate L-2 norm of Matrix*Vec operation for each matrix in F_dbdp and each
      // vec3 miller index across all exp_ids
      for (jj = 0; jj < F_dbdp.size(); ++jj) {
        for (unsigned int exp_id_range = 0; exp_id_range < data_range.size();
             ++exp_id_range) {
          for (ii = data_offset[exp_id_range];
               ii < data_offset[exp_id_range] + data_range[exp_id_range];
               ++ii) {
            vec3<double> v3(ref_miller[ii]);
            mul_norm.push_back((F_dbdp[jj] * v3).length());
          }
        }
      }

      // Sum the total offsets to determine each range for splitting of mul_norm to
      // nref param values with F_dbdp.size() number of partitions
      int sum_offsets = std::accumulate(data_range.begin(), data_range.end(), 0);
      scitbx::af::shared<int> nref_each_param;
      for (ii = 0; ii < F_dbdp.size(); ++ii) {
        // Count the number of elements greater than 0 from the L-2 norms.
        nref_each_param.push_back(
          std::count_if(mul_norm.begin() + ii * sum_offsets,
                        mul_norm.begin() + (ii + 1) * sum_offsets,
                        gtZero));
      }

      // Return the value of the smallest element
      scitbx::af::shared<int>::iterator it =
        std::min_element(nref_each_param.begin(), nref_each_param.end());
      result = *(nref_each_param.begin() + std::distance(nref_each_param.begin(), it));
    }
    int result;
    static bool gtZero(double x) {
      return ((x > 0.0) == 1);
    }  // Used for comparison operation
  };

  struct surpl_iter {
    template <typename IntType>
    surpl_iter(scitbx::af::shared<IntType> ref_ids,
               const scitbx::af::shared<std::size_t> exp_ids) {
      // Expose pointers for experiment flex arrays
      const std::size_t* ptr_exp_ids = exp_ids.begin();

      std::size_t exp_it;  // ref_id;
      std::pair<const IntType*, const IntType*> refIDRange;

      unsigned int data_count = 0;
      result = 0;

      // Perform subselection of data ranges
      for (exp_it = 0; exp_it < exp_ids.size(); ++exp_it) {
        refIDRange =
          std::equal_range(ref_ids.begin(), ref_ids.end(), ptr_exp_ids[exp_it]);
        data_count = refIDRange.second - refIDRange.first;
        result += data_count;
      }
    }
    int result;
  };

  void export_pg_surpl_iter() {
    typedef return_value_policy<return_by_value> rbv;
    // templated constructor for int and size_t flex arrays of id
    class_<dials::refinement::boost_python::pg_surpl_iter>("pg_surpl_iter", no_init)
      .def(init<scitbx::af::shared<std::size_t>,
                const scitbx::af::shared<std::size_t>,
                const boost::python::list,
                const scitbx::af::shared<std::size_t>,
                const int>((boost::python::arg("ref_ids"),
                            boost::python::arg("panel_id"),
                            boost::python::arg("pnl_ids"),
                            boost::python::arg("exp_ids"),
                            boost::python::arg("cutoff"))))
      .def(init<scitbx::af::shared<int>,
                const scitbx::af::shared<std::size_t>,
                const boost::python::list,
                const scitbx::af::shared<std::size_t>,
                const int>((boost::python::arg("ref_ids"),
                            boost::python::arg("panel_id"),
                            boost::python::arg("pnl_ids"),
                            boost::python::arg("exp_ids"),
                            boost::python::arg("cutoff"))))
      .add_property(
        "result",
        make_getter(&dials::refinement::boost_python::pg_surpl_iter::result, rbv()));
  }

  void export_uc_surpl_iter() {
    typedef return_value_policy<return_by_value> rbv;
    class_<dials::refinement::boost_python::uc_surpl_iter>("uc_surpl_iter", no_init)
      .def(init<const scitbx::af::shared<std::size_t>,
                const scitbx::af::shared<cctbx::miller::index<> >,
                const scitbx::af::shared<std::size_t>,
                const scitbx::af::shared<mat3<double> > >(
        (boost::python::arg("ref_ids"),
         boost::python::arg("ref_miller"),
         boost::python::arg("exp_ids"),
         boost::python::arg("F_dbdp"))))
      .def(init<const scitbx::af::shared<int>,
                const scitbx::af::shared<cctbx::miller::index<> >,
                const scitbx::af::shared<std::size_t>,
                const scitbx::af::shared<mat3<double> > >(
        (boost::python::arg("ref_ids"),
         boost::python::arg("ref_miller"),
         boost::python::arg("exp_ids"),
         boost::python::arg("F_dbdp"))))
      .add_property(
        "result",
        make_getter(&dials::refinement::boost_python::uc_surpl_iter::result, rbv()));
  }

  void export_surpl_iter() {
    typedef return_value_policy<return_by_value> rbv;
    class_<dials::refinement::boost_python::surpl_iter>("surpl_iter", no_init)
      .def(
        init<scitbx::af::shared<std::size_t>, const scitbx::af::shared<std::size_t> >(
          (boost::python::arg("ref_ids"), boost::python::arg("exp_ids"))))
      .def(init<scitbx::af::shared<int>, const scitbx::af::shared<std::size_t> >(
        (boost::python::arg("ref_ids"), boost::python::arg("exp_ids"))))
      .add_property(
        "result",
        make_getter(&dials::refinement::boost_python::surpl_iter::result, rbv()));
  }
}}}  // namespace dials::refinement::boost_python


/*
 *  Copyright (C) 2018 Naomasa Matsubayashi
 *  Licensed under MIT license, see file LICENSE in this source tree.
 */
#ifndef UWG_TYPES_HPP
#define UWG_TYPES_HPP
#include <functional>
#include <memory>
#include <boost/container/static_vector.hpp>
#include <boost/container/flat_map.hpp>
#include <boost/asio.hpp>
#include <boost/asio/system_timer.hpp>
#include <uwg/defs.hpp>
#include <uwg/peer_identifier.hpp>
#include <uwg/window.hpp>

namespace uwg {
  enum class result_t {
    OK,
    FORBIDDEN,
    DISCONNECTED,
    NOT_IMPLEMENTED
  };
  using wg_key_type = boost::container::static_vector< unsigned char, wg_key_len >;
  using wg_tai64n_type = boost::container::static_vector< unsigned char, wg_tai64n_len >;
  using timer_type = boost::asio::system_timer;
  using kxcb_type = std::function< void( result_t ) >;
  struct kx_state {
    kx_state( boost::asio::io_service &ios ) : timer( new timer_type( ios ) ), cb( []( result_t ){} ) {}
    wg_key_type self_ephemeral_private;
    wg_key_type self_ephemeral_public;
    wg_key_type chain_key;
    wg_key_type hash_key;
    wg_key_type cookie;
    wg_key_type mac1;
    std::shared_ptr< timer_type > timer;
    kxcb_type cb;
  };
  using kx_states = boost::container::flat_map< peer_identifier, kx_state >;
  struct responder_state {
    responder_state() : cookie_since( 0 ) {}
    wg_key_type cookie;
    time_t cookie_since;
  };
  using responder_states = boost::container::flat_map< peer_identifier, responder_state >;
  struct key_state {
    key_state() : since( 0 ), tx_count( 0 ) {}
    key_state( boost::asio::io_service &ios ) : timer( new timer_type( ios ) ), since( 0 ), tx_count( 0 ) {}
    wg_key_type send_key;
    wg_key_type receive_key;
    std::shared_ptr< timer_type > timer;
    time_t since;
    size_t tx_count;
    uint32_t remote_kxid;
    window_state window;
  };
  using key_ring_t = boost::container::flat_map< peer_identifier, std::shared_ptr< key_state > >;
  struct peer_state {
    peer_state() : timestamp( wg_tai64n_len, 0 ) {}
    peer_state( const wg_tai64n_type &ts ) : timestamp( ts.begin(), ts.end() ) {}
    wg_tai64n_type timestamp;
  };
  struct key_comp {
    bool operator()( const wg_key_type &l, const  wg_key_type &r ) const {
      const auto min_size = std::min( l.size(), r.size() );
      const auto data_diff = memcmp( l.data(), r.data(), min_size );
      if( data_diff == 0 ) return l.size() < r.size();
      else return data_diff < 0;
    }
  };
  using peer_states = boost::container::flat_map< wg_key_type, peer_state, key_comp >;
}

#endif


#include <boost/archive/iterators/base64_from_binary.hpp>
#include <boost/archive/iterators/binary_from_base64.hpp>
#include <boost/archive/iterators/ostream_iterator.hpp>
#include <boost/archive/iterators/transform_width.hpp>

#include <elle/Exception.hh>
#include <elle/assert.hh>
#include <elle/format/base64.hh>
#include <elle/log.hh>

ELLE_LOG_COMPONENT("elle.format.base64")

namespace elle
{
  namespace format
  {
    namespace base64
    {
      using namespace boost::archive::iterators;

      size_t
      encoded_size(ConstWeakBuffer input)
      {
        return (signed(input.size()) + 2) / 3 * 4;
      }

      Buffer
      encode(ConstWeakBuffer input)
      {
        ELLE_TRACE_SCOPE("encode %s", input);
        size_t size = encoded_size(input);
        ELLE_DEBUG("previsional size: %s", size);
        Buffer res;
        res.capacity(size);
        IOStream stream(res.ostreambuf());
        Stream base64_stream(stream);
        base64_stream.write(reinterpret_cast<char const*>(input.contents()),
                            input.size());
        return res;
      }

      size_t
      decoded_size(ConstWeakBuffer encoded)
      {
        size_t size = encoded.size();
        int padding = 0;
        if (size >= 1 && encoded.contents()[size - 1] == '=')
          ++padding;
        if (size >= 2 && encoded.contents()[size - 2] == '=')
          ++padding;
        return size / 4 * 3 - padding;
      }

      Buffer
      decode(ConstWeakBuffer input)
      {
        ELLE_TRACE_SCOPE("decode %s", input);
        size_t size = decoded_size(input);
        ELLE_DEBUG("previsional size: %s", size);
        Buffer res(size);
        IOStream stream(input.istreambuf());
        Stream base64_stream(stream);
        base64_stream.read(reinterpret_cast<char*>(res.mutable_contents()),
                           size);
        ELLE_ASSERT_EQ(base64_stream.gcount(), signed(size));
        ELLE_ASSERT(
          (base64_stream.read(reinterpret_cast<char*>(
                                res.mutable_contents()), 1),
           base64_stream.gcount() == 0));

        return res;
      }

      using encoder = base64_from_binary<transform_width<const char *, 6, 8>>;
      template <typename SourceIterator, typename DestinationIterator>
      void
      _stream_encode(SourceIterator begin,
                     SourceIterator end,
                     DestinationIterator output)
      {
        std::copy(encoder(begin), encoder(end), output);
      }

      using decoder = transform_width<binary_from_base64<char const *>, 8, 6>;
      template <typename SourceIterator, typename DestinationIterator>
      void
      _stream_decode(SourceIterator begin,
                     SourceIterator end,
                     DestinationIterator output)
      {
        std::copy(decoder(begin), decoder(end), output);
      }

      /*-------------.
      | Construction |
      `-------------*/

      StreamBuffer::StreamBuffer(std::iostream& stream):
        _stream(stream),
        _remaining_write(0),
        _remaining_read(0)
      {}

      /*-------.
      | Buffer |
      `-------*/

      WeakBuffer
      StreamBuffer::read_buffer()
      {
        static auto const size = sizeof(this->_buffer_read) / 3 * 4;
        ELLE_TRACE_SCOPE("%s: decode up to %s bytes", *this, size);
        char buffer[size];
        size_t read;
        ELLE_DEBUG("%s: read up to %s bytes from the backend", *this, size)
        {
          this->_stream.read(buffer, size);
          read = this->_stream.gcount();
          ELLE_DUMP("%s: got %s bytes", *this, read);
          if (read > 0)
            ELLE_DUMP("%s: encoded data: %s", *this,
                      elle::WeakBuffer(buffer, read));
        }
        this->_remaining_read = read % 4;
        read = read - this->_remaining_read;
        size_t decoded_size = read / 4 * 3;
        if (read > 0)
        {
          ELLE_DEBUG_SCOPE("%s: decode %s bytes", *this, read);
          bool end = buffer[read - 1] == '=';
          if (end)
          {
            read -= 4;
            decoded_size -= 3;
          }
          _stream_decode(buffer, buffer + read, this->_buffer_read);
          if (end)
          {
            ELLE_DEBUG_SCOPE("%s: decode last 4 bytes with padding", *this);
            ELLE_DUMP("%s: remaining bytes: \"%s\"", *this,
                      std::string(buffer + read, 4));
            auto source = decoder(buffer + read);
            auto destination = this->_buffer_read + (read / 4 * 3);
            if (buffer[read + 2] != '=')
            {
              *(destination++) = *(source++);
              ++decoded_size;
            }
            *(destination++) = *(source++);
            ++decoded_size;
          }
          if (decoded_size > 0)
            ELLE_DUMP("%s: decoded data: %s", *this,
                      elle::WeakBuffer(this->_buffer_read, decoded_size));
        }
        if (this->_remaining_read > 0)
        {
          ELLE_DEBUG_SCOPE("%s: store %s remaining bytes",
                           *this, this->_remaining_read);
        }
        return WeakBuffer(this->_buffer_read, decoded_size);
      }

      WeakBuffer
      StreamBuffer::write_buffer()
      {
        return WeakBuffer(this->_buffer_write + this->_remaining_write,
                          sizeof(_buffer_write) - this->_remaining_write);
      }

      void
      StreamBuffer::flush(Size size)
      {
        size += this->_remaining_write;
        ELLE_TRACE_SCOPE("%s: encode %s bytes", *this, size);
        this->_remaining_write = size % 3;
        size = size - size % 3;
        if (size > 0)
        {
          ELLE_DEBUG_SCOPE("%s: encode %s bytes to the backend",
                           *this, size);
          _stream_encode(this->_buffer_write,
                         this->_buffer_write + size,
                         std::ostreambuf_iterator<char>(_stream));
        }
        if (size && this->_remaining_write > 0)
        {
          ELLE_DEBUG_SCOPE("%s: push back %s bytes in the buffer",
                           *this, this->_remaining_write);
          memcpy(this->_buffer_write,
                 this->_buffer_write + size,
                 this->_remaining_write);
        }
      }

      void
      StreamBuffer::finalize()
      {
        if (this->_remaining_write > 0)
        {
          ELLE_DEBUG_SCOPE("%s: encode last %s remaining bytes",
                           *this, this->_remaining_write);
          ELLE_ASSERT_LT(this->_remaining_write, 3);
          _stream_encode(
            this->_buffer_write,
            this->_buffer_write + this->_remaining_write,
            std::ostreambuf_iterator<char>(this->_stream));
          switch (this->_remaining_write)
          {
            case 1:
              this->_stream << "==";
              break;
            case 2:
              this->_stream << "=";
              break;
            default:
              elle::unreachable();
          }
        }
      }

      /*----------.
      | Printable |
      `----------*/

      void
      StreamBuffer::print(std::ostream& output) const
      {
        output << "base64::StreamBuffer";
      }

      Stream::Stream(std::iostream& underlying):
        IOStream(this->_buffer = new StreamBuffer(underlying)),
        _underlying(underlying)
      {}

      Stream::~Stream()
      {
        ELLE_TRACE_SCOPE("%s: finalize encoding", *this->_buffer);
        this->_buffer->sync();
        this->_buffer->finalize();
      }
    }
  }
}


#pragma once

#include <DXFeed.h>
#include <fmt/format.h>

#include <algorithm>
#include <boost/multi_index/identity.hpp>
#include <boost/multi_index/member.hpp>
#include <boost/multi_index/ordered_index.hpp>
#include <boost/multi_index/random_access_index.hpp>
#include <boost/multi_index_container.hpp>
#include <cassert>
#include <cmath>
#include <deque>
#include <functional>
#include <limits>
#include <memory>
#include <mutex>
#include <set>
#include <string>
#include <utility>
#include <vector>

#include "StringConverter.hpp"

namespace dxf {

struct OrderData {
  dxf_long_t index = 0;
  double price = std::numeric_limits<double>::quiet_NaN();
  double size = std::numeric_limits<double>::quiet_NaN();
  dxf_long_t time = 0;
  dxf_order_side_t side = dxf_osd_undefined;
};

struct PriceLevel {
  double price = std::numeric_limits<double>::quiet_NaN();
  double size = std::numeric_limits<double>::quiet_NaN();
  std::int64_t time = 0;

  friend bool operator<(const PriceLevel& a, const PriceLevel& b) {
    if (std::isnan(b.price)) return true;
    if (std::isnan(a.price)) return false;

    return a.price < b.price;
  }
};

struct PriceLevelChanges {
  std::vector<PriceLevel> asks{};
  std::vector<PriceLevel> bids{};
};

struct PriceLevelChangesSet {
  PriceLevelChanges additions{};
  PriceLevelChanges updates{};
  PriceLevelChanges removals{};
};

namespace bmi = boost::multi_index;

using PriceLevelContainer = bmi::multi_index_container<
  PriceLevel,
  bmi::indexed_by<bmi::random_access<>, bmi::ordered_unique<bmi::member<PriceLevel, double, &PriceLevel::price>>>>;

class PriceLevelBook final {
  dxf_snapshot_t snapshot_;
  std::string symbol_;
  std::string source_;
  std::size_t levelsNumber_;
  PriceLevelContainer asks_;
  PriceLevelContainer bids_;
  std::unordered_map<dxf_long_t, OrderData> orderDataSnapshot_;
  bool isValid_;
  std::mutex mutex_;

  std::function<void(const PriceLevelChanges&)> onNewBook_;
  std::function<void(const PriceLevelChanges&)> onBookUpdate_;
  std::function<void(const PriceLevelChangesSet&)> onIncrementalChange_;

  static bool isZeroPriceLevel(const PriceLevel& pl) {
    return std::abs(pl.size) < std::numeric_limits<double>::epsilon();
  };

  static bool areEqualPrices(double price1, double price2) {
    return std::abs(price1 - price2) < std::numeric_limits<double>::epsilon();
  }

  PriceLevelBook(std::string symbol, std::string source, std::size_t levelsNumber = 0)
      : snapshot_{nullptr},
        symbol_{std::move(symbol)},
        source_{std::move(source)},
        levelsNumber_{levelsNumber},
        asks_{},
        bids_{},
        orderDataSnapshot_{},
        isValid_{false},
        mutex_{} {}

  // Process the tx\snapshot data, converts it to PL changes. Also, changes the orderDataSnapshot_
  PriceLevelChanges convertToUpdates(const dxf_snapshot_data_ptr_t snapshotData) {
    assert(snapshotData->records_count != 0);
    assert(snapshotData->event_type != dx_eid_order);

    std::set<PriceLevel> askUpdates{};
    std::set<PriceLevel> bidUpdates{};

    auto orders = reinterpret_cast<const dxf_order_t*>(snapshotData->records);

    auto isOrderRemoval = [](const dxf_order_t& o) {
      return (o.event_flags & dxf_ef_remove_event) != 0 || o.size == 0 || std::isnan(o.size);
    };

    auto processOrderAddition = [&bidUpdates, &askUpdates](const dxf_order_t& order) {
      auto& updatesSide = order.side == dxf_osd_buy ? bidUpdates : askUpdates;
      auto priceLevelChange = PriceLevel{order.price, order.size, order.time};
      auto foundPriceLevel = updatesSide.find(priceLevelChange);

      if (foundPriceLevel != updatesSide.end()) {
        priceLevelChange.size = foundPriceLevel->size + priceLevelChange.size;
        updatesSide.erase(foundPriceLevel);
      }

      updatesSide.insert(priceLevelChange);
    };

    auto processOrderRemoval = [&bidUpdates, &askUpdates](const dxf_order_t& order, const OrderData& foundOrderData) {
      auto& updatesSide = foundOrderData.side == dxf_osd_buy ? bidUpdates : askUpdates;
      auto priceLevelChange = PriceLevel{foundOrderData.price, -foundOrderData.size, order.time};
      auto foundPriceLevel = updatesSide.find(priceLevelChange);

      if (foundPriceLevel != updatesSide.end()) {
        priceLevelChange.size = foundPriceLevel->size + priceLevelChange.size;
        updatesSide.erase(foundPriceLevel);
      }

      if (!isZeroPriceLevel(priceLevelChange)) {
        updatesSide.insert(priceLevelChange);
      }
    };

    for (std::size_t i = 0; i < snapshotData->records_count; i++) {
      auto order = orders[i];

      fmt::print("O:ind={},pr={},sz={},sd={}\n", order.index, order.price, order.size,
                 order.side == dxf_osd_buy ? "buy" : "sell");

      auto removal = isOrderRemoval(order);
      auto foundOrderDataIt = orderDataSnapshot_.find(order.index);

      if (foundOrderDataIt == orderDataSnapshot_.end()) {
        if (removal) {
          continue;
        }

        processOrderAddition(order);
        orderDataSnapshot_[order.index] = OrderData{order.index, order.price, order.size, order.time, order.side};
      } else {
        const auto& foundOrderData = foundOrderDataIt->second;

        if (removal) {
          processOrderRemoval(order, foundOrderData);
          orderDataSnapshot_.erase(foundOrderData.index);
        } else {
          if (order.side != foundOrderData.side) {
            processOrderRemoval(order, foundOrderData);
          }

          processOrderAddition(order);
          orderDataSnapshot_[foundOrderData.index] =
            OrderData{order.index, order.price, order.size, order.time, order.side};
        }
      }
    }

    return {std::vector<PriceLevel>{askUpdates.begin(), askUpdates.end()},
            std::vector<PriceLevel>{bidUpdates.rbegin(), bidUpdates.rend()}};
  }

  // TODO: price levels number
  PriceLevelChangesSet applyUpdates(const PriceLevelChanges& priceLevelUpdates) {
    PriceLevelChanges additions{};
    PriceLevelChanges updates{};
    PriceLevelChanges removals{};

    // We generate lists of additions, updates, removals
    for (const auto& updateAsk : priceLevelUpdates.asks) {
      auto found = std::lower_bound(asks_.begin(), asks_.end(), updateAsk);

      if (found == asks_.end() || !areEqualPrices(found->price, updateAsk.price)) {
        additions.asks.push_back(updateAsk);
      } else {
        auto newPriceLevelChange = *found;

        newPriceLevelChange.size += updateAsk.size;
        newPriceLevelChange.time = updateAsk.time;

        if (isZeroPriceLevel(newPriceLevelChange)) {
          removals.asks.push_back(*found);
        } else {
          updates.asks.push_back(newPriceLevelChange);
        }
      }
    }

    for (const auto& updateBid : priceLevelUpdates.bids) {
      auto found = std::lower_bound(bids_.begin(), bids_.end(), updateBid);

      if (found == bids_.end() || !areEqualPrices(found->price, updateBid.price)) {
        additions.bids.push_back(updateBid);
      } else {
        auto newPriceLevelChange = *found;

        newPriceLevelChange.size += updateBid.size;
        newPriceLevelChange.time = updateBid.time;

        if (isZeroPriceLevel(newPriceLevelChange)) {
          removals.bids.push_back(*found);
        } else {
          updates.bids.push_back(newPriceLevelChange);
        }
      }
    }

    std::set<PriceLevel> askRemovals{};
    std::set<PriceLevel> bidRemovals{};
    std::set<PriceLevel> askAdditions{};
    std::set<PriceLevel> bidAdditions{};
    std::set<PriceLevel> askUpdates{};
    std::set<PriceLevel> bidUpdates{};

    for (const auto& askRemoval : removals.asks) {
      if (asks_.empty()) continue;

      // Determine what will be the removal given the number of price levels.
      if (levelsNumber_ == 0 || asks_.size() <= levelsNumber_ || askRemoval.price < asks_[levelsNumber_].price) {
        askRemovals.insert(askRemoval);
      }

      // Determine what will be the shift in price levels after removal.
      if (levelsNumber_ != 0 && asks_.size() > levelsNumber_ && askRemoval.price < asks_[levelsNumber_].price) {
        askAdditions.insert(asks_[levelsNumber_]);
      }

      // remove price level by price
      asks_.get<1>().erase(askRemoval.price);
    }

    for (const auto& askAddition : additions.asks) {
      // We determine what will be the addition of the price level, taking into account the possible quantity.
      if (levelsNumber_ == 0 || asks_.size() < levelsNumber_ || askAddition.price < asks_[levelsNumber_ - 1].price) {
        askAdditions.insert(askAddition);
      }

      // We determine what will be the shift after adding
      if (levelsNumber_ != 0 && asks_.size() >= levelsNumber_ && askAddition.price < asks_[levelsNumber_ - 1].price) {
        const auto& toRemove = asks_[levelsNumber_ - 1];

        // We take into account the possibility that the previously added price level will be deleted.
        if (askAdditions.contains(toRemove)) {
          askAdditions.erase(toRemove);
        } else {
          askRemovals.insert(toRemove);
        }
      }

      asks_.get<1>().insert(askAddition);
    }

    for (const auto& askUpdate : updates.asks) {
      if (levelsNumber_ == 0 || asks_.get<1>().count(askUpdate.price) > 0) {
        askUpdates.insert(askUpdate);
      }

      asks_.get<1>().erase(askUpdate.price);
      asks_.get<1>().insert(askUpdate);
    }

    for (const auto& bidRemoval : removals.bids) {
      if (bids_.empty()) continue;

      // Determine what will be the removal given the number of price levels.
      if (levelsNumber_ == 0 || bids_.size() <= levelsNumber_ ||
          bidRemoval.price > bids_[bids_.size() - 1 - levelsNumber_].price) {
        bidRemovals.insert(bidRemoval);
      }

      // Determine what will be the shift in price levels after removal.
      if (levelsNumber_ != 0 && bids_.size() > levelsNumber_ &&
          bidRemoval.price > bids_[bids_.size() - 1 - levelsNumber_].price) {
        bidAdditions.insert(bids_[bids_.size() - 1 - levelsNumber_]);
      }

      // remove price level by price
      bids_.get<1>().erase(bidRemoval.price);
    }

    for (const auto& bidAddition : additions.bids) {
      // We determine what will be the addition of the price level, taking into account the possible quantity.
      if (levelsNumber_ == 0 || bids_.size() < levelsNumber_ ||
          bidAddition.price > bids_[bids_.size() - levelsNumber_].price) {
        bidAdditions.insert(bidAddition);
      }

      // We determine what will be the shift after adding
      if (levelsNumber_ != 0 && bids_.size() >= levelsNumber_ &&
          bidAddition.price > bids_[bids_.size() - levelsNumber_].price) {
        const auto& toRemove = bids_[bids_.size() - levelsNumber_];

        // We take into account the possibility that the previously added price level will be deleted.
        if (bidAdditions.contains(toRemove)) {
          bidAdditions.erase(toRemove);
        } else {
          bidRemovals.insert(toRemove);
        }
      }

      bids_.get<1>().insert(bidAddition);
    }

    for (const auto& bidUpdate : updates.bids) {
      if (levelsNumber_ == 0 || bids_.get<1>().count(bidUpdate.price) > 0) {
        bidUpdates.insert(bidUpdate);
      }

      bids_.get<1>().erase(bidUpdate.price);
      bids_.get<1>().insert(bidUpdate);
    }

    return {PriceLevelChanges{std::vector<PriceLevel>{askAdditions.begin(), askAdditions.end()},
                              std::vector<PriceLevel>{bidAdditions.rbegin(), bidAdditions.rend()}},
            PriceLevelChanges{std::vector<PriceLevel>{askUpdates.begin(), askUpdates.end()},
                              std::vector<PriceLevel>{bidUpdates.rbegin(), bidUpdates.rend()}},
            PriceLevelChanges{std::vector<PriceLevel>{askRemovals.begin(), askRemovals.end()},
                              std::vector<PriceLevel>{bidRemovals.rbegin(), bidRemovals.rend()}}};
  }

  [[nodiscard]] std::vector<PriceLevel> getAsks() const {
    return {asks_.begin(),
            (levelsNumber_ == 0 || asks_.size() <= levelsNumber_) ? asks_.end() : asks_.begin() + levelsNumber_};
  }

  [[nodiscard]] std::vector<PriceLevel> getBids() const {
    return {bids_.rbegin(),
            (levelsNumber_ == 0 || bids_.size() <= levelsNumber_) ? bids_.rend() : bids_.rbegin() + levelsNumber_};
  }

 public:
  // TODO: move to another thread
  void processSnapshotData(const dxf_snapshot_data_ptr_t snapshotData, int newSnapshot) {
    std::lock_guard<std::mutex> lk(mutex_);

    auto newSnap = newSnapshot != 0;

    if (newSnap) {
      asks_.clear();
      bids_.clear();
      orderDataSnapshot_.clear();
    }

    if (snapshotData->records_count == 0) {
      if (newSnap && onNewBook_) {
        onNewBook_({});
      }

      return;
    }

    auto updates = convertToUpdates(snapshotData);
    auto resultingChangesSet = applyUpdates(updates);

    if (newSnap) {
      if (onNewBook_) {
        onNewBook_(PriceLevelChanges{getAsks(), getBids()});
      }
    } else {
      if (onIncrementalChange_) {
        onIncrementalChange_(resultingChangesSet);
      }

      if (onBookUpdate_) {
        onBookUpdate_(PriceLevelChanges{getAsks(), getBids()});
      }
    }
  }

  ~PriceLevelBook() {
    if (isValid_) {
      dxf_close_price_level_book(snapshot_);
    }
  }

  static std::unique_ptr<PriceLevelBook> create(dxf_connection_t connection, const std::string& symbol,
                                                const std::string& source, std::size_t levelsNumber) {
    auto plb = std::unique_ptr<PriceLevelBook>(new PriceLevelBook(symbol, source, levelsNumber));
    auto wSymbol = StringConverter::utf8ToWString(symbol);
    dxf_snapshot_t snapshot = nullptr;

    dxf_create_order_snapshot(connection, wSymbol.c_str(), source.c_str(), 0, &snapshot);

    dxf_attach_snapshot_inc_listener(
      snapshot,
      [](const dxf_snapshot_data_ptr_t snapshot_data, int new_snapshot, void* user_data) {
        static_cast<PriceLevelBook*>(user_data)->processSnapshotData(snapshot_data, new_snapshot);
      },
      plb.get());
    plb->isValid_ = true;

    return plb;
  }

  void setOnNewBook(std::function<void(const PriceLevelChanges&)> onNewBookHandler) {
    onNewBook_ = std::move(onNewBookHandler);
  }

  void setOnBookUpdate(std::function<void(const PriceLevelChanges&)> onBookUpdateHandler) {
    onBookUpdate_ = std::move(onBookUpdateHandler);
  }

  void setOnIncrementalChange(std::function<void(const PriceLevelChangesSet&)> onIncrementalChangeHandler) {
    onIncrementalChange_ = std::move(onIncrementalChangeHandler);
  }
};

}  // namespace dxf