//
// Created by david on 2020-12-20.
//

#include "class_quantum_gates.h"
#include <Eigen/Core>
#include <general/nmspc_iter.h>
#include <general/nmspc_tensor_extra.h>
#include <math/linalg/tensor.h>
#include <math/num.h>
#include <set>
#include <tools/common/log.h>
#include <unsupported/Eigen/MatrixFunctions>
#include <utility>

template<typename T>
std::vector<T> subset(const std::vector<T> &vec, size_t idx_start, size_t num) {
    if(idx_start + num > vec.size())
        throw std::range_error(fmt::format("Vector subset start {} num {} out of range for vector of size {}", idx_start, num, vec.size()));
    auto vec_bgn = vec.begin() + static_cast<long>(idx_start);
    auto vec_end = vec_bgn + static_cast<long>(num);
    return std::vector<T>(vec_bgn, vec_end);
}

template<auto N, typename T, auto M>
std::array<T, N> subset(const std::array<T, M> &arr, size_t idx_start) {
    if(idx_start + N > arr.size())
        throw std::range_error(fmt::format("Vector subset start {} num {} out of range for vector of size {}", idx_start, N, arr.size()));
    auto             vec_bgn = arr.begin() + idx_start;
    auto             vec_end = vec_bgn + N;
    std::array<T, N> res{};
    std::copy(vec_bgn, vec_end, res.begin());
    return res;
}

template<typename T>
T prod(const std::vector<T> &vec, size_t idx_start, size_t num, T seed = 1) {
    auto sub = subset(vec, idx_start, num);
    return std::accumulate(sub.begin(), sub.end(), std::multiplies(seed));
}

template<auto N, typename T, auto M>
T prod(const std::array<T, M> &arr, size_t idx_start, T seed = 1) {
    auto sub = subset<N>(arr, idx_start);
    T    acc = seed;
    for(auto &s : sub) acc *= s;
    return acc;
}

template<typename T>
std::vector<T> concat(const std::vector<T> &v1, const std::vector<T> &v2) {
    auto res = v1;
    res.insert(res.end(), v2.begin(), v2.end());
    return res;
}

template<typename T, auto N, auto M>
auto concat(const std::array<T, N> &a1, const std::array<T, M> &a2) {
    std::array<T, N + M> result{};
    std::copy(a1.cbegin(), a1.cend(), result.begin());
    std::copy(a2.cbegin(), a2.cend(), result.begin() + N);
    return result;
}

template<typename T, auto N>
auto repeat(const std::array<T, N> &a) {
    return concat(a, a);
}

auto group(const std::vector<long> &dim, const std::vector<size_t> &pattern) {
    std::vector<long> res(pattern.size());
    size_t            dim_offset = 0;
    for(size_t i = 0; i < pattern.size(); i++) {
        long product = 1;
        for(size_t j = 0; j < pattern[i]; j++) product *= dim[dim_offset + j];
        res[i] = product;
        dim_offset += pattern[i];
    }
    return res;
}

template<auto N, auto M>
auto group(const std::array<long, N> &dim, const std::array<size_t, M> &pattern) {
    std::array<long, M> res{};
    size_t              dim_offset = 0;
    for(size_t i = 0; i < M; i++) {
        long product = 1;
        for(size_t j = 0; j < pattern[i]; j++) product *= dim[dim_offset + j];
        res[i] = product;
        dim_offset += pattern[i];
    }
    return res;
}

template<typename T1, typename T2>
void erase(std::vector<T1> &vec, T2 val) {
    auto it = vec.begin();
    while(it != vec.end()) {
        if(*it == static_cast<T1>(val))
            it = vec.erase(it);
        else
            ++it;
    }
}

Eigen::Tensor<qm::Scalar, 2> contract_a(const Eigen::Tensor<qm::Scalar, 2> &m, const Eigen::Tensor<qm::Scalar, 2> &ud, const std::array<long, 4> &shp_mid4,
                                        const std::array<long, 4> &shp_udn4, const std::array<long, 6> &shf6, const Textra::idxlistpair<1> &idx1,
                                        const Textra::idxlistpair<2> &idx2, const std::array<long, 2> &dim2) {
    return ud.reshape(shp_udn4)
        .contract(m.reshape(shp_mid4), idx1)
        .contract(ud.conjugate().shuffle(Textra::array2{1, 0}).reshape(shp_udn4), idx2)
        .shuffle(shf6)
        .reshape(dim2);
}

Eigen::Tensor<qm::Scalar, 2> contract_b(const Eigen::Tensor<qm::Scalar, 2> &m, const Eigen::Tensor<qm::Scalar, 2> &ud, const std::array<long, 2> &shp_udn2,
                                        const std::array<long, 4> &shp_udn4, const Textra::idxlistpair<1> &idx1, const Textra::idxlistpair<2> &idx2) {
    return ud.reshape(shp_udn4).contract(m, idx1).contract(ud.conjugate().shuffle(Textra::array2{1, 0}).reshape(shp_udn4), idx2).reshape(shp_udn2);
}

Eigen::Tensor<qm::Scalar, 2> contract_c(const Eigen::Tensor<qm::Scalar, 2> &m, const Eigen::Tensor<qm::Scalar, 2> &ud, const std::array<long, 6> &shp_mid6,
                                        const Textra::idxlistpair<1> &idx_up, const Textra::idxlistpair<1> &idx_dn, const std::array<long, 6> &shf6,
                                        const std::array<long, 2> &dim2) {
    return ud.contract(m.reshape(shp_mid6), idx_up).contract(ud.conjugate().shuffle(Textra::array2{1, 0}), idx_dn).shuffle(shf6).reshape(dim2);
}

Eigen::Tensor<qm::Scalar, 2> contract_d(const Eigen::Tensor<qm::Scalar, 2> &m, const Eigen::Tensor<qm::Scalar, 2> &ud, const std::array<long, 4> &shp_mid4,
                                        const Textra::idxlistpair<1> &idx_up, const Textra::idxlistpair<1> &idx_dn, const std::array<long, 4> &shf4,
                                        const std::array<long, 2> &dim2) {
    return ud.contract(m.reshape(shp_mid4), idx_up).contract(ud.conjugate().shuffle(Textra::array2{1, 0}), idx_dn).shuffle(shf4).reshape(dim2);
}

Eigen::Tensor<qm::Scalar, 2> qm::Gate::exp_internal(const Eigen::Tensor<Scalar, 2> &op_, Scalar alpha) const {
    /* Note fore flbit:
     *  Let h = op(i,i), i.e. h are the diagonal entries in op
     *  Let alpha = -i * delta be purely imaginary (since delta is real). So delta = imag(-alpha)
     *  Now notice that
     *       exp( -i * delta * h )
     *  can become imprecise when delta is large, since delta and h are real and h ~O(1)
     *
     *  Instead, in the flbit case, we can exploit that h is real to compute
     *       exp(-i * mod(delta * real(h), 2*pi))
     *  which is equivalent but with a much smaller exponent.
     *
     *  Remember to do the modulo separately on each diagonal entry h!
     */
    auto op_map = Textra::MatrixMap(op_);
    if(op_map.isDiagonal() and op_map.imag().isZero() and std::real(alpha) == 0) {
        auto minus_i = std::complex<double>(0, -1);

        auto diag =
            op_map.diagonal()
                .unaryViewExpr([&alpha, &minus_i](const Scalar &h) { return std::exp(minus_i * std::fmod(std::imag(-alpha) * std::real(h), 2.0 * M_PI)); })
                .asDiagonal();
        return Textra::TensorMap(diag.toDenseMatrix());
        //        std::cout << "alpha: " << alpha << std::endl;
        //        std::cout << "op:\n" << op << std::endl;
        //        std::cout << "old:\n" << (alpha * Textra::MatrixMap(op_)).exp() << std::endl;
        //        if(not Textra::MatrixMap(op).isApprox((alpha * Textra::MatrixMap(op_)).exp())){
        //            throw std::runtime_error("Mismatch");
        //        }
    } else {
        return Textra::TensorMap((alpha * Textra::MatrixMap(op_)).exp().eval());
    }
}

qm::Gate::Gate(const Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic> &op_, std::vector<size_t> pos_, std::vector<long> dim_)
    : pos(std::move(pos_)), dim(std::move(dim_)) {
    auto dim_prod = std::accumulate(std::begin(dim), std::end(dim), 1l, std::multiplies<>());
    if(dim_prod != op_.rows() or dim_prod != op_.cols())
        throw std::logic_error(fmt::format("dim {} not compatible with matrix dimensions {} x {}", dim, op_.rows(), op_.cols()));
    if(pos.size() != dim.size()) throw std::logic_error(fmt::format("pos.size() {} != dim.size() {}", pos, dim));
    op = Textra::TensorMap(op_);
}

qm::Gate::Gate(const Eigen::Tensor<Scalar, 2> &op_, std::vector<size_t> pos_, std::vector<long> dim_) : op(op_), pos(std::move(pos_)), dim(std::move(dim_)) {
    auto dim_prod = std::accumulate(std::begin(dim), std::end(dim), 1, std::multiplies<>());
    if(dim_prod != op_.dimension(0) or dim_prod != op_.dimension(1))
        throw std::logic_error(fmt::format("dim {} not compatible with matrix dimensions {} x {}", dim, op_.dimension(0), op_.dimension(1)));
    if(pos.size() != dim.size()) throw std::logic_error(fmt::format("pos.size() {} != dim.size() {}", pos, dim));
}
qm::Gate::Gate(const Eigen::Tensor<Scalar, 2> &op_, std::vector<size_t> pos_, std::vector<long> dim_, Scalar alpha)
    : pos(std::move(pos_)), dim(std::move(dim_)) {
    auto dim_prod = std::accumulate(std::begin(dim), std::end(dim), 1, std::multiplies<>());
    if(dim_prod != op_.dimension(0) or dim_prod != op_.dimension(1))
        throw std::logic_error(fmt::format("dim {} not compatible with matrix dimensions {} x {}", dim, op_.dimension(0), op_.dimension(1)));
    if(pos.size() != dim.size()) throw std::logic_error(fmt::format("pos.size() {} != dim.size() {}", pos, dim));
    op = exp_internal(op_, alpha);
}

void                          qm::Gate::exp_inplace(Scalar alpha) { op = exp_internal(op, alpha); }
qm::Gate                      qm::Gate::exp(Scalar alpha) const { return Gate(op, pos, dim, alpha); }
bool                          qm::Gate::isUnitary(double prec) const { return Textra::MatrixMap(op).isUnitary(prec); }
Eigen::Tensor<qm::Scalar, 2> &qm::Gate::adjoint() const {
    if(adj) return adj.value();
    adj = op.conjugate().shuffle(Textra::array2{1, 0});
    return adj.value();
}

template<auto rank>
std::array<long, rank> qm::Gate::shape() const {
    if constexpr(rank == 1) { return {op.dimension(0) * op.dimension(1)}; }
    if constexpr(rank == 2) {
        return {op.dimension(0), op.dimension(1)};
    } else {
        if(rank == 2 * dim.size()) {
            std::array<long, rank> dims{};
            for(size_t i = 0; i < dims.size(); i++) dims[i] = dim[i % dim.size()];
            return dims;
        } else {
            throw std::range_error(fmt::format("Can't compute shape of rank {} for gate with pos {} and dim {}", rank, pos, dim));
        }
    }
}
template std::array<long, 1>  qm::Gate::shape() const;
template std::array<long, 2>  qm::Gate::shape() const;
template std::array<long, 4>  qm::Gate::shape() const;
template std::array<long, 6>  qm::Gate::shape() const;
template std::array<long, 8>  qm::Gate::shape() const;
template std::array<long, 10> qm::Gate::shape() const;
template std::array<long, 12> qm::Gate::shape() const;
template std::array<long, 14> qm::Gate::shape() const;
template std::array<long, 16> qm::Gate::shape() const;

