//==================================================================================================
/*!
  @file

    @copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_FUNCTION_EXPM1_HPP_INCLUDED
#define BOOST_SIMD_FUNCTION_EXPM1_HPP_INCLUDED

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

  /*!
    @ingroup group-exponential
    This function object returns the exponential of its argument minus one:\f$e^{x}-1\f$

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

    @par Notes

    - result is accurate even for @c x of small modulus

    @par Decorators

     - std_ for floating entries calls @c std::expm1

    @see exp

    @par Example:

      @snippet expm1.cpp expm1

    @par Possible output:

      @snippet expm1.txt expm1

  **/
  IEEEValue expm1(IEEEValue const& x);
} }
#endif

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

#endif


#pragma once

#include <vector>
#include <tuple>
#include <array>
#include <Eigen/Dense>
#include <iod/symbols.hh>
#include <vpp/core/image2d.hh>
#include <vpp/core/vector.hh>
#include <vpp/core/make_array.hh>

namespace vpp
{

  template <typename T, typename U>
  void euclide_distance_transform(image2d<T>& input, image2d<U>& sedt)
  {
    image2d<vshort2> R(input.domain(), _border = 1);
    fill_with_border(R, vshort2{0,0});

    auto forward4 = [] () { return make_array(vint2{-1, -1}, vint2{-1, 0}, vint2{-1, 1}, vint2{0, -1}); };
    auto backward4 = [] () { return make_array(vint2{1, 1}, vint2{1, 0}, vint2{1, -1}, vint2{0, 1}); };

    fill_with_border(sedt, input.nrows() + input.ncols());
    pixel_wise(input, sedt) | [] (auto& i, auto& s) { if (i == 0) s = 0; };

    auto run = [&] (auto neighborhood, auto col_direction,
                    auto row_direction1, auto row_direction2, auto spn) {

      row_wise(sedt, R)(col_direction, _no_threads) | [&] (auto sedt_row, auto R_row)
      {
        // Forward pass
        pixel_wise(relative_access(sedt_row), relative_access(R_row))(row_direction1, _no_threads)
        | [&] (auto sedt_nbh, auto R_nbh)
        {
          vint2 min_rel_coord = neighborhood()[0];
          int min_dist = INT_MAX;
          for (vint2 nc : neighborhood())
          {
            int d = sedt_nbh(nc) + 2 * (std::abs(R_nbh(nc)[0] * nc[0]) +
                                        std::abs(R_nbh(nc)[1] * nc[1]))
              + nc.cwiseAbs().sum();

            if (d < min_dist)
            {
              min_dist = d;
              min_rel_coord = nc;
            }
          }

          if (min_dist < sedt_nbh(0, 0))
          {
            R_nbh(0, 0) = (R_nbh(min_rel_coord) + min_rel_coord.cast<short>()).template cast<short>();
            sedt_nbh(0, 0) = min_dist;
          }
        };

        // Backward pass
        pixel_wise(relative_access(sedt_row), relative_access(R_row))
        (row_direction2, _no_threads) | [&] (auto sedt_nbh, auto R_nbh)
        {
          int d = sedt_nbh(spn()) + 2 * std::abs(R_nbh(spn())[1]) + 1;
          if (d < sedt_nbh(0, 0))
          {
            sedt_nbh(0, 0) = d;
            R_nbh(0, 0) = (R_nbh(spn()) + spn().template cast<short>()).template cast<short>();
          }
        };

      };

    };

    run(forward4, _top_to_bottom, _left_to_right, _right_to_left, [] () { return vint2{0, 1}; });
    run(backward4, _bottom_to_top, _right_to_left, _left_to_right, [] () { return vint2{0, -1}; });

    // pixel_wise(sedt) | [] (auto& p) { p/=100; };
  }

  template <unsigned N, typename F>
  void loop_unroll(F f, std::enable_if_t<N == 0>* = 0) { f(N); }

  template <unsigned N, typename F>
  void loop_unroll(F f, std::enable_if_t<N != 0>* = 0) { f(N); loop_unroll<N-1>(f); }

  template <typename T, typename U, typename F, typename FW, typename B, typename BW, int WS = 3>
  void generic_incremental_distance_transform(image2d<T>& input, image2d<U>& sedt,
                                              F forward,
                                              FW forward_ws,
                                              B backward,
                                              BW backward_ws,
                                              std::integral_constant<int, WS> = std::integral_constant<int, WS>())
  {
    fill_with_border(sedt, input.nrows() + input.ncols());
    pixel_wise(input, sedt) | [] (auto& i, auto& s) { if (i == 0) s = 0; };

    auto run = [&] (auto neighb, auto ws,
                    auto col_direction,
                    auto row_direction) {
      pixel_wise(relative_access(sedt))(col_direction, row_direction, _no_threads) | [neighb, ws] (auto sn) {
        int min_dist = sn(0,0);

        auto nbh = neighb();

        // if (neighb().size() < 6)
        auto it = [&] (int i) {
          min_dist = std::min(min_dist, ws()[i] + sn(neighb()[i]));
        };

        typedef decltype(nbh) NBH;
        loop_unroll<std::tuple_size<NBH>::value - 1>(it);
        sn(0,0) = min_dist;
      };
    };

    run(forward, forward_ws, _top_to_bottom, _left_to_right);
    run(backward, backward_ws, _bottom_to_top, _right_to_left);
  }

  const auto d4_distance_transform = [] (auto& a, auto& b) {
    generic_incremental_distance_transform(a, b,
                                 [] () { return make_array(vint2{-1, 0}, vint2{0, -1}); },
                                 [] () { return make_array(1, 1); },
                                 [] () { return make_array(vint2{1, 0}, vint2{0, 1}); },
                                 [] () { return make_array(1, 1); });
  };

  const auto d8_distance_transform = [] (auto& a, auto& b) {
    generic_incremental_distance_transform(a, b,
                                 [] () { return make_array(vint2{-1, -1}, vint2{-1, 0}, vint2{-1, 1}, vint2{0, -1}); },
                                 [] () { return make_array(1,1,1,1); },
                                 [] () { return make_array(vint2{1, 1}, vint2{1, 0}, vint2{1, -1}, vint2{0, 1}); },
                                           [] () { return make_array(1, 1, 1, 1); });
  };

  const auto d3_4_distance_transform = [] (auto& a, auto& b) {
    generic_incremental_distance_transform(a, b,
                                           [] () { return make_array(vint2{-1, -1}, vint2{-1, 0}, vint2{-1, 1}, vint2{0, -1}); },
                                           [] () { return make_array(4,3,4,3); },
                                           [] () { return make_array(vint2{1, 1}, vint2{1, 0}, vint2{1, -1}, vint2{0, 1}); },
                                           [] () { return make_array(4, 3, 4, 3); }, std::integral_constant<int, 5>());
  };

  const auto d5_7_11_distance_transform = [] (auto& a, auto& b) {
    generic_incremental_distance_transform(a, b,

                                           [] () { return make_array(vint2{-2, -1}, vint2{-2, 1}, vint2{-1, -2}, vint2{-1, -1}, vint2{-1, 0}, vint2{-1, 1}, vint2{-1, 2}, vint2{0, -1}); },
                                           [] () { return make_array(11,11,11,7,5,7,11,5); },
                                           [] () { return make_array(vint2{0, 1}, vint2{1, -2}, vint2{1, -1}, vint2{1, 0}, vint2{1, 1}, vint2{1, 2}, vint2{2, -1}, vint2{2, 1}); },
                                           [] () { return make_array(5,11,7,5,7,11,11,11); },
                                           std::integral_constant<int, 5>());
  };

}


#include "mpc.h"
#include "param.h"
#include "assert.h"
#include "primes.h"
#include <NTL/mat_ZZ_p.h>
#include <NTL/mat_ZZ.h>
#include <NTL/ZZ.h>
#include <cmath>
#include <iostream>
#include <fstream>
#include <sstream>

using namespace NTL;
using namespace std;

bool MPCEnv::Initialize(int pid, vector< pair<int, int> > &pairs) {
  cout << "Initializing MPC environment" << endl;

  /* Set base prime for the finite field */
  ZZ base_p = conv<ZZ>(Param::BASE_P.c_str());
  ZZ_p::init(base_p);

  this->pid = pid;
  this->clock_start = chrono::steady_clock::now();
  debug = false;

  if (!SetupChannels(pairs)) {
    cout << "MPCEnv::Initialize: failed to initialize communication channels" << endl;
    return false;
  }

  if (!SetupPRGs(pairs)) {
    cout << "MPCEnv::Initialize: failed to initialize PRGs" << endl;
    return false;
  }

  SetSeed(prg.find(pid)->second);
  cur_prg_pid = pid;

  primes.SetLength(3);
  primes[0] = ZZ_p::modulus();

  bool found1 = false;
  bool found2 = false;
  long thres1 = Param::NBIT_K / 2;
  long thres2 = ((long) ceil(sqrt((double) NumBits(ZZ_p::modulus())))) + 1;

  long ind = -1;
  long maxind = sizeof(PRIME_LIST) / sizeof(PRIME_LIST[0]);
  while ((!found1 || !found2) && ++ind < maxind) {
    long p = PRIME_LIST[ind];
    if (!found1 && p > thres1) {
      found1 = true;
      primes[1] = ZZ(p);
    }
    if (!found2 && p > thres2) {
      found2 = true;
      primes[2] = ZZ(p);
    }
  }

  if (!found1 || !found2) {
    cout << "Failed to find suitable small primes" << endl;
    return false;
  }

  cout << "Small base primes selected: " << primes[1] << "(>" << thres1 << ") and "
       << primes[2] << "(>" << thres2 << ")" << endl;

  ZZ_bytes.SetLength(primes.length());
  ZZ_bits.SetLength(primes.length());
  ZZ_per_buf.SetLength(primes.length());
  for (int i = 0; i < primes.length(); i++) {
    ZZ_bytes[i] = NumBytes(primes[i]);
    ZZ_bits[i] = NumBits(primes[i]);
    ZZ_per_buf[i] = (uint64_t) (Param::MPC_BUF_SIZE / ZZ_bytes[i]);
  }

  assert(ZZ_bytes[0] <= Param::MPC_BUF_SIZE); // buffer should contain at least one ZZ_p

  pstate.push("");

  buf = (unsigned char *) malloc(Param::MPC_BUF_SIZE + GCM_AUTH_TAG_LEN);
  if (buf == NULL) {
    cout << "Fail to allocate MPC buffer" << endl;
    exit(1);
  } else {
    cout << "Allocated MPC buffer of size " << Param::MPC_BUF_SIZE << endl;
  }

  cout << "Number of bytes per ZZ_p: " << ZZ_bytes[0] << endl;

  cout << "Setting up lookup tables" << endl;

  table_cache.SetLength(3);
  table_type_ZZ.SetLength(3);
  table_field_index.SetLength(3);
  lagrange_cache.SetLength(3);

  Mat<ZZ_p> table;

  // Table 0
  table.SetDims(1, 2);
  if (pid > 0) {
    table[0][0] = 1;
    table[0][1] = 0;
  }
  table_type_ZZ[0] = true;
  table_cache[0] = table;
  table_field_index[0] = 2;

  // Table 1
  int half_len = Param::NBIT_K / 2;
  table.SetDims(2, half_len + 1);
  if (pid > 0) {
    for (int i = 0; i < half_len + 1; i++) {
      if (i == 0) {
        table[0][i] = 1;
        table[1][i] = 1;
      } else {
        table[0][i] = table[0][i - 1] * 2;
        table[1][i] = table[1][i - 1] * 4;
      }
    }

  }
  table_type_ZZ[1] = true;
  table_cache[1] = table;
  table_field_index[1] = 1;

  // Table 2: parameters (intercept, slope) for piecewise-linear approximation of
  //          negative log-sigmoid function
  table.SetDims(2, 64);
  if (pid > 0) {
    ifstream ifs;
    ifs.open("sigmoid_approx.txt");
    if (!ifs.is_open()) {
      cout << "Error opening sigmoid_approx.txt" << endl;
      clear(table);
    }
    for (int i = 0; i < table.NumCols(); i++) {
      double intercept, slope;
      ifs >> intercept >> slope;

      ZZ_p fp_intercept, fp_slope;
      DoubleToFP(fp_intercept, intercept, Param::NBIT_K, Param::NBIT_F);
      DoubleToFP(fp_slope, slope, Param::NBIT_K, Param::NBIT_F);

      table[0][i] = fp_intercept;
      table[1][i] = fp_slope;
    }
    ifs.close();
  }
  table_type_ZZ[2] = false;
  table_cache[2] = table;
  table_field_index[2] = 0;

  cout << "Generating lagrange cache" << endl;

  for (int cid = 0; cid < table_cache.length(); cid++) {
    long nrow = table_cache[cid].NumRows();
    long ncol = table_cache[cid].NumCols();
    bool index_by_ZZ = table_type_ZZ[cid];
    if (index_by_ZZ) {
      lagrange_cache[cid].SetDims(nrow, 2 * ncol);
    } else {
      lagrange_cache[cid].SetDims(nrow, ncol);
    }

    if (pid > 0) {
      cout << "Lagrange interpolation for Table " << cid << " ... ";
      for (int i = 0; i < nrow; i++) {
        Vec<long> x;
        Vec<ZZ_p> y;
        if (index_by_ZZ) {
          x.SetLength(2 * ncol);
          y.SetLength(2 * ncol);
        } else {
          x.SetLength(ncol);
          y.SetLength(ncol);
        }
        for (int j = 0; j < ncol; j++) {
          x[j] = j + 1;
          y[j] = table_cache[cid][i][j];
          if (index_by_ZZ) {
            x[j + ncol] = x[j] + conv<long>(primes[table_field_index[cid]]);
            y[j + ncol] = table_cache[cid][i][j];
          }
        }

        lagrange_interp(lagrange_cache[cid][i], x, y);
      }
      cout << "done" << endl;
    }
  }

  return true;
}

bool MPCEnv::SetupChannels(vector< pair<int, int> > &pairs) {
  for (int i = 0; i < pairs.size(); i++) {
    int p1 = pairs[i].first;
    int p2 = pairs[i].second;

    if (p1 != pid && p2 != pid) {
      continue;
    }

    int port = 8000;
    if (p1 == 0 && p2 == 1) {
      port = Param::PORT_P0_P1;
    } else if (p1 == 0 && p2 == 2) {
      port = Param::PORT_P0_P2;
    } else if (p1 == 1 && p2 == 2) {
      port = Param::PORT_P1_P2;
    } else if (p1 == 1 && p2 == 3) {
      port = Param::PORT_P1_P3;
    } else if (p1 == 2 && p2 == 3) {
      port = Param::PORT_P2_P3;
    }

    ostringstream oss;
    oss << Param::KEY_PATH << "P" << p1 << "_P" << p2 << ".key";
    string key_file = oss.str();

    int pother = p1 + p2 - pid;
    sockets.insert(map<int, CSocket>::value_type(pother, CSocket()));

    if (p1 == pid) {
      if (!OpenChannel(sockets[pother], port)) {
        cout << "Failed to connect with P" << pother << endl;
        return false;
      }
    } else {
      string ip_addr;
      if (pother == 0) {
        ip_addr = Param::IP_ADDR_P0;
      } else if (pother == 1) {
        ip_addr = Param::IP_ADDR_P1;
      } else if (pother == 2) {
        ip_addr = Param::IP_ADDR_P2;
      }

      if (!Connect(sockets[pother], ip_addr.c_str(), port)) {
        cout << "Failed to connect with P" << pother << endl;
        return false;
      }
    }

    if (!sockets[pother].SetKey(key_file)) {
      cout << "Failed to establish a secure channel with P" << pother << endl;
      return false;
    }

    cout << "Established a secure channel with P" << pother << endl;
  }

  cout << "Network setup complete" << endl;
  return true;
}

