#ifndef BATCHACCESS_H
 #define BATCHACCESS_H

#include "propagate.h"

/***********************************************************
# v はアクセスしたい配列．長さ U
# idx は v のアクセスしたい要素の添え字の配列を 1 進数表現にしたもの．
# idx の値の最大値は len(v) 未満 
def BatchAccessUnary(v, idx):
  U = len(v)
  N = len(idx)
  sigma = StableSort(idx)
  n = sum(idx)
  X = v + ([0] * (N-U))
  Y = AppPerm(X, sigma)
  Z = Propagate(vneg(idx), Y)
  W = AppInvPerm(Z, sigma)
  return W[U:]
***********************************************************/
_ BatchAccessUnary(_ v, _ idx)
{
  if (_party >  2) return NULL;
  int U = len(v);
  int N = len(idx);
  _ idx_tmp = _shrink(idx, 2);
  _ sigma = StableSort2(idx_tmp);
  printf("sigma "); _print(sigma);
//  _ n = sum(idx); // 使ってない?
  _ zeros = _const(N-U, 0, order(v));
  _ X = _concat(v, zeros);
  _ Y = AppPerm(X, sigma);
  _ nidx = vneg(idx);
  //printf("nidx "); _print(nidx);
  _move_(nidx, _shrink(nidx,2));
  _ Z = Propagate2(nidx, Y);
  //_ Z = Propagate(nidx, Y);
  _ W = AppInvPerm(Z, sigma);
  _ ans = _slice(W, U, N);

  _free(sigma);
//  _free(n);
  _free(X);
  _free(Y);
  _free(nidx);
  _free(Z);
  _free(W);
  _free(zeros);

  return ans;
}


/************************************************
def Unary(x, U):
  N = len(x)
  i = 0
  X = [0] * (N+U)
  print('N', N, 'U', U)
  while i < U:
    X[i] = (i, 0, i) # (x, b, pos)
    i += 1
  i = 0
  while i < N:
    X[i+U] = (x[i], 1, i+U)
    i += 1
  print('X1', X)
  X.sort() # log N rounds
  print('X2', X)
  B = [b for (x, b, i) in X]
##  sigma = [i for (x, b, i) in X]
##  print(B)
  return B
************************************************/
_ Unary(_ x, int U)
{
  if (_party >  2) return NULL;
  int N = len(x);
  share_t q = order(x);
  _ X = _const(N+U, 0, q);
//  _ Y = _const(N+U, 0, N+U+1);
  int d = blog(N+U+1)+1;
  _ Y = _const(N+U, 0, 1<<d);
//  _ Z = _const(N+U, 0, N+U+1);
  for (int i=0; i<U; i++) {
    _setpublic(X, i, i);
    _setpublic(Y, i, 0);
//    _setpublic(Z, i, i);
  }
  for (int i=0; i<N; i++) {
    _setshare(X, i+U, x, i);
    _setpublic(Y, i+U, 1);
//    _setpublic(Z, i+U, i+U);
  }
  _pair tmp = share_radix_sort(X);
  _ sigma = _move(tmp.y);
  _ ans = AppPerm(Y, sigma);
  _free(sigma);
  _free(X);
  _free(Y);
//  _free(Z);
  _free(tmp.x);

  return ans;

}

_pair Unary2(_ x, int U)
{
  if (_party >  2) {
    _pair ans = {NULL, NULL};
    return ans;
  }
//  printf("Unary2 x "); _print(x);
  int N = len(x);
  share_t q = order(x);
  _ X = _const(N+U, 0, q);
  int d = blog(N+U+1)+1;
  _ Y = _const(N+U, 0, 1<<d);
  for (int i=0; i<U; i++) {
    _setpublic(X, i, i);
    _setpublic(Y, i, 0);
  }
  for (int i=0; i<N; i++) {
    _setshare(X, i+U, x, i);
    _setpublic(Y, i+U, 1);
  }
//  printf("X "); _print(X);
  _pair tmp = share_radix_sort(X);
  _pair ans;
//  printf("tmp.x "); _print(tmp.x);
//  printf("tmp.y "); _print(tmp.y);
  ans.x = AppPerm(Y, tmp.y);
//  printf("ans.x "); _print(ans.x);

  _ sigma = StableSort(ans.x);
//  printf("sigma "); _print(sigma);
  _ qq = AppInvPerm(tmp.y, sigma);
//  printf("qq "); _print(qq);
  _ rho = _slice(qq, U, U+N);
  for (int i=0; i<N; i++) {
    _addpublic(rho, i, -U);
  }
//  printf("rho "); _print(rho);
  ans.y = rho;

  _free(X);
  _free(Y);
//  _free(Z);
  _free(tmp.x);
  _free(tmp.y);
  _free(sigma);
  _free(qq);

  return ans;

}

