#ifndef SHARE_BITS_H
 #define SHARE_BITS_H

/*
  ビット分解されたシェアに対する基本的な操作
*/

typedef struct bits {
  int d; // 桁数
  _* a;
}* _bits;

typedef struct {
  _bits x;
  _ y;
} share_pair_bits;
#define _pair_bits share_pair_bits



// 定数 x をビット分解してシェアの配列にする
_bits share_bits_const(int d, int n, share_t q, share_t x) {
    share_t *x_bits;
    NEWA(x_bits, share_t, d);
    for (int i = 0; i < d; ++i) {
        x_bits[i] = (x>>i) & 1;
    }

    NEWT(_bits, b);
    b->d = d;
    NEWA(b->a, share_array, d);
    for (int j = 0; j < d; ++j) {
        b->a[j] = share_const(n, 0, q);
        //pa_iter_new(b->a[j]->A);
    }

#if 0
    for (int i = 0; i < n; ++i) {
        for (int j = 0; j < d; ++j) {
            share_setpublic(b->a[j], i, x_bits[j]); 
        }
    }
#else
    if (_party <= 1) {
      for (int j = 0; j < d; ++j) {
        pa_iter itr = pa_iter_new(b->a[j]->A);
        for (int i = 0; i < n; ++i) {
          //pa_iter_set(b->a[j]->A, x_bits[j]);
          pa_iter_set(itr, x_bits[j]);
        }
        //pa_iter_flush(b->a[j]->A);
        pa_iter_flush(itr);
      }
    }
#endif
    free(x_bits);

    return b;
}

// 平文の配列 A をビット分解してシェアにする
_bits share_bits_new(int d, int n, share_t q, share_t *A) {
    share_t *A_bits;

    NEWT(_bits, ans);
    ans->d = d;
    NEWA(ans->a, share_array, n);
    if (_party <= 0) {
        NEWA(A_bits, share_t, n);

        for (int j = 0; j < d; ++j) {
            for (int i = 0; i < n; ++i) {
                A_bits[i] = (A[i]>>j) & 1;
            }
            ans->a[j] = share_new(n, q, A_bits);
        }

        free(A_bits);
    } else if (1 <= _party && _party <= 2) {
        printf("share_bits_new: A_bits is not initialized for party %d\n", _party);
        for (int j = 0; j < d; ++j) {
            ans->a[j] = share_new(n, q, A_bits); // A_bits が初期化されていない
        }
    }

    return ans;
}


void _free_bits(_bits a)
{
//  if (_party >  2) return;
  for (int i=0; i<a->d; i++){
    _free(a->a[i]);
  }
  free(a->a);
  free(a);
}

////////////////////////////////////////////
// a[i] := b[j]
////////////////////////////////////////////
static void share_setshare_bits(_bits a, int i, _bits b, int j)
{
  if (_party >  2) return;
  if (i < 0 || i >= a->a[0]->n) {
    printf("share_setshare_bits a: n %d i %d\n", a->a[0]->n, i);
    exit(1);
  }
  if (j < 0 || j >= b->a[0]->n) {
    printf("share_setshare_bits b: n %d j %d\n", b->a[0]->n, j);
    exit(1);
  }
  if (a->a[0]->q != b->a[0]->q) {
    printf("share_setshare_bits a->q %d b->q %d\n", (int)a->a[0]->q, (int)b->a[0]->q);
    exit(1);
  }
  if (a->d != b->d) {
    printf("share_setshare_bits a->d %d b->d %d\n", (int)a->d, (int)b->d);
    exit(1);
  }
  for (int k=0; k<a->d; k++) {
    pa_set(a->a[k]->A,i, pa_get(b->a[k]->A,j));
  }
}
#define _setshare_bits share_setshare_bits

void share_setshares_bits(_bits x, int si, int ei, _bits y, int j) {
    if (x->d != y->d) {
        printf("share_setshres_bits x->d: %d    y->d: %d\n", x->d, y->d);
        exit(1);
    }
    if (order(x->a[0]) != order(y->a[0])) {
        printf("share_setshraes_bits: order(x->a[0]): %d    order(y->a[0]): %d\n", order(x->a[0]), order(y->a[0]));
        exit(1);
    }
    if (si > ei || si < 0 || ei > len(x->a[0])) {
        printf("share_setshares_bits si: %d ei: %d\n", si, ei);
        exit(1);
    }
    if (j + ei - si > len(y->a[0])) {
        printf("share_setshre_bits  si: %d  ei: %d  j: %d   len(y->a[0]): %d\n", si, ei, j, len(y->a[0]));
        exit(1);
    }

    for (int i = 0; i < ei - si; ++i) {
        share_setshare_bits(x, si + i, y, j + i);
    }
}