bool MPCEnv::SetupPRGs(vector< pair<int, int> > &pairs) {
  int key_len = NTL_PRG_KEYLEN; // from NTL
  unsigned char key[NTL_PRG_KEYLEN + GCM_AUTH_TAG_LEN];

  /* Internal PRG */
  int bytes = randread(key, key_len);
  if (bytes != key_len) {
    cout << "Failed to generate an internal PRG key" << endl;
    return false;
  }

  prg.insert(map<int, RandomStream>::value_type(pid, NewRandomStream(key)));

  /* Global PRG */
  ifstream ifs;
  string key_file = Param::KEY_PATH + "global.key";
  ifs.open(key_file.c_str(), ios::binary);
  if (!ifs.is_open()) {
    cout << "Failed to open global PRG key file: " << key_file << endl;
    return false;
  }

  ifs.read((char *)key, PRF_KEY_BYTES);
  if (ifs.gcount() != PRF_KEY_BYTES) {
    cout << "Failed to read " << PRF_KEY_BYTES << " bytes from global key file: " << key_file << endl;
    return false;
  }
  ifs.close();

  AESStream aes(key);
  aes.get(key, key_len);

  prg.insert(map<int, RandomStream>::value_type(-1, NewRandomStream(key)));

  /* Shared PRG (pairwise) */
  for (int i = 0; i < pairs.size(); i++) {
    int p1 = pairs[i].first;
    int p2 = pairs[i].second;

    if (p1 != pid && p2 != pid) {
      continue;
    }

    int pother = p1 + p2 - pid;

    if (p1 == pid) {
      bytes = randread(key, key_len);
      if (bytes != key_len) {
        cout << "Failed to generate a shared PRG key" << endl;
        return false;
      }

      prg.insert(map<int, RandomStream>::value_type(pother, NewRandomStream(key)));
      sockets[pother].SendSecure(key, key_len);
    } else {
      sockets[pother].ReceiveSecure(key, key_len);
      prg.insert(map<int, RandomStream>::value_type(pother, NewRandomStream(key)));
    }

    cout << "Shared PRG with P" << pother << " initialized" << endl;
  }

  cout << "PRG setup complete" << endl;
  return true;
}

void MPCEnv::CleanUp() {
  cout << "Closing sockets ... ";
  for (map<int, CSocket>::iterator it = sockets.begin(); it != sockets.end(); ++it) {
    CloseChannel(it->second);
  }
  cout << "done." << endl;
}

void MPCEnv::ProfilerResetTimer() {
  if (!Param::PROFILER) return;

  vector<uint64_t> stat(5, 0);

  chrono::time_point<chrono::steady_clock> clock_end = chrono::steady_clock::now();

  stat[0] = chrono::duration_cast<chrono::milliseconds>(clock_end - clock_start).count();

  int ind = 1;
  for (int p = 0; p < 3; p++) {
    if (p == pid) continue;
    map<int, CSocket>::iterator it = sockets.find(p);
    stat[ind] = it->second.GetBytesSent();
    stat[ind+1] = it->second.GetBytesReceived();
    it->second.ResetStats();
    ind += 2;
  }

  string state = pstate.top();
  if (state != "") {
    map<string, int>::iterator it = ptable_index.find(state);
    if (it != ptable_index.end()) {
      for (int i = 0; i < stat.size(); i++) {
        ptable[it->second].second[i] += stat[i];
      }
    } else {
      ptable.push_back(make_pair(state, stat));
      ptable_index[state] = ptable.size() - 1;
    }
  }

  clock_start = clock_end;
}

void MPCEnv::ProfilerPushState(string desc) {
  if (!Param::PROFILER) return;

  assert(desc != "");

  ProfilerResetTimer();

  string full_desc;
  if (pstate.top() == "") {
    full_desc = desc;
  } else {
    full_desc = pstate.top() + ">" + desc;
  }

  pstate.push(full_desc);
}

void MPCEnv::ProfilerPopState(bool write) {
  if (!Param::PROFILER) return;

  assert(pstate.top() != "");

  ProfilerResetTimer();

  pstate.pop();

  if (write) {
    ProfilerWriteToFile();
  }
}

void MPCEnv::ProfilerWriteToFile() {
  if (!Param::PROFILER) return;

  ProfilerResetTimer();

  // Open log file
  logfs.open(Param::LOG_FILE.c_str());
  if (!logfs.is_open()) {
    cout << "Fail to open the log file: " << Param::LOG_FILE << endl;
    exit(1);
  }

  ostringstream oss;

  // log file header
  oss << "Desc\tTime(ms)";
  for (int p = 0; p < 3; p++) {
    if (p == pid) continue;
    oss << "\tTo_" << p;
    oss << "\tFrom_" << p;
  }
  oss << endl;

  for (int i = 0; i < ptable.size(); i++) {
    oss << ptable[i].first;
    for (int j = 0; j < ptable[i].second.size(); j++) {
      oss << "\t" << ptable[i].second[j];
    }
    oss << endl;
  }

  logfs << oss.str();
  logfs.flush();
  logfs.close();
}

void MPCEnv::ParallelLogisticRegression(Vec<ZZ_p>& b0, Mat<ZZ_p>& bv, Vec<ZZ_p>& bx,
                                Mat<ZZ_p>& xr, Mat<ZZ_p>& xm,
                                Mat<ZZ_p>& vr, Mat<ZZ_p>& vm,
                                Vec<ZZ_p>& yr, Vec<ZZ_p>& ym,
                                int max_iter) {
  cout << "ParallelLogisticRegression" << endl;

  size_t n = vr.NumCols();
  size_t p = vr.NumRows();
  size_t c = xr.NumRows();
  assert(vm.NumRows() == p); assert(vm.NumCols() == n);
  assert(xm.NumRows() == c); assert(xm.NumCols() == n);
  assert(xr.NumCols() == n);
  assert(yr.length() == n);
  assert(ym.length() == n);

  // Initialize
  Init(b0, c);
  Init(bv, c, p);
  Init(bx, c);

  // Flip y
  Vec<ZZ_p> yneg_r = -yr;
  Vec<ZZ_p> yneg_m = -ym;
  if (pid > 0) {
    for (int i = 0; i < n; i++) {
      yneg_r[i] += 1;
    }
  }

  Vec<ZZ_p> yneg = yneg_m;
  if (pid == 1) {
    for (int i = 0; i < n; i++) {
      yneg[i] += yneg_r[i];
    }
  }

  ZZ_p fp_memory;
  DoubleToFP(fp_memory, 0.5, Param::NBIT_K, Param::NBIT_F);

  ZZ_p fp_n_inv;
  DoubleToFP(fp_n_inv, 1 / ((double) n), Param::NBIT_K, Param::NBIT_F);

  double eta = 0.3;

  ZZ_p fp_one;
  DoubleToFP(fp_one, 1, Param::NBIT_K, Param::NBIT_F);

  Vec<ZZ_p> step0;
  Mat<ZZ_p> stepv;
  Vec<ZZ_p> stepx;
  Init(step0, c);
  Init(stepv, c, p);
  Init(stepx, c);

  int nbatch = 10;
  int batch_size = (n + nbatch - 1) / nbatch;

  Mat<ZZ_p> xr_batch, xm_batch;
  Mat<ZZ_p> vr_batch, vm_batch;
  Vec<ZZ_p> yn_batch, ynr_batch, ynm_batch;

  // Gradient descent (with momentum)
  for (int it = 0; it < max_iter; it++) {
    cout << "Iter " << it << endl;

    int batch_index = it % nbatch;
    int start_ind = batch_size * batch_index;
    int end_ind = start_ind + batch_size; // exclusive
    if (end_ind > n) {
      end_ind = n;
    }
    int cur_bsize = end_ind - start_ind;

    // Construct mini-batch
    xr_batch.SetDims(c, cur_bsize);
    xm_batch.SetDims(c, cur_bsize);
    vr_batch.SetDims(p, cur_bsize);
    vm_batch.SetDims(p, cur_bsize);
    yn_batch.SetLength(cur_bsize);
    ynr_batch.SetLength(cur_bsize);
    ynm_batch.SetLength(cur_bsize);
    for (int j = 0; j < c; j++) {
      for (int i = 0; i < cur_bsize; i++) {
        xr_batch[j][i] = xr[j][start_ind + i];
        xm_batch[j][i] = xm[j][start_ind + i];
      }
    }
    for (int j = 0; j < p; j++) {
      for (int i = 0; i < cur_bsize; i++) {
        vr_batch[j][i] = vr[j][start_ind + i];
        vm_batch[j][i] = vm[j][start_ind + i];
      }
    }
    for (int i = 0; i < cur_bsize; i++) {
      yn_batch[i] = yneg[start_ind + i];
      ynr_batch[i] = yneg_r[start_ind + i];
      ynm_batch[i] = yneg_m[start_ind + i];
    }

    ZZ_p fp_bsize_inv;
    DoubleToFP(fp_bsize_inv, eta * (1 / ((double) cur_bsize)), Param::NBIT_K, Param::NBIT_F);

    Mat<ZZ_p> bvr, bvm;
    BeaverPartition(bvr, bvm, bv);

    Vec<ZZ_p> bxr, bxm;
    BeaverPartition(bxr, bxm, bx);

    Mat<ZZ_p> h;
    Init(h, c, cur_bsize);

    BeaverMult(h, bvr, bvm, vr_batch, vm_batch);
    for (int j = 0; j < c; j++) {
      Vec<ZZ_p> xrvec = fp_one * xr_batch[j];
      Vec<ZZ_p> xmvec = fp_one * xm_batch[j];
      BeaverMult(h[j], xrvec, xmvec, bxr[j], bxm[j]);
    }
    BeaverReconstruct(h);
    Trunc(h);

    for (int j = 0; j < c; j++) {
      AddScalar(h[j], b0[j]);
    }

    Vec<ZZ_p> hvec;
    Reshape(hvec, h);

    Vec<ZZ_p> svec, s_grad_vec;
    NegLogSigmoid(svec, s_grad_vec, hvec);
    hvec.kill();

    Mat<ZZ_p> s, s_grad;
    Reshape(s, svec, c, cur_bsize);
    Reshape(s_grad, s_grad_vec, c, cur_bsize);
    svec.kill();
    s_grad_vec.kill();

    // Compute gradient
    Vec<ZZ_p> d0;
    Mat<ZZ_p> dv;
    Vec<ZZ_p> dx;
    Init(d0, c);
    Init(dv, c, p);
    Init(dx, c);

    for (int j = 0; j < c; j++) {
      s_grad[j] += yn_batch * fp_one;
      d0[j] = Sum(s_grad[j]);
    }

    Mat<ZZ_p> s_grad_r, s_grad_m;
    BeaverPartition(s_grad_r, s_grad_m, s_grad);

    for (int j = 0; j < c; j++) {
      BeaverInnerProd(dx[j], xr_batch[j], xm_batch[j], s_grad_r[j], s_grad_m[j]);
    }
    BeaverReconstruct(dx);

    transpose(vr_batch, vr_batch);
    transpose(vm_batch, vm_batch);
    BeaverMult(dv, s_grad_r, s_grad_m, vr_batch, vm_batch);
    BeaverReconstruct(dv);
    Trunc(dv);

    step0 = step0 * fp_memory - d0 * fp_bsize_inv;
    stepv = stepv * fp_memory - dv * fp_bsize_inv;
    stepx = stepx * fp_memory - dx * fp_bsize_inv;
    Trunc(step0);
    Trunc(stepv);
    Trunc(stepx);

    b0 = b0 + step0;
    bv = bv + stepv;
    bx = bx + stepx;
  }
}

void MPCEnv::NegLogSigmoid(Vec<ZZ_p>& b, Vec<ZZ_p>& b_grad, Vec<ZZ_p>& a) {
  size_t n = a.length();

  int depth = 6;

  Vec<ZZ_p> cur = a; // copy

  Vec<ZZ_p> a_ind;
  a_ind.SetLength(a.length());
  clear(a_ind);

  double step = 4;

  for (int i = 0; i < depth; i++) {
    Vec<ZZ_p> cur_sign;
    IsPositive(cur_sign, cur);

    ZZ_p index_step(1 << (depth - 1 - i));

    for (int j = 0; j < n; j++) {
      a_ind[j] += cur_sign[j] * index_step;
    }

    cur_sign *= 2;
    if (pid == 1) {
      for (int j = 0; j < n; j++) {
        cur_sign[j] -= 1;
      }
    }

    ZZ_p step_fp;
    DoubleToFP(step_fp, step, Param::NBIT_K, Param::NBIT_F);

    for (int j = 0; j < n; j++) {
      cur[j] -= step_fp * cur_sign[j];
    }

    step /= 2;
  }

  // Make indices 1-based
  if (pid == 1) {
    for (int j = 0; j < n; j++) {
      a_ind[j]++;
    }
  }

  // Fetch piecewise linear approx parameters
  Mat<ZZ_p> param;
  TableLookup(param, a_ind, 2);

  MultElem(b, param[1], a);
  Trunc(b);

  if (pid > 0) {
    for (int j = 0; j < n; j++) {
      b[j] += param[0][j];
    }
  }

  b_grad = param[1];
}

void MPCEnv::InnerProd(Vec<ZZ_p>& c, Mat<ZZ_p>& a) {
  if (debug) cout << "InnerProd: " << a.NumRows() << ", " << a.NumCols() << endl;

  Mat<ZZ_p> ar, am;
  BeaverPartition(ar, am, a);

  Init(c, a.NumRows());
  for (int i = 0; i < a.NumRows(); i++) {
    BeaverInnerProd(c[i], ar[i], am[i]);
  }

  BeaverReconstruct(c);
}
void MPCEnv::InnerProd(ZZ_p& c, Vec<ZZ_p>& a) {
  if (debug) cout << "InnerProd: " << a.length() << endl;

  Vec<ZZ_p> ar, am;
  BeaverPartition(ar, am, a);
  BeaverInnerProd(c, ar, am);
  BeaverReconstruct(c);
}