std::vector<size_t> qm::Gate::idx() const {
    std::vector<size_t> idx(pos.size() * 2);
    std::iota(idx.begin(), idx.end(), 0);
    return idx;
}

std::vector<size_t> qm::Gate::idx_up() const {
    std::vector<size_t> idx(pos.size());
    std::iota(idx.begin(), idx.end(), 0);
    return idx;
}

std::vector<size_t> qm::Gate::idx_dn() const {
    std::vector<size_t> idx(pos.size());
    std::iota(idx.begin(), idx.end(), pos.size());
    return idx;
}

std::vector<size_t> qm::Gate::idx(const std::vector<size_t> &pos_) const {
    std::vector<size_t> idx;
    for(const auto &p_ : pos_)
        for(const auto &[i, p] : iter::enumerate(pos))
            if(p == p_) {
                idx.emplace_back(i);
                idx.emplace_back(i + pos.size());
            }
    return idx;
}

std::vector<size_t> qm::Gate::idx_up(const std::vector<size_t> &pos_) const {
    std::vector<size_t> idx;
    for(const auto &p_ : pos_)
        for(const auto &[i, p] : iter::enumerate(pos))
            if(p == p_) { idx.emplace_back(i); }
    return idx;
}

std::vector<size_t> qm::Gate::idx_dn(const std::vector<size_t> &pos_) const {
    std::vector<size_t> idx;
    for(const auto &p_ : pos_)
        for(const auto &[i, p] : iter::enumerate(pos))
            if(p == p_) { idx.emplace_back(i + pos.size()); }
    return idx;
}

qm::Gate qm::Gate::insert(const Gate &other) const { return qm::insert(*this, other); }
qm::Gate qm::Gate::connect_above(const Gate &other) const { return qm::connect(other, *this); }
qm::Gate qm::Gate::connect_under(const Gate &other) const { return qm::connect(*this, other); }

template<auto N>
qm::Gate qm::Gate::trace(const std::array<Eigen::IndexPair<Eigen::Index>, N> &idxpair) const {
    return qm::trace(*this, idxpair);
}
template qm::Gate qm::Gate::trace(const std::array<Eigen::IndexPair<Eigen::Index>, 1> &idxpairs) const;
template qm::Gate qm::Gate::trace(const std::array<Eigen::IndexPair<Eigen::Index>, 2> &idxpairs) const;

qm::Gate   qm::Gate::trace_idx(const std::vector<long> &idx_) const { return qm::trace_idx(*this, idx_); }
qm::Gate   qm::Gate::trace_pos(const std::vector<size_t> &pos_) const { return qm::trace_pos(*this, pos_); }
qm::Gate   qm::Gate::trace_pos(size_t pos_) const { return qm::trace_pos(*this, pos_); }
qm::Scalar qm::Gate::trace() const { return qm::trace(*this); }

std::vector<std::vector<size_t>> qm::get_gate_sequence(const std::vector<qm::Gate> &layer) {
    std::vector<std::vector<size_t>> gate_sequence;
    size_t                           gate_size = layer.front().pos.size();
    size_t                           pos_max   = layer.back().pos.back();
    for(size_t offset = 0; offset < gate_size; offset++) {
        if(offset + gate_size > pos_max + 1) break;
        auto off_idx = num::range<size_t>(offset, pos_max - gate_size + 2, gate_size);
        if(num::mod<size_t>(offset, 2) == 1) std::reverse(off_idx.begin(), off_idx.end()); // If odd, reverse the sequence
        gate_sequence.emplace_back(off_idx);
    }
    return gate_sequence;
}

template<iter::order o>
std::vector<std::vector<size_t>> qm::get_lightcone(const std::vector<std::vector<qm::Gate>> &layers, size_t pos) {
    std::vector<std::vector<size_t>> cone;
    cone.emplace_back(std::vector<size_t>{pos});
    for(const auto &[idx_layer, layer] : iter::enumerate<o>(layers)) {
        if(layer.empty()) continue;
        auto gate_sequence = get_gate_sequence(layer);
        for(const auto &[idx_sublayer, seq] : iter::enumerate<o>(gate_sequence)) {
            std::set<size_t> match;
            // The cone width is strictly increasing.
            // So we immediately add the sites from the previous layer.
            match.insert(cone.back().begin(), cone.back().end());
            for(const auto &[idx_seq, pos_gate] : iter::enumerate<o>(seq)) {
                // Now we try to find some more matching sites in the new layer.
                auto &              u = layer[pos_gate];
                std::vector<size_t> pos_isect;
                std::set_intersection(cone.back().begin(), cone.back().end(), u.pos.begin(), u.pos.end(), back_inserter(pos_isect));
                if(not pos_isect.empty()) match.insert(u.pos.begin(), u.pos.end()); // Add positions from u if the gate connects to the current cone
            }
            if(match.empty())
                cone.emplace_back(cone.back());
            else
                cone.emplace_back(std::vector<size_t>(match.begin(), match.end()));
        }
    }
    if constexpr(o == iter::order::rev) std::reverse(cone.begin(), cone.end());
    return cone;
}

template std::vector<std::vector<size_t>> qm::get_lightcone<iter::order::def>(const std::vector<std::vector<qm::Gate>> &layers, size_t pos);
template std::vector<std::vector<size_t>> qm::get_lightcone<iter::order::rev>(const std::vector<std::vector<qm::Gate>> &layers, size_t pos);

std::vector<std::vector<size_t>> qm::get_lightcone_intersection(const std::vector<std::vector<qm::Gate>> &unitary_layers, size_t pos_tau, size_t pos_sig) {
    auto tau_cone = qm::get_lightcone<iter::order::def>(unitary_layers, pos_tau);
    auto sig_cone = qm::get_lightcone<iter::order::rev>(unitary_layers, pos_sig); // This cone is upside down!
    if(tau_cone.size() != sig_cone.size()) throw std::runtime_error("tau and sig cones should have equal size!");

    std::vector<std::vector<size_t>> int_cone;

    // Find the intersection between tau and sig cones
    for(size_t idx_sublayer = 0; idx_sublayer < tau_cone.size(); idx_sublayer++) {
        auto &              tau_sublayer = tau_cone[idx_sublayer];
        auto &              sig_sublayer = sig_cone[idx_sublayer];
        std::vector<size_t> pos_isect;
        std::set_intersection(tau_sublayer.begin(), tau_sublayer.end(), sig_sublayer.begin(), sig_sublayer.end(), back_inserter(pos_isect));
        int_cone.emplace_back(pos_isect);
    }
    bool deb = tools::log->level() == spdlog::level::trace;
    if(deb) {
        fmt::print("Lightcones\n");
        auto tau_pic = get_lightcone_picture(unitary_layers, tau_cone, "tau");
        auto sig_pic = get_lightcone_picture(unitary_layers, sig_cone, "sig");
        auto int_pic = get_lightcone_picture(unitary_layers, int_cone, "int");
        for(const auto &c : iter::reverse(tau_pic)) fmt::print("{}\n", c);
        for(const auto &c : iter::reverse(sig_pic)) fmt::print("{}\n", c);
        for(const auto &c : iter::reverse(int_pic)) fmt::print("{}\n", c);
    }

    return int_cone;
}

std::vector<std::string> qm::get_lightcone_picture(const std::vector<std::vector<qm::Gate>> &layers, const std::vector<std::vector<size_t>> &cone,
                                                   std::string_view tag, size_t pw, std::string_view sep) {
    std::vector<std::string> pic;
    if(not layers.empty() and not cone.empty()) {
        size_t sw = sep.size();
        size_t tw = static_cast<size_t>(tag.size()) + 5;                                       // Tag width (brackets, number and colon)
        size_t mw = static_cast<size_t>(layers.front().back().pos.back() + 1) * (pw + sw) + 1; // max cone width
        pic       = std::vector<std::string>(cone.size(), fmt::format("{0:^{1}}", " ", tw + mw));
        for(const auto &[i, c] : iter::enumerate(cone)) {
            pic[i].replace(0, tw, fmt::format("{}[{:^2}]:", tag, i));
            for(const auto &[j, p] : iter::enumerate(c)) pic[i].replace(tw + p * (pw + sw), pw, fmt::format("{0:>{1}}{2}", p, pw, sep));
        }
    }
    return pic;
}