_bits share_reconstruct_bits_channel(_bits x, int channel) {
    if (_party > 2) {
        return NULL;
    }
    if (x == NULL) {
        return NULL;
    }

    int n = len(x->a[0]);
    int d = x->d;
    share_t q = order(x->a[0]);
    _bits ans = share_bits_const(d, n, q, 0);
    if (_party == 0) {
        share_setshares_bits(ans, 0, n, x, 0);
    }
    else {
        for (int k = 0; k < d; ++k) {   // TODO: ここを並列化
            mpc_exchange_channel(x->a[k]->A->B, ans->a[k]->A->B, pa_size(x->a[k]->A), channel);
            vadd_(ans->a[k], x->a[k]);
        }
    }

    return ans;
}



/////////////////////////////////////////////////
// ビット分解されている v に対し v[i] := x を行う
/////////////////////////////////////////////////
void _setpublic_bits(_bits v, int i, share_t x)
{
  if (_party >  2) return;
  for (int d=0; d<v->d; d++) {
    _setpublic(v->a[d], i, (x>>d) & 1);
  }
}

int len_bits(_bits b)
{
  return len(b->a[0]);
}

int order_bits(_bits b)
{
  return order(b->a[0]);
}

int depth_bits(_bits b)
{
  return b->d;
}

static _bits share_const_bits(int n, share_t v, share_t q, int d)
{
  if (_party >  2) return NULL;
  NEWT(_bits, ans);
  ans->d = d;
  NEWA(ans->a, _, d); 
  for (int i=0; i<d; i++) {
    ans->a[i] = _const(n, 0, q);
  }

  for (int i=0; i<n; i++) {
    _setpublic_bits(ans, i, v);
  }
  return ans;
}
#define _const_bits share_const_bits

static _bits share_const_bits_3party(int n, share_t v, share_t q, int d)
{
  if (_party >  3) return NULL;
  NEWT(_bits, ans);
  ans->d = d;
  NEWA(ans->a, _, d); 
  for (int i=0; i<d; i++) {
    ans->a[i] = share_const_type(n, 0, q, SHARE_T_SHAMIR);
  }

  for (int i=0; i<n; i++) {
    _setpublic_bits(ans, i, v);
  }
  return ans;
}
#define _const_bits_3party share_const_bits_3party


static _bits share_slice_bits(_bits a, int start, int end)
{
  if (_party >  2) return NULL;
  if (start < 0) start = a->a[0]->n + start;
  if (end <= 0) end = a->a[0]->n + end;
  if (start < 0 || start > a->a[0]->n) {
    printf("share_slice_bits n %d start %d\n", a->a[0]->n, start);
  }
  if (end < 0 || end > a->a[0]->n) {
    printf("share_slice_bits n %d end %d\n", a->a[0]->n, end);
  }
  _bits ans = _const_bits(end-start, 0, a->a[0]->q, a->d);
  for (int i=0; i<ans->a[0]->n; i++) _setshare_bits(ans, i, a, start+i);
  return ans;
}
#define _slice_bits share_slice_bits

static share_t debug_get_bits(_bits a, int i)
{
  _bits tmp = share_slice_bits(a, i, i+1);
  //_bits tmp2 = share_reconstruct_bits_channel(tmp, 0);
  share_t ans = 0;
  for (int d = tmp->d-1; d >= 0; d--) {
    ans <<= 1;
    ans += debug_get(tmp->a[d], 0);
  }
  _free_bits(tmp);
  return ans;
}

_ B2A_channel(_ a, share_t q, int channel); // in func.h
#ifndef B2A
 #define B2A(b, q) B2A_channel(b, q, 0)
#endif





//////////////////////////////////////////////////
// ビットごとのシェアから加法的シェアに変換
// 各桁のシェアの位数は元の位数と等しいと仮定
//（つまり繰り上がりはそのまま扱えば良い）
//////////////////////////////////////////////////
_ _B2A_bits(_bits b)
{
  if (_party >  2) return NULL;
  // 元の位数が 2 だとダメ
  if (b->a[0]->q == 2) {
    printf("B2A_bits: warning q = %d\n", b->a[0]->q);
  }
  _ ans = _const(b->a[0]->n, 0, b->a[0]->q);
//  printf("B2A: not supported yet\n");
//  exit(1);
  for (int k=b->d-1; k>=0; k--) {
    smul_(2, ans);
    vadd_(ans, b->a[k]);
  }
  return ans;
}