void MPCEnv::Householder(Vec<ZZ_p>& v, Vec<ZZ_p>& x) {
  if (debug) cout << "Householder: " << x.length() << endl;

  int n = x.length();

  Vec<ZZ_p> xr, xm;
  BeaverPartition(xr, xm, x);

  Vec<ZZ_p> xdot;
  Init(xdot, 1);
  BeaverInnerProd(xdot[0], xr, xm);
  BeaverReconstruct(xdot);
  Trunc(xdot);

  Vec<ZZ_p> xnorm, dummy;
  FPSqrt(xnorm, dummy, xdot);

  Vec<ZZ_p> x1;
  x1.SetLength(1);
  x1[0] = x[0];

  Vec<ZZ_p> x1sign;
  IsPositive(x1sign, x1);

  x1sign *= 2;
  if (pid == 1) {
    x1sign[0] -= 1;
  }

  Vec<ZZ_p> shift;
  MultElem(shift, xnorm, x1sign);

  ZZ_p sr, sm;
  BeaverPartition(sr, sm, shift[0]);

  ZZ_p dot_shift(0);
  BeaverMult(dot_shift, xr[0], xm[0], sr, sm);
  BeaverReconstruct(dot_shift);
  Trunc(dot_shift);

  Vec<ZZ_p> vdot;
  vdot.SetLength(1);
  if (pid > 0) {
    vdot[0] = 2 * (xdot[0] + dot_shift);
  }

  Vec<ZZ_p> vnorm_inv;
  FPSqrt(dummy, vnorm_inv, vdot);

  ZZ_p invr, invm;
  BeaverPartition(invr, invm, vnorm_inv[0]);

  Vec<ZZ_p> vr, vm;
  if (pid > 0) {
    vr = xr;
    vr[0] += sr;
  } else {
    vr.SetLength(n);
  }
  vm = xm;
  vm[0] += sm;

  Init(v, n);
  BeaverMult(v, vr, vm, invr, invm);
  BeaverReconstruct(v);
  Trunc(v);
}

void MPCEnv::QRFactSquare(Mat<ZZ_p>& Q, Mat<ZZ_p>& R, Mat<ZZ_p>& A) {
  if (debug) cout << "QRFactSquare: " << A.NumRows() << ", " << A.NumCols() << endl;

  assert(A.NumRows() == A.NumCols());

  int n = A.NumRows();
  R.SetDims(n, n);
  if (pid > 0) {
    clear(R);
  }

  Mat<ZZ_p> Ap;
  if (pid == 0) {
    Ap.SetDims(n, n);
  } else {
    Ap = A;
  }

  ZZ_p one;
  DoubleToFP(one, 1, Param::NBIT_K, Param::NBIT_F);

  for (int i = 0; i < n - 1; i++) {
    Mat<ZZ_p> v;
    v.SetDims(1, Ap.NumCols());
    Householder(v[0], Ap[0]);

    Mat<ZZ_p> vt;
    if (pid == 0) {
      vt.SetDims(Ap.NumCols(), 1);
    } else {
      transpose(vt, v);
    }

    Mat<ZZ_p> P;
    MultMat(P, vt, v);
    Trunc(P);
    if (pid > 0) {
      P *= -2;
      if (pid == 1) {
        for (int j = 0; j < P.NumCols(); j++) {
          P[j][j] += one;
        }
      }
    }

    Mat<ZZ_p> B;
    if (i == 0) {
      Q = P;
      MultMat(B, Ap, P);
      Trunc(B);
    } else {
      Mat<ZZ_p> Qsub;
      Qsub.SetDims(n - i, n);
      if (pid > 0) {
        for (int j = 0; j < n - i; j++) {
          Qsub[j] = Q[j+i];
        }
      }

      Vec< Mat<ZZ_p> > left;
      Vec< Mat<ZZ_p> > right;
      left.SetLength(2);
      right.SetLength(2);
      left[0] = P;
      right[0] = Qsub;
      left[1] = Ap;
      right[1] = P;

      Vec< Mat<ZZ_p> > prod;
      MultMatParallel(prod, left, right);
      // TODO: parallelize Trunc
      Trunc(prod[0]);
      Trunc(prod[1]);

      if (pid > 0) {
        for (int j = 0; j < n - i; j++) {
          Q[j+i] = prod[0][j];
        }
        B = prod[1];
      } else {
        B.SetDims(n - i, n - i);
      }
    }

    if (pid > 0) {
      for (int j = 0; j < n - i; j++) {
        R[i+j][i] = B[j][0];
      }
      if (i == n - 2) {
        R[n-1][n-1] = B[1][1];
      }

      Ap.SetDims(n - i - 1, n - i - 1);
      for (int j = 0; j < n - i - 1; j++) {
        for (int k = 0; k < n - i - 1; k++) {
          Ap[j][k] = B[j+1][k+1];
        }
      }
    } else {
      Ap.SetDims(n - i - 1, n - i - 1);
    }
  }
}

void MPCEnv::OrthonormalBasis(Mat<ZZ_p>& Q, Mat<ZZ_p>& A) {
  if (debug) cout << "OrthonormalBasis: " << A.NumRows() << ", " << A.NumCols() << endl;

  assert(A.NumCols() >= A.NumRows());

  int c = A.NumRows();
  int n = A.NumCols();

  Vec< Vec<ZZ_p> > v_list;
  v_list.SetLength(c);

  Mat<ZZ_p> Ap;
  if (pid == 0) {
    Ap.SetDims(c, n);
  } else {
    Ap = A;
  }

  ZZ_p one;
  DoubleToFP(one, 1, Param::NBIT_K, Param::NBIT_F);

  for (int i = 0; i < c; i++) {
    Mat<ZZ_p> v;
    v.SetDims(1, Ap.NumCols());
    Householder(v[0], Ap[0]);

    if (pid == 0) {
      v_list[i].SetLength(Ap.NumCols());
    } else {
      v_list[i] = v[0];
    }

    Mat<ZZ_p> vt;
    if (pid == 0) {
      vt.SetDims(Ap.NumCols(), 1);
    } else {
      transpose(vt, v);
    }

    Mat<ZZ_p> Apv;
    MultMat(Apv, Ap, vt);
    Trunc(Apv);

    Mat<ZZ_p> B;
    MultMat(B, Apv, v);
    Trunc(B);
    if (pid > 0) {
      B *= -2;
      B += Ap;
    }

    Ap.SetDims(B.NumRows() - 1, B.NumCols() - 1);
    if (pid > 0) {
      for (int j = 0; j < B.NumRows() - 1; j++) {
        for (int k = 0; k < B.NumCols() - 1; k++) {
          Ap[j][k] = B[j+1][k+1];
        }
      }
    }
  }

  Q.SetDims(c, n);
  if (pid > 0) {
    clear(Q);
    if (pid == 1) {
      for (int i = 0; i < c; i++) {
        Q[i][i] = one;
      }
    }
  }

  for (int i = c - 1; i >= 0; i--) {
    Mat<ZZ_p> v;
    v.SetDims(1, v_list[i].length());
    if (pid > 0) {
      v[0] = v_list[i];
    }

    Mat<ZZ_p> vt;
    if (pid == 0) {
      vt.SetDims(v.NumCols(), 1);
    } else {
      transpose(vt, v);
    }

    Mat<ZZ_p> Qsub;
    Qsub.SetDims(c, n - i);
    if (pid > 0) {
      for (int j = 0; j < c; j++) {
        for (int k = 0; k < n - i; k++) {
          Qsub[j][k] = Q[j][k+i];
        }
      }
    }

    Mat<ZZ_p> Qv;
    MultMat(Qv, Qsub, vt);
    Trunc(Qv);

    Mat<ZZ_p> Qvv;
    MultMat(Qvv, Qv, v);
    Trunc(Qvv);
    if (pid > 0) {
      Qvv *= -2;
    }

    if (pid > 0) {
      for (int j = 0; j < c; j++) {
        for (int k = 0; k < n - i; k++) {
          Q[j][k+i] += Qvv[j][k];
        }
      }
    }
  }
}

void MPCEnv::Tridiag(Mat<ZZ_p>& T, Mat<ZZ_p>& Q, Mat<ZZ_p>& A) {
  if (debug) cout << "Tridiag: " << A.NumRows() << ", " << A.NumCols() << endl;

  assert(A.NumRows() == A.NumCols());
  assert(A.NumRows() > 2);

  int n = A.NumRows();

  ZZ_p one;
  DoubleToFP(one, 1, Param::NBIT_K, Param::NBIT_F);

  Q.SetDims(n, n);
  T.SetDims(n, n);
  if (pid > 0) {
    clear(Q);
    clear(T);
    if (pid == 1) {
      for (int i = 0; i < n; i++) {
        Q[i][i] = one;
      }
    }
  }

  Mat<ZZ_p> Ap;
  if (pid == 0) {
    Ap.SetDims(n, n);
  } else {
    Ap = A;
  }

  for (int i = 0; i < n - 2; i++) {
    Vec<ZZ_p> x;
    x.SetLength(Ap.NumCols() - 1);
    if (pid > 0) {
      for (int j = 0; j < Ap.NumCols() - 1; j++) {
        x[j] = Ap[0][j+1];
      }
    }

    Mat<ZZ_p> v;
    v.SetDims(1, x.length());
    Householder(v[0], x);

    Mat<ZZ_p> vt;
    if (pid == 0) {
      vt.SetDims(x.length(), 1);
    } else {
      transpose(vt, v);
    }

    Mat<ZZ_p> vv;
    MultMat(vv, vt, v);
    Trunc(vv);

    Mat<ZZ_p> P;
    P.SetDims(Ap.NumCols(), Ap.NumCols());
    if (pid > 0) {
      P[0][0] = (pid == 1) ? one : ZZ_p(0);
      for (int j = 1; j < Ap.NumCols(); j++) {
        for (int k = 1; k < Ap.NumCols(); k++) {
          P[j][k] = -2 * vv[j-1][k-1];
          if (pid == 1 && j == k) {
            P[j][k] += one;
          }
        }
      }
    }

    // TODO: parallelize? (minor improvement)
    Mat<ZZ_p> PAp;
    MultMat(PAp, P, Ap);
    Trunc(PAp);

    Mat<ZZ_p> B;
    MultMat(B, PAp, P);
    Trunc(B);

    Mat<ZZ_p> Qsub;
    Qsub.SetDims(n, n - i);
    if (pid > 0) {
      for (int j = 0; j < n; j++) {
        for (int k = 0; k < n - i; k++) {
          Qsub[j][k] = Q[j][k+i];
        }
      }
    }

    MultMat(Qsub, Qsub, P);
    Trunc(Qsub);
    if (pid > 0) {
      for (int j = 0; j < n; j++) {
        for (int k = 0; k < n - i; k++) {
          Q[j][k+i] = Qsub[j][k];
        }
      }
    }

    if (pid > 0) {
      T[i][i] = B[0][0];
      T[i+1][i] = B[1][0];
      T[i][i+1] = B[0][1];
      if (i == n - 3) {
        T[i+1][i+1] = B[1][1];
        T[i+1][i+2] = B[1][2];
        T[i+2][i+1] = B[2][1];
        T[i+2][i+2] = B[2][2];
      }
    }

    Ap.SetDims(B.NumRows() - 1, B.NumCols() - 1);
    if (pid > 0) {
      for (int j = 0; j < B.NumRows() - 1; j++) {
        for (int k = 0; k < B.NumCols() - 1; k++) {
          Ap[j][k] = B[j+1][k+1];
        }
      }
    }
  }
}

void MPCEnv::EigenDecomp(Mat<ZZ_p>& V, Vec<ZZ_p>& L, Mat<ZZ_p>& A) {
  if (debug) cout << "EigenDecomp: " << A.NumRows() << ", " << A.NumCols() << endl;

  assert(A.NumRows() == A.NumCols());
  int n = A.NumRows();

  L.SetLength(n);
  clear(L);

  Mat<ZZ_p> Ap, Q;
  Tridiag(Ap, Q, A);

  if (pid == 0) {
    V.SetDims(n, n);
  } else {
    transpose(V, Q);
  }

  for (int i = n - 1; i >= 1; i--) {
    cout << "EigenDecomp: " << i << "-th eigenvalue" << endl;
    for (int it = 0; it < Param::ITER_PER_EVAL; it++) {
      ZZ_p shift = Ap[i][i];
      if (pid > 0) {
        for (int j = 0; j < Ap.NumCols(); j++) {
          Ap[j][j] -= shift;
        }
      }

      Mat<ZZ_p> R;
      QRFactSquare(Q, R, Ap);

      MultMat(Ap, Q, R);
      Trunc(Ap);

      if (pid > 0) {
        for (int j = 0; j < Ap.NumCols(); j++) {
          Ap[j][j] += shift;
        }
      }

      Mat<ZZ_p> Vsub;
      Vsub.SetDims(i + 1, n);
      if (pid > 0) {
        for (int j = 0; j < i + 1; j++) {
          Vsub[j] = V[j];
        }
      }

      MultMat(Vsub, Q, Vsub);
      Trunc(Vsub);

      if (pid > 0) {
        for (int j = 0; j < i + 1; j++) {
          V[j] = Vsub[j];
        }
      }
    }

    L[i] = Ap[i][i];
    if (i == 1) {
      L[0] = Ap[0][0];
    }

    Mat<ZZ_p> Ap_copy = Ap;
    Ap.SetDims(i, i);
    if (pid > 0) {
      for (int j = 0; j < i; j++) {
        for (int k = 0; k < i; k++) {
          Ap[j][k] = Ap_copy[j][k];
        }
      }
    }
  }
  cout << "EigenDecomp: complete" << endl;
}

void MPCEnv::LessThanBitsAux(Vec<ZZ>& c, Mat<ZZ>& a, Mat<ZZ>& b, int public_flag, int fid) {
  if (debug) cout << "LessThanBitsAux: " << a.NumRows() << ", " << a.NumCols() << endl;

  assert(a.NumRows() == b.NumRows());
  assert(a.NumCols() == b.NumCols());

  int n = a.NumRows();
  int L = a.NumCols();

  /* Calculate XOR */
  Mat<ZZ> x;
  x.SetDims(n, L);

  if (public_flag == 0) {
    MultElem(x, a, b, fid);
    if (pid > 0) {
      x = a + b - 2 * x;
      Mod(x, fid);
    }
  } else if (pid > 0) {
    mul_elem(x, a, b);
    x = a + b - 2 * x;
    if (pid == 2) {
      x -= (public_flag == 1) ? a : b;
    }
    Mod(x, fid);
  }

  Mat<ZZ> f;
  PrefixOr(f, x, fid);
  x.kill();

  if (pid > 0) {
    for (int i = 0; i < n; i++) {
      for (int j = L - 1; j >= 1; j--) {
        f[i][j] -= f[i][j - 1];
      }
    }
    Mod(f, fid);
  }

  if (public_flag == 2) {
    c.SetLength(n);

    if (pid > 0) {
      for (int i = 0; i < n; i++) {
        c[i] = 0;
        for (int j = 0; j < L; j++) {
          c[i] += f[i][j] * b[i][j];
        }
      }
      Mod(c, fid);
    }
  } else {
    //TODO: optimize
    Vec< Mat<ZZ> > f_arr, b_arr;
    f_arr.SetLength(n);
    b_arr.SetLength(n);
    for (int i = 0; i < n; i++) {
      f_arr[i].SetDims(1, L);
      b_arr[i].SetDims(L, 1);
    }

    if (pid > 0) {
      for (int i = 0; i < n; i++) {
        f_arr[i][0] = f[i];
        for (int j = 0; j < L; j++) {
          b_arr[i][j][0] = b[i][j];
        }
      }
    }

    Vec< Mat<ZZ> > c_arr;
    MultMatParallel(c_arr, f_arr, b_arr, fid);

    c.SetLength(n);
    if (pid > 0) {
      for (int i = 0; i < n; i++) {
        c[i] = c_arr[i][0][0];
      }
    }
  }
}