qm::Gate qm::insert(const qm::Gate &middle_gate, const qm::Gate &updown_gate) {
    std::vector<size_t> pos_isect; // locations that intersect: both on middle and updown gates
    std::vector<size_t> pos_nsect; // locations that do not intersect: not in both on middle and updown gates
    std::set_intersection(middle_gate.pos.begin(), middle_gate.pos.end(), updown_gate.pos.begin(), updown_gate.pos.end(), back_inserter(pos_isect));
    std::set_symmetric_difference(middle_gate.pos.begin(), middle_gate.pos.end(), updown_gate.pos.begin(), updown_gate.pos.end(), back_inserter(pos_nsect));
    if(pos_isect.empty()) return middle_gate;
    bool inc = std::includes(middle_gate.pos.begin(), middle_gate.pos.end(), updown_gate.pos.begin(), updown_gate.pos.end());

    auto shp_udn4 = updown_gate.shape<4>();
    auto shp_udn2 = updown_gate.shape<2>();
    if(not pos_isect.empty() and pos_nsect.empty() and inc) {
        // In this case we stack gates vertically that are equally wide.
        // Should be the simplest case
        Eigen::Tensor<Scalar, 2> op = updown_gate.op.contract(middle_gate.op, Textra::idx({1}, {0}))
                                          .contract(updown_gate.op.conjugate().shuffle(Textra::array2{1, 0}), Textra::idx({1}, {0}));
        return qm::Gate{op, middle_gate.pos, middle_gate.dim};
    }
    if(pos_isect.size() == 1 and pos_nsect.size() == 1 and middle_gate.pos.size() == 1 and updown_gate.pos.size() == 2) {
        // One common location, one two uncommon. Then this connects a 1-site gate with 2-site gates up and down
        Textra::idxlistpair<1> idx1;
        Textra::idxlistpair<2> idx2;
        // Decide if this is connects on the left or right leg
        if(middle_gate.pos.front() == updown_gate.pos.front()) {
            /*  Left insert
             *            0    1              0    1               0    1
             *            |    |              |    |               |    |
             *           [  up  ]            [  up  ]             [  up  ]
             *            |    |              |    |               |    |
             *           (2)   3              |   <2>              |    |
             *           (0)                  |                    |    |
             *            |                   |                    |    |
             *          [mid]       ===>    [mid]        ===>    [mid]  |
             *            |                   |                    |    |
             *            1                  (3)                   |    |
             *            0    1             (0)  <1>              |    |
             *            |    |              |    |               |    |
             *           [  dn  ]            [  dn  ]             [  dn  ]
             *            |    |              |    |               |    |
             *            2    3              2    3               2    3
             *
             */

            idx1 = Textra::idx({2}, {0});
            idx2 = Textra::idx({3, 2}, {0, 1});
        } else {
            /*  Right insert
             *            0    1              0    1               0    1
             *            |    |              |    |               |    |
             *           [  up  ]            [  up  ]             [  up  ]
             *            |    |              |    |               |    |
             *            2   (3)            <2>   |               |    |
             *                (0)                  |               |    |
             *                 |                   |               |    |
             *               [mid]   ===>        [mid]     ===>    |  [mid]
             *                 |                   |               |    |
             *                 1                  (3)              |    |
             *            0    1             <0>  (1)              |    |
             *            |    |              |    |               |    |
             *           [  dn  ]            [  dn  ]             [  dn  ]
             *            |    |              |    |               |    |
             *            2    3              2    3               2    3
             *
             */
            idx1 = Textra::idx({3}, {0});
            idx2 = Textra::idx({2, 3}, {0, 1});
        }
        tools::log->trace("Inserting gate pos {} between gates pos {} | pos_isect {} | pos_nsect {} | inc {} | contract_b", middle_gate.pos, updown_gate.pos,
                          pos_isect, pos_nsect, inc);
        auto op = contract_b(middle_gate.op, updown_gate.op, shp_udn2, shp_udn4, idx1, idx2);
        return qm::Gate{op, updown_gate.pos, updown_gate.dim};
    }
    if(pos_isect.size() == 1 and pos_nsect.size() >= 2 and updown_gate.pos.size() == 2) {
        // One common location, so it must be at the edge.
        std::array<long, 4>         shp_mid4{};
        std::array<long, 2>         dim2{};
        std::array<Eigen::Index, 6> shf6{};
        Textra::idxlistpair<1>      idx1;
        Textra::idxlistpair<2>      idx2;
        std::vector<size_t>         pos;
        std::vector<long>           dim;
        size_t                      merged = pos_nsect.size() - pos_isect.size();
        // Decide if this is connects on the right or left leg
        if(middle_gate.pos.front() == updown_gate.pos.back()) {
            /*  Right insert (Free legs in mid are merged)
             *
             *    0    1                      0    1                        0    1
             *    |    |                      |    |                        |    |
             *   [  up  ]                    [  up  ]                      [  up  ]
             *    |    |                      |    |                        |    |
             *   2    (3)                     |    |                        |    |
             *        (0)    1               <2>   |     3                  |    |     2                                       0   1   2              0
             *         |     |                     |     |                  |    |     |                                       |   |   |              |
             *        [  mid  ]       ===>        [  mid  ]     ===>        |   [  mid  ]       shuffle({0,1,2,4,5,3})  =    [   gate    ]  =   [   gate    ]
             *         |     |                     |     |                  |    |     |                                       |   |   |              |
             *         2     3                    (4)    5                  |    |     3                                       3   4   5              1
             *    0    1                     <0>  (1)                       |    |
             *    |    |                      |    |                        |    |
             *   [  dn  ]                    [  dn  ]                      [  dn  ]
             *    |    |                      |    |                        |    |
             *    2    3                      2    3                        4    5
             *
             */

            if(pos_nsect.size() == 1) shp_mid4 = group(middle_gate.shape<2>(), std::array<size_t, 4>{1, merged, 1, merged});
            if(pos_nsect.size() == 2) shp_mid4 = group(middle_gate.shape<4>(), std::array<size_t, 4>{1, merged, 1, merged});
            if(pos_nsect.size() == 3) shp_mid4 = group(middle_gate.shape<6>(), std::array<size_t, 4>{1, merged, 1, merged});
            if(pos_nsect.size() == 4) shp_mid4 = group(middle_gate.shape<8>(), std::array<size_t, 4>{1, merged, 1, merged});
            if(pos_nsect.size() == 5) shp_mid4 = group(middle_gate.shape<10>(), std::array<size_t, 4>{1, merged, 1, merged});
            if(pos_nsect.size() == 6) shp_mid4 = group(middle_gate.shape<12>(), std::array<size_t, 4>{1, merged, 1, merged});
            if(pos_nsect.size() == 7) shp_mid4 = group(middle_gate.shape<14>(), std::array<size_t, 4>{1, merged, 1, merged});
            if(pos_nsect.size() == 8) shp_mid4 = group(middle_gate.shape<16>(), std::array<size_t, 4>{1, merged, 1, merged});
            if(pos_nsect.size() == 9) shp_mid4 = group(middle_gate.shape<18>(), std::array<size_t, 4>{1, merged, 1, merged});
            if(pos_nsect.size() == 10) shp_mid4 = group(middle_gate.shape<20>(), std::array<size_t, 4>{1, merged, 1, merged});
            if(pos_nsect.size() == 11) shp_mid4 = group(middle_gate.shape<22>(), std::array<size_t, 4>{1, merged, 1, merged});
            if(pos_nsect.size() == 12) shp_mid4 = group(middle_gate.shape<24>(), std::array<size_t, 4>{1, merged, 1, merged});
            if(pos_nsect.size() == 13) shp_mid4 = group(middle_gate.shape<26>(), std::array<size_t, 4>{1, merged, 1, merged});
            if(pos_nsect.size() == 14) shp_mid4 = group(middle_gate.shape<28>(), std::array<size_t, 4>{1, merged, 1, merged});
            if(pos_nsect.size() == 15) shp_mid4 = group(middle_gate.shape<30>(), std::array<size_t, 4>{1, merged, 1, merged});
            if(pos_nsect.size() == 16) shp_mid4 = group(middle_gate.shape<32>(), std::array<size_t, 4>{1, merged, 1, merged});
            idx1 = Textra::idx({3}, {0});
            idx2 = Textra::idx({2, 4}, {0, 1});
            shf6 = std::array<Eigen::Index, 6>{0, 1, 2, 4, 5, 3};
            pos  = concat(updown_gate.pos, subset(middle_gate.pos, 1, merged));
            dim  = concat(updown_gate.dim, subset(middle_gate.dim, 1, merged));
            dim2 = repeat(std::array<long, 1>{shp_mid4[1] * shp_udn2[0]});
        } else {
            /*  Left insert (Free legs in mid are merged)
             *
             *           0    1                      0    1                        0    1
             *           |    |                      |    |                        |    |
             *          [  up  ]                    [  up  ]                      [  up  ]
             *           |    |                      |    |                        |    |
             *          (2)   3                      |   <2>                       |    |
             *     0    (1)                    3     |                       2     |    |                                  0   1   2              0
             *     |     |                     |     |                       |     |    |                                  |   |   |              |
             *    [  mid  ]       ===>        [  mid  ]           ===>      [  mid  ]   |   shuffle({2,0,1,3,4,5})  =    [   gate    ]  =   [   gate    ]
             *     |     |                     |     |                       |     |    |                                  |   |   |              |
             *     2     3                     4    (5)                      3     |    |                                  3   4   5              1
             *           0    1                     (0)   <1>                      |    |
             *           |    |                      |    |                        |    |
             *          [  dn  ]                    [  dn  ]                      [  dn  ]
             *           |    |                      |    |                        |    |
             *           2    3                      2    3                        4    5
             *
             */
            if(pos_nsect.size() == 1) shp_mid4 = group(middle_gate.shape<2>(), std::array<size_t, 4>{merged, 1, merged, 1});
            if(pos_nsect.size() == 2) shp_mid4 = group(middle_gate.shape<4>(), std::array<size_t, 4>{merged, 1, merged, 1});
            if(pos_nsect.size() == 3) shp_mid4 = group(middle_gate.shape<6>(), std::array<size_t, 4>{merged, 1, merged, 1});
            if(pos_nsect.size() == 4) shp_mid4 = group(middle_gate.shape<8>(), std::array<size_t, 4>{merged, 1, merged, 1});
            if(pos_nsect.size() == 5) shp_mid4 = group(middle_gate.shape<10>(), std::array<size_t, 4>{merged, 1, merged, 1});
            if(pos_nsect.size() == 6) shp_mid4 = group(middle_gate.shape<12>(), std::array<size_t, 4>{merged, 1, merged, 1});
            if(pos_nsect.size() == 7) shp_mid4 = group(middle_gate.shape<14>(), std::array<size_t, 4>{merged, 1, merged, 1});
            if(pos_nsect.size() == 8) shp_mid4 = group(middle_gate.shape<16>(), std::array<size_t, 4>{merged, 1, merged, 1});
            if(pos_nsect.size() == 9) shp_mid4 = group(middle_gate.shape<18>(), std::array<size_t, 4>{merged, 1, merged, 1});
            if(pos_nsect.size() == 10) shp_mid4 = group(middle_gate.shape<20>(), std::array<size_t, 4>{merged, 1, merged, 1});
            if(pos_nsect.size() == 11) shp_mid4 = group(middle_gate.shape<22>(), std::array<size_t, 4>{merged, 1, merged, 1});
            if(pos_nsect.size() == 12) shp_mid4 = group(middle_gate.shape<24>(), std::array<size_t, 4>{merged, 1, merged, 1});
            if(pos_nsect.size() == 13) shp_mid4 = group(middle_gate.shape<26>(), std::array<size_t, 4>{merged, 1, merged, 1});
            if(pos_nsect.size() == 14) shp_mid4 = group(middle_gate.shape<28>(), std::array<size_t, 4>{merged, 1, merged, 1});
            if(pos_nsect.size() == 15) shp_mid4 = group(middle_gate.shape<30>(), std::array<size_t, 4>{merged, 1, merged, 1});
            if(pos_nsect.size() == 16) shp_mid4 = group(middle_gate.shape<32>(), std::array<size_t, 4>{merged, 1, merged, 1});
            idx1 = Textra::idx({2}, {1});
            idx2 = Textra::idx({5, 2}, {0, 1});
            shf6 = std::array<Eigen::Index, 6>{2, 0, 1, 3, 4, 5};
            pos  = concat(subset(middle_gate.pos, 0, merged), updown_gate.pos);
            dim  = concat(subset(middle_gate.dim, 0, merged), updown_gate.dim);
            dim2 = repeat(std::array<long, 1>{shp_mid4[0] * shp_udn2[0]});
        }
        tools::log->trace("Inserting gate pos {} between gates pos {} | pos_isect {} | pos_nsect {} | inc {} | contract_a", middle_gate.pos, updown_gate.pos,
                          pos_isect, pos_nsect, inc);
        auto op = contract_a(middle_gate.op, updown_gate.op, shp_mid4, shp_udn4, shf6, idx1, idx2, dim2);
        return qm::Gate{op, pos, dim};
    }
    if(pos_isect.size() == 2 and pos_nsect.size() >= 1 and updown_gate.pos.size() == 2) {
        // All the updown locations get contracted
        size_t                      usize = updown_gate.pos.size();
        size_t                      msize = middle_gate.pos.size();
        std::array<long, 4>         shp_mid4{};
        std::array<long, 6>         shp_mid6{};
        std::array<long, 2>         dim2{};
        std::array<Eigen::Index, 6> shf6{};
        std::array<Eigen::Index, 4> shf4{};
        Textra::idxlistpair<1>      idx_up, idx_dn;
        Eigen::Tensor<Scalar, 2>    op;
        auto                        offset =
            static_cast<size_t>(std::distance(middle_gate.pos.begin(), find(middle_gate.pos.begin(), middle_gate.pos.end(), updown_gate.pos.front())));
        size_t offmin = 0;
        size_t offmax = middle_gate.pos.size() - updown_gate.pos.size();
        size_t merged = offmax;
        // Decide if this is connects on the left, right or somewhere in the center
        if(offset == offmin or offset == offmax) {
            if(offset == 0) {
                /*  Insert at offset 0
                 *            0                   0                 0
                 *            |                   |                 |
                 *           [up]                [up]              [up]
                 *            |                   |                 |
                 *           (1)                  |                 |
                 *           (0)  1               |   1             |   1
                 *            |   |               |   |             |   |
                 *          [  mid  ]   ===>    [  mid  ]         [  mid  ]     ===> shuffle({0,1,3,2})
                 *           |   |               |   |             |   |
                 *           2   3              (2)   3            |   2
                 *           0                  (0)                |
                 *           |                   |                 |
                 *          [dn]                [dn]              [dn]
                 *           |                   |                 |
                 *           1                   1                 3
                 *
                 */
                if(msize == 2) shp_mid4 = group(middle_gate.shape<4>(), std::array<size_t, 4>{usize, merged, usize, merged});
                if(msize == 3) shp_mid4 = group(middle_gate.shape<6>(), std::array<size_t, 4>{usize, merged, usize, merged});
                if(msize == 4) shp_mid4 = group(middle_gate.shape<8>(), std::array<size_t, 4>{usize, merged, usize, merged});
                if(msize == 5) shp_mid4 = group(middle_gate.shape<10>(), std::array<size_t, 4>{usize, merged, usize, merged});
                if(msize == 6) shp_mid4 = group(middle_gate.shape<12>(), std::array<size_t, 4>{usize, merged, usize, merged});
                if(msize == 7) shp_mid4 = group(middle_gate.shape<14>(), std::array<size_t, 4>{usize, merged, usize, merged});
                if(msize == 8) shp_mid4 = group(middle_gate.shape<16>(), std::array<size_t, 4>{usize, merged, usize, merged});
                if(msize == 9) shp_mid4 = group(middle_gate.shape<18>(), std::array<size_t, 4>{usize, merged, usize, merged});
                if(msize == 10) shp_mid4 = group(middle_gate.shape<20>(), std::array<size_t, 4>{usize, merged, usize, merged});
                if(msize == 11) shp_mid4 = group(middle_gate.shape<22>(), std::array<size_t, 4>{usize, merged, usize, merged});
                if(msize == 12) shp_mid4 = group(middle_gate.shape<24>(), std::array<size_t, 4>{usize, merged, usize, merged});
                if(msize == 13) shp_mid4 = group(middle_gate.shape<26>(), std::array<size_t, 4>{usize, merged, usize, merged});
                if(msize == 14) shp_mid4 = group(middle_gate.shape<28>(), std::array<size_t, 4>{usize, merged, usize, merged});
                if(msize == 15) shp_mid4 = group(middle_gate.shape<30>(), std::array<size_t, 4>{usize, merged, usize, merged});
                if(msize == 16) shp_mid4 = group(middle_gate.shape<32>(), std::array<size_t, 4>{usize, merged, usize, merged});
                dim2   = middle_gate.shape<2>();
                idx_up = Textra::idx({1}, {0});
                idx_dn = Textra::idx({2}, {0});
                shf4   = Textra::array4{0, 1, 3, 2};
                tools::log->trace("Inserting gate pos {} between gates pos {} | pos_isect {} | pos_nsect {} | inc {} | contract_d", middle_gate.pos,
                                  updown_gate.pos, pos_isect, pos_nsect, inc);
                op = contract_d(middle_gate.op, updown_gate.op, shp_mid4, idx_up, idx_dn, shf4, dim2);
            } else if(offset == offmax) {
                /*  Insert at offmax
                 *            0                   0                 0
                 *            |                   |                 |
                 *           [up]                [up]              [up]
                 *            |                   |                 |
                 *           (1)                  |                 |
                 *        0  (1)              1   |             1   |
                 *        |   |               |   |             |   |
                 *      [  mid  ]   ===>    [  mid  ]         [  mid  ]     ===> shuffle({1,0,2,3})
                 *        |   |               |   |             |   |
                 *        2   3               2  (3)            2   |
                 *            0                  (0)                |
                 *            |                   |                 |
                 *           [dn]                [dn]              [dn]
                 *            |                   |                 |
                 *            1                   1                 3
                 *
                 */

                if(msize == 2) shp_mid4 = group(middle_gate.shape<4>(), std::array<size_t, 4>{merged, usize, merged, usize});
                if(msize == 3) shp_mid4 = group(middle_gate.shape<6>(), std::array<size_t, 4>{merged, usize, merged, usize});
                if(msize == 4) shp_mid4 = group(middle_gate.shape<8>(), std::array<size_t, 4>{merged, usize, merged, usize});
                if(msize == 5) shp_mid4 = group(middle_gate.shape<10>(), std::array<size_t, 4>{merged, usize, merged, usize});
                if(msize == 6) shp_mid4 = group(middle_gate.shape<12>(), std::array<size_t, 4>{merged, usize, merged, usize});
                if(msize == 7) shp_mid4 = group(middle_gate.shape<14>(), std::array<size_t, 4>{merged, usize, merged, usize});
                if(msize == 8) shp_mid4 = group(middle_gate.shape<16>(), std::array<size_t, 4>{merged, usize, merged, usize});
                if(msize == 9) shp_mid4 = group(middle_gate.shape<18>(), std::array<size_t, 4>{merged, usize, merged, usize});
                if(msize == 10) shp_mid4 = group(middle_gate.shape<20>(), std::array<size_t, 4>{merged, usize, merged, usize});
                if(msize == 11) shp_mid4 = group(middle_gate.shape<22>(), std::array<size_t, 4>{merged, usize, merged, usize});
                if(msize == 12) shp_mid4 = group(middle_gate.shape<24>(), std::array<size_t, 4>{merged, usize, merged, usize});
                if(msize == 13) shp_mid4 = group(middle_gate.shape<26>(), std::array<size_t, 4>{merged, usize, merged, usize});
                if(msize == 14) shp_mid4 = group(middle_gate.shape<28>(), std::array<size_t, 4>{merged, usize, merged, usize});
                if(msize == 15) shp_mid4 = group(middle_gate.shape<30>(), std::array<size_t, 4>{merged, usize, merged, usize});
                if(msize == 16) shp_mid4 = group(middle_gate.shape<32>(), std::array<size_t, 4>{merged, usize, merged, usize});
                dim2   = middle_gate.shape<2>();
                idx_up = Textra::idx({1}, {1});
                idx_dn = Textra::idx({3}, {0});
                shf4   = Textra::array4{1, 0, 2, 3};
            }
            tools::log->trace("Inserting gate pos {} between gates pos {} | pos_isect {} | pos_nsect {} | inc {} | contract_d", middle_gate.pos,
                              updown_gate.pos, pos_isect, pos_nsect, inc);
            op = contract_d(middle_gate.op, updown_gate.op, shp_mid4, idx_up, idx_dn, shf4, dim2);
        } else {
            /*  Insert at offmax
             *           0                     0                      0
             *           |                     |                      |
             *          [up]                  [up]                   [up]
             *           |                     |                      |
             *          (1)                    |                      |
             *       0  (1)  2             1   |   2              1   |   2
             *       |   |   |             |   |   |              |   |   |
             *      [   mid   ]    ===>   [   mid   ]    ===>    [   mid   ]       ===> shuffle({1,0,2,3,5,4})
             *       |   |   |             |   |   |              |   |   |
             *       3   4   5             3  (4)  5              3   |   4
             *           0                    (0)                     |
             *           |                     |                      |
             *          [dn]                  [dn]                   [dn]
             *           |                     |                      |
             *           1                     1                      5
             *
             */

            if(msize == 4) shp_mid6 = group(middle_gate.shape<8>(), repeat(std::array<size_t, 3>{offset, usize, msize - usize - offset}));
            if(msize == 5) shp_mid6 = group(middle_gate.shape<10>(), repeat(std::array<size_t, 3>{offset, usize, msize - usize - offset}));
            if(msize == 6) shp_mid6 = group(middle_gate.shape<12>(), repeat(std::array<size_t, 3>{offset, usize, msize - usize - offset}));
            if(msize == 7) shp_mid6 = group(middle_gate.shape<14>(), repeat(std::array<size_t, 3>{offset, usize, msize - usize - offset}));
            if(msize == 8) shp_mid6 = group(middle_gate.shape<16>(), repeat(std::array<size_t, 3>{offset, usize, msize - usize - offset}));
            if(msize == 9) shp_mid6 = group(middle_gate.shape<18>(), repeat(std::array<size_t, 3>{offset, usize, msize - usize - offset}));
            if(msize == 10) shp_mid6 = group(middle_gate.shape<20>(), repeat(std::array<size_t, 3>{offset, usize, msize - usize - offset}));
            if(msize == 11) shp_mid6 = group(middle_gate.shape<22>(), repeat(std::array<size_t, 3>{offset, usize, msize - usize - offset}));
            if(msize == 12) shp_mid6 = group(middle_gate.shape<24>(), repeat(std::array<size_t, 3>{offset, usize, msize - usize - offset}));
            if(msize == 13) shp_mid6 = group(middle_gate.shape<26>(), repeat(std::array<size_t, 3>{offset, usize, msize - usize - offset}));
            if(msize == 14) shp_mid6 = group(middle_gate.shape<28>(), repeat(std::array<size_t, 3>{offset, usize, msize - usize - offset}));
            if(msize == 15) shp_mid6 = group(middle_gate.shape<30>(), repeat(std::array<size_t, 3>{offset, usize, msize - usize - offset}));
            if(msize == 16) shp_mid6 = group(middle_gate.shape<32>(), repeat(std::array<size_t, 3>{offset, usize, msize - usize - offset}));
            dim2   = middle_gate.shape<2>();
            idx_up = Textra::idx({1}, {1});
            idx_dn = Textra::idx({4}, {0});
            shf6   = Textra::array6{1, 0, 2, 3, 5, 4};
            tools::log->trace("Inserting gate pos {} between gates pos {} | pos_isect {} | pos_nsect {} | inc {} | contract_c", middle_gate.pos,
                              updown_gate.pos, pos_isect, pos_nsect, inc);
            op = contract_c(middle_gate.op, updown_gate.op, shp_mid6, idx_up, idx_dn, shf6, dim2);
        }
        if(op.size() == 0) throw std::logic_error("No op computed!");
        return qm::Gate{op, middle_gate.pos, middle_gate.dim};
    }

    throw std::runtime_error(fmt::format("Insert case not implemented: middle pos {} | updown pos {} | pos_isect {} | pos_nsect {}", middle_gate.pos,
                                         updown_gate.pos, pos_isect, pos_nsect));
}