static _ B2A_GF(_bits x, share_t irr_poly)
{
  share_t q = 1 << x->d;
  int n = len(x->a[0]);
  //int n = x->a[0]->n;
  _ ans = share_const_type(n, 0, q, x->a[0]->type);
  ans->type = x->a[0]->type;
  ans->irr_poly = irr_poly;
  for (int i=0; i<ans->n; i++) {
    share_t z;
    z = 0;
    for (int j=x->d-1; j>=0; j--) {
    //  z <<= 1;
    //  z += share_getraw(x->a[j], i);
      share_t t = share_getraw(x->a[j], i);
      z ^= GF_mul(t, (1<<j), irr_poly);
    }
    share_setraw(ans, i, z);
  }
  return ans;
}

void _print_bits(_bits a)
{
//  if (_party >  2) return;
  for (int i=0; i<a->d; i++) {
    printf("i = %d ", i); _print(a->a[i]);
  }
}

char *share_print_bits_str(_bits a)
{
  char **bufs;
  NEWA(bufs, char*, a->d);
  int total = 0;
  for (int i=0; i<a->d; i++) {
    bufs[i] = share_print_str(a->a[i]);
    total += strlen(bufs[i]);
  }
  char *buf;
  NEWA(buf, char, total + 20*a->d);
  total = 0;
  int s;
  for (int i=0; i<a->d; i++) {
    s = sprintf(buf+total, "i = %d %s\n", i, bufs[i]);
    total += s;
  }
  buf[total] = 0;
  buf = realloc(buf, total+1);
  for (int i=0; i<a->d; i++) {
    free(bufs[i]);
  }
  free(bufs);
  return buf;
}


_bits _dup_bits(_bits a)
{
  if (_party >  2) return NULL;
  NEWT(_bits, ans);
  ans->d = a->d;
  NEWA(ans->a, _, a->d); 
  for (int i=0; i<a->d; i++) {
    ans->a[i] = _dup(a->a[i]);
  }
  return ans;
}

_bits _vconcat_bits(_bits low, _bits high)
{
  if (_party >  2) return NULL;
  NEWT(_bits, ans);
  ans->d = low->d + high->d;
  NEWA(ans->a, _, ans->d); 
  for (int i=0; i<low->d; i++) {
    ans->a[i] = _dup(low->a[i]);
  }
  for (int i=0; i<high->d; i++) {
    ans->a[low->d + i] = _dup(high->a[i]);
  }
  return ans;
}

_bits _concat_bits(_bits first, _bits second)
{
  if (_party >  2) return NULL;
  int d = first->d;
  if (d != second->d) {
    printf("concat_bits: depth %d %d\n", d, second->d);
    exit(1);
  }
  NEWT(_bits, ans);
  ans->d = d;
  NEWA(ans->a, _, d); 
  for (int i=0; i<d; i++) {
    ans->a[i] = _concat(first->a[i], second->a[i]);
  }
  return ans;
}

static void _print_debug_bits(_bits a)
{
  for (int i=0; i<a->d; i++) {
    _ tmp = _reconstruct(a->a[i]);
    printf("debug "); _print(tmp);
    _free(tmp);
  }
}

// この関数はaes.cの中からコピーしたものなので，二重定義に気をつける．
static void share_xor_bits(_bits x, int i, _bits y, int j)
{
  for (int d=0; d<x->d; d++) {
    share_t xtmp, ytmp;
    xtmp = share_getraw(x->a[d], i);
    ytmp = share_getraw(y->a[d], j);
    xtmp ^= ytmp;
    share_setraw(x->a[d], i, xtmp);
  }
}

static void share_xor_bits2(_bits x, int i, _bits y, int j) {
    for (int d = 0; d < min(x->d, y->d); ++d) {
        share_t xtmp, ytmp;
        xtmp = share_getraw(x->a[d], i);
        ytmp = share_getraw(y->a[d], j);
        xtmp ^= ytmp;
        share_setraw(x->a[d], i, xtmp);
    }
}

void bits_free(_bits x) {
    for (int i = 0; i < x->d; ++i)
        share_free(x->a[i]);
    free(x->a);
    free(x);
}

void _move_bits(_bits x, _bits y) {
    if (x == NULL)  return;
    if (y == NULL)  return;
    for (int i = 0; i < x->d; ++i) {
        share_free(x->a[i]);
    }
    free(x->a);
    *x = *y;
    free(y);
}

