//---------------------------------------------------------------------------//
// Copyright (c) 2018-2020 Mikhail Komarov <nemo@nil.foundation>
// Copyright (c) 2020 Nikita Kaskov <nbering@nil.foundation>
//
// MIT License
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.
//---------------------------------------------------------------------------//

#ifndef CRYPTO3_R1CS_GG_PPZKSNARK_IPP2_PROVE_HPP
#define CRYPTO3_R1CS_GG_PPZKSNARK_IPP2_PROVE_HPP

#include <algorithm>
#include <memory>
#include <vector>
#include <tuple>

#include <boost/iterator/zip_iterator.hpp>

#include <nil/crypto3/detail/pack_numeric.hpp>

#include <nil/crypto3/hash/algorithm/hash.hpp>
#include <nil/crypto3/hash/sha2.hpp>

#include <nil/crypto3/algebra/multiexp/multiexp.hpp>
#include <nil/crypto3/algebra/multiexp/policies.hpp>

#include <nil/crypto3/zk/snark/schemes/ppzksnark/r1cs_gg_ppzksnark/ipp2/proof.hpp>
#include <nil/crypto3/zk/snark/schemes/ppzksnark/r1cs_gg_ppzksnark/ipp2/srs.hpp>

namespace nil {
    namespace crypto3 {
        namespace zk {
            namespace snark {
                /// Returns the vector used for the linear combination fo the inner pairing product
                /// between A and B for the Groth16 aggregation: A^r * B. It is required as it
                /// is not enough to simply prove the ipp of A*B, we need a random linear
                /// combination of those.
                template<typename FieldType>
                std::vector<typename FieldType::value_type>
                    structured_scalar_power(std::size_t num, const typename FieldType::value_type &s) {
                    std::vector<typename FieldType::value_type> powers = {FieldType::value_type::one()};
                    for (int i = 1; i < num; i++) {
                        powers.emplace_back(powers[i - 1] * s);
                    }
                    return powers;
                }

                /// compress is similar to commit::{V,W}KEY::compress: it modifies the `vec`
                /// vector by setting the value at index $i:0 -> split$  $vec[i] = vec[i] +
                /// vec[i+split]^scaler$. The `vec` vector is half of its size after this call.
                template<typename CurveAffine>
                void compress(std::vector<typename CurveAffine::value_type> &vec, std::size_t split,
                              const typename CurveAffine::scalar &scaler) {
                    std::vector<typename CurveAffine::value_type> left = {vec.begin(), vec.begin() + split},
                                                                  right = {vec.begin() + split, vec.end()};
                    std::for_each(boost::make_zip_iterator(std::make_tuple(left.begin(), right.begin())),
                                  boost::make_zip_iterator(std::make_tuple(left.end(), right.end())),
                                  [&](const std::tuple<const typename CurveAffine::value_type &,
                                                       const typename CurveAffine::value_type &> &t) {
                                      auto x = std::get<1>(t).to_projective() * scaler;
                                      x += std::get<0>(t);
                                      std::get<0>(t) = x.to_affine();
                                  });

                    vec.resize(left.size());
                }

                /// Aggregate `n` zkSnark proofs, where `n` must be a power of two.
                template<typename CurveType, typename InputProofIterator, typename Hash = hashes::sha2<256>>
                typename std::enable_if<std::is_same<typename std::iterator_traits<InputProofIterator>::value_type,
                                                     r1cs_gg_ppzksnark_aggregate_proof<CurveType>>::value,
                                        r1cs_gg_ppzksnark_aggregate_proof<CurveType>>::type
                    aggregate_proofs(const r1cs_gg_ppzksnark_proving_srs<CurveType> &srs, InputProofIterator first,
                                     InputProofIterator last) {
                    std::size_t size = std::distance(first, last);
                    BOOST_ASSERT((size & (size - 1)) == 0);
                    BOOST_ASSERT(srs.valid(size));

                    std::vector<typename CurveType::g1_type::value_type> a, c;
                    std::vector<typename CurveType::g2_type::value_type> b;

                    // We first commit to A B and C - these commitments are what the verifier
                    // will use later to verify the TIPP and MIPP proofs
                    while (first != last) {
                        a.emplace_back(*first.g_A);
                        b.emplace_back(*first.g_B);
                        c.emplace_back(*first.g_C);
                        ++first;
                    }

                    // A and B are committed together in this scheme
                    // we need to take the reference so the macro doesn't consume the value
                    // first
                    r1cs_gg_ppzksnark_ipp2_commitment_output<CurveType> com_ab =
                        r1cs_gg_ppzksnark_ipp2_commitment<CurveType>::pair(srs.vkey, srs.wkey, a.begin(), a.end(),
                                                                           b.begin(), b.end());
                    r1cs_gg_ppzksnark_ipp2_commitment_output<CurveType> com_c =
                        r1cs_gg_ppzksnark_ipp2_commitment<CurveType>::pair(srs.vkey, c.begin(), c.end());

                    // Random linear combination of proofs
                    std::size_t counter_nonce = 1;
                    std::array<std::uint8_t, sizeof(std::size_t)> counter_nonce_bytes;
                    detail::pack<stream_endian::big_byte_big_bit>({counter_nonce}, counter_nonce_bytes);
                    accumulator_set<Hash> acc;

                    hash<Hash>(counter_nonce_bytes, acc);
                    hash<Hash>(std::get<0>(com_ab), acc);
                    hash<Hash>(std::get<1>(com_ab), acc);
                    hash<Hash>(std::get<0>(com_c), acc);
                    hash<Hash>(std::get<1>(com_c), acc);

                    typename Hash::digest_type d = accumulators::extract::hash<Hash>(acc);
                    typename CurveType::scalar_field_type::value_type r;
                    detail::pack(d, r.data);
                    r = r.inversed();

                    // r, r^2, r^3, r^4 ...
                    std::vector<typename CurveType::scalar_field_type::value_type> r_vec =
                        structured_scalar_power<typename CurveType::scalar_field_type>(std::distance(first, last), r);
                    // r^-1, r^-2, r^-3
                    std::vector<typename CurveType::scalar_field_type::value_type> r_inv;
                    for (const typename CurveType::scalar_field_type::value_type &ri : r_vec) {
                        r_inv.emplace_back(ri.inversed());
                    }

                    // B^{r}
                    std::vector<typename CurveType::scalar_field_type::value_type> b_r;
                    std::for_each(boost::make_zip_iterator(std::make_tuple(b.begin(), r_vec.begin())),
                                  boost::make_zip_iterator(std::make_tuple(b.end(), r_vec.end())),
                                  [&](const std::tuple<const typename CurveType::g2_type::value_type &,
                                                       const typename CurveType::scalar_field_type::value_type &> &t) {
                                      b_r.emplace_back((std::get<0>(t).to_projective() * std::get<1>(t)).to_affine());
                                  });

                    // w^{r^{-1}}
                    auto wkey_r_inv = srs.wkey.scale(r_inv);

                    // we prove tipp and mipp using the same recursive loop
                    tipp_mipp_proof<CurveType> proof =
                        prove_tipp_mipp(srs, wkey_r_inv, a.begin(), a.end(), b_r.begin(), b_r.end(), c.begin(), c.end(),
                                        r_vec.begin(), r_vec.end());
                    // compute A * B^r for the verifier
                    auto ip_ab = algebra::pair<CurveType>(a, b_r);
                    // compute C^r for the verifier
                    auto agg_c = algebra::multiexp<algebra::policies::multiexp_method_bos_coster>(c, r_vec);

                    // debug assert
                    auto computed_com_ab = r1cs_gg_ppzksnark_ipp2_commitment<CurveType>::pair(
                        srs.vkey, wkey_r_inv, a.begin(), a.end(), b_r.begin(), b_r.end());
                    BOOST_ASSERT(com_ab == computed_com_ab);

                    return {com_ab, com_c, ip_ab, agg_c, proof};
                }

                /// Proves a TIPP relation between A and B as well as a MIPP relation with C and
                /// r. Commitment keys must be of size of A, B and C. In the context of Groth16
                /// aggregation, we have that B = B^r and wkey is scaled by r^{-1}. The
                /// commitment key v is used to commit to A and C recursively in GIPA such that
                /// only one KZG proof is needed for v. In the original paper version, since the
                /// challenges of GIPA would be different, two KZG proofs would be needed.
                template<typename CurveType, typename InputG1Iterator, typename InputG2Iterator,
                         typename InputScalarIterator, typename Hash = hashes::sha2<256>>
                tipp_mipp_proof<CurveType> prove_tipp_mipp(const r1cs_gg_ppzksnark_proving_srs<CurveType> &srs,
                                                           const r1cs_gg_ppzksnark_ipp2_wkey<CurveType> &wkey,
                                                           InputG1Iterator afirst, InputG1Iterator alast,
                                                           InputG2Iterator bfirst, InputG2Iterator blast,
                                                           InputG1Iterator cfirst, InputG1Iterator clast,
                                                           InputScalarIterator rfirst, InputScalarIterator rlast) {
                    std::size_t asize = std::distance(afirst, alast);
                    std::size_t bsize = std::distance(bfirst, blast);

                    BOOST_ASSERT((asize & (asize - 1)) == 0 || asize == bsize);
                    typename std::iterator_traits<InputScalarIterator>::value_type r_shift = *rfirst + 1;

                    // Run GIPA
                    std::tuple<gipa_proof<CurveType>, std::vector<typename CurveType::scalar_field_type::value_type>,
                               std::vector<typename CurveType::scalar_field_type::value_type>>
                        gtm =
                            gipa_tipp_mipp(afirst, alast, bfirst, blast, cfirst, clast, rfirst, rlast, srs.vkey, wkey);

                    // Prove final commitment keys are wellformed
                    // we reverse the transcript so the polynomial in kzg opening is constructed
                    // correctly - the formula indicates x_{l-j}. Also for deriving KZG
                    // challenge point, input must be the last challenge.
                    std::reverse(std::get<1>(gtm).begin(), std::get<1>(gtm).end());
                    std::reverse(std::get<2>(gtm).begin(), std::get<2>(gtm).end());
                    typename std::iterator_traits<InputScalarIterator>::value_type r_inverse = r_shift.inverse();

                    // KZG challenge point
                    std::size_t counter_nonce = 1;
                    std::array<std::uint8_t, sizeof(std::size_t)> counter_nonce_bytes;
                    detail::pack<stream_endian::big_byte_big_bit>({counter_nonce}, counter_nonce_bytes);
                    accumulator_set<Hash> acc;

                    hash<Hash>(counter_nonce_bytes, acc);
                    hash<Hash>(*std::get<1>(gtm).begin(), acc);
                    hash<Hash>(std::get<0>(std::get<0>(gtm).final_vkey), acc);
                    hash<Hash>(std::get<1>(std::get<0>(gtm).final_vkey), acc);
                    hash<Hash>(std::get<0>(std::get<0>(gtm).final_wkey), acc);
                    hash<Hash>(std::get<1>(std::get<0>(gtm).final_wkey), acc);

                    typename Hash::digest_type d = accumulators::extract::hash<Hash>(acc);
                    typename CurveType::scalar_field_type::value_type z;
                    multiprecision::import_bits(z.data, d);
                    z = z.inversed();

                    // Complete KZG proofs
                    kzg_opening<typename CurveType::g2_type> vkey_opening = prove_commitment_key_kzg_opening(
                        srs.h_alpha_powers_table, srs.h_beta_powers_table, srs.n, std::get<2>(gtm),
                        CurveType::scalar_field_type::value_type::one(), z);
                    kzg_opening<typename CurveType::g1_type> wkey_opening = prove_commitment_key_kzg_opening(
                        srs.g_alpha_powers_table, srs.g_beta_powers_table, srs.n, std::get<1>(gtm), r_inverse, z);

                    return {std::get<0>(gtm), vkey_opening, wkey_opening};
                }