qm::Gate qm::connect(const qm::Gate &dn_gate, const qm::Gate &up_gate) {
    std::vector<size_t> pos_isect; // locations that intersect: both on middle and updown gates
    std::vector<size_t> pos_nsect; // locations that do not intersect: not in both on middle and updown gates
    std::set_intersection(dn_gate.pos.begin(), dn_gate.pos.end(), up_gate.pos.begin(), up_gate.pos.end(), back_inserter(pos_isect));
    std::set_symmetric_difference(dn_gate.pos.begin(), dn_gate.pos.end(), up_gate.pos.begin(), up_gate.pos.end(), back_inserter(pos_nsect));
    bool inc = std::includes(dn_gate.pos.begin(), dn_gate.pos.end(), up_gate.pos.begin(), up_gate.pos.end());
    if(pos_isect.empty()) return dn_gate;

    tools::log->trace("Connecting dn {} | up {}", dn_gate.pos, up_gate.pos);

    if(not pos_isect.empty() and pos_nsect.empty() and inc) {
        // This case is a vertical stack. Should be the simplest case
        Eigen::Tensor<Scalar, 2> op = up_gate.op.contract(dn_gate.op, Textra::idx({1}, {0}));
        return qm::Gate{op, dn_gate.pos, dn_gate.dim};
    }
    if(not pos_isect.empty() and not pos_nsect.empty() and inc) {
        std::array<long, 6> dn_shp6{};
        std::array<long, 4> dn_shp4{};
        std::vector<size_t> pos;
        std::vector<long>   dim;
        auto                dn_size  = dn_gate.pos.size();
        auto                up_size  = up_gate.pos.size();
        size_t              dn_merge = dn_size - pos_isect.size();
        size_t              up_merge = up_size;
        size_t              offset   = up_gate.pos.front() - dn_gate.pos.front();
        // Decide if this is connects on the right or right leg
        if(dn_gate.pos.back() == pos_isect.back()) {
            /*  Right connection
             *
             *            |    |
             *            |  [up]
             *            |   |
             *          [  dn  ]
             *           |    |
             *
             */
            if(dn_size == 2) dn_shp4 = group(dn_gate.shape<4>(), repeat(std::array<size_t, 2>{dn_merge, up_merge}));
            if(dn_size == 3) dn_shp4 = group(dn_gate.shape<6>(), repeat(std::array<size_t, 2>{dn_merge, up_merge}));
            if(dn_size == 4) dn_shp4 = group(dn_gate.shape<8>(), repeat(std::array<size_t, 2>{dn_merge, up_merge}));
            if(dn_size == 5) dn_shp4 = group(dn_gate.shape<10>(), repeat(std::array<size_t, 2>{dn_merge, up_merge}));
            if(dn_size == 6) dn_shp4 = group(dn_gate.shape<12>(), repeat(std::array<size_t, 2>{dn_merge, up_merge}));
            if(dn_size == 7) dn_shp4 = group(dn_gate.shape<14>(), repeat(std::array<size_t, 2>{dn_merge, up_merge}));
            if(dn_size == 8) dn_shp4 = group(dn_gate.shape<16>(), repeat(std::array<size_t, 2>{dn_merge, up_merge}));
            auto                     dim2 = dn_gate.shape<2>();
            Eigen::Tensor<Scalar, 2> op =
                up_gate.op.contract(dn_gate.op.reshape(dn_shp4), Textra::idx({1}, {1})).shuffle(Textra::array4{1, 0, 2, 3}).reshape(dim2);
            return qm::Gate{op, dn_gate.pos, dn_gate.dim};
        } else if(dn_gate.pos.front() == pos_isect.front()) {
            /*  Left connection
             *
             *            |   |
             *          [up]  |
             *           |    |
             *          [  dn  ]
             *           |    |
             */
            if(dn_size == 2) dn_shp4 = group(dn_gate.shape<4>(), repeat(std::array<size_t, 2>{up_merge, dn_merge}));
            if(dn_size == 3) dn_shp4 = group(dn_gate.shape<6>(), repeat(std::array<size_t, 2>{up_merge, dn_merge}));
            if(dn_size == 4) dn_shp4 = group(dn_gate.shape<8>(), repeat(std::array<size_t, 2>{up_merge, dn_merge}));
            if(dn_size == 5) dn_shp4 = group(dn_gate.shape<10>(), repeat(std::array<size_t, 2>{up_merge, dn_merge}));
            if(dn_size == 6) dn_shp4 = group(dn_gate.shape<12>(), repeat(std::array<size_t, 2>{up_merge, dn_merge}));
            if(dn_size == 7) dn_shp4 = group(dn_gate.shape<14>(), repeat(std::array<size_t, 2>{up_merge, dn_merge}));
            if(dn_size == 8) dn_shp4 = group(dn_gate.shape<16>(), repeat(std::array<size_t, 2>{up_merge, dn_merge}));
            auto                     dim2 = dn_gate.shape<2>();
            Eigen::Tensor<Scalar, 2> op   = up_gate.op.contract(dn_gate.op.reshape(dn_shp4), Textra::idx({1}, {0})).reshape(dim2);
            return qm::Gate{op, dn_gate.pos, dn_gate.dim};
        } else {
            /*  Center connection
             *
             *           |    |   |
             *           |  [up]  |
             *           |   |    |
             *          [    dn    ]
             *           |    |   |
             */
            if(dn_size == 3) dn_shp6 = group(dn_gate.shape<6>(), repeat(std::array<size_t, 3>{offset, up_size, dn_size - up_size - offset}));
            if(dn_size == 4) dn_shp6 = group(dn_gate.shape<8>(), repeat(std::array<size_t, 3>{offset, up_size, dn_size - up_size - offset}));
            if(dn_size == 5) dn_shp6 = group(dn_gate.shape<10>(), repeat(std::array<size_t, 3>{offset, up_size, dn_size - up_size - offset}));
            if(dn_size == 6) dn_shp6 = group(dn_gate.shape<12>(), repeat(std::array<size_t, 3>{offset, up_size, dn_size - up_size - offset}));
            if(dn_size == 7) dn_shp6 = group(dn_gate.shape<14>(), repeat(std::array<size_t, 3>{offset, up_size, dn_size - up_size - offset}));
            if(dn_size == 8) dn_shp6 = group(dn_gate.shape<16>(), repeat(std::array<size_t, 3>{offset, up_size, dn_size - up_size - offset}));
            auto                     dim2 = dn_gate.shape<2>();
            Eigen::Tensor<Scalar, 2> op =
                up_gate.op.contract(dn_gate.op.reshape(dn_shp6), Textra::idx({1}, {1})).shuffle(Textra::array6{1, 0, 2, 3, 4, 5}).reshape(dim2);
            return qm::Gate{op, dn_gate.pos, dn_gate.dim};
        }
    }
    if(pos_isect.size() == 1 and pos_nsect.size() >= 2) {
        // One common location, and one uncommon. Then this connects two 2-site gates
        std::array<long, 4> dn_shp4{};
        std::array<long, 4> up_shp4{};
        std::vector<size_t> pos;
        std::vector<long>   dim;
        auto                dn_size  = dn_gate.pos.size();
        auto                up_size  = up_gate.pos.size();
        size_t              dn_merge = dn_size - pos_isect.size();
        size_t              up_merge = up_size - pos_isect.size();
        // Decide if this is connects on the right or right leg
        if(dn_gate.pos.front() == up_gate.pos.back()) {
            /*  Right connection
             *
             *      |    |     |           0   1   2              0
             *     [  up  ]    |           |   |   |              |
             *      |    |     |   =     [   gate    ]  =   [   gate    ]
             *      |   [  dn  ]           |   |   |              |
             *      |    |     |           3   4   5              1
             *
             */
            if(dn_size == 2) dn_shp4 = group(dn_gate.shape<4>(), repeat(std::array<size_t, 2>{1, dn_merge}));
            if(dn_size == 3) dn_shp4 = group(dn_gate.shape<6>(), repeat(std::array<size_t, 2>{1, dn_merge}));
            if(dn_size == 4) dn_shp4 = group(dn_gate.shape<8>(), repeat(std::array<size_t, 2>{1, dn_merge}));
            if(dn_size == 5) dn_shp4 = group(dn_gate.shape<10>(), repeat(std::array<size_t, 2>{1, dn_merge}));
            if(dn_size == 6) dn_shp4 = group(dn_gate.shape<12>(), repeat(std::array<size_t, 2>{1, dn_merge}));
            if(dn_size == 7) dn_shp4 = group(dn_gate.shape<14>(), repeat(std::array<size_t, 2>{1, dn_merge}));
            if(dn_size == 8) dn_shp4 = group(dn_gate.shape<16>(), repeat(std::array<size_t, 2>{1, dn_merge}));
            if(up_size == 2) up_shp4 = group(dn_gate.shape<4>(), repeat(std::array<size_t, 2>{up_merge, 1}));
            if(up_size == 3) up_shp4 = group(dn_gate.shape<6>(), repeat(std::array<size_t, 2>{up_merge, 1}));
            if(up_size == 4) up_shp4 = group(dn_gate.shape<8>(), repeat(std::array<size_t, 2>{up_merge, 1}));
            if(up_size == 5) up_shp4 = group(dn_gate.shape<10>(), repeat(std::array<size_t, 2>{up_merge, 1}));
            if(up_size == 6) up_shp4 = group(dn_gate.shape<12>(), repeat(std::array<size_t, 2>{up_merge, 1}));
            if(up_size == 7) up_shp4 = group(dn_gate.shape<14>(), repeat(std::array<size_t, 2>{up_merge, 1}));
            if(up_size == 8) up_shp4 = group(dn_gate.shape<16>(), repeat(std::array<size_t, 2>{up_merge, 1}));
            auto                     dim2 = repeat(std::array<long, 1>{up_shp4[0] * up_shp4[1] * dn_shp4[1]});
            Eigen::Tensor<Scalar, 2> op   = dn_gate.op.reshape(dn_shp4)
                                              .contract(up_gate.op.reshape(up_shp4), Textra::idx({0}, {3}))
                                              .shuffle(Textra::array6{3, 4, 0, 5, 1, 2})
                                              .reshape(dim2);
            pos = concat(up_gate.pos, subset(dn_gate.pos, 1, dn_merge));
            dim = concat(up_gate.dim, subset(dn_gate.dim, 1, dn_merge));
            return qm::Gate{op, pos, dim};
        } else {
            /*  Left connection
             *
             *     |     |    |             0   1   2              0
             *     |    [  up  ]            |   |   |              |
             *     |     |    |      =    [   gate    ]  =   [   gate    ]
             *    [  dn  ]    |             |   |   |              |
             *     |     |    |             3   4   5              1
             */
            if(dn_size == 2) dn_shp4 = group(dn_gate.shape<4>(), repeat(std::array<size_t, 2>{dn_merge, 1}));
            if(dn_size == 3) dn_shp4 = group(dn_gate.shape<6>(), repeat(std::array<size_t, 2>{dn_merge, 1}));
            if(dn_size == 4) dn_shp4 = group(dn_gate.shape<8>(), repeat(std::array<size_t, 2>{dn_merge, 1}));
            if(dn_size == 5) dn_shp4 = group(dn_gate.shape<10>(), repeat(std::array<size_t, 2>{dn_merge, 1}));
            if(dn_size == 6) dn_shp4 = group(dn_gate.shape<12>(), repeat(std::array<size_t, 2>{dn_merge, 1}));
            if(dn_size == 7) dn_shp4 = group(dn_gate.shape<14>(), repeat(std::array<size_t, 2>{dn_merge, 1}));
            if(dn_size == 8) dn_shp4 = group(dn_gate.shape<16>(), repeat(std::array<size_t, 2>{dn_merge, 1}));
            if(up_size == 2) up_shp4 = group(dn_gate.shape<4>(), repeat(std::array<size_t, 2>{1, up_merge}));
            if(up_size == 3) up_shp4 = group(dn_gate.shape<6>(), repeat(std::array<size_t, 2>{1, up_merge}));
            if(up_size == 4) up_shp4 = group(dn_gate.shape<8>(), repeat(std::array<size_t, 2>{1, up_merge}));
            if(up_size == 5) up_shp4 = group(dn_gate.shape<10>(), repeat(std::array<size_t, 2>{1, up_merge}));
            if(up_size == 6) up_shp4 = group(dn_gate.shape<12>(), repeat(std::array<size_t, 2>{1, up_merge}));
            if(up_size == 7) up_shp4 = group(dn_gate.shape<14>(), repeat(std::array<size_t, 2>{1, up_merge}));
            if(up_size == 8) up_shp4 = group(dn_gate.shape<16>(), repeat(std::array<size_t, 2>{1, up_merge}));
            auto                     dim2 = repeat(std::array<long, 1>{dn_shp4[0] * up_shp4[0] * up_shp4[1]});
            Eigen::Tensor<Scalar, 2> op   = dn_gate.op.reshape(dn_shp4)
                                              .contract(up_gate.op.reshape(up_shp4), Textra::idx({1}, {2}))
                                              .shuffle(Textra::array6{0, 3, 4, 1, 2, 5})
                                              .reshape(dim2);
            pos = concat(subset(dn_gate.pos, 0, dn_merge), up_gate.pos);
            dim = concat(subset(dn_gate.dim, 0, dn_merge), up_gate.dim);
            return qm::Gate{op, pos, dim};
        }
    }

    throw std::runtime_error(
        fmt::format("Connect case not implemented: dn pos {} | up pos {} | pos_isect {} | pos_nsect {}", dn_gate.pos, up_gate.pos, pos_isect, pos_nsect));
}