_bits vsub_bits(_bits x, _bits y) {
    if (len(x->a[0]) != len(y->a[0])) {
        printf("vsub_bits: len(x->a[0]) %d  len(y->a[0]) %d\n", len(x->a[0]), len(y->a[0]));
        exit(1);
    }
    if (x->d != y->d) {
        printf("vsub_bits: x->d %d  y->d %d\n", x->d, y->d);
        exit(1);
    }
    if (order(x->a[0]) != 2) {
        printf("total_xor_bits order(x->a[0]): %d\n", order(x->a[0]));
        exit(1);
    }
    if (order(y->a[0]) != 2) {
        printf("total_xor_bits order(y->a[0]): %d\n", order(y->a[0]));
        exit(1);
    }
    
    int n = len(x->a[0]);
    int d = x->d;
    _bits z = share_const_bits(n, 0, 2, d);
    share_setshares_bits(z, 0, n, x, 0);
    for (int i = 0; i < n; ++i) {
        share_xor_bits(z, i, y, i);
    }

    return z;
}

void vsub_bits_(_bits x, _bits y) {
    _bits tmp = vsub_bits(x, y);
    _move_bits(x, tmp);
}

_bits total_xor_bits(_bits x) {
    // _print_bits(x);
    if (order(x->a[0]) != 2) {
        printf("total_xor_bits order(x->a[0]): %d\n", order(x->a[0]));
        exit(1);
    }
    int n = len(x->a[0]);
    int d = x->d;

    _bits y = share_const_bits(1, 0, 2, d);
    // printf("y\n");
    // _print_bits(y);
    // printf("x\n");
    // _print_bits(x);
    for (int i = 0; i < n; ++i) {
        share_xor_bits(y, 0, x, i);
    }

    return y;
}

// TODO: 並列化？
_bits random_bits(int n, int d, share_t q) {
    if (_party > 2) {
        return NULL;
    }

    NEWT(_bits, ans);
    ans->d = d;
    NEWA(ans->a, share_array, d);
    share_t *A;
    NEWA(A, share_t, d);
    
    for (int i = 0; i < d; ++i) {
        if (_party <= 0) {
            for (int j = 0; j < n; ++j) {
                A[j] = RANDOM0(q);
            }
        }
        ans->a[i] = share_new(n, q, A);
    }

    free(A);

    return ans;
}

share_array serialize_bits(_bits x) {
    int d = x->d;
    int n = len(x->a[0]);
    share_t q = order(x->a[0]);
    share_array ans = share_const(d * n, 0, q);
    for (int i = 0; i < n; ++i) {
        for (int j = 0; j < d; ++j) {
            share_setshare(ans, i * d + j, x->a[j], i);
        }
    }
    return ans;
}

//_ IfThenElse_channel(_ f, _ a, _ b, int channel);

_bits IfThenElse_bits_channel(share_array e, _bits x, _bits y, int channel) {
    int n = len(x->a[0]);
    int q = order(x->a[0]);

    share_array serialized_x = serialize_bits(x);
    share_array serialized_y = serialize_bits(y);
    share_array extended_e = extend_share_array(x->d, e);

    share_array serialized_ans = IfThenElse_channel(extended_e, serialized_x, serialized_y, channel);

    _bits ans = share_bits_const(x->d, n, q, 0);
    for (int i = 0; i < n; ++i) {
        for (int k = 0; k < x->d; ++k) {
            share_setshare(ans->a[k], i, serialized_ans, x->d * i + k);
        }
    }

    share_free(serialized_x);
    share_free(serialized_y);
    share_free(extended_e);
    share_free(serialized_ans);

    // _print_bits(ans);

    return ans;
}

_ IfThen_b(_ c, _ x) {
    if (c->q != 2) {
        printf("IfThen_b: c->q = %d\n", c->q);
        exit(1);
    }
    _ c_ = B2A(c, x->q);
    _ ans = vmul(x, c_);
    _free(c_);

    return ans;
}


// この関数はtreefunction2.hからコピーしたものなので，二重定義に気をつける
_ IfThenElse2(_ c, _ x, _ y) {
    if (c->q != 2) {
        printf("IfThenElse2: c->q = %d\n", c->q);
        exit(1);
    }
    else if (x->q != y->q) {
        printf("IfThenElse2: x->q = %d y->q = %d\n", x->q, y->q);
        exit(1);
    }
    else if (x->n != y->n) {
        printf("IfThenElse2: x->n = %d y->n = %d\n", x->n, y->n);
        exit(1);
    }
    _ c_ = B2A(c, x->q);
    _ dif = vsub(x, y);
    vmul_(dif, c_);
    _ ans = _dup(y);
    vadd_(ans, dif);
    _free(c_);   _free(dif);

    return ans;
}