                /*
                 * @brief gipa_tipp_mipp peforms the recursion of the GIPA protocol for TIPP and MIPP.
                 * It returns a proof containing all intermdiate committed values, as well as
                 * the challenges generated necessary to do the polynomial commitment proof
                 * later in TIPP.
                 * @param wkey scaled key w^r^-1
                 */
                template<typename CurveType, typename InputG1Iterator, typename InputG2Iterator,
                         typename InputScalarIterator, typename Hash = hashes::sha2<256>>
                std::tuple<gipa_proof<CurveType>, std::vector<typename CurveType::scalar_field_type::value_type>,
                           std::vector<typename CurveType::scalar_field_type::value_type>>
                    gipa_tipp_mipp(InputG1Iterator afirst, InputG1Iterator alast, InputG2Iterator bfirst,
                                   InputG2Iterator blast, InputG1Iterator cfirst, InputG1Iterator clast,
                                   InputScalarIterator rfirst, InputScalarIterator rlast,
                                   const r1cs_gg_ppzksnark_ipp2_vkey<CurveType> &vkey,
                                   const r1cs_gg_ppzksnark_ipp2_wkey<CurveType> &wkey) {

                    std::vector<typename std::iterator_traits<InputG1Iterator>::value_type> m_a = {afirst, alast},
                                                                                m_c = {cfirst, clast};
                    std::vector<typename std::iterator_traits<InputG2Iterator>::value_type> m_b = {bfirst, blast};
                    std::vector<typename std::iterator_traits<InputScalarIterator>::value_type> m_r = {rfirst, rlast};

                    r1cs_gg_ppzksnark_ipp2_vkey<CurveType> vkey = vkey;
                    r1cs_gg_ppzksnark_ipp2_wkey<CurveType> wkey = wkey;

                    // storing the values for including in the proof
                    std::vector<std::tuple<typename r1cs_gg_ppzksnark_ipp2_commitment<CurveType>::output_type,
                                           typename r1cs_gg_ppzksnark_ipp2_commitment<CurveType>::output_type>>
                        comms_ab;
                    std::vector<std::tuple<typename r1cs_gg_ppzksnark_ipp2_commitment<CurveType>::output_type,
                                           typename r1cs_gg_ppzksnark_ipp2_commitment<CurveType>::output_type>>
                        comms_c;
                    std::vector<typename CurveType::pairing::gt_type::value_type> z_ab;
                    std::vector<typename std::iterator_traits<InputG1Iterator>::value_type> z_c;
                    std::vector<typename CurveType::scalar_field_type::value_type> challenges, challenges_inv;

                    while (m_a.size() > 1) {
                        // recursive step
                        // Recurse with problem of half size
                        std::size_t split = m_a.size() / 2;

                        // TIPP ///
                        std::vector<typename std::iterator_traits<InputG1Iterator>::value_type>
                            a_left = {m_a.begin(), m_a.begin() + split},
                            a_right = {m_a.begin() + split, m_a.end()};
                        std::vector<typename std::iterator_traits<InputG2Iterator>::value_type>
                            b_left = {m_b.begin(), m_b.begin() + split},
                            b_right = {m_b.begin() + split, m_b.end()};
                        // MIPP ///
                        // c[:n']   c[n':]
                        std::vector<typename std::iterator_traits<InputG1Iterator>::value_type>
                            c_left = {m_c.begin() + m_c.begin() + split},
                            c_right = {m_c.begin() + split, m_c.end()};
                        // r[:n']   r[:n']
                        std::vector<typename std::iterator_traits<InputScalarIterator>::value_type>
                            r_left = {m_r.begin(), m_r.begin() + split},
                            r_right = {m_r.begin() + split, m_r.end()};

                        r1cs_gg_ppzksnark_ipp2_vkey<CurveType> vk_left = {{vkey.a.begin(), vkey.a.begin() + split},
                                                                          {vkey.b.begin(), vkey.b.begin() + split}},
                                                               vk_right = {{vkey.a.begin() + split, vkey.a.end()},
                                                                           {vkey.b.begin() + split, vkey.b.end()}};

                        r1cs_gg_ppzksnark_ipp2_wkey<CurveType> wk_left = {{wkey.a.begin(), wkey.a.begin() + split},
                                                                          {wkey.b.begin(), wkey.b.begin() + split}},
                                                               wk_right = {{wkey.a.begin() + split, wkey.a.end()},
                                                                           {wkey.b.begin() + split, wkey.b.end()}};

                        // See section 3.3 for paper version with equivalent names
                        typename r1cs_gg_ppzksnark_ipp2_commitment<CurveType>::output_type tab_l =
                            r1cs_gg_ppzksnark_ipp2_commitment<CurveType>::pair(vk_left, wk_right, a_right, b_left);
                        typename r1cs_gg_ppzksnark_ipp2_commitment<CurveType>::output_type tab_r =
                            r1cs_gg_ppzksnark_ipp2_commitment<CurveType>::pair(vk_right, wk_left, a_left, b_right);

                        // TIPP part
                        typename CurveType::pairing::gt_type::value_type zab_l =
                            algebra::pair<CurveType>(a_right, b_left);
                        typename CurveType::pairing::gt_type::value_type zab_r =
                            algebra::pair<CurveType>(a_left, b_right);

                        // MIPP part
                        // z_l = c[n':] ^ r[:n']
                        typename std::iterator_traits<InputG1Iterator>::value_type zc_l =
                            algebra::multiexp<algebra::policies::multiexp_method_bos_coster>(c_right, r_left);
                        // Z_r = c[:n'] ^ r[n':]
                        typename std::iterator_traits<InputG1Iterator>::value_type zc_r =
                            algebra::multiexp<algebra::policies::multiexp_method_bos_coster>(c_left, r_right);
                        // u_l = c[n':] * v[:n']
                        typename r1cs_gg_ppzksnark_ipp2_commitment<CurveType>::output_type tuc_l =
                            r1cs_gg_ppzksnark_ipp2_commitment<CurveType>::pair(vk_left, c_right);
                        // u_r = c[:n'] * v[n':]
                        typename r1cs_gg_ppzksnark_ipp2_commitment<CurveType>::output_type tuc_r =
                            r1cs_gg_ppzksnark_ipp2_commitment<CurveType>::pair(vk_right, c_left);

                        // Fiat-Shamir challenge
                        typename CurveType::scalar_field_type::value_type default_transcript =
                            CurveType::scalar_field_type::value_type::zero();
                        auto transcript = challenges.empty() ? default_transcript : *(challenges.end() - 1);

                        // combine both TIPP and MIPP transcript
                        std::size_t counter_nonce = 1;
                        std::array<std::uint8_t, sizeof(std::size_t)> counter_nonce_bytes;
                        detail::pack<stream_endian::big_byte_big_bit>({counter_nonce}, counter_nonce_bytes);
                        accumulator_set<Hash> acc;

                        hash<Hash>(counter_nonce_bytes, acc);
                        hash<Hash>(transcript, acc);
                        hash<Hash>(std::get<0>(tab_l), acc);
                        hash<Hash>(std::get<1>(tab_l), acc);
                        hash<Hash>(std::get<0>(tab_r), acc);
                        hash<Hash>(std::get<1>(tab_r), acc);
                        hash<Hash>(zab_l, acc);
                        hash<Hash>(zab_r, acc);
                        hash<Hash>(zc_l, acc);
                        hash<Hash>(zc_r, acc);
                        hash<Hash>(std::get<0>(tuc_l), acc);
                        hash<Hash>(std::get<1>(tuc_l), acc);
                        hash<Hash>(std::get<0>(tuc_r), acc);
                        hash<Hash>(std::get<1>(tuc_r), acc);

                        typename hashes::sha2<256>::digest_type d = accumulators::extract::hash<hashes::sha2<256>>(acc);
                        typename CurveType::scalar_field_type::value_type c_inv;
                        detail::pack(d, c_inv.data);
                        c_inv = c_inv.inversed();

                        // Optimization for multiexponentiation to rescale G2 elements with
                        // 128-bit challenge Swap 'c' and 'c_inv' since can't control bit size
                        // of c_inv
                        typename CurveType::scalar_field_type::value_type c = c_inv.inversed();

                        // Set up values for next step of recursion
                        // A[:n'] + A[n':] ^ x
                        compress(m_a, split, c);
                        // B[:n'] + B[n':] ^ x^-1
                        compress(m_b, split, c_inv);

                        // c[:n'] + c[n':]^x
                        compress(m_c, split, c);
                        std::for_each(
                            boost::make_zip_iterator(std::make_tuple(r_left.begin(), r_right.begin())),
                            boost::make_zip_iterator(std::make_tuple(r_left.end(), r_right.end())),
                            [&](const std::tuple<typename std::iterator_traits<InputScalarIterator>::value_type &,
                                                 typename std::iterator_traits<InputScalarIterator>::value_type &> &t) {
                                // r[:n'] + r[n':]^x^-1
                                std::get<1>(t) *= c_inv;
                                std::get<0>(t) += std::get<1>(t);
                            });
                        std::size_t len = r_left.size();
                        m_r.resize(len);    // shrink to new size

                        // v_left + v_right^x^-1
                        vkey = vk_left.compress(vk_right, c_inv);
                        // w_left + w_right^x
                        wkey = wk_left.compress(wk_right, c);

                        comms_ab.emplace_back({tab_l, tab_r});
                        comms_c.emplace_back({tuc_l, tuc_r});
                        z_ab.emplace_back({zab_l, zab_r});
                        z_c.emplace_back({zc_l, zc_r});
                        challenges.emplace_back(c);
                        challenges_inv.emplace_back(c_inv);
                    }

                    BOOST_ASSERT(m_a.size() == 1 && m_b.size() == 1);
                    BOOST_ASSERT(m_c.size() == 1 && m_r.size() == 1);
                    BOOST_ASSERT(vkey.a.size() == 1 && vkey.b.size() == 1);
                    BOOST_ASSERT(wkey.a.size() == 1 && wkey.b.size() == 1);

                    return std::make_tuple({a.size(), comms_ab, comms_c, z_ab, z_c, m_a[0], m_b[0], m_c[0], m_r[0],
                                            *vkey.begin(), *wkey.begin()},
                                           challenges, challenges_inv);
                }

                /// It returns the evaluation of the polynomial $\prod (1 + x_{l-j}(rX)^{2j}$ at
                /// the point z, where transcript contains the reversed order of all challenges (the x).
                /// The challenges must be in reversed order for the correct evaluation of the
                /// polynomial in O(logn)
                template<typename FieldType, typename InputFieldValueIterator>
                typename std::enable_if<std::is_same<typename std::iterator_traits<InputFieldIterator>::value_type,
                                                     typename FieldType::value_type>::value,
                                        typename FieldType::value_type>::type
                    polynomial_evaluation_product_form_from_transcript(InputFieldValueIterator transcript_first,
                                                                       InputFieldValueIterator transcript_last,
                                                                       const typename FieldType::value_type &z,
                                                                       const typename FieldType::value_type &r_shift) {
                    // this is the term (rz) that will get squared at each step to produce the
                    // $(rz)^{2j}$ of the formula
                    typename FieldType::value_type power_zr = z;
                    power_zr *= r_shift;

                    // 0 iteration
                    typename FieldType::value_type res =
                        FieldType::value_type::one() + (*transcript_first * power_zr);
                    power_zr *= power_zr;
                    ++transcript_first;

                    // the rest
                    while (transcript_first != transcript_last) {
                        res *= FieldType::value_type::one() + (*transcript_first * power_zr);
                        power_zr *= power_zr;
                        ++transcript_first;
                    }

                    return res;
                }