qm::Gate qm::trace_idx(const qm::Gate &gate, const std::vector<long> &idx) {
    if(idx.size() == 2) return qm::trace(gate, Textra::idx({idx[0]}, {idx[1]}));
    // idx is already pair-ordered, so when we call Textra::idx we need to undo that order
    // which explains the weird indexing in the following case
    if(idx.size() == 4) return qm::trace(gate, Textra::idx({idx[0], idx[2]}, {idx[1], idx[3]}));
    throw std::runtime_error(fmt::format("Tracing {} indices is not implemented", idx.size()));
}

qm::Gate qm::trace_pos(const qm::Gate &gate, const std::vector<size_t> &pos) {
    if(pos.size() <= 2) {
        auto              idx = gate.idx(pos);
        std::vector<long> idx_long(idx.begin(), idx.end());
        return qm::trace_idx(gate, idx_long);
    }
    qm::Gate tmp = gate;
    for(auto &p : pos) { tmp = tmp.trace_pos(std::vector<size_t>{p}); }
    return tmp;
}

qm::Gate qm::trace_pos(const qm::Gate &gate, size_t pos) { return qm::trace_pos(gate, gate.idx(std::vector<size_t>{pos})); }

qm::Scalar qm::trace(const qm::Gate &gate) {
    qm::Gate t = qm::trace_pos(gate, gate.pos);
    if(not t.pos.empty()) throw std::logic_error(fmt::format("Gate should be empty after tracing all positions. Got pos: {}", t.pos));
    if(t.op.dimension(0) * t.op.dimension(1) != 1)
        throw std::logic_error(fmt::format("Gate should have scalar op after tracing all positions. Got dims: {}", t.op.dimensions()));
    return t.op(0);
}

