import math, random 

# sample implementation of Elias-Fano code
# Use
#  ef_build(a,u,verbose)
# to compute the Elisa-Fano representation 
# of the integers in a (must be increasing an non -negative). 
# u is an upper bound to the max of a, if not given we set u=max(a)+1  
# 
# For testing purposes:
#   create_set(n,u)
# create a set containing n randomly selected 
# increasing integers in [0,u)


# return the bitarrays (represented as strings)
# L and H associated to the integers in a[]
def ef_build(a,u = -1, verbose=False):
  assert len(a)>0
  if u<0:
    u = 1 + max(a)
  else: 
    assert u > max(a), "Invalid u parameter"   
  n = len(a)
  log_u = bits(u)
  ell = int(math.ceil(math.log(u/n,2)))
  # number of necessary buckets 
  buckets = math.ceil(u/(2**ell))
  print(f"n={n}, u={u}, log(u)={log_u}, ell={ell}, #buckets={buckets}")
  # for debug purposes print binary representation of input
  if verbose:
    print("Binary representation:\n  ",list2bin(a,log_u))
  # compute L bitvector with ell less significant bits
  # of each entry
  Llist = [ dec2bin(i%(2**ell),ell) for i in a]
  if verbose:
    print("Less significant bits:\n  ",Llist) 
  # compute L
  L = "".join(Llist)
  assert(len(L)==n*ell)
  # compute bucket id for each value
  Hvalues = [i >> ell for i in a] # i//(2**ell)
  if verbose:
    print("Bucket ids:\n  ", Hvalues)
  # build H using array count() method
  H = ""
  for i in range(buckets):
    H += Hvalues.count(i)*"1" + "0"
  assert(len(H) == n+buckets)
  return L, H, ell

# access the i-th element of the original array
# the index i is 1-based
def access(L,H,ell,i):
  assert i>0, "Invalid index"
  lowbits = L[(i-1)*ell:i*ell]
  bucket = select(H,1,i) - i
  return bucket * 2**ell + bin2dec(lowbits)


# return the pair (i,v) where v is the i-th (1-based)
# element of the set and v is the smallest value
# in the set which is >= x
# return (-1,-1) is the set contains no value >= x
def nextGEQ(L,H,ell,x,verbose=True):
  assert x>=0
  bx = x >> ell         # bucket that would contain x 
  if bx>0: 
    j =select(H,0,bx)+1 # first bit of bucket bx>0
  else:
    j = 1               # bucket 0 starts at position 1 
  if(j<=0):             # j is invalid, ie x>= u
    return (-1,-1)      # there is no value => x   
  if verbose:  
    print(f"The first bit of bucket {bx} is in pos {j}") 
  # scan the elements of bucket bx, if any
  while H[j-1]=='1':
    i = j - bx       # there are i-1 1's in H[1..j]
    v = (bx<< ell) + bin2dec(L[(i-1)*ell:i*ell])
    if v>=x:         # found element >= x 
      return (i ,v)  # return its position and its value 
    j += 1
  # now we have reached the end of bucket bx
  assert(H[j-1]=='0')
  i = j - bx            # there are i-1 1's in H[1..j] 
  next1 = select(H,1,i) # position of the i-th 1  
  if next1 == -1:       # there no i-th element
    return (-1, -1)     # x is larger than all set elements 
  # the desired value is the one in position i+1        
  return i, access(L,H,ell,i) 



# --- auxiliary functions

# select function for a binary string
# return the index 1-based of the ith occurrence of d
def select(s,d,i):
  assert i>0, "Invalid index"
  assert d==0 or d==1, "Invalid bit"
  if d==0: c='0'
  else: c='1'
  # check there are at least i c's
  tot = s.count(c)
  if i>tot:
    return -1
  found = 0  
  for j in range(len(s)):
    if s[j]==c:
      found += 1
      if found==i:
        return j+1 # +1 because we are 1-based 
  assert False,"Something is wrong"
  

# create a sorted sequence of n values in [0,u)
def create_set(n,u):
  a = random.sample(range(u),n)
  a.sort()
  return a
  
# test nextGEQ() function on a random set  
def test_set(n,u):
  y = create_set(n,u)
  print(y)
  L,H,ell = ef_build(y,False)
  for j in range(0,y[0]+1):
    ij,vj = nextGEQ(L,H,ell,j,False)
    assert ij== 1
    assert vj == y[0]       
  for i in range(n-1):
    for j in range(y[i]+1,y[i+1]+1):
      ij,vj = nextGEQ(L,H,ell,j,False)
      assert ij== i+2, f"j={j} a"
      assert vj == y[i+1], f"j={j} b"
  for j in range(y[-1]+1,y[-1]+10):
    ij,vj = nextGEQ(L,H,ell,j,False)
    assert ij== -1
    assert vj == -1     
  print("Done!")


# convert n to binary using at least d bits  
def dec2bin(n,d=8):
  assert n>=0
  return f"{n:0{d}b}"

# convert a list of integers to binary using at least d bits
def list2bin(lis,d=8):
  return [dec2bin(i,d) for i in lis]

# convert from binary 
def bin2dec(s):
    return int(s, 2)



# number of bits required to encode k different values
def bits(k):
    assert k > 0
    if k == 1:
        return 0
    bits = 1
    values = 2  # 1 bit -> 2 values
    while values < k:
        bits += 1
        values *= 2
    return bits



# --- list creation for experiments

# give a list of pairs of positive integers 
# (i1,n1) (i2,n2) ... (ik,nk)
# creates a random integer array with n_h copies of i_h 
def random_ia(lis):
  a = []
  for (i,n) in lis:
    a += n*[i]
  random.shuffle(a)  # permutes a's elements in place
  return a

# creates a random array with  x^{-2} distribution
# with values between 1 and nmax 
def random_list2(n,nmax):
  a = [ (i, int(1+n/(i))) for i in range(1,nmax+1) ]
  return random_ia(a)