void MPCEnv::LessThan(Vec<ZZ_p>& c, Vec<ZZ_p>& a, Vec<ZZ_p>& b) {
  Vec<ZZ_p> a_cpy;
  a_cpy.SetLength(a.length());
  if (pid > 0) {
    for (int i = 0; i < a.length(); i++) {
      a_cpy[i] = a[i] - b[i];
    }
  }

  // a - b >= 0?
  IsPositive(c, a_cpy);
  FlipBit(c);
}

void MPCEnv::LessThanPublic(Vec<ZZ_p>& c, Vec<ZZ_p>& a, ZZ_p bpub) {
  Vec<ZZ_p> a_cpy;
  a_cpy.SetLength(a.length());
  if (pid > 0) {
    for (int i = 0; i < a.length(); i++) {
      a_cpy[i] = a[i];
      if (pid == 1) {
        a_cpy[i] -= bpub;
      }
    }
  }

  // a - b >= 0?
  IsPositive(c, a_cpy);
  FlipBit(c);
}

// Failure probability of 1 / BASE_P
// Base field index 2
void MPCEnv::IsPositive(Vec<ZZ_p>& b, Vec<ZZ_p>& a) {
  if (debug) cout << "IsPositive: " << a.length() << endl;

  int n = a.length();
  int nbits = ZZ_bits[0];
  int fid = 2;

  Vec<ZZ_p> r;
  Mat<ZZ> r_bits;
  if (pid == 0) {
    RandVec(r, n);
    NumToBits(r_bits, r, nbits);

    SwitchSeed(1);
    Vec<ZZ_p> r_mask;
    Mat<ZZ> r_bits_mask;
    RandVec(r_mask, n);
    RandMat(r_bits_mask, n, nbits, fid);
    RestoreSeed();

    r -= r_mask;
    r_bits -= r_bits_mask;
    Mod(r_bits, fid);

    SendVec(r, 2);
    SendMat(r_bits, 2, fid);
  } else if (pid == 2) {
    ReceiveVec(r, 0, n);
    ReceiveMat(r_bits, 0, n, nbits, fid);
  } else {
    SwitchSeed(0);
    RandVec(r, n);
    RandMat(r_bits, n, nbits, fid);
    RestoreSeed();
  }

  Vec<ZZ_p> c;
  if (pid == 0) {
    c.SetLength(n);
  } else {
    c = 2 * a + r;
  }

  RevealSym(c);

  Mat<ZZ> c_bits;
  if (pid == 0) {
    c_bits.SetDims(n, nbits);
  } else {
    NumToBits(c_bits, c, nbits);
  }

  // Incorrect result if r = 0, which happens with probaility 1 / BASE_P
  Vec<ZZ> no_overflow;
  LessThanBitsPublic(no_overflow, r_bits, c_bits, fid);

  Vec<ZZ> c_xor_r;
  c_xor_r.SetLength(n);
  if (pid > 0) {
    for (int i = 0; i < n; i++) {
      c_xor_r[i] = r_bits[i][nbits-1] - 2 * c_bits[i][nbits-1] * r_bits[i][nbits-1];
      if (pid == 1) {
        c_xor_r[i] += c_bits[i][nbits-1];
      }
    }
    Mod(c_xor_r, fid);
  }

  Vec<ZZ> lsb;
  MultElem(lsb, c_xor_r, no_overflow, fid);
  if (pid > 0) {
    lsb *= 2;
    for (int i = 0; i < n; i++) {
      lsb[i] -= no_overflow[i] + c_xor_r[i];
      if (pid == 1) {
        lsb[i] += 1;
      }
    }
    Mod(lsb, fid);
  }

  // 0, 1 -> 1, 2
  if (pid == 1) {
    for (int i = 0; i < n; i++) {
      lsb[i] += 1;
    }
    Mod(lsb, fid);
  }

  Mat<ZZ_p> b_mat;
  TableLookup(b_mat, lsb, 0, fid);

  b = b_mat[0];
}

void MPCEnv::FlipBit(Vec<ZZ_p>& b, Vec<ZZ_p>& a) {
  if (debug) cout << "FlipBit: " << a.length() << endl;
  if (pid == 0) {
    b.SetLength(a.length());
  } else {
    b = -a;
  }

  if (pid == 1) {
    for (int i = 0; i < b.length(); i++) {
      b[i] += 1;
    }
  }
}

// Assumes Param::NBIT_K - NBIT_F is even
void MPCEnv::FPSqrt(Vec<ZZ_p>& b, Vec<ZZ_p>& b_inv, Vec<ZZ_p>& a) {
  if (debug) cout << "FPSqrt: " << a.length() << endl;

  int n = a.length();

  if (n > Param::DIV_MAX_N) {
    int nbatch = ceil(n / ((double) Param::DIV_MAX_N));
    b.SetLength(n);
    b_inv.SetLength(n);
    for (int i = 0; i < nbatch; i++) {
      cout << "FPSqrt on large vector: " << i + 1 << "/" << nbatch << endl;
      int start = Param::DIV_MAX_N * i;
      int end = start + Param::DIV_MAX_N;
      if (end > n) {
        end = n;
      }
      int batch_size = end - start;
      Vec<ZZ_p> a_copy;
      a_copy.SetLength(batch_size);
      for (int j = 0; j < batch_size; j++) {
        a_copy[j] = a[start + j];
      }
      Vec<ZZ_p> b_copy, b_inv_copy;
      FPSqrt(b_copy, b_inv_copy, a_copy);
      for (int j = 0; j < batch_size; j++) {
        b[start + j] = b_copy[j];
        b_inv[start + j] = b_inv_copy[j];
      }
    }
    return;
  }

  // TODO: Currently using the same # iter as division -- possibly need to update
  int niter = 2 * ceil(log2(((double) Param::NBIT_K) / 3.5));

  /* Initial approximation: 1 / sqrt(a_scaled) ~= 2.9581 - 4 * a_scaled + 2 * a_scaled^2 */
  Vec<ZZ_p> s, s_sqrt;
  NormalizerEvenExp(s, s_sqrt, a);

  Vec<ZZ_p> a_scaled;
  MultElem(a_scaled, a, s);
  Trunc(a_scaled, Param::NBIT_K, Param::NBIT_K - Param::NBIT_F);

  Vec<ZZ_p> a_scaled_sq;
  MultElem(a_scaled_sq, a_scaled, a_scaled);
  Trunc(a_scaled_sq);

  Vec<ZZ_p> scaled_est;
  if (pid == 0) {
    scaled_est.SetLength(n);
  } else {
    scaled_est = - 4 * a_scaled + 2 * a_scaled_sq;
    if (pid == 1) {
      ZZ_p coeff;
      DoubleToFP(coeff, 2.9581, Param::NBIT_K, Param::NBIT_F);
      for (int i = 0; i < n; i++) {
        scaled_est[i] += coeff;
      }
    }
  }

  Vec< Mat<ZZ_p> > h_and_g;
  h_and_g.SetLength(2);
  h_and_g[0].SetDims(1, n);
  h_and_g[1].SetDims(1, n);

  MultElem(h_and_g[0][0], scaled_est, s_sqrt);
  // Our scaled initial approximation (scaled_est) has bit length <= NBIT_F + 2
  // and s_sqrt is at most NBIT_K/2 bits, so their product is at most NBIT_K/2 + NBIT_F + 2
  Trunc(h_and_g[0], Param::NBIT_K/2 + Param::NBIT_F + 2, ((Param::NBIT_K - Param::NBIT_F) / 2) + 1);

  h_and_g[1][0] = h_and_g[0][0] * 2;
  MultElem(h_and_g[1][0], h_and_g[1][0], a);
  Trunc(h_and_g[1]);

  ZZ_p onepointfive;
  DoubleToFP(onepointfive, 1.5, Param::NBIT_K, Param::NBIT_F);

  for (int it = 0; it < niter; it++) {
    Mat<ZZ_p> r;
    MultElem(r, h_and_g[0], h_and_g[1]);
    Trunc(r);
    r = -r;
    if (pid == 1) {
      for (int i = 0; i < n; i++) {
        r[0][i] += onepointfive;
      }
    }

    Vec< Mat<ZZ_p> > r_dup;
    r_dup.SetLength(2);
    r_dup[0] = r;
    r_dup[1] = r;

    MultElemParallel(h_and_g, h_and_g, r_dup);
    // TODO: write a version of Trunc with parallel processing
    Trunc(h_and_g[0]);
    Trunc(h_and_g[1]);
  }

  b_inv = 2 * h_and_g[0][0];
  b = h_and_g[1][0];
}

void MPCEnv::FPDiv(Vec<ZZ_p>& c, Vec<ZZ_p>& a, Vec<ZZ_p>& b) {
  if (debug) cout << "FPDiv: " << a.length() << endl;

  assert(a.length() == b.length());

  int n = a.length();
  if (n > Param::DIV_MAX_N) {
    int nbatch = ceil(n / ((double) Param::DIV_MAX_N));
    c.SetLength(n);
    for (int i = 0; i < nbatch; i++) {
      int start = Param::DIV_MAX_N * i;
      int end = start + Param::DIV_MAX_N;
      if (end > n) {
        end = n;
      }
      int batch_size = end - start;

      cout << "FPDiv on large vector: " << i + 1 << "/" << nbatch << ", n = " << batch_size << endl;

      Vec<ZZ_p> a_copy, b_copy;
      a_copy.SetLength(batch_size);
      b_copy.SetLength(batch_size);
      for (int j = 0; j < batch_size; j++) {
        a_copy[j] = a[start + j];
        b_copy[j] = b[start + j];
      }
      Vec<ZZ_p> c_copy;
      FPDiv(c_copy, a_copy, b_copy);
      for (int j = 0; j < batch_size; j++) {
        c[start + j] = c_copy[j];
      }
    }
    return;
  }

  int niter = 2 * ceil(log2(((double) Param::NBIT_K) / 3.5)) + 1;

  /* Initial approximation: 1 / x_scaled ~= 5.9430 - 10 * x_scaled + 5 * x_scaled^2 */
  Vec<ZZ_p> s, s_sqrt;
  NormalizerEvenExp(s, s_sqrt, b);

  Vec<ZZ_p> b_scaled;
  MultElem(b_scaled, b, s);
  Trunc(b_scaled, Param::NBIT_K, Param::NBIT_K - Param::NBIT_F);

  Vec<ZZ_p> b_scaled_sq;
  MultElem(b_scaled_sq, b_scaled, b_scaled);
  Trunc(b_scaled_sq);

  Vec<ZZ_p> scaled_est;
  if (pid == 0) {
    scaled_est.SetLength(n);
  } else {
    scaled_est = - 10 * b_scaled + 5 * b_scaled_sq;
    if (pid == 1) {
      ZZ_p coeff;
      DoubleToFP(coeff, 5.9430, Param::NBIT_K, Param::NBIT_F);
      AddScalar(scaled_est, coeff);
    }
  }

  Vec<ZZ_p> w;
  MultElem(w, scaled_est, s);
  // scaled_est has bit length <= NBIT_F + 2, and s has bit length <= NBIT_K
  // so the bit length of w is at most NBIT_K + NBIT_F + 2
  Trunc(w, Param::NBIT_K + Param::NBIT_F + 2, Param::NBIT_K - Param::NBIT_F);

  Vec<ZZ_p> x;
  MultElem(x, w, b);
  Trunc(x);

  ZZ_p one;
  IntToFP(one, 1, Param::NBIT_K, Param::NBIT_F);

  x *= -1;
  if (pid == 1) {
    for (int i = 0; i < x.length(); i++) {
      x[i] += one;
    }
  }

  Vec<ZZ_p> y;
  MultElem(y, a, w);
  Trunc(y);

  for (int i = 0; i < niter; i++) {
    Vec<ZZ_p> xr, xm, yr, ym;
    BeaverPartition(xr, xm, x);
    BeaverPartition(yr, ym, y);

    Vec<ZZ_p> xpr = xr;
    if (pid > 0) {
      AddScalar(xpr, one);
    }

    Init(x, n);
    Init(y, n);

    BeaverMultElem(y, yr, ym, xpr, xm);
    BeaverMultElem(x, xr, xm, xr, xm);
    BeaverReconstruct(x);
    BeaverReconstruct(y);

    Trunc(x);
    Trunc(y);
  }

  if (pid == 1) {
    for (int i = 0; i < x.length(); i++) {
      x[i] += one;
    }
  }

  MultElem(c, y, x);
  Trunc(c);
}

void MPCEnv::Trunc(Mat<ZZ_p>& a, int k, int m) {
  if (debug) cout << "Trunc: " << a.NumRows() << ", " << a.NumCols() << endl;

  Mat<ZZ_p> r;
  Mat<ZZ_p> r_low;
  if (pid == 0) {
    RandMatBits(r, a.NumRows(), a.NumCols(), k + Param::NBIT_V);

    r_low.SetDims(a.NumRows(), a.NumCols());
    for (int i = 0; i < a.NumRows(); i++) {
      for (int j = 0; j < a.NumCols(); j++) {
        r_low[i][j] = conv<ZZ_p>(trunc_ZZ(rep(r[i][j]), m));
      }
    }

    Mat<ZZ_p> r_mask;
    Mat<ZZ_p> r_low_mask;
    SwitchSeed(1);
    RandMat(r_mask, a.NumRows(), a.NumCols());
    RandMat(r_low_mask, a.NumRows(), a.NumCols());
    RestoreSeed();

    r -= r_mask;
    r_low -= r_low_mask;

    SendMat(r, 2);
    SendMat(r_low, 2);
  } else if (pid == 2) {
    ReceiveMat(r, 0, a.NumRows(), a.NumCols());
    ReceiveMat(r_low, 0, a.NumRows(), a.NumCols());
  } else {
    SwitchSeed(0);
    RandMat(r, a.NumRows(), a.NumCols());
    RandMat(r_low, a.NumRows(), a.NumCols());
    RestoreSeed();
  }

  Mat<ZZ_p> c;
  if (pid > 0) {
    c = a + r;
  } else {
    c.SetDims(a.NumRows(), a.NumCols());
  }

  RevealSym(c);

  Mat<ZZ_p> c_low;
  c_low.SetDims(a.NumRows(), a.NumCols());
  if (pid > 0) {
    for (int i = 0; i < a.NumRows(); i++) {
      for (int j = 0; j < a.NumCols(); j++) {
        c_low[i][j] = conv<ZZ_p>(trunc_ZZ(rep(c[i][j]), m));
      }
    }
  }

  if (pid > 0) {
    a += r_low;
    if (pid == 1) {
      a -= c_low;
    }

    ZZ_p twoinvm;
    map<int, ZZ_p>::iterator it = invpow_cache.find(m);
    if (it == invpow_cache.end()) {
      ZZ_p two(2);
      ZZ_p twoinv;
      inv(twoinv, two);
      power(twoinvm, twoinv, m);
      invpow_cache[m] = twoinvm;
    } else {
      twoinvm = it->second;
    }

    a *= twoinvm;
  }
}