template<auto N>
qm::Gate qm::trace(const qm::Gate &gate, const std::array<Eigen::IndexPair<Eigen::Index>, N> &idxpairs) {
    if constexpr(N == 1) tools::log->trace("Tracing pos {} | dim {} | index pair [{},{}]", gate.pos, gate.dim, idxpairs[0].first, idxpairs[0].second);
    if constexpr(N == 2)
        tools::log->trace("Tracing gate pos {} | dim {} | index pairs [{},{}][{},{}]", gate.pos, gate.dim, idxpairs[0].first, idxpairs[0].second,
                          idxpairs[1].first, idxpairs[1].second);

    // Compute the remaining indices positions and dimensions
    auto idx = gate.idx(); // Twice as long as pos and dim! It has the top and bottom indices
    auto pos = gate.pos;
    auto dim = gate.dim;
    for(const auto &pair : idxpairs) { // Idx pairs has the indices of the legs we remove
        auto it1 = std::find(idx.begin(), idx.end(), std::max(pair.first, pair.second));
        if(it1 != idx.end()) {
            auto idx_rm1 = std::distance(idx.begin(), it1);
            if(idx_rm1 < static_cast<long>(pos.size())) pos.erase(pos.begin() + idx_rm1);
            if(idx_rm1 < static_cast<long>(dim.size())) dim.erase(dim.begin() + idx_rm1);
            idx.erase(it1);
        }
        auto it2 = std::find(idx.begin(), idx.end(), std::min(pair.first, pair.second));
        if(it2 != idx.end()) {
            auto idx_rm2 = std::distance(idx.begin(), it2);
            if(idx_rm2 < static_cast<long>(pos.size())) pos.erase(pos.begin() + idx_rm2);
            if(idx_rm2 < static_cast<long>(dim.size())) dim.erase(dim.begin() + idx_rm2);
            idx.erase(it2);
        }
    }
    if(idx.size() != 2 * (gate.pos.size() - N)) throw std::logic_error("Wrong number of indices removed");

    // Assuming the gate is symmetric, we can compute the rank2 dimensions for the storage of the gate op
    long dim_prod = std::accumulate(std::begin(dim), std::end(dim), 1, std::multiplies<>()); // Product of all dimensions of the remaining top legs of the gate
    std::array<long, 2> dim2{dim_prod, dim_prod};

    // Shorthand tensors
    using T2  = Eigen::Tensor<Scalar, 2>;
    using T4  = Eigen::Tensor<Scalar, 4>;
    using T6  = Eigen::Tensor<Scalar, 6>;
    using T8  = Eigen::Tensor<Scalar, 8>;
    using T10 = Eigen::Tensor<Scalar, 10>;
    using T12 = Eigen::Tensor<Scalar, 12>;
    using T14 = Eigen::Tensor<Scalar, 14>;
    using T16 = Eigen::Tensor<Scalar, 16>;
    using T18 = Eigen::Tensor<Scalar, 18>;
    using T20 = Eigen::Tensor<Scalar, 20>;
    using T22 = Eigen::Tensor<Scalar, 22>;
    using T24 = Eigen::Tensor<Scalar, 24>;
    using T26 = Eigen::Tensor<Scalar, 26>;
    using T28 = Eigen::Tensor<Scalar, 28>;
    using T30 = Eigen::Tensor<Scalar, 30>;
    using T32 = Eigen::Tensor<Scalar, 32>;
    Eigen::Tensor<Scalar, 2> op_traced;

    // Trace
    if(gate.dim.size() == 1) {
        if constexpr(N == 1) { // dim size 1 is special! It can't take 2 index pairs
            op_traced = linalg::tensor::trace(static_cast<T2>(gate.op.reshape(gate.shape<2>())), idxpairs).reshape(dim2);
        } else {
            throw std::runtime_error(fmt::format("Can't trace {} index pairs on gate with pos {}", N, gate.pos));
        }
    }
    /* clang-format off */
    else if(gate.dim.size() == 2 ) op_traced = linalg::tensor::trace(static_cast<T4>(gate.op.reshape(gate.shape<4>())), idxpairs).reshape(dim2);
    else if(gate.dim.size() == 3 ) op_traced = linalg::tensor::trace(static_cast<T6>(gate.op.reshape(gate.shape<6>())), idxpairs).reshape(dim2);
    else if(gate.dim.size() == 4 ) op_traced = linalg::tensor::trace(static_cast<T8>(gate.op.reshape(gate.shape<8>())), idxpairs).reshape(dim2);
    else if(gate.dim.size() == 5 ) op_traced = linalg::tensor::trace(static_cast<T10>(gate.op.reshape(gate.shape<10>())), idxpairs).reshape(dim2);
    else if(gate.dim.size() == 6 ) op_traced = linalg::tensor::trace(static_cast<T12>(gate.op.reshape(gate.shape<12>())), idxpairs).reshape(dim2);
    else if(gate.dim.size() == 7 ) op_traced = linalg::tensor::trace(static_cast<T14>(gate.op.reshape(gate.shape<14>())), idxpairs).reshape(dim2);
    else if(gate.dim.size() == 8 ) op_traced = linalg::tensor::trace(static_cast<T16>(gate.op.reshape(gate.shape<16>())), idxpairs).reshape(dim2);
    else if(gate.dim.size() == 9 ) op_traced = linalg::tensor::trace(static_cast<T18>(gate.op.reshape(gate.shape<18>())), idxpairs).reshape(dim2);
    else if(gate.dim.size() == 10) op_traced = linalg::tensor::trace(static_cast<T20>(gate.op.reshape(gate.shape<20>())), idxpairs).reshape(dim2);
    else if(gate.dim.size() == 11) op_traced = linalg::tensor::trace(static_cast<T22>(gate.op.reshape(gate.shape<22>())), idxpairs).reshape(dim2);
    else if(gate.dim.size() == 12) op_traced = linalg::tensor::trace(static_cast<T24>(gate.op.reshape(gate.shape<24>())), idxpairs).reshape(dim2);
    else if(gate.dim.size() == 13) op_traced = linalg::tensor::trace(static_cast<T26>(gate.op.reshape(gate.shape<26>())), idxpairs).reshape(dim2);
    else if(gate.dim.size() == 14) op_traced = linalg::tensor::trace(static_cast<T28>(gate.op.reshape(gate.shape<28>())), idxpairs).reshape(dim2);
    else if(gate.dim.size() == 15) op_traced = linalg::tensor::trace(static_cast<T30>(gate.op.reshape(gate.shape<30>())), idxpairs).reshape(dim2);
    else if(gate.dim.size() == 16) op_traced = linalg::tensor::trace(static_cast<T32>(gate.op.reshape(gate.shape<32>())), idxpairs).reshape(dim2);
    /* clang-format on */
    else
        throw std::runtime_error(fmt::format("Trace not implemented: N == {} | dim.size() == {}", N, gate.dim.size()));

    return Gate{op_traced, pos, dim};
}

template qm::Gate qm::trace(const qm::Gate &gate, const std::array<Eigen::IndexPair<Eigen::Index>, 1> &idxpairs);
template qm::Gate qm::trace(const qm::Gate &gate, const std::array<Eigen::IndexPair<Eigen::Index>, 2> &idxpairs);


#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/io.hpp>
#include <boost/numeric/ublas/lu.hpp>
#include "BodyModel.hpp"
#include "Generator.hpp"
#include "utils/body.hpp"
#include "utils/matrix_helpers.hpp"

const float BodyModel::COM_HEIGHT = 200;
const float BodyModel::MOTION_DT = 1 / 100.0f;
const float BodyModel::FOOT_LENGTH = 160;
const float BodyModel::L_RATE = 0.5;

BodyModel::BodyModel() : filAccX(0.0),
                         filAccY(0.0),
                         filteredTotalPressure(0.0),
                         filHighZMPF(0.0),
                         filLowZMPF(0.0),
                         filZMPL(0.0)
{
   stepCounter = 0;
   sagittalAcceleration = 0.0;
   forwardR = forwardL = 0.0;

   isLeftPhase = false;
   isOnFront = false;
   lastIsLeftPhase = false;

   // real body model
   isStopped = false;

   //Used to AutoCalibrate Footsensors
	fsLfr = fsLfl = fsLrr = fsLrl = 0.1;                   // keeps maximum foot sensor readings
	fsRfr = fsRfl = fsRrr = fsRrl = 0.1;                   // small value deemed close to zero initially
}