_bits IfThenElse2_bits(share_array e, _bits x, _bits y) {
    int n = len(x->a[0]);
    int q = order(x->a[0]);

    share_array serialized_x = serialize_bits(x);
    share_array serialized_y = serialize_bits(y);
    share_array extended_e = extend_share_array(x->d, e);

    share_array serialized_ans = IfThenElse2(extended_e, serialized_x, serialized_y);

    _bits ans = share_bits_const(x->d, n, q, 0);
    for (int i = 0; i < n; ++i) {
        for (int k = 0; k < x->d; ++k) {
            share_setshare(ans->a[k], i, serialized_ans, x->d * i + k);
        }
    }

    share_free(serialized_x);
    share_free(serialized_y);
    share_free(extended_e);
    share_free(serialized_ans);

    // _print_bits(ans);

    return ans;
}

_bits IfThen_b_bits(share_array e, _bits x) {
    int n = len(x->a[0]);
    int q = order(x->a[0]);

    share_array serialized_x = serialize_bits(x);
    share_array extended_e = extend_share_array(x->d, e);

    share_array serialized_ans = IfThen_b(extended_e, serialized_x);

    _bits ans = share_bits_const(x->d, n, q, 0);
    for (int i = 0; i < n; ++i) {
        for (int k = 0; k < x->d; ++k) {
            share_setshare(ans->a[k], i, serialized_ans, x->d * i + k);
        }
    }

    share_free(serialized_x);
    share_free(extended_e);
    share_free(serialized_ans);

    // _print_bits(ans);

    return ans;
}


///////////////////////////////////////////////////////////////////////
// x の連続する l 個の要素を各桁とする値を n/l 個作る．
///////////////////////////////////////////////////////////////////////
_bits vextend_share_array(int l, share_array x) 
{
  int n = len(x);
  int m = n/l;
  if (m*l != n) {
    printf("vextend_share_array: n = %d l = %d\n", n, l);
    exit(1);
  }
  share_t q = order(x);
  _bits ans = share_const_bits(m, 0, q, l);
  pa_iter *itr_ans;
  NEWA(itr_ans, pa_iter, l);
  for (int i=0; i<l; i++) {
    itr_ans[i] = pa_iter_new(ans->a[i]->A);
  }
  pa_iter itr_x = pa_iter_new(x->A);
  for (int i = 0; i < m; ++i) {
    for (int j = 0; j < l; j++) {
      pa_iter_set(itr_ans[j], pa_iter_get(itr_x));
    }
  }
  for (int i=0; i<l; i++) {
    pa_iter_flush(itr_ans[i]);
  }
  free(itr_ans);
  pa_iter_free(itr_x);
    
  return ans;
}

_bits extended_total_xor_bits(int l, _bits x) {
    int n = len(x->a[0]);
    int d = x->d;
    int q = order(x->a[0]);
    if (n % l != 0) {
        printf("extended_total_xor_bits n = %d  l = %d  n %% l = %d\n", n, l, n % l);
        exit(1);
    }
    _bits ans = share_bits_const(d, l, q, 0);
    for (int i = 0; i < n / l; ++i) {
        for (int j = 0; j < l; ++j) {
            share_xor_bits(ans, j, x, i * l + j);
        }
    }

    return ans;
}


////////////////////////////////////////////
// シェアを深さ 1 のビット分解したものに変換
// 元の変数は以後使用できない
////////////////////////////////////////////
_bits share_to_bits(_ a)
{
  NEWT(_bits, ans);
  NEWA(ans->a, _, 1);
  ans->d = 1;
  ans->a[0] = a;

  return ans;
}

_ bits_to_share(_bits b)
{
  _ ans = _B2A_bits(b);
  _free_bits(b);
  return ans;
}

_ bits_shrink(_bits b)
{
  int d = b->d;
  int n = len_bits(b);
  _ ans = _const(d * n, 0, order_bits(b));
  pa_iter *itr_b;
  NEWA(itr_b, pa_iter, d);
  for (int i=0; i<d; i++) {
    itr_b[i] = pa_iter_new(b->a[i]->A);
  }
  pa_iter itr_ans = pa_iter_new(ans->A);
  for (int i=0; i<n; i++) {
    for (int j=0; j<d; j++) {
      pa_iter_set(itr_ans, pa_iter_get(itr_b[j]));
    }
  }
  pa_iter_flush(itr_ans);
  for (int i=0; i<d; i++) {
    pa_iter_free(itr_b[i]);
  }
  free(itr_b);
  _free_bits(b);

  return ans;
}

#endif // SHARE_BITS_H