void MPCEnv::PrefixOr(Mat<ZZ>& b, Mat<ZZ>& a, int fid) {
  if (debug) cout << "PrefixOr: " << a.NumRows() << ", " << a.NumCols() << endl;

  int n = a.NumRows();

  /* Find next largest squared integer */
  int L = (int) ceil(sqrt((double) a.NumCols()));
  int L2 = L * L;

  assert(primes[fid] > L + 1);

  /* Zero-pad to L2 bits */
  Mat<ZZ> a_padded;
  a_padded.SetDims(n, L2);

  if (pid > 0) {
    for (int i = 0; i < n; i++) {
      for (int j = 0; j < L2; j++) {
        if (j < L2 - a.NumCols())
          a_padded[i][j] = 0;
        else
          a_padded[i][j] = a[i][j - L2 + a.NumCols()];
      }
    }
  }

  Reshape(a_padded, n * L, L);

  Vec<ZZ> x;
  FanInOr(x, a_padded, fid);

  Mat<ZZ> xpre;
  xpre.SetDims(n * L, L);

  if (pid > 0) {
    for (int i = 0; i < n; i++) {
      for (int j = 0; j < L; j++) {
        int xpi = L * i + j;
        for (int k = 0; k < L; k++) {
          xpre[xpi][k] = (k <= j) ? x[L * i + k] : ZZ(0);
        }
      }
    }
  }

  Vec<ZZ> y;
  FanInOr(y, xpre, fid);
  xpre.kill();

  Vec< Mat<ZZ> > f; // f is a concatenation of n 1-by-L matrices
  f.SetLength(n);
  for (int i = 0; i < n; i++) {
    f[i].SetDims(1, L);
  }

  if (pid > 0) {
    for (int i = 0; i < n; i++) {
      for (int j = 0; j < L; j++) {
        if (j == 0) {
          f[i][0][j] = x[L * i];
        } else {
          f[i][0][j] = y[L * i + j] - y[L * i + j - 1];
        }
      }
      Mod(f[i], fid);
    }
  }
  x.kill();

  Vec< Mat<ZZ> > tmp;
  tmp.SetLength(n);
  for (int i = 0; i < n; i++) {
    tmp[i].SetDims(L, L);
  }

  if (pid > 0) {
    for (int i = 0; i < n; i++) {
      for (int j = 0; j < L; j++) {
        tmp[i][j] = a_padded[L * i + j];
      }
    }
  }
  a_padded.kill();

  Vec< Mat<ZZ> > c;
  MultMatParallel(c, f, tmp, fid); // c is a concatenation of n 1-by-L matrices
  tmp.kill();

  Mat<ZZ> cpre;
  cpre.SetDims(n * L, L);
  if (pid > 0) {
    for (int i = 0; i < n; i++) {
      for (int j = 0; j < L; j++) {
        int cpi = L * i + j;
        for (int k = 0; k < L; k++) {
          cpre[cpi][k] = (k <= j) ? c[i][0][k] : ZZ(0);
        }
      }
    }
  }
  c.kill();

  Vec<ZZ> bdot_vec;
  FanInOr(bdot_vec, cpre, fid);
  cpre.kill();

  Vec< Mat<ZZ> > bdot;
  bdot.SetLength(n);
  for (int i = 0; i < n; i++) {
    bdot[i].SetDims(1, L);
  }

  if (pid > 0) {
    for (int i = 0; i < n; i++) {
      for (int j = 0; j < L; j++) {
        bdot[i][0][j] = bdot_vec[L * i + j];
      }
    }
  }
  bdot_vec.kill();

  for (int i = 0; i < n; i++) {
    Reshape(f[i], L, 1);
  }

  Vec< Mat<ZZ> > s;
  MultMatParallel(s, f, bdot, fid);
  bdot.kill();

  b.SetDims(n, a.NumCols());
  if (pid > 0) {
    for (int i = 0; i < n; i++) {
      for (int j = 0; j < a.NumCols(); j++) {
        int j_pad = L2 - a.NumCols() + j;

        int il = (int) (j_pad / L);
        int jl = j_pad - il * L;

        b[i][j] = s[i][il][jl] + y[L * i + il] - f[i][il][0];
      }
    }
  }
  Mod(b, fid);
  s.kill();
  y.kill();
  f.kill();
}

void MPCEnv::FanInOr(Vec<ZZ>& b, Mat<ZZ>& a, int fid) {
  if (debug) cout << "FanInOr: " << a.NumRows() << ", " << a.NumCols() << endl;

  int n = a.NumRows();
  int d = a.NumCols();

  Vec<ZZ> a_sum;
  a_sum.SetLength(n);

  if (pid > 0) {
    for (int i = 0; i < n; i++) {
      a_sum[i] = (pid == 1) ? 1 : 0;
      for (int j = 0; j < d; j++) {
        a_sum[i] += a[i][j];
      }
    }
    Mod(a_sum, fid);
  }

  Mat<ZZ> coeff;
  coeff.SetDims(1, d + 1);
  pair<int, int> key = make_pair(d + 1, fid);
  if (or_lagrange_cache.find(key) != or_lagrange_cache.end()) {
    coeff[0] = or_lagrange_cache[key];
  } else {
    Vec<ZZ> y;
    y.SetLength(d + 1);
    for (int i = 0; i < d + 1; i++) {
      y[i] = (i == 0) ? 0 : 1;
    }
    lagrange_interp_simple(coeff[0], y, fid); // OR function
    or_lagrange_cache[key] = coeff[0];
  }

  Mat<ZZ> bmat;
  EvaluatePoly(bmat, a_sum, coeff, fid);
  b = bmat[0];
}

void MPCEnv::ShareRandomBits(Vec<ZZ_p>& r, Mat<ZZ>& rbits, int k, int n, int fid) {
  if (debug) cout << "ShareRandomBits: " << n << endl;

  if (pid == 0) {
    RandVecBits(r, n, k + Param::NBIT_V);
    NumToBits(rbits, r, k);

    Vec<ZZ_p> r_mask;
    Mat<ZZ> rbits_mask;

    SwitchSeed(1);
    RandVec(r_mask, n);
    RandMat(rbits_mask, n, k, fid);
    RestoreSeed();

    r -= r_mask;
    rbits -= rbits_mask;
    Mod(rbits, fid);

    SendVec(r, 2);
    SendMat(rbits, 2, fid);
  } else if (pid == 2) {
    ReceiveVec(r, 0, n);
    ReceiveMat(rbits, 0, n, k, fid);
  } else {
    SwitchSeed(0);
    RandVec(r, n);
    RandMat(rbits, n, k, fid);
    RestoreSeed();
  }
}

void MPCEnv::TableLookup(Mat<ZZ_p>& b, Vec<ZZ_p>& a, int table_id) {
  if (debug) cout << "TableLookup: " << a.length() << endl;

  assert(!table_type_ZZ[table_id]);

  EvaluatePoly(b, a, lagrange_cache[table_id]);
}

void MPCEnv::TableLookup(Mat<ZZ_p>& b, Vec<ZZ>& a, int table_id, int fid) {
  if (debug) cout << "TableLookup: " << a.length() << endl;

  assert(table_type_ZZ[table_id]);
  assert(table_field_index[table_id] == fid);

  int s = table_cache[table_id].NumCols();
  int n = a.length();

  Vec<ZZ_p> a_exp;
  a_exp.SetLength(n);
  if (pid > 0) {
    for (int i = 0; i < n; i++) {
      a_exp[i] = conv<ZZ_p>(a[i]);
    }
  }

  if (debug) cout << "Evaluating polynomial" << endl;
  if (debug) cout << s << ", " << lagrange_cache[table_id].NumCols() << endl;

  EvaluatePoly(b, a_exp, lagrange_cache[table_id]);
}

// Base field index 1
void MPCEnv::NormalizerEvenExp(Vec<ZZ_p>& b, Vec<ZZ_p>& b_sqrt, Vec<ZZ_p>& a) {
  if (debug) cout << "NormalizerEvenExp: " << a.length() << endl;

  int n = a.length();
  int fid = 1;

  Vec<ZZ_p> r;
  Mat<ZZ> rbits;
  ShareRandomBits(r, rbits, Param::NBIT_K, n, fid);

  Vec<ZZ_p> e;
  if (pid == 0) {
    e.SetLength(n);
  } else {
    e = a + r;
  }
  r.kill();

  RevealSym(e);

  Mat<ZZ> ebits;
  if (pid == 0) {
    ebits.SetDims(n, Param::NBIT_K);
  } else {
    NumToBits(ebits, e, Param::NBIT_K);
  }
  e.kill();

  Vec<ZZ> c;
  LessThanBitsPublic(c, rbits, ebits, fid);
  if (pid > 0) {
    c = -c;
    if (pid == 1) {
      for (int i = 0; i < n; i++) {
        c[i] += 1;
      }
    }
    Mod(c, fid);
  }

  Mat<ZZ> ep;
  ep.SetDims(n, Param::NBIT_K + 1);
  if (pid > 0) {
    for (int i = 0; i < n; i++) {
      ep[i][0] = c[i];
      for (int j = 1; j < Param::NBIT_K + 1; j++) {
        ep[i][j] = (1 - 2 * ebits[i][j-1]) * rbits[i][j-1];
        if (pid == 1) {
          ep[i][j] += ebits[i][j-1];
        }
      }
    }
    Mod(ep, fid);
  }
  c.kill();

  Mat<ZZ> E;
  PrefixOr(E, ep, fid);
  ep.kill();

  Mat<ZZ> tpneg;
  tpneg.SetDims(n, Param::NBIT_K);
  if (pid > 0) {
    for (int i = 0; i < n; i++) {
      for (int j = 0; j < Param::NBIT_K; j++) {
        tpneg[i][j] = E[i][j] - rbits[i][j] * (1 - ebits[i][j]);
      }
    }
    Mod(tpneg, fid);
  }
  E.kill();

  Mat<ZZ> Tneg;
  PrefixOr(Tneg, tpneg, fid);
  tpneg.kill();

  int half_len = Param::NBIT_K / 2;

  Mat<ZZ> efir, rfir;
  efir.SetDims(n, Param::NBIT_K);
  rfir.SetDims(n, Param::NBIT_K);
  if (pid > 0) {
    mul_elem(efir, ebits, Tneg);
    Mod(efir, fid);
  }
  MultElem(rfir, rbits, Tneg, fid);
  ebits.kill();
  rbits.kill();

  Vec<ZZ> double_flag;
  LessThanBits(double_flag, efir, rfir, fid);
  efir.kill();
  rfir.kill();

  Mat<ZZ> odd_bits, even_bits;
  odd_bits.SetDims(n, half_len);
  even_bits.SetDims(n, half_len);
  if (pid > 0) {
    for (int i = 0; i < n; i++) {
      for (int j = 0; j < half_len; j++) {
        odd_bits[i][j] = (pid == 1) ? (1 - Tneg[i][2*j+1]) : -Tneg[i][2*j+1];
        if ((2 * j + 2) < Param::NBIT_K) {
          even_bits[i][j] = (pid == 1) ? (1 - Tneg[i][2*j+2]) : -Tneg[i][2*j+2];
        } else {
          even_bits[i][j] = 0;
        }
      }
    }
    Mod(odd_bits, fid);
    Mod(even_bits, fid);
  }
  Tneg.kill();

  Vec<ZZ> odd_bit_sum, even_bit_sum;
  Init(odd_bit_sum, n);
  Init(even_bit_sum, n);
  for (int i = 0; i < n; i++) {
    for (int j = 0; j < half_len; j++) {
      odd_bit_sum[i] += odd_bits[i][j];
      even_bit_sum[i] += even_bits[i][j];
    }
    if (pid == 1) {
      odd_bit_sum[i] += 1;
      even_bit_sum[i] += 1;
    }
  }
  Mod(odd_bit_sum, fid);
  Mod(even_bit_sum, fid);
  odd_bits.kill();
  even_bits.kill();

  // If double_flag = true, then use odd_bits, otherwise use even_bits

  Vec<ZZ> diff;
  if (pid == 0) {
    diff.SetLength(n);
  } else {
    diff = odd_bit_sum - even_bit_sum;
    Mod(diff, fid);
  }
  MultElem(diff, double_flag, diff, fid);
  double_flag.kill();

  Vec<ZZ> chosen_bit_sum;
  if (pid == 0) {
    chosen_bit_sum.SetLength(n);
  } else {
    chosen_bit_sum = even_bit_sum + diff;
    Mod(chosen_bit_sum, fid);
  }
  odd_bit_sum.kill();
  even_bit_sum.kill();
  diff.kill();

  Mat<ZZ_p> b_mat;
  TableLookup(b_mat, chosen_bit_sum, 1, fid);

  if (pid > 0) {
    b_sqrt = b_mat[0];
    b = b_mat[1];
  } else {
    b_sqrt.SetLength(n);
    b.SetLength(n);
  }
}

void MPCEnv::ReadFromFile(ZZ_p& a, ifstream& ifs) {
  Vec<ZZ_p> avec;
  if (pid > 0) {
    Read(avec, ifs, 1);
    a = avec[0];
  }
}
void MPCEnv::ReadFromFile(Vec<ZZ_p>& a, ifstream& ifs, int n) {
  if (pid > 0) {
    Read(a, ifs, n);
  } else {
    a.SetLength(n);
  }
}
void MPCEnv::ReadFromFile(Mat<ZZ_p>& a, ifstream& ifs, int nrow, int ncol) {
  if (pid > 0) {
    Read(a, ifs, nrow, ncol);
  } else {
    a.SetDims(nrow, ncol);
  }
}
void MPCEnv::WriteToFile(ZZ_p& a, fstream& ofs) {
  if (pid > 0) {
    Vec<ZZ_p> avec;
    avec.SetLength(1);
    avec[0] = a;
    Write(avec, ofs);
  }
}
void MPCEnv::WriteToFile(Vec<ZZ_p>& a, fstream& ofs) {
  if (pid > 0) {
    Write(a, ofs);
  }
}
void MPCEnv::WriteToFile(Mat<ZZ_p>& a, fstream& ofs) {
  if (pid > 0) {
    Write(a, ofs);
  }
}