                // Compute the coefficients of the polynomial $\prod_{j=0}^{l-1} (1 + x_{l-j}(rX)^{2j})$
                // It does this in logarithmic time directly; here is an example with 2
                // challenges:
                //
                //     We wish to compute $(1+x_1ra)(1+x_0(ra)^2) = 1 +  x_1ra + x_0(ra)^2 + x_0x_1(ra)^3$
                //     Algorithm: $c_{-1} = [1]$; $c_j = c_{i-1} \| (x_{l-j} * c_{i-1})$; $r = r*r$
                //     $c_0 = c_{-1} \| (x_1 * r * c_{-1}) = [1] \| [rx_1] = [1, rx_1]$, $r = r^2$
                //     $c_1 = c_0 \| (x_0 * r^2c_0) = [1, rx_1] \| [x_0r^2, x_0x_1r^3] = [1, x_1r, x_0r^2, x_0x_1r^3]$
                //     which is equivalent to $f(a) = 1 + x_1ra + x_0(ra)^2 + x_0x_1r^2a^3$
                //
                // This method expects the coefficients in reverse order so transcript[i] =
                // x_{l-j}.
                template<typename FieldType, typename InputFieldValueIterator>
                typename std::enable_if<std::is_same<typename std::iterator_traits<InputFieldValueIterator>::value_type,
                                                     typename FieldType::value_type>::value,
                                        std::vector<typename FieldType::value_type>>::type
                    polynomial_coefficients_from_transcript(InputFieldValueIterator transcript_first,
                                                            InputFieldValueIterator transcript_last,
                                                            const typename FieldType::value_type &r_shift) {
                    std::vector<typename FieldType::value_type> coefficients = {FieldType::value_type::one()};
                    typename FieldType::value_type power_2_r = r_shift;

                    while (transcript_first != transcript_last) {
                        std::size_t n = coefficients.size();
                        for (int j = 0; j < n; j++) {
                            coefficients.emplace_back(coefficients[j] * (*transcript_first * power_2_r));
                        }
                        power_2_r *= power_2_r;

                        ++transcript_first;
                    }

                    return coefficients;
                }

                /// Returns the KZG opening proof for the given commitment key. Specifically, it
                /// returns $g^{f(alpha) - f(z) / (alpha - z)}$ for $a$ and $b$.
                template<typename CurveAffine, typename InputScalarIterator>
                kzg_opening<CurveAffine> prove_commitment_key_kzg_opening(
                    MultiscalarPrecomp<CurveAffine> &srs_powers_alpha_table,
                    MultiscalarPrecomp<CurveAffine> &srs_powers_beta_table, std::size_t srs_powers_len,
                    InputScalarIterator transcript_first, InputScalarIterator transcript_last,
                    const typename std::iterator_traits<InputScalarIterator>::value_type &r_shift,
                    const typename std::iterator_traits<InputScalarIterator>::value_type &kzg_challenge) {
                    // f_v
                    DensePolynomial vkey_poly(
                        polynomial_coefficients_from_transcript(transcript_first, transcript_last, r_shift));

                    BOOST_ASSERT_MSG(srs_powers_len != vkey_poly.coeffs().size(), "Malformed SRS");
                    // f_v(z)
                    std::vector<typename std::iterator_traits<InputScalarIterator>::value_type> vkey_poly_z =
                        polynomial_evaluation_product_form_from_transcript(transcript_first, transcript_last,
                                                                           kzg_challenge, r_shift);

                    typename std::iterator_traits<InputScalarIterator>::value_type neg_kzg_challenge =
                        kzg_challenge.negate();

                    // f_v(X) - f_v(z) / (X - z)
                    DensePolynomial quotient_polynomial =
                        (vkey_poly - DensePolynomial(vkey_poly_z)) /
                        (DensePolynomial({neg_kzg_challenge,
                                          std::iterator_traits<InputScalarIterator>::value_type::one()}));

                    std::vector<typename std::iterator_traits<InputScalarIterator>::value_type>
                        quotient_polynomial_coeffs = quotient_polynomial.into_coeffs();

                    // multiexponentiation inner_product, inlined to optimize
                    std::size_t quotient_polynomial_coeffs_len = quotient_polynomial_coeffs.size();
                    auto getter = [&](std::size_t i) -> typename CurveAffine::scalar_field_type::value_type {
                        return i >= quotient_polynomial_coeffs_len ?
                                   std::iterator_traits<InputScalarIterator>::value_type::zero() :
                                   quotient_polynomial_coeffs[i];
                    };
                }
            }    // namespace snark
        }        // namespace zk
    }            // namespace crypto3
}    // namespace nil

#endif    // CRYPTO3_R1CS_GG_PPZKSNARK_TYPES_POLICY_HPP


#include "trajopt/problem_description.hpp"
#include "trajopt/common.hpp"
#include <boost/foreach.hpp>
#include "utils/logging.hpp"
#include "sco/expr_ops.hpp"
#include "trajopt/kinematic_constraints.hpp"
#include "trajopt/belief_constraints.hpp"
#include "trajopt/collision_avoidance.hpp"
#include "trajopt/rave_utils.hpp"
#include "trajopt/plot_callback.hpp"
#include "trajopt/rave_utils.hpp"
#include "utils/eigen_conversions.hpp"
#include "utils/eigen_slicing.hpp"
#include <boost/algorithm/string.hpp>
using namespace Json;
using namespace std;
using namespace OpenRAVE;
using namespace trajopt;
using namespace util;

namespace {

bool gRegisteredMakers = false;
void RegisterMakers() {

	CostInfo::RegisterMaker("pose", &PoseCostInfo::create);
	CostInfo::RegisterMaker("joint_pos", &JointPosCostInfo::create);
	CostInfo::RegisterMaker("joint_vel", &JointVelCostInfo::create);
	CostInfo::RegisterMaker("collision", &CollisionCostInfo::create);
	CostInfo::RegisterMaker("continuous_collision", &ContinuousCollisionCostInfo::create);

	CntInfo::RegisterMaker("joint", &JointConstraintInfo::create);
	CntInfo::RegisterMaker("pose", &PoseCntInfo::create);
	CntInfo::RegisterMaker("collision", &CollisionCntInfo::create);
	CntInfo::RegisterMaker("continuous_collision", &ContinuousCollisionCntInfo::create);
	CntInfo::RegisterMaker("cart_vel", &CartVelCntInfo::create);

	// belief costs and controls
	CostInfo::RegisterMaker("control", &ControlCostInfo::create);
	CostInfo::RegisterMaker("covariance", &CovarianceCostInfo::create);

	CntInfo::RegisterMaker("control", &ControlCntInfo::create);

	gRegisteredMakers = true;
}

BeliefRobotAndDOFPtr RADFromName(const string& name, RobotBasePtr robot) {
	if (name == "active") {
		return BeliefRobotAndDOFPtr(new BeliefRobotAndDOF(robot, robot->GetActiveDOFIndices(), robot->GetAffineDOF(), robot->GetAffineRotationAxis()));
	}
	vector<int> dof_inds;
	int affinedofs = 0;
	Vector rotationaxis(0,0,1);
	vector<string> components;
	boost::split(components, name, boost::is_any_of("+"));
	for (int i=0; i < components.size(); ++i) {
		std::string& component = components[i];
		if (RobotBase::ManipulatorPtr manip = GetManipulatorByName(*robot, component)) {
			vector<int> inds = manip->GetArmIndices();
			dof_inds.insert(dof_inds.end(), inds.begin(), inds.end());
		}
		else if (component == "base") {
			affinedofs |= DOF_X | DOF_Y | DOF_RotationAxis;
		}
		else if (component == "base_point") {
			affinedofs |= DOF_X | DOF_Y;
		}
		else if (KinBody::JointPtr joint = robot->GetJoint(component)) {
			dof_inds.push_back(joint->GetDOFIndex());
		}
		else PRINT_AND_THROW( boost::format("error in reading manip description: %s must be a manipulator, link, or 'base'")%component );
	}
	return BeliefRobotAndDOFPtr(new BeliefRobotAndDOF(robot, dof_inds, affinedofs, rotationaxis));
}

BoolVec toMask(const VectorXd& x) {
	BoolVec out(x.size());
	for (int i=0; i < x.size(); ++i) out[i] = (x[i] > 0);
	return out;
}

bool allClose(const VectorXd& a, const VectorXd& b) {
	return (a-b).array().abs().maxCoeff() < 1e-4;
}

}

namespace Json { //funny thing with two-phase lookup

void fromJson(const Json::Value& v, Vector3d& x) {
	vector<double> vx;
	fromJsonArray(v, vx, 3);
	x = Vector3d(vx[0], vx[1], vx[2]);
}
void fromJson(const Json::Value& v, Vector4d& x) {
	vector<double> vx;
	fromJsonArray(v, vx, 4);
	x = Vector4d(vx[0], vx[1], vx[2], vx[3]);
}
template <class T>
inline void fromJson(const Json::Value& v, Eigen::Matrix<T,Eigen::Dynamic,Eigen::Dynamic>& m) {
	int nRows = v.size();
	if (nRows != 0) {
		int nCols = v[0].size();
		m.resize(nRows, nCols);
		for (int i=0; i<nRows; i++) {
			if (v[i].size() != nCols)
				PRINT_AND_THROW(boost::format("matrix with variable number of cols not supported. expected %i cols at the %ith row, but got %i cols\n")%nCols%i%v[i].size());
			std::vector<T> row;
			fromJsonArray(v[i], row, nCols);
			for (int j=0; j<nCols; j++) m(i,j) = row[j];
		}
	} else {
		m.resize(0,0);
	}
}

}