void BodyModel::update(Odometry *odometry,
                       const SensorValues &sensors) {
   // Various Center of Pressure calcs to help stabilise walk (ZMP)
   stepCounter++;
   // Update maximum sensed foot-sensor readings
	float temp = sensors.sensors[Sensors::LFoot_FSR_FrontLeft]; if(fsLfl<temp and fsLfl<5.0) fsLfl = temp;
	temp = sensors.sensors[Sensors::LFoot_FSR_FrontLeft];       if(fsLfl<temp and fsLfl<5.0) fsLfl = temp;
	temp = sensors.sensors[Sensors::LFoot_FSR_FrontRight];      if(fsLfr<temp and fsLfr<5.0) fsLfr = temp;
	temp = sensors.sensors[Sensors::LFoot_FSR_RearLeft];        if(fsLrl<temp and fsLrl<5.0) fsLrl = temp;
	temp = sensors.sensors[Sensors::LFoot_FSR_RearRight];       if(fsLrr<temp and fsLrr<5.0) fsLrr = temp;
	temp = sensors.sensors[Sensors::RFoot_FSR_FrontLeft];       if(fsRfl<temp and fsRfl<5.0) fsRfl = temp;
	temp = sensors.sensors[Sensors::RFoot_FSR_FrontRight];      if(fsRfr<temp and fsRfr<5.0) fsRfr = temp;
	temp = sensors.sensors[Sensors::RFoot_FSR_RearLeft];        if(fsRrl<temp and fsRrl<5.0) fsRrl = temp;
	temp = sensors.sensors[Sensors::RFoot_FSR_RearRight];       if(fsRrr<temp and fsRrr<5.0) fsRrr = temp;
	lastZMPL = ZMPL;
	ZMPL = 0;
	// Calculate ZMPL (left-right) used to eg switch support foot in Walk2014
	float pressureL =
	  +sensors.sensors[Sensors::LFoot_FSR_FrontLeft]/fsLfl
	  + sensors.sensors[Sensors::LFoot_FSR_FrontRight]/fsLfr
	  + sensors.sensors[Sensors::LFoot_FSR_RearLeft]/fsLrl
	  + sensors.sensors[Sensors::LFoot_FSR_RearRight]/fsLrr;
	float pressureR =
	  +sensors.sensors[Sensors::RFoot_FSR_FrontLeft]/fsRfl
	  + sensors.sensors[Sensors::RFoot_FSR_FrontRight]/fsRfr
	  + sensors.sensors[Sensors::RFoot_FSR_RearLeft]/fsRrl
	  + sensors.sensors[Sensors::RFoot_FSR_RearRight]/fsRrr;
	float totalPressure = pressureL + pressureR;
	if (ABS(totalPressure) > 0.000001f) {
			ZMPL =
			(  .080 * sensors.sensors[Sensors::LFoot_FSR_FrontLeft]/fsLfl
			+ .030 * sensors.sensors[Sensors::LFoot_FSR_FrontRight]/fsLfr
			+ .080 * sensors.sensors[Sensors::LFoot_FSR_RearLeft]/fsLrl
			+ .030 * sensors.sensors[Sensors::LFoot_FSR_RearRight]/fsLrr
			- .030 * sensors.sensors[Sensors::RFoot_FSR_FrontLeft]/fsRfl
			- .080 * sensors.sensors[Sensors::RFoot_FSR_FrontRight]/fsRfr
			- .030 * sensors.sensors[Sensors::RFoot_FSR_RearLeft]/fsRrl
			- .080 * sensors.sensors[Sensors::RFoot_FSR_RearRight]/fsRrr) / totalPressure;
		}

   isOnFront = isOnFrontOfFoot(sensors);
   processUpdate(odometry, sensors);
   observationUpdate(odometry, sensors);
}

bool BodyModel::isFootOnGround(const SensorValues &sensors) {
   float leftFrontL = sensors.sensors[Sensors::LFoot_FSR_FrontLeft];
   float leftFrontR = sensors.sensors[Sensors::LFoot_FSR_FrontRight];
   float rightFrontL = sensors.sensors[Sensors::RFoot_FSR_FrontLeft];
   float rightFrontR = sensors.sensors[Sensors::RFoot_FSR_FrontRight];
   float leftRearL  = sensors.sensors[Sensors::LFoot_FSR_RearLeft];
   float leftRearR  = sensors.sensors[Sensors::LFoot_FSR_RearRight];
   float rightRearL = sensors.sensors[Sensors::RFoot_FSR_RearLeft];
   float rightRearR = sensors.sensors[Sensors::RFoot_FSR_RearRight];
   if (isLeftPhase) {
      return rightFrontL > 0.01 && rightFrontR > 0.01 &&
             rightRearL > 0.01 && rightRearR > 0.01;
   }
   return leftFrontL > 0.01 &&  leftFrontR > 0.01 &&
          leftRearL > 0.01 && leftRearR > 0.01;
}

bool BodyModel::isOnFrontOfFoot(const SensorValues &sensors) {
   float leftFront = sensors.sensors[Sensors::LFoot_FSR_FrontLeft] +
                     sensors.sensors[Sensors::LFoot_FSR_FrontRight];
   float rightFront = sensors.sensors[Sensors::RFoot_FSR_FrontLeft] +
                      sensors.sensors[Sensors::RFoot_FSR_FrontRight];
   float leftRear  = sensors.sensors[Sensors::LFoot_FSR_RearLeft] +
                     sensors.sensors[Sensors::LFoot_FSR_RearRight];
   float rightRear = sensors.sensors[Sensors::RFoot_FSR_RearLeft] +
                     sensors.sensors[Sensors::RFoot_FSR_RearRight];
   if (isLeftPhase) {
      return rightFront > rightRear;
   }
   return leftFront > leftRear;
}

float BodyModel::getFootZMP(bool isLeft, const SensorValues &sensors) {
   if (walkCycle.isDoubleSupportPhase()) return 0;
   float ZMPF = 0;
   float totalPressure = 0;
   if (isLeft) {
      totalPressure += sensors.sensors[Sensors::LFoot_FSR_FrontLeft]
                         + sensors.sensors[Sensors::LFoot_FSR_FrontRight]
                         + sensors.sensors[Sensors::LFoot_FSR_RearLeft]
                         + sensors.sensors[Sensors::LFoot_FSR_RearRight];
   } else {
      totalPressure += sensors.sensors[Sensors::RFoot_FSR_FrontLeft]
                         + sensors.sensors[Sensors::RFoot_FSR_FrontRight]
                         + sensors.sensors[Sensors::RFoot_FSR_RearLeft]
                         + sensors.sensors[Sensors::RFoot_FSR_RearRight];
   }
   if (isLeft) {
      ZMPF =
            (75.0f) * sensors.sensors[Sensors::LFoot_FSR_FrontLeft]
          + (-45.0f) * sensors.sensors[Sensors::LFoot_FSR_RearLeft]
          + (75.0f) * sensors.sensors[Sensors::LFoot_FSR_FrontRight]
          + (-45.0f) * sensors.sensors[Sensors::LFoot_FSR_RearRight];
   } else {
      ZMPF =
           (75.0f) * sensors.sensors[Sensors::RFoot_FSR_FrontRight]
         + (-45.0f) * sensors.sensors[Sensors::RFoot_FSR_RearRight]
         + (75.0f) * sensors.sensors[Sensors::RFoot_FSR_FrontLeft]
         + (-45.0f) * sensors.sensors[Sensors::RFoot_FSR_RearLeft];
   }
   if (totalPressure == 0) return  0;
   ZMPF /= totalPressure;
   return ZMPF;
}

float BodyModel::getHorizontalFootZMP(bool isLeft, const SensorValues &sensors) {
   float ZMPF = 0;
   float totalPressure = 0;
   if (isLeft) {
      totalPressure += sensors.sensors[Sensors::LFoot_FSR_FrontLeft]
                         + sensors.sensors[Sensors::LFoot_FSR_FrontRight]
                         + sensors.sensors[Sensors::LFoot_FSR_RearLeft]
                         + sensors.sensors[Sensors::LFoot_FSR_RearRight];
   } else {
      totalPressure += sensors.sensors[Sensors::RFoot_FSR_FrontLeft]
                         + sensors.sensors[Sensors::RFoot_FSR_FrontRight]
                         + sensors.sensors[Sensors::RFoot_FSR_RearLeft]
                         + sensors.sensors[Sensors::RFoot_FSR_RearRight];
   }
   if (isLeft) {
      ZMPF =
            (20.0f) * sensors.sensors[Sensors::LFoot_FSR_FrontLeft]
          + (20.0f) * sensors.sensors[Sensors::LFoot_FSR_RearLeft]
          + (-20.0f) * sensors.sensors[Sensors::LFoot_FSR_FrontRight]
          + (-20.0f) * sensors.sensors[Sensors::LFoot_FSR_RearRight];
   } else {
      ZMPF =
           (-20.0f) * sensors.sensors[Sensors::RFoot_FSR_FrontRight]
         + (-20.0f) * sensors.sensors[Sensors::RFoot_FSR_RearRight]
         + (20.0f) * sensors.sensors[Sensors::RFoot_FSR_FrontLeft]
         + (20.0f) * sensors.sensors[Sensors::RFoot_FSR_RearLeft];
   }
   if (totalPressure == 0) return  0;
   ZMPF /= totalPressure;
   return ZMPF;
}

void BodyModel::simulationUpdate() {}