void MPCEnv::Write(Vec<ZZ_p>& a, fstream& ofs) {
  Mat<ZZ_p> a_copy;
  a_copy.SetDims(1, a.length());
  a_copy[0] = a;
  Write(a_copy, ofs);
}

void MPCEnv::Read(Vec<ZZ_p>& a, ifstream& ifs, int n) {
  Mat<ZZ_p> tmp;
  Read(tmp, ifs, 1, n);
  a = tmp[0];
}

void MPCEnv::ReadWithFilter(Vec<ZZ_p>& a, ifstream& ifs, Vec<ZZ_p>& filt) {
  assert(ifs.is_open());
  a.SetLength(filt.length());

  unsigned char *buf_ptr = buf;
  uint64_t stored_in_buf = 0;

  for (int i = 0; i < filt.length(); i++) {
    if (filt[i] != 1) {
      uint64_t count = 0;
      int k = i;
      while (k < filt.length() && filt[k] != 1) {
        k++;
        count++;
      }
      ifs.ignore(count * ZZ_bytes[0]);
      i += count - 1;
    } else {
      if (stored_in_buf == 0) {
        uint64_t count = 0;
        int k = i;
        while (k < filt.length() && filt[k] == 1) {
          k++;
          count++;
        }

        if (count > ZZ_per_buf[0]) {
          count = ZZ_per_buf[0];
        }

        ifs.read((char *)buf, count * ZZ_bytes[0]);
        stored_in_buf += count;
        buf_ptr = buf;
      }

      a[i] = conv<ZZ_p>(ZZFromBytes(buf_ptr, ZZ_bytes[0]));
      buf_ptr += ZZ_bytes[0];
      stored_in_buf--;
    }
  }
}

void MPCEnv::Write(Mat<ZZ_p>& a, fstream& ofs) {
  assert(ofs.is_open());

  unsigned char *buf_ptr = buf;
  uint64_t stored_in_buf = 0;
  for (int i = 0; i < a.NumRows(); i++) {
    for (int j = 0; j < a.NumCols(); j++) {
      if (stored_in_buf == ZZ_per_buf[0]) {
        ofs.write((const char *)buf, ZZ_bytes[0] * stored_in_buf);
        stored_in_buf = 0;
        buf_ptr = buf;
      }

      BytesFromZZ(buf_ptr, rep(a[i][j]), ZZ_bytes[0]);
      stored_in_buf++;
      buf_ptr += ZZ_bytes[0];
    }
  }

  if (stored_in_buf > 0) {
    ofs.write((const char *)buf, ZZ_bytes[0] * stored_in_buf);
  }
}

void MPCEnv::SkipData(ifstream& ifs, int n) {
  if (pid > 0) {
    assert(ifs.is_open());
    ifs.ignore(n * ZZ_bytes[0]);
  }
}

void MPCEnv::SkipData(ifstream& ifs, int nrows, int ncols) {
  if (pid > 0) {
    assert(ifs.is_open());
    for (int i = 0; i < nrows; i++) {
      ifs.ignore(ncols * ZZ_bytes[0]);
    }
  }
}

void MPCEnv::Read(Mat<ZZ_p>& a, ifstream& ifs, int nrows, int ncols) {
  assert(ifs.is_open());

  a.SetDims(nrows, ncols);
  unsigned char *buf_ptr = buf;
  uint64_t stored_in_buf = 0;
  uint64_t remaining = nrows * ncols;
  for (int i = 0; i < a.NumRows(); i++) {
    for (int j = 0; j < a.NumCols(); j++) {
      if (stored_in_buf == 0) {
        uint64_t count;
        if (remaining < ZZ_per_buf[0]) {
          count = remaining;
        } else {
          count = ZZ_per_buf[0];
        }
        ifs.read((char *)buf, count * ZZ_bytes[0]);
        stored_in_buf += count;
        remaining -= count;
        buf_ptr = buf;
      }

      a[i][j] = conv<ZZ_p>(ZZFromBytes(buf_ptr, ZZ_bytes[0]));
      buf_ptr += ZZ_bytes[0];
      stored_in_buf--;
    }
  }
}

void MPCEnv::SendInt(int num, int to_pid) {
  cout << "SendInt called: num(" << num << "), to_pid(" << to_pid << ")" << endl;
  *((int *)buf) = num;
  sockets.find(to_pid)->second.Send(buf, sizeof(int));
}

int MPCEnv::ReceiveInt(int from_pid) {
  cout << "ReceiveInt called: from_pid(" << from_pid << ")" << endl;
  sockets.find(from_pid)->second.Receive(buf, sizeof(int));
  return *((int *)buf);
}

void MPCEnv::SendBool(bool flag, int to_pid) {
  cout << "SendBool called: flag(" << flag << "), to_pid(" << to_pid << ")" << endl;
  *((bool *)buf) = flag;
  sockets.find(to_pid)->second.Send(buf, sizeof(bool));
}

bool MPCEnv::ReceiveBool(int from_pid) {
  cout << "ReceiveBool called: from_pid(" << from_pid << ")" << endl;
  sockets.find(from_pid)->second.Receive(buf, sizeof(bool));
  return *((bool *)buf);
}

void MPCEnv::SwitchSeed(int pid) {
  prg.find(cur_prg_pid)->second = GetCurrentRandomStream();
  SetSeed(prg.find(pid)->second);
  cur_prg_pid = pid;
}

void MPCEnv::ExportSeed(fstream& ofs, int pid) {
  assert(ofs.is_open());

  RandomStream rs = prg.find(pid)->second;
  rs.serialize(buf);

  ofs.write((const char *)buf, RandomStream::numBytes());
}

void MPCEnv::ExportSeed(fstream& ofs) {
  assert(ofs.is_open());

  RandomStream rs = GetCurrentRandomStream();
  rs.serialize(buf);

  ofs.write((const char *)buf, RandomStream::numBytes());
}

void MPCEnv::ImportSeed(int newid, ifstream& ifs) {
  assert(ifs.is_open());

  ifs.read((char *)buf, RandomStream::numBytes());

  RandomStream rs((const unsigned char *)buf, true);

  pair<map<int,RandomStream>::iterator,bool> ret;
  ret = prg.insert(pair<int, RandomStream>(newid, rs));
  if (!ret.second) { // ID exists already
    ret.first->second = rs;
  }
}

void MPCEnv::BeaverReadFromFile(Mat<ZZ_p>& ar, Mat<ZZ_p>& am, ifstream& ifs, int nrow, int ncol) {
  if (pid > 0) {
    Read(ar, ifs, nrow, ncol);
  } else {
    ar.SetDims(nrow, ncol);
  }
  Read(am, ifs, nrow, ncol);
}

void MPCEnv::BeaverReadFromFile(Vec<ZZ_p>& ar, Vec<ZZ_p>& am, ifstream& ifs, int n) {
  if (pid > 0) {
    Read(ar, ifs, n);
  } else {
    ar.SetLength(n);
  }
  Read(am, ifs, n);
}

void MPCEnv::BeaverReadFromFileWithFilter(Vec<ZZ_p>& ar, Vec<ZZ_p>& am, ifstream& ifs, Vec<ZZ_p>& filt) {
  if (pid > 0) {
    ReadWithFilter(ar, ifs, filt);
  } else {
    ar.SetLength(filt.length());
  }
  ReadWithFilter(am, ifs, filt);
}

void MPCEnv::BeaverWriteToFile(Vec<ZZ_p>& ar, Vec<ZZ_p>& am, fstream& ofs) {
  if (pid > 0) {
    Write(ar, ofs);
  }
  Write(am, ofs);
}

void MPCEnv::BeaverWriteToFile(Mat<ZZ_p>& ar, Mat<ZZ_p>& am, fstream& ofs) {
  if (pid > 0) {
    Write(ar, ofs);
  }
  Write(am, ofs);
}

void MPCEnv::BeaverMultElem(Vec<ZZ_p>& ab, Vec<ZZ_p>& ar, Vec<ZZ_p>& am, Vec<ZZ_p>& br, Vec<ZZ_p>& bm, int fid) {
  if (pid == 0) {
    Vec<ZZ_p> ambm;
    mul_elem(ambm, am, bm);
    ab += ambm;
  } else {

    ZZ_pContext context;
    context.save();

    NTL_GEXEC_RANGE(ab.length() > Param::PAR_THRES, ab.length(), first, last)

    context.restore();

    for (int i = first; i < last; i++) {
      ab[i] += ar[i] * bm[i];
      ab[i] += am[i] * br[i];
      if (pid == 1) {
        ab[i] += ar[i] * br[i];
      }
    }

    NTL_GEXEC_RANGE_END
  }
}

void MPCEnv::BeaverMult(Mat<ZZ_p>& ab, Mat<ZZ_p>& ar, Mat<ZZ_p>& am, Mat<ZZ_p>& br, Mat<ZZ_p>& bm, bool elem_wise, int fid) {
  if (pid == 0) {
    Mat<ZZ_p> ambm;
    if (elem_wise) {
      mul_elem(ambm, am, bm);
    } else {
      mul(ambm, am, bm);
    }
    ab += ambm;
  } else {
    if (elem_wise) {

      ZZ_pContext context;
      context.save();

      NTL_GEXEC_RANGE(ab.NumRows() > Param::PAR_THRES, ab.NumRows(), first, last)

      context.restore();

      for (int i = first; i < last; i++) {
        for (int j = 0; j < ab.NumCols(); j++) {
          ab[i][j] += ar[i][j] * bm[i][j];
          ab[i][j] += am[i][j] * br[i][j];
          if (pid == 1) {
            ab[i][j] += ar[i][j] * br[i][j];
          }
        }
      }

      NTL_GEXEC_RANGE_END

    } else {
      ab += ar * bm;
      ab += am * br;
      if (pid == 1) {
        ab += ar * br;
      }
    }
  }
}

void MPCEnv::BeaverMultElem(Vec<ZZ>& ab, Vec<ZZ>& ar, Vec<ZZ>& am, Vec<ZZ>& br, Vec<ZZ>& bm, int fid) {
  if (pid == 0) {
    Vec<ZZ> ambm;
    mul_elem(ambm, am, bm);
    ab += ambm;
  } else {
    NTL_GEXEC_RANGE(ab.length() > Param::PAR_THRES, ab.length(), first, last)

    for (int i = first; i < last; i++) {
      ab[i] += ar[i] * bm[i];
      ab[i] += am[i] * br[i];
      if (pid == 1) {
        ab[i] += ar[i] * br[i];
      }
    }

    NTL_GEXEC_RANGE_END
  }

  Mod(ab, fid);
}

void MPCEnv::BeaverMult(Mat<ZZ>& ab, Mat<ZZ>& ar, Mat<ZZ>& am, Mat<ZZ>& br, Mat<ZZ>& bm, bool elem_wise, int fid) {
  if (pid == 0) {
    Mat<ZZ> ambm;
    if (elem_wise) {
      mul_elem(ambm, am, bm);
    } else {
      mul(ambm, am, bm);
    }
    ab += ambm;
  } else {
    if (elem_wise) {
      NTL_GEXEC_RANGE(ab.NumRows() > Param::PAR_THRES, ab.NumRows(), first, last)

      for (int i = first; i < last; i++) {
        for (int j = 0; j < ab.NumCols(); j++) {
          ab[i][j] += ar[i][j] * bm[i][j];
          ab[i][j] += am[i][j] * br[i][j];
          if (pid == 1) {
            ab[i][j] += ar[i][j] * br[i][j];
          }
        }
      }

      NTL_GEXEC_RANGE_END

    } else {
      ab += ar * bm;
      ab += am * br;
      if (pid == 1) {
        ab += ar * br;
      }
    }
  }

  Mod(ab, fid);
}


// Copyright  (C)  2007  Ruben Smits <ruben dot smits at mech dot kuleuven dot be>

// Version: 1.0
// Author: Ruben Smits <ruben dot smits at mech dot kuleuven dot be>
// Maintainer: Ruben Smits <ruben dot smits at mech dot kuleuven dot be>
// URL: http://www.orocos.org/kdl

// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; either
// version 2.1 of the License, or (at your option) any later version.

// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.

// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA

#include "articulatedbodyinertia.hpp"

#include <Eigen/Core>

namespace KDL{

  ArticulatedBodyInertia::ArticulatedBodyInertia(const RigidBodyInertia& rbi)
    {
        this->M=Eigen::Matrix3d::Identity()*rbi.m;
        this->I=Eigen::Map<const Eigen::Matrix3d>(rbi.I.data);
        this->H << 0,-rbi.h[2],rbi.h[1],
            rbi.h[2],0,-rbi.h[0],
            -rbi.h[1],rbi.h[0],0;
    }

    ArticulatedBodyInertia::ArticulatedBodyInertia(double m, const Vector& c, const RotationalInertia& Ic)
    {
        *this = RigidBodyInertia(m,c,Ic);
    }

  ArticulatedBodyInertia::ArticulatedBodyInertia(const Eigen::Matrix3d& M, const Eigen::Matrix3d& H, const Eigen::Matrix3d& I)
    {
        this->M=M;
        this->I=I;
        this->H=H;
    }

    ArticulatedBodyInertia operator*(double a,const ArticulatedBodyInertia& I){
        return ArticulatedBodyInertia(a*I.M,a*I.H,a*I.I);
    }

    ArticulatedBodyInertia operator+(const ArticulatedBodyInertia& Ia, const ArticulatedBodyInertia& Ib){
        return ArticulatedBodyInertia(Ia.M+Ib.M,Ia.H+Ib.H,Ia.I+Ib.I);
    }

    ArticulatedBodyInertia operator+(const RigidBodyInertia& Ia, const ArticulatedBodyInertia& Ib){
        return ArticulatedBodyInertia(Ia)+Ib;
    }
    ArticulatedBodyInertia operator-(const ArticulatedBodyInertia& Ia, const ArticulatedBodyInertia& Ib){
        return ArticulatedBodyInertia(Ia.M-Ib.M,Ia.H-Ib.H,Ia.I-Ib.I);
    }

    ArticulatedBodyInertia operator-(const RigidBodyInertia& Ia, const ArticulatedBodyInertia& Ib){
        return ArticulatedBodyInertia(Ia)-Ib;
    }

    Wrench operator*(const ArticulatedBodyInertia& I,const Twist& t){
        Wrench result;
        Eigen::Vector3d::Map(result.force.data)=I.M*Eigen::Vector3d::Map(t.vel.data)+I.H.transpose()*Eigen::Vector3d::Map(t.rot.data);
        Eigen::Vector3d::Map(result.torque.data)=I.I*Eigen::Vector3d::Map(t.rot.data)+I.H*Eigen::Vector3d::Map(t.vel.data);
        return result;
    }