namespace trajopt {

TRAJOPT_API ProblemConstructionInfo* gPCI;

void BasicInfo::fromJson(const Json::Value& v) {
	childFromJson(v, start_fixed, "start_fixed", true);
	childFromJson(v, n_steps, "n_steps");
	childFromJson(v, manip, "manip");
	childFromJson(v, robot, "robot", string(""));
	childFromJson(v, dofs_fixed, "dofs_fixed", IntVec());
	childFromJson(v, belief_space, "belief_space", false);
}

////
void fromJson(const Json::Value& v, CostInfoPtr& cost) {
	string type;
	childFromJson(v, type, "type");
	cost = CostInfo::fromName(type);
	if (!cost) PRINT_AND_THROW( boost::format("failed to construct cost named %s")%type );
	cost->fromJson(v);
	childFromJson(v, cost->name, "name", type);
}
CostInfoPtr CostInfo::fromName(const string& type) {
	if (!gRegisteredMakers) RegisterMakers();
	if (name2maker.find(type) != name2maker.end()) {
		return (*name2maker[type])();
	}
	else {
		RAVELOG_ERROR("There is no cost of type%s\n", type.c_str());
		return CostInfoPtr();
	}
}
map<string, CostInfo::MakerFunc> CostInfo::name2maker;

void CostInfo::RegisterMaker(const std::string& type, MakerFunc f) {
	name2maker[type] = f;
}

////
//// almost copied

void fromJson(const Json::Value& v, CntInfoPtr& cnt) {
	string type;
	childFromJson(v, type, "type");
	LOG_DEBUG("reading constraint: %s", type.c_str());
	cnt = CntInfo::fromName(type);
	if (!cnt) PRINT_AND_THROW( boost::format("failed to construct constraint named %s")%type );
	cnt->fromJson(v);
	childFromJson(v, cnt->name, "name", type);
}
CntInfoPtr CntInfo::fromName(const string& type) {
	if (!gRegisteredMakers) RegisterMakers();
	if (name2maker.find(type) != name2maker.end()) {
		return (*name2maker[type])();
	}
	else {
		RAVELOG_ERROR("There is no constraint of type%s\n", type.c_str());
		return CntInfoPtr();
	}
}
map<string,CntInfo::MakerFunc> CntInfo::name2maker;

void InitInfo::fromJson(const Json::Value& v) {
	string type_str;
	childFromJson(v, type_str, "type");
	int n_steps = gPCI->basic_info.n_steps;
	int n_dof = gPCI->rad->GetDOF();
	int b_dim = gPCI->rad->GetBDim(), u_dim = gPCI->rad->GetUDim();

	bool belief_space = gPCI->basic_info.belief_space;
	MatrixXd rt_Sigma0;
	if (belief_space && (type_str == "stationary" || type_str == "given_traj")) {
		FAIL_IF_FALSE(v.isMember("initial_rt_sigma"));
		const Value& rt_Sigma0_v = v["initial_rt_sigma"];
		Json::fromJson(rt_Sigma0_v, rt_Sigma0);
		if (rt_Sigma0.rows()!=n_dof) PRINT_AND_THROW("initial square root of sigma has wrong number of rows");
		if (rt_Sigma0.cols()!=n_dof) PRINT_AND_THROW("initial square root of sigma has wrong number of cols");
	}

	if (type_str == "stationary") {
		if (!belief_space) {
			data = toVectorXd(gPCI->rad->GetDOFValues()).transpose().replicate(n_steps, 1);
		} else {
			data.resize(n_steps, b_dim+u_dim);
			VectorXd theta;
			//gPCI->rad->composeBelief(toVectorXd(gPCI->rad->GetDOFValues()), rt_Sigma0, theta);
			VectorXd x;
			const DblVec& xvec = gPCI->rad->GetDOFValues();
			for(int i = 0; i < xvec.size(); ++i) { x[i] = xvec[i]; }

			//cout << x << endl;
			//cout << rt_Sigma0 << endl;
			//cout << theta << endl;

			gPCI->rad->composeBelief(x, rt_Sigma0, theta);
			for (int i=0; i < n_steps; i++) {
				data.block(i,0,1,b_dim) = theta.transpose();
				if (i != (n_steps-1)) {
					VectorXd u = VectorXd::Zero(u_dim);
					data.block(i,b_dim,1,u_dim) = u.transpose();
					theta = gPCI->rad->BeliefDynamics(theta,u);
				} else {
					data.block(i,b_dim,1,u_dim) = VectorXd::Zero(u_dim).transpose();
				}
			}
		}
	}
	else if (type_str == "given_traj") {
		FAIL_IF_FALSE(v.isMember("data"));
		const Value& data_v = v["data"];
		TrajArray x_data;
		Json::fromJson(data_v, x_data);
		if (x_data.rows()!=n_steps) PRINT_AND_THROW("x data has wrong number of rows");

		if (!belief_space) {
			if (x_data.cols()!=n_dof) PRINT_AND_THROW("x data has wrong number of cols");
			data = x_data;
		} else {
			if (x_data.cols() == b_dim+u_dim) {
				data = x_data;
			} else if (x_data.cols() == n_dof) {
				data.resize(n_steps, b_dim+u_dim);
				VectorXd theta;
				gPCI->rad->composeBelief(x_data.row(0).transpose(), rt_Sigma0, theta);
				for (int i=0; i < n_steps; i++) {
					data.block(i,0,1,b_dim) = theta.transpose();
					if (i != (n_steps-1)) {
						VectorXd u = x_data.row(i+1).transpose() - x_data.row(i).transpose();
						data.block(i,b_dim,1,u_dim) = u.transpose();
						theta = gPCI->rad->BeliefDynamics(theta,u);
					} else {
						data.block(i,b_dim,1,u_dim) = VectorXd::Zero(u_dim).transpose();
					}
				}
			} else {
				PRINT_AND_THROW("x data has wrong number of cols");
			}
		}
	}
	else if (type_str == "straight_line") {
		FAIL_IF_FALSE(!belief_space);
		FAIL_IF_FALSE(v.isMember("endpoint"));
		DblVec endpoint;
		childFromJson(v, endpoint, "endpoint");
		if (endpoint.size() != n_dof) {
			PRINT_AND_THROW(boost::format("wrong number of dof values in initialization. expected %i got %j")%n_dof%endpoint.size());
		}
		data = TrajArray(n_steps, n_dof);
		DblVec start = gPCI->rad->GetDOFValues();
		for (int idof = 0; idof < n_dof; ++idof) {
			data.col(idof) = VectorXd::LinSpaced(n_steps, start[idof], endpoint[idof]);
		}
	}
}

void ProblemConstructionInfo::fromJson(const Value& v) {
	childFromJson(v, basic_info, "basic_info");

	RobotBasePtr robot = (basic_info.robot=="") ? GetRobot(*env) : GetRobotByName(*env, basic_info.robot);
	if (!robot) {
		PRINT_AND_THROW("couldn't get robot");
	}
	rad = RADFromName(basic_info.manip, robot);
	if (!rad) {
		PRINT_AND_THROW( boost::format("couldn't get manip %s")%basic_info.manip );
	}

	gPCI = this;
	if (v.isMember("costs")) fromJsonArray(v["costs"], cost_infos);
	if (v.isMember("constraints")) fromJsonArray(v["constraints"], cnt_infos);

	childFromJson(v, init_info, "init_info");
	gPCI = NULL;

}
void CntInfo::RegisterMaker(const std::string& type, MakerFunc f) {
	name2maker[type] = f;
}

TrajOptResult::TrajOptResult(OptResults& opt, TrajOptProb& prob) :
				  cost_vals(opt.cost_vals),
				  cnt_viols(opt.cnt_viols) {
	BOOST_FOREACH(const CostPtr& cost, prob.getCosts()) {
		cost_names.push_back(cost->name());
	}
	BOOST_FOREACH(const ConstraintPtr& cnt, prob.getConstraints()) {
		cnt_names.push_back(cnt->name());
	}
	traj = getTraj(opt.x, prob.GetVars());
}

Vector3d endEffectorPosition(BeliefRobotAndDOFPtr brad, VectorXd dofs) {
	Vector3d eetrans;
	brad->ForwardKinematics(dofs, eetrans);

	double c6 = cos((double)dofs[6]);
	double s6 = sin((double)dofs[6]);

	eetrans[0] = c6*eetrans[0] - s6*eetrans[1];
	eetrans[0] = s6*eetrans[0] + c6*eetrans[1];
	eetrans[2] = eetrans[2] + 0.16;

	return eetrans;
}

void renderSigmaPts(BeliefRobotAndDOFPtr rad, const MatrixXd& sigma_pts, const osg::Vec4f& colorvec, vector<GraphHandlePtr>& handles, OSGViewerPtr viewer) {
	vector<KinBody::LinkPtr> links;
	vector<int> joint_inds;
	rad->GetAffectedLinks(links, true, joint_inds);

	// render sigma points
	//	for (int j=0; j<sigma_pts.cols(); j++) {
	//		rad->SetDOFValues(toDblVec(sigma_pts.col(j)));
	//		handles.push_back(viewer->PlotKinBody(rad->GetRobot()));
	////		if (j==0) SetColor(handles.back(), osg::Vec4f(0,0,1,1));
	////		else //SetColor(handles.back(), osg::Vec4f(0,0,1,1));
	//		SetColor(handles.back(), colorvec);
	//	}

	// render convex hulls of sigma points
	vector<DblVec> dofvals(sigma_pts.cols());
	for (int i=0; i<sigma_pts.cols(); i++)
		dofvals[i] = toDblVec(sigma_pts.col(i));
	//cc->SetContactDistance(100);
	boost::shared_ptr<CollisionChecker> cc = CollisionChecker::GetOrCreate(*rad->GetRobot()->GetEnv());
	cc->PlotCastHull(*rad, links, dofvals, handles);
}

void renderTrajectory(BeliefRobotAndDOFPtr brad, OSGViewerPtr viewer, const TrajArray& traj, vector<GraphHandlePtr>& handles, int sample_rate = 1) {
	//	viewer->GetEnv()->Load("/home/alex/rll/trajopt/data/barrett_sensor_wall.env.xml");
	//	KinBodyPtr wall = viewer->GetEnv()->GetKinBody("wall");
	//	handles.push_back(viewer->PlotKinBody(wall));
	//	SetColor(handles.back(), osg::Vec4f(0.8,0.8,0,1));
	//	SetTransparency(handles.back(), 0.1);
	//	viewer->GetEnv()->Remove(wall);

	int n_skipped = sample_rate;
	int n_dof = brad->GetDOF();
	int b_dim = brad->GetBDim();
	Vector3d last_ee_pos = endEffectorPosition(brad, traj.block(0,0,1,n_dof).transpose());
	for (int i=0; i<traj.rows(); i++) {
		if ((n_skipped < sample_rate) && ((endEffectorPosition(brad, traj.block(i,0,1,n_dof).transpose())-last_ee_pos).norm()<0.1) && (i != (traj.rows()-1))) {
			n_skipped++;
			continue;
		}
		n_skipped = 0;
		last_ee_pos = endEffectorPosition(brad, traj.block(i,0,1,n_dof).transpose());

		brad->SetDOFValues(toDblVec(traj.block(i,0,1,n_dof).transpose()));
		handles.push_back(viewer->PlotKinBody(brad->GetRobot()));
		if (i==0 || i==(traj.rows()-1)) renderSigmaPts(brad, brad->sigmaPoints(traj.block(i,0,1,b_dim).transpose()), osg::Vec4f(0,1,0,0.2), handles, viewer);
		viewer->Idle();
	}

	viewer->Idle();
	viewer->Idle();
	viewer->Idle();
	viewer->Idle();
}

bool isTrajectoryInCollision(CollisionCheckerPtr cc, TrajArray traj, BeliefRobotAndDOFPtr rad) {
	vector<Collision> collisions;
	cc->DiscreteCheckTrajectory(traj, rad, collisions);
	for (int i=0; i<collisions.size(); i++) {
		if (collisions[i].distance < 0) return true;
	}
	return false;
}

double calcTrajectoryCost(BeliefRobotAndDOFPtr brad, const TrajArray& traj) {
	int x_dim = brad->GetXDim(), b_dim = brad->GetBDim(), u_dim = brad->GetUDim();
	VectorXd theta(b_dim);
	VectorXd u(u_dim);
	VectorXd x(x_dim);
	MatrixXd rt_Sigma(x_dim,x_dim);
	MatrixXd Sigma(x_dim,x_dim);
	MatrixXd Q = MatrixXd::Identity(x_dim,x_dim)*10;
	MatrixXd R = MatrixXd::Identity(u_dim,u_dim)*0.1;
	double total_cost = 0;
	for (int i=0; i<traj.rows(); i++) {
		theta = traj.block(i,0,1,b_dim).transpose();
		u = traj.block(i,b_dim,1,u_dim).transpose();
		brad->decomposeBelief(theta, x, rt_Sigma);
		total_cost += ((MatrixXd) (Q*rt_Sigma*rt_Sigma.transpose())).trace() + u.transpose()*R*u;
	}
	return total_cost;
}

VectorXd SimulateAndReplan(const Json::Value& root, OpenRAVE::EnvironmentBasePtr env, bool sigma_pts_scale, bool interactive) {
	ProblemConstructionInfo pci(env);
	pci.fromJson(root);

	TrajArray& traj = pci.init_info.data;
	TrajArray init_traj = traj;
	BeliefRobotAndDOFPtr brad = pci.rad;
	brad->SetSigmaPointsScale(sigma_pts_scale);
	int n_steps = pci.basic_info.n_steps;
	int x_dim = brad->GetXDim(), b_dim = brad->GetBDim(), u_dim = brad->GetUDim(), q_dim = brad->GetQDim(), r_dim = brad->GetRDim();
	TrajArray plan_traj = TrajArray::Zero(n_steps, b_dim+u_dim);
	TrajArray exec_mpc_traj = TrajArray::Zero(n_steps, b_dim+u_dim);
	TrajArray exec_mpc_gt_traj = TrajArray::Zero(n_steps, x_dim);
	TrajArray exec_open_traj = TrajArray::Zero(n_steps, b_dim+u_dim);
	TrajArray exec_open_gt_traj = TrajArray::Zero(n_steps, x_dim);
	//	TrajArray exec_mpc_traj(n_steps, b_dim);
	//	TrajArray plan_traj(n_steps, b_dim+u_dim);
	//	TrajArray exec_open_traj(n_steps, b_dim);
	double mpc_trans_err, open_trans_err;

	VectorXd theta_init = traj.block(0,0,1,b_dim).transpose();
	VectorXd x_init;
	MatrixXd rt_Sigma_init;
	brad->decomposeBelief(theta_init, x_init, rt_Sigma_init);
	x_init = brad->mvnrnd(x_init, rt_Sigma_init * rt_Sigma_init.transpose());
	brad->composeBelief(x_init, rt_Sigma_init, theta_init);

	MatrixXd qq(q_dim, n_steps);
	for (int j=0; j<n_steps; j++) {
		qq.col(j) = brad->VectorXdRand(q_dim);
	}
	MatrixXd rr(r_dim, n_steps);
	for (int j=0; j<n_steps; j++) {
		rr.col(j) = brad->VectorXdRand(r_dim);
	}

	VectorXd x;
	VectorXd x_gt;
	MatrixXd rt_Sigma;
	VectorXd theta;
	Vector3d trans_gt, trans_est;

	struct timeval startTimeStruct;
	gettimeofday(&startTimeStruct, NULL);
	unsigned long int startTime = startTimeStruct.tv_sec*(long unsigned int)(1e6) + startTimeStruct.tv_usec;

	// MPC
	x = x_init;
	x_gt = x_init;
	rt_Sigma = rt_Sigma_init;
	theta = theta_init;
	int i=0;
	do {
		gPCI = &pci;
		if (root.isMember("costs")) fromJsonArray(root["costs"], pci.cost_infos);
		if (root.isMember("constraints")) fromJsonArray(root["constraints"], pci.cnt_infos);
		gPCI = NULL;

		TrajOptProbPtr prob = ConstructProblem(pci);
		TrajOptResultPtr result = OptimizeProblem(prob, interactive);
		if (i==0) plan_traj = result->traj;
		VectorXd u = result->traj.block(0,b_dim,1,u_dim).transpose();
		exec_mpc_traj.block(i,0,1,b_dim) = theta.transpose();
		exec_mpc_traj.block(i,b_dim,1,u_dim) = u.transpose();

		// simulate dynamics and observe
		VectorXd q = qq.col(i);
		VectorXd r = rr.col(i);
		x_gt = brad->Dynamics(x_gt,u,q);
		VectorXd z = brad->Observe(x_gt,r);
		exec_mpc_gt_traj.row(i) = x_gt.transpose();

		brad->ekfUpdate(u, x, rt_Sigma, x, rt_Sigma, true, z);

		DblVec lower, upper;
		brad->GetDOFLimits(lower, upper);
		for (int j=0; j<x.size(); j++) {
			if (x(j) < lower[j]) x(j) = lower[j] + STEP;
			if (x(j) > upper[j]) x(j) = upper[j] - STEP;
		}

		brad->composeBelief(x, rt_Sigma, theta);

		brad->SetDOFValues(toDblVec(x));
		cout << "setting robot with dofs values " << endl;
		cout << x.transpose() << endl;
		cout << "actual robot dofs values are " << endl;
		cout << toVectorXd(brad->GetDOFValues()).transpose() << endl;
		traj = result->traj.bottomRows(result->traj.rows()-1);
		traj.block(0,0,1,b_dim) = theta.transpose();
		pci.basic_info.n_steps--;

		cout << "-------------------------------------------" << endl;

		i++;
	} while (pci.basic_info.n_steps > 1);
	exec_mpc_traj.block(n_steps-1,0,1,b_dim) = theta.transpose();
	exec_mpc_traj.block(n_steps-1,b_dim,1,u_dim) = VectorXd::Zero(u_dim).transpose();
	exec_mpc_gt_traj.row(n_steps-1) = x_gt.transpose();
	mpc_trans_err = (endEffectorPosition(brad, x) - endEffectorPosition(brad, x_gt)).norm();

	gettimeofday(&startTimeStruct, NULL);
	unsigned long int curTime = startTimeStruct.tv_sec*(long unsigned int)(1e6) + startTimeStruct.tv_usec;

	cout << "Total MPC time (s): " << (1e-6) * (curTime - startTime) << endl;

	// non MPC
	x = x_init;
	x_gt = x_init;
	rt_Sigma = rt_Sigma_init;
	theta = theta_init;
	for (int i=0; i<n_steps-1; i++) {
		VectorXd u = plan_traj.block(i,b_dim,1,u_dim).transpose();

		exec_open_traj.block(i,0,1,b_dim) = theta.transpose();
		exec_open_traj.block(i,b_dim,1,u_dim) = u.transpose();

		VectorXd q = qq.col(i);
		VectorXd r = rr.col(i);
		x_gt = brad->Dynamics(x_gt,u,q);
		VectorXd z = brad->Observe(x_gt,r);
		exec_open_gt_traj.row(i) = x_gt.transpose();

		brad->ekfUpdate(u, x, rt_Sigma, x, rt_Sigma, true, z);
		brad->composeBelief(x, rt_Sigma, theta);
	}
	exec_open_traj.block(n_steps-1,0,1,b_dim) = theta.transpose();
	exec_open_traj.block(n_steps-1,b_dim,1,u_dim) = VectorXd::Zero(u_dim).transpose();
	exec_open_gt_traj.row(n_steps-1) = x_gt.transpose();
	open_trans_err = (endEffectorPosition(brad, x) - endEffectorPosition(brad, x_gt)).norm();

	cout << "-------------------------------------------" << endl;
	cout << "cost exec_mpc_traj " << calcTrajectoryCost(brad, exec_mpc_traj) << "\t" << mpc_trans_err << endl;
	cout << "cost exec_open_traj " << calcTrajectoryCost(brad, exec_open_traj) << "\t" << open_trans_err << endl;
	cout << "cost plan_traj " << calcTrajectoryCost(brad, plan_traj) << endl;

	boost::shared_ptr<CollisionChecker> cc = CollisionChecker::GetOrCreate(*brad->GetRobot()->GetEnv());
	cout << isTrajectoryInCollision(cc, exec_mpc_gt_traj.leftCols(x_dim), brad) << endl;
	cout << isTrajectoryInCollision(cc, exec_open_gt_traj.leftCols(x_dim), brad) << endl;
	cout << isTrajectoryInCollision(cc, plan_traj.leftCols(x_dim), brad) << endl;

	VectorXd stats(8);
	stats << calcTrajectoryCost(brad, exec_mpc_traj), calcTrajectoryCost(brad, exec_open_traj), calcTrajectoryCost(brad, plan_traj), mpc_trans_err, open_trans_err,
			isTrajectoryInCollision(cc, exec_mpc_gt_traj.leftCols(x_dim), brad), isTrajectoryInCollision(cc, exec_open_gt_traj.leftCols(x_dim), brad),
			isTrajectoryInCollision(cc, plan_traj.leftCols(x_dim), brad);

	OSGViewerPtr viewer = OSGViewer::GetOrCreate(brad->GetRobot()->GetEnv());
	vector<GraphHandlePtr> handles;

	bool inter = false;
	if (inter) {
		bool gen_figs = false;
		while (true) {
			cout << "executed MPC" << endl;
			if (gen_figs) dynamic_cast<OSGViewer::EventHandler*>(viewer->GetViewer().getCameraManipulator())->setTransformation(osg::Vec3d(0,0,4), osg::Vec3d(0,0,0), osg::Vec3d(0,1,0));
			renderTrajectory(brad, viewer, exec_mpc_traj, handles, 3);
			if (gen_figs) {
				viewer->Idle();
				dynamic_cast<OSGViewer::EventHandler*>(viewer->GetViewer().getCameraManipulator())->setTransformation(osg::Vec3d(4,0,4), osg::Vec3d(0,0,0), osg::Vec3d(0,0,1));
				viewer->Idle();
			}
			handles.clear();

			cout << "executed open loop" << endl;
			if (gen_figs) dynamic_cast<OSGViewer::EventHandler*>(viewer->GetViewer().getCameraManipulator())->setTransformation(osg::Vec3d(0,0,4), osg::Vec3d(0,0,0), osg::Vec3d(0,1,0));
			renderTrajectory(brad, viewer, exec_open_traj, handles, 3);
			if (gen_figs) {
				viewer->Idle();
				dynamic_cast<OSGViewer::EventHandler*>(viewer->GetViewer().getCameraManipulator())->setTransformation(osg::Vec3d(4,0,4), osg::Vec3d(0,0,0), osg::Vec3d(0,0,1));
				viewer->Idle();
			}
			handles.clear();

			cout << "planned" << endl;
			if (gen_figs) dynamic_cast<OSGViewer::EventHandler*>(viewer->GetViewer().getCameraManipulator())->setTransformation(osg::Vec3d(0,0,4), osg::Vec3d(0,0,0), osg::Vec3d(0,1,0));
			renderTrajectory(brad, viewer, plan_traj, handles, 3);
			if (gen_figs) {
				viewer->Idle();
				dynamic_cast<OSGViewer::EventHandler*>(viewer->GetViewer().getCameraManipulator())->setTransformation(osg::Vec3d(4,0,4), osg::Vec3d(0,0,0), osg::Vec3d(0,0,1));
				viewer->Idle();
			}
			handles.clear();

			cout << "RRT" << endl;
			if (gen_figs) dynamic_cast<OSGViewer::EventHandler*>(viewer->GetViewer().getCameraManipulator())->setTransformation(osg::Vec3d(0,0,4), osg::Vec3d(0,0,0), osg::Vec3d(0,1,0));
			renderTrajectory(brad, viewer, init_traj, handles, 3);
			if (gen_figs) {
				viewer->Idle();
				dynamic_cast<OSGViewer::EventHandler*>(viewer->GetViewer().getCameraManipulator())->setTransformation(osg::Vec3d(4,0,4), osg::Vec3d(0,0,0), osg::Vec3d(0,0,1));
				viewer->Idle();
			}
			handles.clear();
		}
	}
	return stats;
}

TrajOptResultPtr OptimizeProblem(TrajOptProbPtr prob, bool plot) {
	RobotBase::RobotStateSaver saver = prob->GetRAD()->Save();
	BasicTrustRegionSQP opt(prob);
	opt.max_iter_ = 100;
	opt.min_approx_improve_frac_ = .001;
	opt.merit_error_coeff_ = 20;
	opt.max_merit_coeff_increases_ = 10;

	if (plot) opt.addCallback(PlotCallback(*prob));
	//  opt.addCallback(boost::bind(&PlotCosts, boost::ref(prob->getCosts()),boost::ref(*prob->GetRAD()), boost::ref(prob->GetVars()), _1));
	opt.initialize(trajToDblVec(prob->GetInitTraj()));

	struct timeval startTimeStruct;
	gettimeofday(&startTimeStruct, NULL);
	unsigned long int startTime = startTimeStruct.tv_sec*(long unsigned int)(1e6) + startTimeStruct.tv_usec;

	opt.optimize();

	gettimeofday(&startTimeStruct, NULL);
	unsigned long int curTime = startTimeStruct.tv_sec*(long unsigned int)(1e6) + startTimeStruct.tv_usec;

	cout << "Total optimization time (s): " << (1e-6) * (curTime - startTime) << endl;
	TrajOptResultPtr result(new TrajOptResult(opt.results(), *prob));

	return result;
}

TrajOptProbPtr ConstructProblem(const ProblemConstructionInfo& pci) {
	TrajOptProbPtr prob(new TrajOptProb());
	const BasicInfo& bi = pci.basic_info;
	prob->belief_space = bi.belief_space;
	int n_steps = bi.n_steps;

	prob->m_rad = pci.rad;
	int n_dof = prob->m_rad->GetDOF();
	int x_dim = prob->m_rad->GetXDim(), b_dim = prob->m_rad->GetBDim(), u_dim = prob->m_rad->GetUDim();

	DblVec lower, upper;
	prob->m_rad->GetDOFLimits(lower, upper);
	vector<double> vlower, vupper;
	vector<string> names;
	for (int i=0; i < n_steps; ++i) {
		vlower.insert(vlower.end(), lower.data(), lower.data()+lower.size());
		vupper.insert(vupper.end(), upper.data(), upper.data()+upper.size());
		for (unsigned j=0; j < n_dof; ++j) {
			names.push_back( (boost::format("j_%i_%i")%i%j).str() );
		}
		if (bi.belief_space) {
			for (unsigned jj=0; jj< n_dof; ++jj) {
				for (unsigned ii=jj; ii < n_dof; ++ii) {
					names.push_back( (boost::format("cov_%i_%i_%i")%i%ii%jj).str() );
					vlower.push_back(-INFINITY);
					vupper.push_back(INFINITY);
				}
			}
			for (unsigned j=0; j < n_dof; ++j) {
				names.push_back( (boost::format("u_%i_%i")%i%j).str() );
				vlower.push_back(-INFINITY);
				vupper.push_back(INFINITY);
			}
		}
	}
	prob->createVariables(names, vlower, vupper);

	if (bi.belief_space)
		prob->m_traj_vars = VarArray(n_steps, b_dim + u_dim, prob->vars_.data());
	else
		prob->m_traj_vars = VarArray(n_steps, n_dof, prob->vars_.data());

	DblVec cur_dofvals = prob->m_rad->GetDOFValues();

	if (bi.start_fixed) {
		if (pci.init_info.data.rows() > 0 && !allClose(toVectorXd(cur_dofvals), pci.init_info.data.block(0,0,1,n_dof).transpose())) {
			cout << toVectorXd(cur_dofvals).transpose() << endl;
			cout << pci.init_info.data.block(0,0,1,n_dof) << endl;
			LOG_WARN("robot dof values don't match initialization. I don't know what you want me to use for the dof values");
		}
		int n_fixed_terms = n_dof;
		if (bi.belief_space) n_fixed_terms = b_dim;
		for (int j=0; j < n_fixed_terms; ++j) {
			prob->addLinearConstr(exprSub(AffExpr(prob->m_traj_vars(0,j)), pci.init_info.data(0,j)), EQ);
		}
	}

	if (!bi.dofs_fixed.empty()) {
		BOOST_FOREACH(const int& dof_ind, bi.dofs_fixed) {
			for (int i=1; i < prob->GetNumSteps(); ++i) {
				prob->addLinearConstr(exprSub(AffExpr(prob->m_traj_vars(i,dof_ind)), AffExpr(prob->m_traj_vars(0,dof_ind))), EQ);
			}
		}
	}

	BOOST_FOREACH(const CostInfoPtr& ci, pci.cost_infos) {
		ci->hatch(*prob);
	}
	BOOST_FOREACH(const CntInfoPtr& ci, pci.cnt_infos) {
		ci->hatch(*prob);
	}

	if (bi.belief_space) {
		for (int i=0; i < n_steps-1; i++) {
			VarVector theta0_vars = prob->m_traj_vars.block(0,0,n_steps,b_dim).row(i);
			VarVector theta1_vars = prob->m_traj_vars.block(0,0,n_steps,b_dim).row(i+1);
			VarVector u_vars = prob->m_traj_vars.block(0,b_dim,n_steps,u_dim).row(i);
			prob->addConstr(ConstraintPtr(new BeliefDynamicsConstraint(theta0_vars, theta1_vars, u_vars, prob->GetRAD())));
		}
	}

	prob->SetInitTraj(pci.init_info.data);

	return prob;

}
TrajOptProbPtr ConstructProblem(const Json::Value& root, OpenRAVE::EnvironmentBasePtr env) {
	ProblemConstructionInfo pci(env);
	pci.fromJson(root);
	return ConstructProblem(pci);
}

TrajOptProb::TrajOptProb(int n_steps, BeliefRobotAndDOFPtr rad) : m_rad(rad) {
	DblVec lower, upper;
	m_rad->GetDOFLimits(lower, upper);
	int n_dof = m_rad->GetDOF();
	vector<double> vlower, vupper;
	vector<string> names;
	for (int i=0; i < n_steps; ++i) {
		vlower.insert(vlower.end(), lower.data(), lower.data()+lower.size());
		vupper.insert(vupper.end(), upper.data(), upper.data()+upper.size());
		for (unsigned j=0; j < n_dof; ++j) {
			names.push_back( (boost::format("j_%i_%i")%i%j).str() );
		}
	}
	createVariables(names, vlower, vupper);
	m_traj_vars = VarArray(n_steps, n_dof, getVars().data());

}

TrajOptProb::TrajOptProb() {
}

void PoseCostInfo::fromJson(const Value& v) {
	FAIL_IF_FALSE(v.isMember("params"));
	const Value& params = v["params"];
	childFromJson(params, timestep, "timestep", gPCI->basic_info.n_steps-1);
	childFromJson(params, xyz,"xyz");
	childFromJson(params, wxyz,"wxyz");
	childFromJson(params, pos_coeffs,"pos_coeffs", (Vector3d)Vector3d::Ones());
	childFromJson(params, rot_coeffs,"rot_coeffs", (Vector3d)Vector3d::Ones());

	string linkstr;
	childFromJson(params, linkstr, "link");
	link = gPCI->rad->GetRobot()->GetLink(linkstr);
	if (!link) {
		PRINT_AND_THROW(boost::format("invalid link name: %s")%linkstr);
	}
}
CostInfoPtr PoseCostInfo::create() {
	return CostInfoPtr(new PoseCostInfo());
}
void PoseCostInfo::hatch(TrajOptProb& prob) {
	prob.addCost(CostPtr(new CartPoseCost(prob.GetVarRow(timestep), toRaveTransform(wxyz, xyz), rot_coeffs, pos_coeffs, prob.GetRAD(), link)));
	prob.getCosts().back()->setName(name);
}

void PoseCntInfo::fromJson(const Value& v) {
	FAIL_IF_FALSE(v.isMember("params"));
	const Value& params = v["params"];
	childFromJson(params, timestep, "timestep", gPCI->basic_info.n_steps-1);
	childFromJson(params, xyz,"xyz");
	childFromJson(params, wxyz,"wxyz");
	childFromJson(params, pos_coeffs,"pos_coeffs", (Vector3d)Vector3d::Ones());
	childFromJson(params, rot_coeffs,"rot_coeffs", (Vector3d)Vector3d::Ones());

	string linkstr;
	childFromJson(params, linkstr, "link");
	link = gPCI->rad->GetRobot()->GetLink(linkstr);
	if (!link) {
		PRINT_AND_THROW(boost::format("invalid link name: %s")%linkstr);
	}
}

void JointPosCostInfo::fromJson(const Value& v) {
	FAIL_IF_FALSE(v.isMember("params"));
	int n_steps = gPCI->basic_info.n_steps;
	const Value& params = v["params"];
	childFromJson(params, vals, "vals");
	childFromJson(params, coeffs, "coeffs");
	if (coeffs.size() == 1) coeffs = DblVec(n_steps, coeffs[0]);

	int n_dof = gPCI->rad->GetDOF();
	if (vals.size() != n_dof) {
		PRINT_AND_THROW( boost::format("wrong number of dof vals. expected %i got %i")%n_dof%vals.size());
	}
	childFromJson(params, timestep, "timestep", gPCI->basic_info.n_steps-1);
}
void JointPosCostInfo::hatch(TrajOptProb& prob) {
	prob.addCost(CostPtr(new JointPosCost(prob.GetVarRow(timestep), toVectorXd(vals), toVectorXd(coeffs))));
	prob.getCosts().back()->setName(name);
}
CostInfoPtr JointPosCostInfo::create() {
	return CostInfoPtr(new JointPosCostInfo());
}

CntInfoPtr PoseCntInfo::create() {
	return CntInfoPtr(new PoseCntInfo());
}
void PoseCntInfo::hatch(TrajOptProb& prob) {
	VectorXd coeffs(6); coeffs << rot_coeffs, pos_coeffs;
	prob.addConstr(ConstraintPtr(new CartPoseConstraint(prob.GetVarRow(timestep), toRaveTransform(wxyz, xyz), prob.GetRAD(), link, coeffs)));
	prob.getEqConstraints().back()->setName(name);
}

void CartVelCntInfo::fromJson(const Value& v) {
	FAIL_IF_FALSE(v.isMember("params"));
	const Value& params = v["params"];
	childFromJson(params, first_step, "first_step");
	childFromJson(params, last_step, "last_step");
	childFromJson(params, distance_limit,"distance_limit");

	FAIL_IF_FALSE((first_step >= 0) && (first_step <= gPCI->basic_info.n_steps-1) && (first_step < last_step));
	FAIL_IF_FALSE((last_step > 0) && (last_step <= gPCI->basic_info.n_steps-1));

	string linkstr;
	childFromJson(params, linkstr, "link");
	link = gPCI->rad->GetRobot()->GetLink(linkstr);
	if (!link) {
		PRINT_AND_THROW( boost::format("invalid link name: %s")%linkstr);
	}
}
CntInfoPtr CartVelCntInfo::create() {
	return CntInfoPtr(new CartVelCntInfo());
}
void CartVelCntInfo::hatch(TrajOptProb& prob) {
	for (int iStep = first_step; iStep < last_step; ++iStep) {
		prob.addConstr(ConstraintPtr(new CartVelConstraint(prob.GetVarRow(iStep), prob.GetVarRow(iStep+1), prob.GetRAD(), link, distance_limit)));
		prob.getIneqConstraints().back()->setName(name);
	}
}

void JointVelCostInfo::fromJson(const Value& v) {
	FAIL_IF_FALSE(v.isMember("params"));
	const Value& params = v["params"];

	childFromJson(params, coeffs,"coeffs");
	int n_dof = gPCI->rad->GetDOF();
	if (coeffs.size() == 1) coeffs = DblVec(n_dof, coeffs[0]);
	else if (coeffs.size() != n_dof) {
		PRINT_AND_THROW( boost::format("wrong number of coeffs. expected %i got %i")%n_dof%coeffs.size());
	}
}
CostInfoPtr JointVelCostInfo::create() {
	return CostInfoPtr(new JointVelCostInfo());
}
void JointVelCostInfo::hatch(TrajOptProb& prob) {
	// belief-alex take the submatrix because we want joint-vel only on the joint variables
	prob.addCost(CostPtr(new JointVelCost(prob.GetVars().block(0,0,prob.GetVars().m_nRow, prob.GetRAD()->GetDOF()), toVectorXd(coeffs))));
	prob.getCosts().back()->setName(name);
}

void CollisionCostInfo::fromJson(const Value& v) {
	FAIL_IF_FALSE(v.isMember("params"));
	const Value& params = v["params"];

	int n_steps = gPCI->basic_info.n_steps;
	childFromJson(params, coeffs,"coeffs");
	if (coeffs.size() == 1) coeffs = DblVec(n_steps, coeffs[0]);
	else if (coeffs.size() != n_steps) {
		PRINT_AND_THROW( boost::format("wrong size: coeffs. expected %i got %i")%n_steps%coeffs.size() );
	}
	childFromJson(params, dist_pen,"dist_pen");
	if (dist_pen.size() == 1) dist_pen = DblVec(n_steps, dist_pen[0]);
	else if (dist_pen.size() != n_steps) {
		PRINT_AND_THROW( boost::format("wrong size: dist_pen. expected %i got %i")%n_steps%dist_pen.size() );
	}
	childFromJson(params, belief_space,"belief_space", false);
}
void CollisionCostInfo::hatch(TrajOptProb& prob) {
	for (int i=0; i < prob.GetNumSteps(); ++i) {
		if (belief_space)
			prob.addCost(CostPtr(new CollisionCost(dist_pen[i], coeffs[i], prob.GetRAD(), prob.GetVars().rblock(i,0,prob.GetRAD()->GetBDim()))));
		else
			prob.addCost(CostPtr(new CollisionCost(dist_pen[i], coeffs[i], prob.GetRAD(), prob.GetVars().rblock(i,0,prob.GetRAD()->GetDOF()))));
		prob.getCosts().back()->setName( (boost::format("%s_%i")%name%i).str() );
	}
	CollisionCheckerPtr cc = CollisionChecker::GetOrCreate(*prob.GetEnv());
	cc->SetContactDistance(*std::max_element(dist_pen.begin(), dist_pen.end()) + .04);
}
CostInfoPtr CollisionCostInfo::create() {
	return CostInfoPtr(new CollisionCostInfo());
}

void ContinuousCollisionCostInfo::fromJson(const Value& v) {
	FAIL_IF_FALSE(v.isMember("params"));
	const Value& params = v["params"];

	int n_steps = gPCI->basic_info.n_steps;
	childFromJson(params, first_step, "first_step", 0);
	childFromJson(params, last_step, "last_step", n_steps-1);
	childFromJson(params, coeffs, "coeffs");
	int n_terms = last_step - first_step;
	if (coeffs.size() == 1) coeffs = DblVec(n_terms, coeffs[0]);
	else if (coeffs.size() != n_terms) {
		PRINT_AND_THROW (boost::format("wrong size: coeffs. expected %i got %i")%n_terms%coeffs.size());
	}
	childFromJson(params, dist_pen,"dist_pen");
	if (dist_pen.size() == 1) dist_pen = DblVec(n_terms, dist_pen[0]);
	else if (dist_pen.size() != n_terms) {
		PRINT_AND_THROW(boost::format("wrong size: dist_pen. expected %i got %i")%n_terms%dist_pen.size());
	}
}
void ContinuousCollisionCostInfo::hatch(TrajOptProb& prob) {
	for (int i=first_step; i < last_step; ++i) {
		prob.addCost(CostPtr(new CollisionCost(dist_pen[i], coeffs[i], prob.GetRAD(), prob.GetVars().rblock(i,0,prob.GetRAD()->GetDOF()), prob.GetVars().rblock(i+1,0,prob.GetRAD()->GetDOF()))));
		prob.getCosts().back()->setName( (boost::format("%s_%i")%name%i).str() );
	}
	CollisionCheckerPtr cc = CollisionChecker::GetOrCreate(*prob.GetEnv());
	cc->SetContactDistance(*std::max_element(dist_pen.begin(), dist_pen.end()) + .04);
}
CostInfoPtr ContinuousCollisionCostInfo::create() {
	return CostInfoPtr(new ContinuousCollisionCostInfo());
}

void CollisionCntInfo::fromJson(const Value& v) {
	FAIL_IF_FALSE(v.isMember("params"));
	const Value& params = v["params"];

	int n_steps = gPCI->basic_info.n_steps;
	childFromJson(params, coeffs,"coeffs");
	if (coeffs.size() == 1) coeffs = DblVec(n_steps, coeffs[0]);
	else if (coeffs.size() != n_steps) {
		PRINT_AND_THROW( boost::format("wrong size: coeffs. expected %i got %i")%n_steps%coeffs.size() );
	}
	childFromJson(params, dist_pen,"dist_pen");
	if (dist_pen.size() == 1) dist_pen = DblVec(n_steps, dist_pen[0]);
	else if (dist_pen.size() != n_steps) {
		PRINT_AND_THROW( boost::format("wrong size: dist_pen. expected %i got %i")%n_steps%dist_pen.size() );
	}
	childFromJson(params, belief_space,"belief_space", false);
}
void CollisionCntInfo::hatch(TrajOptProb& prob) {
	for (int i=0; i < prob.GetNumSteps(); ++i) {
		if (belief_space)
			prob.addConstr(ConstraintPtr(new CollisionConstraint(dist_pen[i], coeffs[i], prob.GetRAD(), prob.GetVars().rblock(i,0,prob.GetRAD()->GetBDim()))));
		else
			prob.addConstr(ConstraintPtr(new CollisionConstraint(dist_pen[i], coeffs[i], prob.GetRAD(), prob.GetVars().rblock(i,0,prob.GetRAD()->GetDOF()))));
		prob.getConstraints().back()->setName( (boost::format("%s_%i")%name%i).str() );
	}
	CollisionCheckerPtr cc = CollisionChecker::GetOrCreate(*prob.GetEnv());
	cc->SetContactDistance(*std::max_element(dist_pen.begin(), dist_pen.end()) + .04);
}
CntInfoPtr CollisionCntInfo::create() {
	return CntInfoPtr(new CollisionCntInfo());
}

void ContinuousCollisionCntInfo::fromJson(const Value& v) {
	FAIL_IF_FALSE(v.isMember("params"));
	const Value& params = v["params"];

	int n_steps = gPCI->basic_info.n_steps;
	childFromJson(params, first_step, "first_step", 0);
	childFromJson(params, last_step, "last_step", n_steps-1);
	childFromJson(params, coeffs, "coeffs");
	int n_terms = last_step - first_step;
	if (coeffs.size() == 1) coeffs = DblVec(n_terms, coeffs[0]);
	else if (coeffs.size() != n_terms) {
		PRINT_AND_THROW (boost::format("wrong size: coeffs. expected %i got %i")%n_terms%coeffs.size());
	}
	childFromJson(params, dist_pen,"dist_pen");
	if (dist_pen.size() == 1) dist_pen = DblVec(n_terms, dist_pen[0]);
	else if (dist_pen.size() != n_terms) {
		PRINT_AND_THROW(boost::format("wrong size: dist_pen. expected %i got %i")%n_terms%dist_pen.size());
	}
}
void ContinuousCollisionCntInfo::hatch(TrajOptProb& prob) {
	for (int i=first_step; i < last_step; ++i) {
		prob.addConstr(ConstraintPtr(new CollisionConstraint(dist_pen[i], coeffs[i], prob.GetRAD(), prob.GetVars().rblock(i,0,prob.GetRAD()->GetDOF()), prob.GetVars().rblock(i+1,0,prob.GetRAD()->GetDOF()))));
		prob.getConstraints().back()->setName( (boost::format("%s_%i")%name%i).str() );
	}
	CollisionCheckerPtr cc = CollisionChecker::GetOrCreate(*prob.GetEnv());
	cc->SetContactDistance(*std::max_element(dist_pen.begin(), dist_pen.end()) + .04);
}
CntInfoPtr ContinuousCollisionCntInfo::create() {
	return CntInfoPtr(new ContinuousCollisionCntInfo());
}

void JointConstraintInfo::fromJson(const Value& v) {
	FAIL_IF_FALSE(v.isMember("params"));
	const Value& params = v["params"];
	childFromJson(params, vals, "vals");

	int n_dof = gPCI->rad->GetDOF();
	if (vals.size() != n_dof) {
		PRINT_AND_THROW( boost::format("wrong number of dof vals. expected %i got %i")%n_dof%vals.size());
	}
	childFromJson(params, timestep, "timestep", gPCI->basic_info.n_steps-1);
}

void JointConstraintInfo::hatch(TrajOptProb& prob) {
	VarVector vars = prob.GetVarRow(timestep);
	int n_dof = vars.size();
	for (int j=0; j < n_dof; ++j) {
		prob.addLinearConstr(exprSub(AffExpr(vars[j]), vals[j]), EQ);
	}
}
CntInfoPtr JointConstraintInfo::create() {
	return CntInfoPtr(new JointConstraintInfo());
}

void ControlCostInfo::fromJson(const Value& v) {
	belief_space = gPCI->basic_info.belief_space;
	if (!belief_space) {
		LOG_WARN("control cost can only be used in belief space. ignoring.");
		return;
	}
	FAIL_IF_FALSE(v.isMember("params"));
	const Value& params = v["params"];

	childFromJson(params, coeffs,"coeffs");
	int n_dof = gPCI->rad->GetDOF();
	if (coeffs.size() == 1) coeffs = DblVec(n_dof, coeffs[0]);
	else if (coeffs.size() != n_dof) {
		PRINT_AND_THROW( boost::format("wrong number of coeffs. expected %i got %i")%n_dof%coeffs.size());
	}
}
CostInfoPtr ControlCostInfo::create() {
	return CostInfoPtr(new ControlCostInfo());
}
void ControlCostInfo::hatch(TrajOptProb& prob) {
	if (!belief_space) return;
	prob.addCost(CostPtr(new ControlCost(prob.GetVars().block(0,prob.GetRAD()->GetBDim(),prob.GetVars().m_nRow-1, prob.GetRAD()->GetDOF()), toVectorXd(coeffs))));
}

void ControlCntInfo::fromJson(const Value& v) {
	belief_space = gPCI->basic_info.belief_space;
	if (!belief_space) {
		LOG_WARN("control constraint can only be used in belief space. ignoring.");
		return;
	}
	FAIL_IF_FALSE(v.isMember("params"));
	const Value& params = v["params"];
	childFromJson(params, u_min, "u_min");
	childFromJson(params, u_max, "u_max");
}
CntInfoPtr ControlCntInfo::create() {
	return CntInfoPtr(new ControlCntInfo());
}
void ControlCntInfo::hatch(TrajOptProb& prob) {
	if (!belief_space) return;
	int b_dim = prob.GetRAD()->GetBDim(), u_dim = prob.GetRAD()->GetUDim();
	for (int i=0; i < prob.GetVars().m_nRow-1; i++) {
		for (int j=b_dim; j < b_dim+u_dim; j++) {
			prob.addLinearConstr(exprSub(AffExpr(prob.GetVars()(i,j)), u_max), INEQ);
			prob.addLinearConstr(exprMult(exprSub(AffExpr(prob.GetVars()(i,j)), u_min), -1), INEQ);
		}
	}
}

void CovarianceCostInfo::fromJson(const Value& v) {
	belief_space = gPCI->basic_info.belief_space;
	if (!belief_space) {
		LOG_WARN("covariance cost can only be used in belief space. ignoring.");
		return;
	}
	FAIL_IF_FALSE(v.isMember("params"));
	const Value& params = v["params"];

	FAIL_IF_FALSE(params.isMember("Q"));
	const Value& Q_array = params["Q"];
	Json::fromJson(Q_array, Q);

	int n_dof = gPCI->rad->GetDOF();
	if (Q.rows()!=n_dof) PRINT_AND_THROW("cost matrix for the covariance has wrong number of rows");
	if (Q.cols()!=n_dof) PRINT_AND_THROW("cost matrix for the covariance has wrong number of cols");
}
CostInfoPtr CovarianceCostInfo::create() {
	return CostInfoPtr(new CovarianceCostInfo());
}
void CovarianceCostInfo::hatch(TrajOptProb& prob) {
	if (!belief_space) return;
	int n_steps = prob.GetVars().m_nRow;
	int b_dim = prob.GetRAD()->GetBDim(), u_dim = prob.GetRAD()->GetUDim();
	for (int i=0; i < n_steps; ++i) {
		VarVector rtSigma_vars = prob.GetVars().rblock(i,u_dim,b_dim-u_dim);
		prob.addCost(CostPtr(new CovarianceCost(rtSigma_vars, Q, prob.GetRAD())));
	}
}

}