void BodyModel::processUpdate(Odometry *odometry, const SensorValues &sensors) {
   float forwardL, forwardR, leftL, leftR, turnLR, liftL, liftR;
   walkCycle.generateWalk(forwardL, forwardR, leftL, leftR, turnLR, liftL, liftR);

   boost::numeric::ublas::matrix<float> b2f =
      kinematics->evaluateDHChain(
         Kinematics::FOOT,
         Kinematics::BODY,
         isLeftPhase ? Kinematics::RIGHT_CHAIN : Kinematics::LEFT_CHAIN);

   boost::numeric::ublas::matrix<float> b2fOther =
      kinematics->evaluateDHChain(
         Kinematics::FOOT,
         Kinematics::BODY,
         isLeftPhase ? Kinematics::LEFT_CHAIN : Kinematics::RIGHT_CHAIN);

   boost::numeric::ublas::matrix<float> n2f =
      kinematics->evaluateDHChain(
         Kinematics::FOOT,
         Kinematics::NECK,
         isLeftPhase ? Kinematics::RIGHT_CHAIN : Kinematics::LEFT_CHAIN);

   boost::numeric::ublas::matrix<float> f2w =
                  kinematics->createFootToWorldTransform(
         isLeftPhase ? Kinematics::RIGHT_CHAIN : Kinematics::LEFT_CHAIN);

   boost::numeric::ublas::matrix<float> b2w = boost::numeric::ublas::prod(f2w, b2f);
   boost::numeric::ublas::matrix<float> n2w = boost::numeric::ublas::prod(f2w, n2f);
   boost::numeric::ublas::matrix<float> origin = vec4<float>(0, 0, 0, 1);
   boost::numeric::ublas::matrix<float> result = boost::numeric::ublas::prod(b2f, origin);
   boost::numeric::ublas::matrix<float> rPend = boost::numeric::ublas::prod(b2w, origin);
   boost::numeric::ublas::matrix<float> neckPend = boost::numeric::ublas::prod(n2w, origin);
   // float deg = atan2(neckPend(0, 0) - rPend(0, 0), neckPend(2, 0) - rPend(2, 0));
//   std::cout << neckPend(0, 0) << " : " << rPend(0, 0) << " " << neckPend(1, 0) << " " << neckPend(2, 0) << " " << RAD2DEG(deg) << std::endl;

   // float alpha = sensors.sensors[Sensors::InertialSensor_AngleY];
   // float beta = deg; // sensors.joints.angles[Joints::LHipPitch];
   // float psi = alpha - beta;
   // float psiKinematic = atan2(rPend(0, 0), COM_HEIGHT);
   // float thau = psi - psiKinematic;

   // Calculate the centre of mass, convert to frame of reference of foot
   // (ie to foot then rotated by body lean)
   boost::numeric::ublas::matrix<float> com = kinematics->evaluateMassChain();
   boost::numeric::ublas::matrix<float> comOther = prod(b2fOther, com); // inserted CoM for other foot - BH
   centreOfMassOther.x = comOther(0, 0);                                // "
   centreOfMassOther.y = comOther(1, 0);                                // "
   centreOfMassOther.z = comOther(2, 0);                                // "
   com = prod(b2f, com);
   centreOfMass.x = com(0, 0);
   centreOfMass.y = com(1, 0);
   centreOfMass.z = com(2, 0);

   //std::cout << "COM x: " << centreOfMass.x << " y: " << centreOfMass.y << " z: " << centreOfMass.z << std::endl;

   //the following is left from 2011 and is not used atm

   // Get the ZMP from the supporting foot
   float zmpFootOffset = 0;
   if (!getWalkCycle().isDoubleSupportPhase()) {
      zmpFootOffset = getFootZMP(!isLeftPhase, sensors);
   }

   //prediction update
   pendulumModel.predictNext(1/100.0, zmpFootOffset);

   //std::cout << "##############" << std::endl;
   //std::cout << "prediction:   " << std::endl;
   //std::cout << "x: " << pendulumModel.x << std::endl;
   //std::cout << "dx: " << pendulumModel.dx << std::endl;
   //std::cout << "accX: " << pendulumModel.accX << std::endl;
   //std::cout << "theta: " << pendulumModel.theta << std::endl;
   //std::cout << "t: " << pendulumModel.walkCycle.t << std::endl;

   //observation update
   float supportFootPosition = result(0, 0);
   //float h = result(2, 0);
   //float deg = atan2(supportFootPosition - zmpFootOffset, h);
   WalkCycle currentWalkCycle = walkCycle;

   /*
   std::cout << "ZMP :" << zmpFootOffset << std::endl;
   std::cout << "observed x:" << supportFootPosition <<  std::endl;
   std::cout << "observed accX:" << sensors.sensors[Sensors::InertialSensor_AccX]
             << std::endl;
   */

   float alpha = 0.2;
   float beta = 0.5;

   pendulumModel.x = (1 - beta) * pendulumModel.x + (beta) * supportFootPosition;
   if (pendulumModel.walkCycle.leftPhase == currentWalkCycle.leftPhase) {
      pendulumModel.dx = (1 - alpha) * pendulumModel.dx + (alpha) * (pendulumModel.x - lastPendulumModel.x) * 100;
   } else { //switching foot
      pendulumModel.x = -pendulumModel.x/1.2;
   }
   pendulumModel.walkCycle = currentWalkCycle;

   lastPendulumModel.x = pendulumModel.x;
   lastPendulumModel.dx = pendulumModel.dx;

   /*
   std::cout << "updated:" << std::endl;
   std::cout << "x: " << pendulumModel.x << std::endl;
   std::cout << "dx: " << pendulumModel.dx << std::endl;
   std::cout << "t: " << pendulumModel.walkCycle.t << std::endl;
   */

   /*foot switched
   if (pendulumModel.walkCycle.leftPhase != currentWalkCycle.leftPhase) {
      if (!currentWalkCycle.leftPhase) {
         // switching to left support phase
         pendulumModel.x = pendulumModel.x + forwardR - forwardL;
      } else {
         pendulumModel.x = pendulumModel.x + forwardL - forwardR;
      }

   }*/
}

void BodyModel::observationUpdate(Odometry *odometry, const SensorValues &sensors) {
   /*
   float supportFootPosition = isLeftPhase ? forwardR : forwardL;
   // supportFootPosition += isOnFront ? FOOT_LENGTH/2 : -FOOT_LENGTH/2;
   float r = sqrt(supportFootPosition * supportFootPosition + COM_HEIGHT * COM_HEIGHT);

   float obsDTheta = lastHatDTheta +
                     sensors.sensors[Sensors::InertialSensor_AccX] * cos(theta) * MOTION_DT / r;
   dTheta += (obsDTheta - dTheta) * 0.5;
   */
}


/**
 * @brief       Battery_model
 * @file        Battery_model.cpp
 * @author      Richard Kirby <rkirby@kspresearch.com>
 * @copyright   Copyright (c) 2018, Swift Engineering Inc.
 * @license     Licensed under the MIT license. See LICENSE for details.
 */

#include "Battery_model.hpp"

#include <boost/math/tools/rational.hpp>

#include <cmath>
#include <cstdint>
#include <cstdio>
#include <limits>

namespace avionics_sim {

bool Battery_model::initialize(const double initial_soc, const double capacity, const uint8_t num_cells, const double C,
                               const double R) {
    // R must be > 0 or divide by zero error, set C = 0 to turn off transient response.
    if (!(R <= 0.0 || C < 0.0)) {
        m_c = C;
        m_r = R;
    }

    m_num_discharge_curves = 0;
    m_initial_soc = initial_soc;
    m_soc = m_initial_soc;  // in Ah
    m_capacity = capacity;  // in Ah
    m_num_cells = num_cells;
    m_low_current_curve_current = std::numeric_limits<double>::infinity();
    m_high_current_curve_current = -1;
    return true;
}

void Battery_model::update_soc(const double current, const double timestep, double *const voltage, double *const soc) {
    // simple coluomb counter
    if (timestep > 0.0) {
        m_soc += current * timestep / 3600;  // Time step is in seconds, need hours for mAhrs
        *voltage = get_voltage(m_soc * 1000, current, timestep);  // curves are in mA hours
        *soc = m_soc;
    } else {
        *voltage = m_last_voltage;
        *soc = m_soc;
    }
}

void Battery_model::update_soc_ocv(const double current, const double timestep, double *const voltage,
                                   double *const soc) {
    // simple coluomb counter
    if (timestep > 0.0 && current > 0.0) {
        m_soc += current * timestep / 3600;  // Time step is in seconds, need hours for mAhrs
        *voltage = get_voltage_ocv(m_soc, current, timestep);
        *soc = m_soc;
    } else {
        *voltage = m_last_voltage_ocv;
        *soc = m_soc;
    }
}

// discharge curves are in mAhrs
double Battery_model::get_voltage(const double soc, const double current_in, const double timestep) {
    double current = current_in;

    // Ensure that there are at least two discharge curves
    if (m_num_discharge_curves < 2) {
        return m_num_cells * m_last_voltage;
    }

    // Ensure current is between the two discharge curves
    if (current > m_high_current_curve_current) {
        current = m_high_current_curve_current;
    }

    if (current < m_low_current_curve_current) {
        current = m_low_current_curve_current;
    }

    // Find the curves that bracket the load
    int8_t low_current_discharge_curve = -1, high_current_discharge_curve = -1;
    double delta1 = std::numeric_limits<double>::infinity(), delta2 = std::numeric_limits<double>::infinity();

    for (uint8_t i = 0 ; i < m_num_discharge_curves ; i++) {
        if (m_discharge_curves[i].load <= current) {
            if (delta1 > (current - m_discharge_curves[i].load)) {
                delta1 = current - m_discharge_curves[i].load;
                low_current_discharge_curve = i;
            }
        }

        if (m_discharge_curves[i].load >= current) {
            if (delta2 > (m_discharge_curves[i].load - current)) {
                delta2 = m_discharge_curves[i].load - current;
                high_current_discharge_curve = i;
            }
        }
    }

    // The low voltage comes from the high discharge current and vice versa
    double voltage_low = compute_voltage(soc, high_current_discharge_curve);
    double voltage_high = compute_voltage(soc, low_current_discharge_curve);

    // If the load is bracketed by the same curve, then the load is equal to that curve otherwise interpolate
    double voltage;

    if (low_current_discharge_curve == high_current_discharge_curve) {
        voltage = voltage_low;
    } else {
        voltage = (
                      (current - m_discharge_curves[low_current_discharge_curve].load) /
                      (
                          m_discharge_curves[high_current_discharge_curve].load
                          - m_discharge_curves[low_current_discharge_curve].load
                      )
                  )
                  * (voltage_high - voltage_low) + voltage_low;
    }

    // add transient response
    if (timestep > 0) {
        voltage = voltage - (m_c / m_r) * (voltage - m_last_voltage) / (timestep);
        m_last_voltage = voltage;
    }

    return m_num_cells * voltage;
}

// discharge curves are in mAhrs
double Battery_model::get_voltage_ocv(const double soc, const double current_in, const double timestep) {
    double current = current_in;

    // get open circuit voltage, OCV uses curve 0
    double voltage = compute_voltage(soc, 0);   // returns voltage per cell

    // compensate for internal resistance drop
    float soc_percent = (m_capacity - soc) / m_capacity;
    double ir = compute_internal_resistance(soc_percent, 1.0);
    voltage = voltage - (ir * current) / m_num_cells;

    // add transient response
    if (timestep > 0) {
        voltage = voltage - (m_c / m_r) * (voltage - m_last_voltage_ocv) / (timestep);
        m_last_voltage_ocv = voltage;
    }

    return  m_num_cells * voltage;
}

bool Battery_model::add_discharge_curves(const discharge_curve &dc) {
    if (m_num_discharge_curves > max_num_discharge_curves) {
        return false;
    }

    m_discharge_curves.push_back(dc);

    if (m_discharge_curves[m_num_discharge_curves].load < m_low_current_curve_current) {
        m_low_current_curve_current = m_discharge_curves[m_num_discharge_curves].load;
        m_index_of_low_current_curve = m_num_discharge_curves;
    }

    if (m_discharge_curves[m_num_discharge_curves].load > m_high_current_curve_current) {
        m_high_current_curve_current = m_discharge_curves[m_num_discharge_curves].load;
        m_index_of_high_current_curve = m_num_discharge_curves;
    }

    m_num_discharge_curves++;
    return true;
}

bool Battery_model::set_internal_resistance_coefficients(const internal_resistance_curve &ir_coeff) {
    if (m_internal_resistance_curve.size() > 0) {
        return false;    // currently only one internal resistance curve is supported, but
    }

    // using a vector allows support for multiple internal resistance curves, which will change with state of health
    m_internal_resistance_curve.push_back(ir_coeff);
    return true;
}

double  Battery_model::compute_voltage(double soc, const uint8_t curve) {
    return boost::math::tools::evaluate_polynomial(m_discharge_curves[curve].c, soc);
}

double Battery_model::compute_internal_resistance(double soc, double soh) {
    // allows adding additional internal resistance curves for different states of health
    // currently only one curve is supported
    const uint8_t curve =  0;

    return (boost::math::tools::evaluate_polynomial(m_internal_resistance_curve[curve].c, soc,
            m_internal_resistance_num_coef) + m_internal_resistance_curve[curve].harness_resistance);
}

double Battery_model::compute_max_power_available(const double v_volts, const double ir_ohms) {
    return v_volts * v_volts / (4 * ir_ohms);
}

double Battery_model::get_soc_percent() {
    return (m_capacity - m_soc) / m_capacity;
}

}  // namespace avionics_sim