    ArticulatedBodyInertia operator*(const Frame& T,const ArticulatedBodyInertia& I){
        Frame X=T.Inverse();
        //mb=ma
        //hb=R*(h-m*r)
        //Ib = R(Ia+r x h x + (h-m*r) x r x)R'
        Eigen::Map<Eigen::Matrix3d> E(X.M.data);
        Eigen::Matrix3d rcross;
        rcross << 0,-X.p[2],X.p[1],
            X.p[2],0,-X.p[0],
            -X.p[1],X.p[0],0;

        Eigen::Matrix3d HrM=I.H-rcross*I.M;
        return ArticulatedBodyInertia(E*I.M*E.transpose(),E*HrM*E.transpose(),E*(I.I-rcross*I.H.transpose()+HrM*rcross)*E.transpose());
    }

    ArticulatedBodyInertia operator*(const Rotation& M,const ArticulatedBodyInertia& I){
        Eigen::Map<const Eigen::Matrix3d> E(M.data);
        return ArticulatedBodyInertia(E.transpose()*I.M*E,E.transpose()*I.H*E,E.transpose()*I.I*E);
    }

    ArticulatedBodyInertia ArticulatedBodyInertia::RefPoint(const Vector& p){
        //mb=ma
        //hb=R*(h-m*r)
        //Ib = R(Ia+r x h x + (h-m*r) x r x)R'
        Eigen::Matrix3d rcross;
        rcross << 0,-p[2],p[1],
            p[2],0,-p[0],
            -p[1],p[0],0;

        Eigen::Matrix3d HrM=this->H-rcross*this->M;
        return ArticulatedBodyInertia(this->M,HrM,this->I-rcross*this->H.transpose()+HrM*rcross);
    }
}//namespace


//==================================================================================================
/*!
  @file
    @copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_FUNCTION_ERFCX_HPP_INCLUDED
#define BOOST_SIMD_FUNCTION_ERFCX_HPP_INCLUDED

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

  /*!
    @ingroup group-euler
     This function object computes the  underflow-compensating (scaled) complementary  error function:
   \f$\displaystyle e^{x^2}\frac{2}{\sqrt\pi}\int_{x}^{\infty} e^{-t^2}\mbox{d}t\f$

    @see erfc, erf

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

    @par Example:

      @snippet erfcx.cpp erfcx

    @par Possible output:

      @snippet erfcx.txt erfcx
  **/
  IEEEValue erfcx(IEEEValue const& x);
} }
#endif

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

#endif


#include <iostream>
#include <vector>
#include <fstream>
#include <thread>
#include <cmath>
#include <chrono>

using namespace std;

#include <boost/timer.hpp>

// for sophus
#include <sophus/se3.hpp>

using Sophus::SE3d;

// for eigen
#include <Eigen/Core>
#include <Eigen/Geometry>

using namespace Eigen;

#include <opencv2/core/core.hpp>
#include <opencv2/imgcodecs/imgcodecs.hpp>

#include "constants.h"
#include "cuda_wrapper.h"
#include "kernels.cuh"
#include "plot.h"

using namespace cv;

/**
 * Dataset from:
 *
 *   http://rpg.ifi.uzh.ch/datasets/remode_test_data.zip
 *
 * */

void plotDepth(const Mat &depth_truth, const Mat &depth_estimate);
void plotCur(const Mat &cur);

bool readDatasetFiles(
    const string &path,
    vector<string> &color_image_files,
    vector<SE3d> &poses,
    cv::Mat &ref_depth
);

void showEpipolarMatch(const Mat &ref, const Mat &curr, const Vector2d &px_ref, const Vector2d &px_curr);

void showEpipolarLine(const Mat &ref, const Mat &curr, const Vector2d &px_ref, const Vector2d &px_min_curr,
                      const Vector2d &px_max_curr);

void evaludateDepth(const Mat &depth_truth, const Mat &depth_estimate);

int main(int argc, char **argv) {
    if (argc != 2) {
        cout << "Usage: dense_mapping path_to_test_dataset" << endl;
        return -1;
    }

    // Read dataset
    vector<string> color_image_files;
    vector<SE3d> poses_TWC;
    Mat ref_depth;
    bool ret = readDatasetFiles(argv[1], color_image_files, poses_TWC, ref_depth);
    if (ret == false) {
        cout << "Reading image files failed!" << endl;
        return -1;
    }
    cout << "read total " << color_image_files.size() << " files." << endl;

    // Initial depth image
    Mat ref = imread(color_image_files[0], 0); // gray-scale image
    SE3d pose_ref_TWC = poses_TWC[0];
    double init_depth = 3.0;
    double init_cov2 = 3.0;
    Mat depth(height, width, CV_64F, init_depth);
    Mat depth_cov2(height, width, CV_64F, init_cov2);

    for (int index = 1; index < color_image_files.size(); index++) {
        cout << "*** loop " << index << " ***" << endl;
        Mat curr = imread(color_image_files[index], 0);
        if (curr.data == nullptr) continue;
        SE3d pose_curr_TWC = poses_TWC[index];
        SE3d pose_T_C_R = pose_curr_TWC.inverse() * pose_ref_TWC;   // T_C_W * T_W_R = T_C_R
        chrono::steady_clock::time_point t1 = chrono::steady_clock::now();
        update_cuda(ref, curr, pose_T_C_R, depth, depth_cov2);
        chrono::steady_clock::time_point t2 = chrono::steady_clock::now();

        auto time_used = chrono::duration_cast<chrono::duration<double>>(t2 - t1);
        std::cout << "Time used: " << time_used.count() << "s\n";

        evaludateDepth(ref_depth, depth);
        plotDepth(ref_depth, depth);
        plotCur(curr);
        // imshow("image", curr);
        // waitKey(1);
    }

    cout << "estimation returns, saving depth map ..." << endl;
    imwrite("depth.png", depth);
    cout << "done." << endl;

    return 0;
}

bool readDatasetFiles(
    const string &path,
    vector<string> &color_image_files,
    std::vector<SE3d> &poses,
    cv::Mat &ref_depth) {
    ifstream fin(path + "/first_200_frames_traj_over_table_input_sequence.txt");
    if (!fin) return false;

    while (!fin.eof()) {
        // 数据格式：图像文件名 tx, ty, tz, qx, qy, qz, qw ，注意是 TWC 而非 TCW
        string image;
        fin >> image;
        double data[7];
        for (double &d:data) fin >> d;

        color_image_files.push_back(path + string("/images/") + image);
        poses.push_back(
            SE3d(Quaterniond(data[6], data[3], data[4], data[5]),
                 Vector3d(data[0], data[1], data[2]))
        );
        if (!fin.good()) break;
    }
    fin.close();

    // load reference depth
    fin.open(path + "/depthmaps/scene_000.depth");
    ref_depth = cv::Mat(height, width, CV_64F);
    if (!fin) return false;
    for (int y = 0; y < height; y++)
        for (int x = 0; x < width; x++) {
            double depth = 0;
            fin >> depth;
            ref_depth.ptr<double>(y)[x] = depth / 100.0;
        }

    return true;
}

void evaludateDepth(const Mat &depth_truth, const Mat &depth_estimate) {
    double ave_depth_error = 0;
    double ave_depth_error_sq = 0;
    int cnt_depth_data = 0;
    for (int y = boarder; y < depth_truth.rows - boarder; y++)
        for (int x = boarder; x < depth_truth.cols - boarder; x++) {
            double error = depth_truth.ptr<double>(y)[x] - depth_estimate.ptr<double>(y)[x];
            ave_depth_error += error;
            ave_depth_error_sq += error * error;
            cnt_depth_data++;
        }
    ave_depth_error /= cnt_depth_data;
    ave_depth_error_sq /= cnt_depth_data;

    cout << "Average squared error = " << ave_depth_error_sq << ", average error: " << ave_depth_error << endl;
}


#include "CalibrationHelpers.h"

#include <iostream>
#include <fstream>
#include <string>
#include <Eigen/Dense>

namespace romocc
{

std::vector<Transform3d> invert_matrices(std::vector<Transform3d> matrices)
{
    std::vector<Transform3d> invertedMatrices;

    for(int i = 0; i<matrices.size(); i++)
    {
        invertedMatrices.push_back(matrices.at(i).inverse());
    }
    return invertedMatrices;
}

double compute_average(std::vector<double> const &v)
{
    double sum = 0;
    for(int i=0; i<v.size(); i++)
        sum += v[i];
    return sum/v.size();
}

double compute_variance(std::vector<double> const &v, double mean)
{
    double sum = 0.0;
    double temp =0.0;
    double var =0.0;

    for ( int j =0; j <= v.size()-1; j++)
    {
        temp = std::pow(v[j]-mean,2);
        sum += temp;
    }

    return var = sum/(v.size()-2);
}

std::vector<double> compute_linspace(double a, double b, int n)
{
    std::vector<double> array;
    double step = (b-a) / (n-1);

    while(a <= b) {
        array.push_back(a);
        a += step;
    }
    return array;
}

double sgn(double val)
{
    return (double(0) < val) - (val < double(0));
}

Transform3d load_calibration_file(std::string filepath){
    auto calMat = Eigen::Affine3d::Identity();
    std::ifstream inFile;
    inFile.open(filepath);

    if(inFile.is_open())
    {
        for(int row = 0; row < calMat.rows(); row++){
            for(int col = 0; col < calMat.cols(); col++)
            {
                double item = 0;
                inFile >> item;
                calMat(row, col) = item;
            }
        }
        inFile.close();
    }
    return calMat;
}

void save_calibration_file(std::string path, Eigen::Affine3d calMat) {
    std::ofstream outFile;
    outFile.open(path);

    if(outFile.is_open()){
        outFile << calMat.matrix();
        outFile.close();
    }
}

}


/**
 *  sim_carAbst.cpp
 *
 *  To simulate the vehicle example using abstraction-based control.
 *
 *  Authors: Yinan Li, Zhibing Sun, Jun Liu
 *  Created: May 27, 2020
 *
 *  Hybrid Systems Group, University of Waterloo.
 */

#include <iostream>
#include <fstream>
#include <sstream>
#include <cstdlib>

#include <string>
#include <math.h>
#include <boost/numeric/odeint.hpp>

#include "src/abstraction.hpp"
#include "src/DBAparser.h"
#include "src/bsolver.hpp"
#include "src/hdf5io.h"

/* define dynamics */
struct car_dynamics {
	rocs::Rn u;
	car_dynamics (const rocs::Rn param): u (param) {}

	void operator() (rocs::Rn &x, rocs::Rn &dxdt, double t) const
	{
	      double alpha = std::atan(std::tan(u[1])/2.0);
	      dxdt[0] = u[0]*std::cos(alpha+x[2])/std::cos(alpha);
	      dxdt[1] = u[0]*std::sin(alpha+x[2])/std::cos(alpha);
	      dxdt[2] = u[0]*std::tan(u[1]);
	}
};

int main(int argc, char *argv[])
{
    /**
     * Default arguments
     */
    std::string specfile{"dba1.txt"};
    double eta[]{0.2, 0.2, 0.2};

    /* Input arguments:
     * sim_abst specfile precision(e.g. 0.2 0.2 0.2)
     */
    if (argc > 2 && argc < 5) {
	std::cout << "Improper number of arguments. Input arguments are:\n";
	std::cout << "./sim_abst specfile precision(e.g. 0.2 0.2 0.2)\n";
	std::exit(1);
    }
    if(argc > 1) {
	specfile = std::string(argv[1]);
    }
    if (argc == 5) {
	for(int i = 2; i < 5; ++i)
	    eta[i-2] = std::atof(argv[i]);
    }
    std::vector<std::string> tokens;
    boost::split(tokens, specfile, boost::is_any_of("."));
    std::string ctlrfile = "controller_" + tokens[0] + "_";
    for(int i = 0; i < 3; ++i) {
	std::stringstream ss;
	ss << std::setprecision(1);
	ss << eta[i];
	ctlrfile += ss.str();
	if (i < 2)
	    ctlrfile += "-";
    }
    ctlrfile += ".h5";

    /**
     * Load specification
     **/
    std::cout << "\nReading the specification...\n";
    rocs::UintSmall nAP = 0, nNodes = 0, q0 = 0;
    std::vector<rocs::UintSmall> acc;
    std::vector<std::vector<rocs::UintSmall>> arrayM;
    if (!rocs::read_spec(specfile, nNodes, nAP, q0, arrayM, acc))
	std::exit(1);
    boost::dynamic_bitset<> isacc(nNodes, false);
    for (rocs::UintSmall i = 0; i < acc.size(); ++i)
	isacc[acc[i]] = true;

    /**
     * Setup the motion planning workspace
     **/
    /* set the state space */
    const int x_dim = 3;
    const double theta = 3.5;
    double xlb[] = {0, 0, -theta};
    double xub[] = {10, 10, theta};
    /* set the control values */
    const int u_dim = 2;
    double ulb[] = {-1.0, -1.0};
    double uub[] = {1.0, 1.0};
    double mu[] = {0.3, 0.3};
    /* Generate grids */
    rocs::grid x_grid(x_dim,eta,xlb,xub);
    x_grid.gridding();
    rocs::grid u_grid(u_dim,mu,ulb,uub);
    std::cout << "Number of discrete states: " << x_grid._nv << "\n";
    std::cout << "Number of discrete inputs: " << u_grid._nv << "\n";

    /**
     * Load global controller
     **/
    std::cout << "\nLoading the controller from " << ctlrfile << "...\n";
    std::vector<long long> w_x0, encode3;
    std::vector<NODE_POST> nts_ctrlr;
    std::vector<CTRL> ctrl;
    std::vector<int> q_prime;
    rocs::h5FileHandler planRdr(ctlrfile, H5F_ACC_RDONLY);
    planRdr.read_discrete_controller(w_x0, encode3, nts_ctrlr, ctrl, q_prime);

    /**
     * Simulation
     **/
    const double tsim = 0.3;  //simulation time step
    /* Open files for writing results and logs */
    std::string simfile = "sim_traj_"+tokens[0]+".txt";
    std::ofstream ctlrWtr(simfile);
    if(!ctlrWtr.is_open())
	ctlrWtr.open(simfile, std::ios::out);
	/********** Logging **********/
    std::ofstream logger;
    logger.open("logs_"+tokens[0]+".txt", std::ios::out);
    /********** Logging **********/

    /* ode solver */
    const double dt = 0.001; //integration step size for odeint
    boost::numeric::odeint::runge_kutta_cash_karp54<rocs::Rn> rk45;

    /* Initial condition */
    rocs::Rn x{3, 2, M_PI/2.};
    rocs::Rn u{0, 0};
    long long x_index = x_grid.val_to_id(x);
    std::vector<long long>::iterator p1;
    p1 = std::lower_bound(w_x0.begin(), w_x0.end(), x_index);
    if(*p1 != x_index) {
	std::cout << x[0] <<  " "  << x[1] << " " << x[2];
	std::cout << " is not in the winning set, change a new initial state." << std::endl;
	return 1;
    }

    float tpat = 0;
    int max_num_achieve_acc=5, max_num_iteration=500; //3000000;
    rocs::UintSmall q;
    int i, j;
    std::vector<CTRL>::iterator p7;
    NODE_POST p5;
    std::cout << "\nLaunching simulation...\n";
    for(q = q0, i = j = 0; i<max_num_achieve_acc && j<max_num_iteration; ++j) {
	/* Determine the control u by the product state (x_index, q) */
	p5 = nts_ctrlr[encode3[x_index]];
	p7 = std::lower_bound(ctrl.begin()+p5.pos, ctrl.begin()+p5.pos+p5.num_a,
			      q, [](const CTRL &item, const int val) {
				     return item.q < val;});
	if(p7->q != q) {
	    std::cout << j << "th iteration: " << "Reach accepting state " << i << " times.\n";
	    std::cout << "Automaton state: " << q << '\n';
	    std::cout << "Position in ctrl list: " << p5.pos << ", number of actions: " << p5.num_a << '\n';
	    std::cout << "Error in ctrl" << std::endl;
	    break;
	}
	u_grid.id_to_val(u, p7->u);

	/* Print to screen */
	std::cout << "# " << j << ":\n";
	std::cout << "x: [" << x[0] <<  ','  << x[1] << ',' << x[2] << "]\n";
	std::cout << "Nominal control: " << p7->u << '('
		  << '[' << u[0] << ',' << u[1] << "])\n";
	std::cout << "current dba state: " << q << "\n";

	/* Write to txt file */
	// ctlrWtr << j << ':';
	for(int d = 0; d < x_dim; ++d) {
	    ctlrWtr << x[d];
	    if(d<x_dim-1) {
		ctlrWtr << ',';
	    }
	}
	ctlrWtr << ';';
	for(int d = 0; d < u_dim; ++d) {
	    ctlrWtr << u[d];
	    if(d<u_dim-1) {
		ctlrWtr << ',';
	    }
	}
	ctlrWtr << '\n';

	/* Integrate the dynamics */
	boost::numeric::odeint::integrate_const(rk45, car_dynamics(u), x, 0.0, tsim, dt);
	if(x[2] > M_PI) //convert angle into [-pi, pi]
	    x[2] -= 2*M_PI;
	if(x[2] < -M_PI)
	    x[2] += 2*M_PI;

	/********** Logging **********/
    logger << "x,q,u=" << x_index << ' ' << q << ' ' << p7->u << "; encode3[x]=" << encode3[x_index];
	logger << "\np5=" << p5.num_a << ' ' << p5.label << ' ' << p5.pos;
	logger << "\np7=";
	for(std::vector<CTRL>::iterator p = ctrl.begin()+p5.pos;
		p!=ctrl.begin()+p5.pos+p5.num_a; ++p) {
		logger << p->q << ' ' << p->u << ';';
	}
	logger << "\n\n";
    /********** Logging **********/

	/* Update automaton state q */
	x_index = x_grid.val_to_id(x);
	q = q_prime[p5.label*nNodes + q];
	if(isacc[q])
	    std::cout << "******** achieve acc \'" << q << "\' " << ++i << " time ********" << std::endl;
    }

    ctlrWtr.close();
	logger.close();

    return 0;
}