// Copyright 2015-2022 The ALMA Project Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.

#pragma once

/// @file
///
/// Code used to build full BZ collision operators as well as the functions and
/// classes used to take and store the exponential of such operator in an
/// efficient way

#include <vector>
#include <tuple>
#include <utilities.hpp>
#include <bulk_hdf5.hpp>
#include <boost/mpi.hpp>
#include <isotopic_scattering.hpp>
#include <Eigen/Dense>
#include <unordered_map>
#if BOOST_VERSION >= 106700
#include <boost/container_hash/hash.hpp>
#else
#include <boost/functional/hash.hpp>
#endif
/// Specialization of std space used for collision_operator.cpp
namespace std {
/// Creating a std template specialization of less functional
template <> struct less<std::array<std::size_t, 3>> {
    bool operator()(const std::array<std::size_t, 3>& a,
                    const std::array<std::size_t, 3>& b) const {
        return std::tuple<std::size_t, std::size_t, std::size_t>(
                   {a[0], a[1], a[2]}) <
               std::tuple<std::size_t, std::size_t, std::size_t>(
                   {b[0], b[1], b[2]});
    }
};
/// Trivial implementation of std::hash for pairs,
/// required to create an unordered_map with pair as key.
template <typename T> struct hash<std::pair<T, T>> {
    std::size_t operator()(const std::pair<T, T>& key) const {
        hash<T> backend;
        std::size_t nruter = 0;
        boost::hash_combine(nruter, backend(key.first));
        boost::hash_combine(nruter, backend(key.second));
        return nruter;
    }
};
/// Trivial implementation of std::hash for arrays,
/// required to create an unordered_set of arrays.
template <typename T, std::size_t S> struct hash<std::array<T, S>> {
    std::size_t operator()(const array<T, S>& key) const {
        hash<T> backend;
        std::size_t nruter = 0;

        for (auto& e : key)
            boost::hash_combine(nruter, backend(e));
        return nruter;
    }
};
} // namespace std

