#ifndef EMP_IKNP_H__
#define EMP_IKNP_H__
#include "emp-ot/co.h"
#include "emp-ot/cot.h"

namespace emp {

/*
 * IKNP OT Extension
 * [REF] Implementation of "Extending oblivious transfers efficiently"
 * https://www.iacr.org/archive/crypto2003/27290145/27290145.pdf
 *
 * [REF] With optimization of "More Efficient Oblivious Transfer and Extensions
 * for Faster Secure Computation" https://eprint.iacr.org/2013/552.pdf [REF]
 * With optimization of "Better Concrete Security for Half-Gates Garbling (in
 * the Multi-Instance Setting)" https://eprint.iacr.org/2019/1168.pdf
 */
template <typename T> class IKNP : public COT<T> {
public:
  using COT<T>::io;
  using COT<T>::Delta;

  OTCO<T> *base_ot = nullptr;
  bool setup = false, *extended_r = nullptr;

  const static int64_t block_size = 1024 * 2;
  block local_out[block_size];
  bool s[128], local_r[256];
  PRG prg, G0[128], G1[128];
  bool malicious = false;
  block k0[128], k1[128];
  IKNP(T *io, bool malicious = false) : malicious(malicious) { this->io = io; }
  ~IKNP() { delete_array_null(extended_r); }

  void setup_send(const bool *in_s = nullptr, block *in_k0 = nullptr) {
    setup = true;
    if (in_s == nullptr)
      prg.random_bool(s, 128);
    else
      memcpy(s, in_s, 128);

    if (in_k0 != nullptr) {
      memcpy(k0, in_k0, 128 * sizeof(block));
    } else {
      this->base_ot = new OTCO<T>(io);
      base_ot->recv(k0, s, 128);
      io->flush();
      delete base_ot;
    }
    for (int64_t i = 0; i < 128; ++i)
      G0[i].reseed(&k0[i]);

    Delta = bool_to_block(s);
  }

  void setup_recv(block *in_k0 = nullptr, block *in_k1 = nullptr) {
    setup = true;
    if (in_k0 != nullptr) {
      memcpy(k0, in_k0, 128 * sizeof(block));
      memcpy(k1, in_k1, 128 * sizeof(block));
    } else {
      this->base_ot = new OTCO<T>(io);
      prg.random_block(k0, 128);
      prg.random_block(k1, 128);
      base_ot->send(k0, k1, 128);
      io->flush();
      delete base_ot;
    }
    for (int64_t i = 0; i < 128; ++i) {
      G0[i].reseed(&k0[i]);
      G1[i].reseed(&k1[i]);
    }
  }
  void send_pre(block *out, int64_t length) {
    if (not setup)
      setup_send();
    int64_t j = 0;
    for (; j < length / block_size; ++j)
      send_pre_block(out + j * block_size, block_size);
    int64_t remain = length % block_size;
    if (remain > 0) {
      send_pre_block(local_out, remain);
      memcpy(out + j * block_size, local_out, sizeof(block) * remain);
    }
    if (malicious)
      send_pre_block(local_out, 256);
  }

  void send_pre_block(block *out, int64_t len) {
    block t[block_size];
    block tmp[block_size];
    int64_t local_block_size = (len + 127) / 128 * 128;
    io->recv_block(tmp, local_block_size);
    for (int64_t i = 0; i < 128; ++i) {
      G0[i].random_data(t + (i * block_size / 128), local_block_size / 8);
      if (s[i])
        xorBlocks_arr(t + (i * block_size / 128), t + (i * block_size / 128),
                      tmp + (i * local_block_size / 128),
                      local_block_size / 128);
    }

    sse_trans((uint8_t *)(out), (uint8_t *)t, 128, block_size);
  }

  void recv_pre(block *out, const bool *r, int64_t length) {
    if (not setup)
      setup_recv();

    block *block_r = new block[(length + 127) / 128];
    for (int64_t i = 0; i < length / 128; ++i)
      block_r[i] = bool_to_block(r + i * 128);
    if (length % 128 != 0) {
      bool tmp_bool_array[128];
      memset(tmp_bool_array, 0, 128);
      int64_t start_point = (length / 128) * 128;
      memcpy(tmp_bool_array, r + start_point, length % 128);
      block_r[length / 128] = bool_to_block(tmp_bool_array);
    }

    int64_t j = 0;
    for (; j < length / block_size; ++j)
      recv_pre_block(out + j * block_size, block_r + (j * block_size / 128),
                     block_size);
    int64_t remain = length % block_size;

    if (remain > 0) {
      recv_pre_block(local_out, block_r + (j * block_size / 128), remain);
      memcpy(out + j * block_size, local_out, sizeof(block) * remain);
    }

    if (malicious) {
      block local_r_block[2];
      prg.random_bool(local_r, 256);
      local_r_block[0] = bool_to_block(local_r);
      local_r_block[1] = bool_to_block(local_r + 128);
      recv_pre_block(local_out, local_r_block, 256);
    }

    delete[] block_r;
  }
  void recv_pre_block(block *out, block *r, int64_t len) {
    block t[block_size];
    block tmp[block_size];
    int64_t local_block_size = (len + 127) / 128 * 128;

    for (int64_t i = 0; i < 128; ++i) {
      G0[i].random_data(t + (i * block_size / 128), local_block_size / 8);
      G1[i].random_data(tmp, local_block_size / 8);
      xorBlocks_arr(tmp, t + (i * block_size / 128), tmp,
                    local_block_size / 128);
      xorBlocks_arr(tmp, r, tmp, local_block_size / 128);
      io->send_data(tmp, local_block_size / 8);
    }

    sse_trans((uint8_t *)(out), (uint8_t *)t, 128, block_size);
  }

  void send_cot(block *data, int64_t length) override {
    send_pre(data, length);
    if (malicious)
      if (!send_check(data, length))
        error("OT Extension check failed");
  }
  void recv_cot(block *data, const bool *b, int64_t length) override {
    recv_pre(data, b, length);
    if (malicious)
      recv_check(data, b, length);
  }

  /*
   *
   * [REF] Implementation of "Actively Secure OT Extension with Optimal
   * Overhead" https://eprint.iacr.org/2015/546.pdf
   */
  bool send_check(block *out, int64_t length) {
    block seed2, x, t[2], q[2], tmp[2];
    block chi[block_size];
    q[0] = q[1] = makeBlock(0, 0);
    io->recv_block(&seed2, 1);
    io->flush();
    PRG chiPRG(&seed2);

    for (int64_t i = 0; i < length / block_size; ++i) {
      chiPRG.random_block(chi, block_size);
      vector_inn_prdt_sum_no_red<block_size>(tmp, chi, out + i * block_size);
      q[0] = q[0] ^ tmp[0];
      q[1] = q[1] ^ tmp[1];
    }
    int64_t remain = length % block_size;
    if (remain != 0) {
      chiPRG.random_block(chi, block_size);
      vector_inn_prdt_sum_no_red(tmp, chi, out + length - remain, remain);
      q[0] = q[0] ^ tmp[0];
      q[1] = q[1] ^ tmp[1];
    }
    {
      chiPRG.random_block(chi, 256);
      vector_inn_prdt_sum_no_red<256>(tmp, chi, local_out);
      q[0] = q[0] ^ tmp[0];
      q[1] = q[1] ^ tmp[1];
    }

    io->recv_block(&x, 1);
    io->recv_block(t, 2);
    mul128(x, Delta, tmp, tmp + 1);
    q[0] = q[0] ^ tmp[0];
    q[1] = q[1] ^ tmp[1];

    return cmpBlock(q, t, 2);
  }
  void recv_check(block *out, const bool *r, int64_t length) {
    block select[2] = {zero_block, all_one_block};
    block seed2, x = makeBlock(0, 0), t[2], tmp[2];
    prg.random_block(&seed2, 1);
    io->send_block(&seed2, 1);
    io->flush();
    block chi[block_size];
    t[0] = t[1] = makeBlock(0, 0);
    PRG chiPRG(&seed2);

    for (int64_t i = 0; i < length / block_size; ++i) {
      chiPRG.random_block(chi, block_size);
      vector_inn_prdt_sum_no_red<block_size>(tmp, chi, out + i * block_size);
      t[0] = t[0] ^ tmp[0];
      t[1] = t[1] ^ tmp[1];
      for (int64_t j = 0; j < block_size; ++j)
        x = x ^ (chi[j] & select[r[i * block_size + j]]);
    }
    int64_t remain = length % block_size;
    if (remain != 0) {
      chiPRG.random_block(chi, block_size);
      vector_inn_prdt_sum_no_red(tmp, chi, out + length - remain, remain);
      t[0] = t[0] ^ tmp[0];
      t[1] = t[1] ^ tmp[1];
      for (int64_t j = 0; j < remain; ++j)
        x = x ^ (chi[j] & select[r[length - remain + j]]);
    }

    {
      chiPRG.random_block(chi, 256);
      vector_inn_prdt_sum_no_red<256>(tmp, chi, local_out);
      t[0] = t[0] ^ tmp[0];
      t[1] = t[1] ^ tmp[1];
      for (int64_t j = 0; j < 256; ++j)
        x = x ^ (chi[j] & select[local_r[j]]);
    }

    io->send_block(&x, 1);
    io->send_block(t, 2);
    io->flush();
  }
};

} // namespace emp
#endif