/*
 *  Distributed under our modified Boost Software License.
 *  Version 1.0 (see accompanying file LICENSE)
 */
/**
 *  @file       Robot.hpp
 *  @author     Lydia Zoghbi, Ari Kupferberg
 *  @copyright  Copyright ARL 2019
 *  @date       10/15/2019
 *  @version    2.0
 *
 *  @brief      Header file for constructing a Robot class.
 *
 */

#ifndef INCLUDE_ROBOT_HPP_
#define INCLUDE_ROBOT_HPP_

#include <vector>
#include <boost/numeric/ublas/matrix.hpp>
#include <boost/math/constants/constants.hpp>
#include <RobotPosition.hpp>
#include <RobotPath.hpp>
#include <Point.hpp>
#include <PathPlanner.hpp>

/**
 *  @brief      Class for creating a Robot object containing useful information about the system
 */

class Robot {
 private:
   int numberLinks = 6;
   Point initialEEPosition;
   std::vector<double> initialJointAngles = {0, 0, 0, 0, 0, 0};
   std::vector<RobotPath> path;
    std::vector<std::vector<double>> dhParams{
     {770, M_PI/2.0, 750, 0},
     {1050, 0, 0, 0},
     {200, M_PI/2.0, 0, 0},
     {0, -M_PI/2.0, 1705, 0},
     {0, M_PI/2.0, 0, 0},
     {0, 0, 325, 0}};

 public:

  /**
   *  @brief    Constructor for class Robot
   *  @param	Point of the robot's end effector position
   *  @return	Instance of robot
   */
   explicit Robot(const Point& startingPos);

  /**
   *  @brief     Compute the intermediate transformation between two DH frames
   *  @param	 DH Table as matrix double
   *  @return	'A' transformation matrix
   */
   boost::numeric::ublas::matrix<double> computeATransform(std::vector<double> dhRow);

  /**
   *  @brief    Computing the forward kinematics for the Robot
   *  @param	Vector of Robot's joint angles as double
   *  @return	Vector of Point objects depicting Robot's joint positions
   */
   std::vector<Point> computeFk(std::vector<double> jointAngles);

  /**
   *  @brief    Computing a set of transformation matrices
   *  @param	Vector of Robot's joint angles
   *  @return	Vector of matrices representing the transformation from each joint frame to the base frame
   */
   std::vector<boost::numeric::ublas::matrix<double>> computeTransformationMatrices(std::vector<double> jointAngles);

  /**
   *  @brief    Computing the geometric Jacobian matrix
   *  @param	Robot's current position for performing a cross product
   *  @param    Vector of transformation matrices
   *  @return	The Jacobian matrix
   */
   boost::numeric::ublas::matrix<double> computeJacobian(RobotPosition robotPosition, std::vector<boost::numeric::ublas::matrix<double>> tTransforms);

  /**
   *  @brief    Computing the inverse kinematics for Robot's generated path
   *  @param	The end effector's target position
   *  @param    The environment to check for collisions
   *  @return	Vector of RobotPosition object containing Points of joint positions
   */
   std::vector<RobotPosition> computeIK(Point targetPoint, Environment environment);

  /**
   *  @brief    Computes cross product of the two input vectors
   *  @param	Vector to compute cross product of (1 of 2)
   *  @param	Vector to compute cross product of (2 of 2)
   *  @return	Vector of the resultant cross product
   */
   boost::numeric::ublas::vector<double> crossProduct(boost::numeric::ublas::vector<double> vector1, boost::numeric::ublas::vector<double> vector2);

  /**
   *  @brief    Computes the Moore-Penrose pseudoinverse
   *  @param	Input matrix to find the pseudoinverse of
   *  @return	Pseudoinverse matrix result
   */
   boost::numeric::ublas::matrix<double> penroseInverseMatrix(boost::numeric::ublas::matrix<double> mat);

};

#endif // INCLUDE_ROBOT_HPP_


///////////////////////////////////////////////////////////////////////////////
// importance_sampling::maximal_finite_sums.hpp                              //
//                                                                           //
//  Copyright 2009 Erwann Rogard. Distributed under the Boost                //
//  Software License, Version 1.0. (See accompanying file                    //
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)         //
///////////////////////////////////////////////////////////////////////////////
#ifndef BOOST_STATISTICS_DETAIL_IMPORTANCE_SAMPLING_WEIGHTS_MAXIMAL_FINITE_SUMS_HPP_ER_2009
#define BOOST_STATISTICS_DETAIL_IMPORTANCE_SAMPLING_WEIGHTS_MAXIMAL_FINITE_SUMS_HPP_ER_2009
#include <boost/iterator/iterator_traits.hpp>
#include <boost/math/special_functions/fpclassify.hpp>

namespace boost{
namespace statistics{
namespace detail{
namespace importance_sampling{

    // Breaks down the summation of *i, i in [b,e) into maximal finite amounts,
    // each of which are copied to the output iterator i
    //
    // Note: Primarily for use by find_scale_to_finite_sum
    template<typename InIt,typename OutIt>
    OutIt maximal_finite_sums(InIt b, InIt e, OutIt i)
    {
        typedef typename iterator_value<InIt>::type value_type;
        value_type sum = static_cast<value_type>(0);
        while(b!=e){
            value_type d = *b;
            if( boost::math::isinf( sum + d ) ){
                *i = sum;
                sum = d;
                ++i;
            }else{
                sum += d;
            }
            ++b;
        }
        return i;
    };

}// importance_weights
}// detail
}// statistics
}// boost

#endif


#include <jfs/differential_ops/grid_diff2d.h>

#include <Eigen/Eigen>

namespace jfs {

template <class SparseMatrix>
JFS_INLINE void gridDiff2D<SparseMatrix>::Laplace(const grid2D* grid, SparseMatrix &dst, unsigned int dims, unsigned int fields)
{
    auto btype = grid->bound_type_;
    auto L = grid->L;
    auto N = grid->N;
    auto D = grid->D;

    typedef Eigen::Triplet<float> T;
    std::vector<T> tripletList;
    tripletList.reserve(N*N*5*dims);

    for (int idx = 0; idx < dims*fields*N*N; idx++)
    {
        int idx_tmp = idx;
        int j = idx_tmp / (N*dims*fields);
        idx_tmp -= j * (N*dims*fields);
        int i = idx_tmp / (dims*fields);
        idx_tmp -= i * (dims*fields);
        int f = idx_tmp / dims;
        idx_tmp -= f * dims;
        int d = idx_tmp;

        int iMat, jMat;

        iMat = N*fields*dims*j + fields*dims*i + dims*f + d;
        jMat = iMat;
        tripletList.push_back(T(iMat,jMat,-4.f));

        int i_tmp = i;
        for (int offset = -1; offset < 2; offset+=2)
        {
            i = i_tmp + offset;
            if ( (i == -1 || i == N) && btype == ZERO)
                continue;
            else if ( i == -1 )
                i = (N-2);
            else if ( i == N )
                i = 1;
            jMat = N*fields*dims*j + fields*dims*i + dims*f + d;
            tripletList.push_back(T(iMat,jMat,1.f));
        }
        i = i_tmp;

        int j_tmp = j;
        for (int offset = -1; offset < 2; offset+=2)
        {
            j = j_tmp + offset;
            if ( (j == -1 || j == N) && btype == ZERO)
                continue;
            else if ( j == -1 )
                j = (N-2);
            else if ( j == N )
                j = 1;
            jMat = N*fields*dims*j + fields*dims*i + dims*f + d;
            tripletList.push_back(T(iMat,jMat,1.f));
        }
        j = j_tmp;
    }

    dst = SparseMatrix(N*N*dims*fields,N*N*dims*fields);
    dst.setFromTriplets(tripletList.begin(), tripletList.end());
    dst = 1.f/(D*D) * dst;
}

template <class SparseMatrix>
JFS_INLINE void gridDiff2D<SparseMatrix>::div(const grid2D* grid, SparseMatrix &dst, unsigned int fields)
{
    auto btype = grid->bound_type_;
    auto L = grid->L;
    auto N = grid->N;
    auto D = grid->D;

    typedef Eigen::Triplet<float> T;
    std::vector<T> tripletList;
    tripletList.reserve(N*N*2*2);

    int dims = 2;

    for (int idx = 0; idx < dims*fields*N*N; idx++)
    {
        int idx_tmp = idx;
        int j = idx_tmp / (N*dims*fields);
        idx_tmp -= j * (N*dims*fields);
        int i = idx_tmp / (dims*fields);
        idx_tmp -= i * (dims*fields);
        int f = idx_tmp / dims;
        idx_tmp -= f * dims;
        int d = idx_tmp;

        int iMat, jMat;

        iMat = N*fields*j + fields*i + f;

        int i_tmp = i;
        for (int offset = -1; offset < 2 && d == 0; offset+=2)
        {
            i = i_tmp + offset;
            if ( (i == -1 || i == N) && btype == ZERO)
                continue;
            else if ( i == -1 )
                i = (N-2);
            else if ( i == N )
                i = 1;
            jMat = N*fields*dims*j + fields*dims*i + dims*f + d;
            tripletList.push_back(T(iMat,jMat,(float) offset));
        }
        i = i_tmp;

        int j_tmp = j;
        for (int offset = -1; offset < 2 && d == 1; offset+=2)
        {
            j = j_tmp + offset;
            if ( (j == -1 || j == N) && btype == ZERO)
                continue;
            else if ( j == -1 )
                j = (N-2);
            else if ( j == N )
                j = 1;
            jMat = N*fields*dims*j + fields*dims*i + dims*f + d;
            tripletList.push_back(T(iMat,jMat,(float) offset));
        }
        j = j_tmp;
    }

    dst = SparseMatrix(N*N*fields,N*N*2*fields);
    dst.setFromTriplets(tripletList.begin(), tripletList.end());
    dst = 1.f/(2*D) * dst;
}

template <class SparseMatrix>
JFS_INLINE void gridDiff2D<SparseMatrix>::grad(const grid2D* grid, SparseMatrix &dst, unsigned int fields)
{
    auto btype = grid->bound_type_;
    auto L = grid->L;
    auto N = grid->N;
    auto D = grid->D;

    typedef Eigen::Triplet<float> T;
    std::vector<T> tripletList;
    tripletList.reserve(N*N*2*2);

    int dims = 2;

    for (int idx = 0; idx < dims*fields*N*N; idx++)
    {
        int idx_tmp = idx;
        int j = idx_tmp / (N*dims*fields);
        idx_tmp -= j * (N*dims*fields);
        int i = idx_tmp / (dims*fields);
        idx_tmp -= i * (dims*fields);
        int f = idx_tmp / dims;
        idx_tmp -= f * dims;
        int d = idx_tmp;

        int iMat, jMat;

        iMat = N*fields*dims*j + fields*dims*i + dims*f + d;

        int i_tmp = i;
        for (int offset = -1; offset < 2 && d == 0; offset+=2)
        {
            i = i_tmp + offset;
            if ( (i == -1 || i == N) && btype == ZERO)
                continue;
            else if ( i == -1 )
                i = (N-2);
            else if ( i == N )
                i = 1;
            jMat = N*fields*j + fields*i + f;
            tripletList.push_back(T(iMat,jMat,(float) offset));
        }
        i = i_tmp;

        int j_tmp = j;
        for (int offset = -1; offset < 2 && d == 1; offset+=2)
        {
            j = j_tmp + offset;
            if ( (j == -1 || j == N) && btype == ZERO)
                continue;
            else if ( j == -1 )
                j = (N-2);
            else if ( j == N )
                j = 1;
            jMat = N*fields*j + fields*i + f;
            tripletList.push_back(T(iMat,jMat,(float) offset));
        }
        j = j_tmp;
    }

    dst = SparseMatrix(N*N*dims*fields,N*N*fields);
    dst.setFromTriplets(tripletList.begin(), tripletList.end());
    dst = 1.f/(2*D) * dst;
}

// explicit instantiation of templates
#ifdef JFS_STATIC
template class gridDiff2D<Eigen::SparseMatrix<float, Eigen::ColMajor>>;
template class gridDiff2D<Eigen::SparseMatrix<float, Eigen::RowMajor>>;
#endif

} // namespace jfs