/// This boost specialization of serialization is required to
/// perform MPI operations of the matrix stored quantities
namespace boost {
namespace serialization {
/// This one is for three phonon elements
template <typename Archive, typename T1, typename T2, typename T3>
void serialize(Archive& ar,
               std::tuple<T1, T2, T3>& t,
               const unsigned int version) {
    ar& std::get<0>(t);
    ar& std::get<1>(t);
    ar& std::get<2>(t);
}

template <typename Archive>
void serialize(Archive& ar,
               std::pair<std::size_t, double>& t,
               const unsigned int version) {
    ar& std::get<0>(t);
    ar& std::get<1>(t);
}

/// Eigen Matrix serialization:
template <class Archive>
void serialize(Archive& ar,
               Eigen::Matrix<double, -1, 1>& t,
               const unsigned int version) {
    Eigen::MatrixXd::Index rows = t.rows();
    Eigen::MatrixXd::Index cols = t.cols();

    ar& rows;
    ar& cols;
    // Because our matrix is dynamic we need to ensure resizing
    if (rows * cols != t.size())
        t.resize(rows, cols);

    ar& boost::serialization::make_array(t.data(), rows * cols);
}

template <class Archive>
void serialize(Archive& ar,
               Eigen::Matrix<double, -1, -1>& t,
               const unsigned int version) {
    Eigen::MatrixXd::Index rows = t.rows();
    Eigen::MatrixXd::Index cols = t.cols();

    ar& rows;
    ar& cols;
    // Because our matrix is dynamic we need to ensure resizing
    if (rows * cols != t.size())
        t.resize(rows, cols);

    ar& boost::serialization::make_array(t.data(), rows * cols);
}

} // namespace serialization
} // namespace boost