//////////////////////////////////////////////////////
// x が xor share の場合
//////////////////////////////////////////////////////
_pair Unary_xor(_ x, int U)
{
  if (_party >  2) {
    _pair ans = {NULL, NULL};
    return ans;
  }
//  printf("Unary2 x "); _print(x);
  int N = len(x);
  share_t q = order(x);
  _ X = _const(N+U, 0, q);
  int d = blog(N+U+1)+1;
  _ Y = _const(N+U, 0, 1<<d);
  for (int i=0; i<U; i++) {
    _setpublic(X, i, i);
    _setpublic(Y, i, 0);
  }
  for (int i=0; i<N; i++) {
    _setshare(X, i+U, x, i);
    _setpublic(Y, i+U, 1);
  }
//  printf("X "); _print(X);
  _pair tmp = share_radix_sort_xor(X);
  _pair ans;
//  printf("tmp.x "); _print(tmp.x);
//  printf("tmp.y "); _print(tmp.y);
  ans.x = AppPerm(Y, tmp.y);
//  printf("ans.x "); _print(ans.x);

  _ sigma = StableSort(ans.x);
//  printf("sigma "); _print(sigma);
  _ qq = AppInvPerm(tmp.y, sigma);
//  printf("qq "); _print(qq);
  _ rho = _slice(qq, U, U+N);
  for (int i=0; i<N; i++) {
    _addpublic(rho, i, -U);
  }
//  printf("rho "); _print(rho);
  ans.y = rho;

  _free(X);
  _free(Y);
//  _free(Z);
  _free(tmp.x);
  _free(tmp.y);
  _free(sigma);
  _free(qq);

  return ans;

}

_pair Unary_bits(_bits x, int U)
{
  if (_party >  2) {
    _pair ans = {NULL, NULL};
    return ans;
  }

  int N = len(x->a[0]);
  share_t q = order(x->a[0]);
  int d = blog(N+U-1)+1;
  _bits X = _const_bits(N+U, 0, 1<<d, x->d);
  _ Y = _const(N+U, 0, 1<<d);
  for (int i=0; i<U; i++) {
    _setpublic_bits(X, i, i);
    _setpublic(Y, i, 0);
  }
  for (int i=0; i<N; i++) {
    _setshare_bits(X, i+U, x, i);
    _setpublic(Y, i+U, 1);
  }
//  printf("X\n"); _print_bits(X);
  _ tmpy = share_radix_sort_bits(X);
  _pair ans;
//  printf("tmpy "); _print(tmpy);
//  _ tmpy_inv = AppInvPerm(Perm_ID2(N+U, order(tmpy)), tmpy);
  _ tmpy_inv = InvPerm(tmpy);
//  printf("tmpy_inv "); _print(tmpy_inv);
  ans.x = AppPerm(Y, tmpy_inv);
//  printf("ans.x "); _print(ans.x);

  _ sigma = StableSort(ans.x);
//  printf("sigma "); _print(sigma);
  _ qq = AppInvPerm(tmpy_inv, sigma);
//  printf("qq "); _print(qq);
  _ rho = _slice(qq, U, U+N);
  for (int i=0; i<N; i++) {
    _addpublic(rho, i, -U);
  }
//  printf("rho "); _print(rho);
  ans.y = rho;

  _free(sigma);
  _free_bits(X);
  _free(Y);
  _free(qq);
  _free(tmpy);
  _free(tmpy_inv);

  return ans;

}