namespace alma {

/// Chunck vectors and scatter them
///@param[in] v2c - vector to chunck and scatter
///@param[in] world - MPI communicator
template <class T>
void chunck_and_scatter(std::vector<T>& v2c, boost::mpi::communicator world) {
    broadcast(world, v2c, 0);
    auto ranges = parrange(0, v2c.size(), world);
    std::vector<T> vproc(0);
    vproc.insert(
        vproc.begin(), v2c.begin() + ranges.first, v2c.begin() + ranges.second);
    std::swap(v2c, vproc);
}

/// Gather data from procs to given rank
///@param[in] v2c - vector to gather from each proc, the gather data is then
/// inserted in the selected process
///@param[in] world - MPI communicator
///@param[in] where - the id-process in which the vector is gathered
template <class T>
void gather_from_procs(std::vector<T>& v2c,
                       boost::mpi::communicator world,
                       std::size_t where) {
    std::vector<std::vector<T>> chunckedv;

    gather(world, v2c, chunckedv, where);
    v2c.clear();
    if ((std::size_t)world.rank() == where) {
        for (auto& chunk : chunckedv)
            std::move(chunk.begin(), chunk.end(), std::back_inserter(v2c));
        v2c.shrink_to_fit();
    }
    else {
        v2c.shrink_to_fit();
    }
}

/// It gives the id-process for a given index
/// @param[in] index - index to look for
/// @param[in] container_size - the size of the container
/// @param[in] world_size - the number of processes in that MPI communicator
inline std::size_t index_proc(std::size_t index,
                              std::size_t container_size,
                              std::size_t world_size) {
    for (std::size_t ip = 0; ip < world_size; ip++) {
        auto limits = alma::my_jobs(container_size, world_size, ip);
        if (index >= limits[0] and index < limits[1])
            return ip;
    }
    return world_size;
}

/// Defining alias for long tuples used to store the full BZ elements to build
/// collision operator
using matel3part = std::tuple<std::size_t, std::size_t, double>;
using matel2part = std::pair<std::size_t, double>;
// template <class T> using std2dvec_<T> = std::vector<std::vector<T>>;

/// This is to build B matrix as defined in
/// [http://dx.doi.org/10.1063/1.4898090] with the inclusion of isotopic
/// scattering terms
///@param[in] grid       - qpoint grid
///@param[in] cell       - crystal structure information
///@param[in] processes  - three phonon processes in irreductible wedge
///@param[in] anhIFC     - Anharmonic force constants
///@param[in] Treference - reference temperature
///@param[in] world      - MPI communicator
Eigen::MatrixXd get_collision_operator_dense(
    const Gamma_grid& grid,
    const Crystal_structure& cell,
    std::vector<alma::Thirdorder_ifcs>& anhIFC,
    double Treference,
    boost::mpi::communicator& world);

/// This calculates the upper bound error for the m-th sized krylov
/// approximation of exponential with a normalized vector
///@param[in] A -  matrix from which the exponential is requested
///@param[in] n -  the krylov subspace size
///@param[in] t -  real scalar constant multiplying the matrix i.e.: exp(t*A)
inline long double error_bound(const Eigen::MatrixXd& A, int n, double t = 1.0) {
    long double At_norm = (t * A).norm();
    return 2.0 * std::pow(At_norm, n) * std::exp(At_norm) / std::tgamma(n);
}

/// Arnoldi algorithm used to compute the basis n-th krylov subspace generated
/// by Ab {b,A**2 * b,...,A**(n-1)* b}
///@param[in]     A - matrix A (see function purpose)
///@param[in]     b - vector b (see function purpose)
///@param[in]     n - krylov subspace order
///@param[in,out] B - Matrix to store the basis (each vector of the basis is
/// stored in different column)
///@param[in,out] H - A in the krylov subspace basis
///@param[in]     t - constant multiplying the A matrix so that the real matrix
/// is A' = t*A
void Arnoldi_algorithm(const Eigen::MatrixXd& A,
                       const Eigen::VectorXd& b,
                       int n,
                       Eigen::MatrixXd& B,
                       Eigen::MatrixXd& H,
                       double t = 1.0);

/// Storage of propagator as histogram
class propagator_H {
private:
    mutable Eigen::VectorXd signed_cumm; /// Signed cummulative function
    double
        pcol_abssum; /// The summ of abs value of original matrix - this is only
                     /// requested to be able to reconstruct P from histograms