/************************************************
def BatchAccess(v, idx):
  I = Unary(idx, len(v))
  print('idx', idx, 'unary', I)
  return BatchAccessUnary(v, I)
************************************************/
// idx は単調増加である必要は無い
_ BatchAccess(_ v, _ idx)
{
  if (_party >  2) return NULL;
//  _ I = Unary(idx, m);
  _pair tmp = Unary2(idx, len(v));
  _ I = tmp.x;
  _ sigma = tmp.y;
//  printf("BatchAccess idx "); _print(idx);
//  printf("BatchAccess I "); _print(I);
//  printf("BatchAccess sigma "); _print(sigma);
  _ ans = BatchAccessUnary(v, I);
//  printf("BatchAccess: ans "); _print(ans); 
//  _pair tmp = share_radix_sort(I);
//  _ ans2 = AppInvPerm(ans, tmp.y);
  _ ans2 = AppInvPerm(ans, sigma); // idx で指定した順番に並び替える
//  printf("BatchAccess: ans2"); _print(ans2); 
  _free(I);
  _free(sigma);
  _free(ans);
  return ans2;
}

_ BatchAccess_xor(_ v, _ idx)
{
  if (_party >  2) return NULL;
//  _ I = Unary(idx, m);
  _pair tmp = Unary_xor(idx, len(v));
  _ I = tmp.x;
  _ sigma = tmp.y;
//  printf("BatchAccess idx "); _print(idx);
//  printf("BatchAccess I "); _print(I);
//  printf("BatchAccess sigma "); _print(sigma);
  _ ans = BatchAccessUnary(v, I);
//  printf("BatchAccess: ans "); _print(ans); 
//  _pair tmp = share_radix_sort(I);
//  _ ans2 = AppInvPerm(ans, tmp.y);
  _ ans2 = AppInvPerm(ans, sigma); // idx で指定した順番に並び替える
//  printf("BatchAccess: ans2"); _print(ans2); 
  _free(I);
  _free(sigma);
  _free(ans);
  return ans2;
}


_bits BatchAccess_bits(_bits v, _ idx)
{
  if (_party >  2) return NULL;
//  _ I = Unary(idx, m);
  _pair tmp = Unary2(idx, len(v->a[0]));
  _ I = tmp.x;
  _ sigma = tmp.y;
//  printf("idx "); _print(idx);
//  printf("unary "); _print(I);
//  printf("sigma "); _print(sigma);
  int d = v->d;
  int n = v->a[0]->n;
  share_t q = v->a[0]->q;
  _bits ans = _const_bits(n, 0, q, d);

//  _pair tmp3 = share_radix_sort(idx);
//  printf("sigma "); _print(sigma);
//  printf("tmp3.y "); _print(tmp3.y);

  for (int i=0; i<d; i++) {
    _ tmp2 = BatchAccessUnary(v->a[i], I);
  //  _move_(ans->a[i], AppInvPerm(tmp2, tmp3.y));
    _move_(ans->a[i], AppInvPerm(tmp2, sigma));
    _free(tmp2);
  }
  _free(sigma);
  _free(I);
  return ans;
}

_bits BatchAccess_bits_bits(_bits v, _bits idx)
{
  if (_party >  2) return NULL;

  _pair tmp = Unary_bits(idx, len(v->a[0]));
  _ I = tmp.x;
  _ sigma = tmp.y;
//  printf("b idx "); _print_bits(idx);
//  printf("b unary "); _print(I);
//  printf("b sigma "); _print(sigma);

  int d = v->d;
  int n = v->a[0]->n;
  share_t q = v->a[0]->q;
  _bits ans;
  ans = _const_bits(n, 0, q, d);

  for (int i=0; i<d; i++) {
    _ tmp2 = BatchAccessUnary(v->a[i], I);
    _move_(ans->a[i], AppInvPerm(tmp2, sigma));
    _free(tmp2);
  }
  _free(I);
  _free(sigma);
  return ans;
}



#endif