    /// Some definitions to serialize the class
    friend class boost::serialization::access;
    template <typename Archive>
    void serialize(Archive& ar, const unsigned int version) {
        ar& signed_cumm;
        ar& pcol_abssum;
    }

    struct abs_less {
        bool operator()(const double& a, const double& b) const {
            return std::abs(a) < std::abs(b);
        }
    };

    /// Adaptation of lower bound to work with signed cumm
    template <class T> std::size_t lower_bound(T value) const {
        auto it = std::lower_bound(this->signed_cumm.data(),
                                   this->signed_cumm.data() + this->signed_cumm.rows(),
                                   value,
                                   abs_less());
        return std::distance(signed_cumm.data(),it);
    }

    friend alma::propagator_H lirp(
        alma::propagator_H& H1,
        alma::propagator_H& H2,
        double T1,
        double T2,
        double Tx);

public:
    /// Empty constructor
    propagator_H();

    /// Constructor
    ///@param[in]  data      - pointer to data to construct the histogram
    propagator_H(Eigen::VectorXd& b);
    ///@param[in] signed_cumm_ - the signed cummulative function
    ///@param[in] pcol_abssum_ - the abs sum of column values
    propagator_H(Eigen::VectorXd& signed_cumm_, double pcol_abssum_)
        : signed_cumm(signed_cumm_), pcol_abssum(pcol_abssum_){};

    /// Destructor is compiler default
    ~propagator_H() = default;
    /// Move constructor
    propagator_H(propagator_H&& B);

    /// Copy constructor
    propagator_H(const propagator_H& B);

    /// Assignement operator
    propagator_H& operator=(const propagator_H& B);

    /// comparison operator is dummy constructed
    bool operator<(const propagator_H& B) const;

    /// Given a random number it selects a mode
    ///@param[in] - R - number [0,1] not [0,1) as the last bit will not be
    /// selected with proper probability
    std::pair<int, std::size_t> get(double R) const;

    /// Returns the vector with the info
    Eigen::VectorXd get_signed_cumm() const;

    /// Returns the abs col summ
    double get_pcol_abssum() const {
        return this->pcol_abssum;
    }

    /// It return the size in bytes of all class
    std::size_t byte_size() const;

    /// It return the number of non-zero elements
    std::size_t size() const;
};

/// Function to interpolate two historgrams at different temperatures:
///@param[in] H1 - histogram at low  temperature
///@param[in] H2 - histogram at high temperature
///@param[in] T1 - low  T
///@param[in] T2 - high T
///@param[in] Tx - T at which the histogram must be interpolated
alma::propagator_H lirp(alma::propagator_H& H1,
                        alma::propagator_H& H2,
                        double T1,
                        double T2,
                        double Tx);

/// Function to interpolate two propagators at different temperatures:
///@param[in] P1 - propagator at low  temperature
///@param[in] P2 - propagator at high temperature
///@param[in] T1 - low  T
///@param[in] T2 - high T
///@param[in] Tx - T at which the propagator must be interpolated
inline Eigen::MatrixXd lirp(Eigen::MatrixXd& P1,
                            Eigen::MatrixXd& P2,
                            double T1,
                            double T2,
                            double Tx) {
    return P1 + (Tx - T1) * (P2 - P1) / (T2 - T1);
}

/// Function to dump P matrix to eigen binary from histogram:
///@param[in] fname  - File name without format extension
///@param[in] data   - reference to P to save
///@param[in] world  - MPI communicator
void save_P(std::string fname,
            Eigen::MatrixXd& data,
            boost::mpi::communicator& world);

/// Function to load P matrix from eigen binary:
///@param[in]     fname  - File name without format extension
///@param[in,out] pmat   - Pmatrix
///@param[in]     world  - MPI communicator
void load_P(std::string fname,
            Eigen::MatrixXd& pmat,
            boost::mpi::communicator& world);

/// Builds up a histogram of propagator matrix.
///@param[in]     A - Collision operator in energy form
///@param[in,out] modes_histo - histogram of propagator matrix
///@param[in]     t - time step
///@param[in]     world - MPI communicator
///@param[in]     eps   - tolerance to select the order of krylov subspace
///@param[in]     print_info - print_info
void build_histogram(Eigen::MatrixXd& A,
                     std::unordered_map<std::size_t, propagator_H>& modes_histo,
                     double t,
                     boost::mpi::communicator& world,
                     double eps = 1.0e-8,
                     bool print_info = false);

/// Builds up propagator matrix.
///@param[in]     A - Collision operator in energy form
///@param[in,out] modes_histo - histogram of propagator matrix
///@param[in]     t - time step
///@param[in]     world - MPI communicator
///@param[in]     eps   - tolerance to select the order of krylov subspace
void build_P(Eigen::MatrixXd& A,
             Eigen::MatrixXd& P,
             double t,
             boost::mpi::communicator& world,
             double eps = 1.0e-8);

} // namespace alma