# CORVIDS implementation, R. Perry, July 2018
#
# reference: https://psyarxiv.com/7shr8
######################################################################
from math import sqrt
from sympy import Matrix, zeros, ones
import diophantine
from util import Hermite, inc, nonneg, avg, table, stats
from plot import plot
Dio = diophantine.Diophantine()
######################################################################
#
# define n, m, v, m_tol, v_tol, vals

v = m_tol = v_tol = 0

# choose one:

if False: # from CORVIDS paper, https://psyarxiv.com/7shr8

  vals = [i+1 for i in range(7)] # 1...7

  n = 20; m = 1.85; s = 0.875094 # skewed data

  # n = 20; m = 3.2; s = 1.43637; v_tol = 0.01 # data with variance range

  # n = 25; m = 2; s = 0.875094 # data with impossible values

  # n = 20; m = 2; s = 0.875094; v_tol = 0.01 # data with variance range and impossible values
  # vals = [i for i in range(2,8)] # 2...7

  v = s*s

elif True: # CATS examples, no variance given

  vals = [i+1 for i in range(5)] # 1...5

  # n = 3; m = 5

  # n = 4; m = 4.8; # actual m = 4.75

  n = 15; m = 4.5; m_tol = 0.05; # actual m = (9*5+4*4+2*3)/15

  # n = 63; m = (17*5+25*4+15*3+5*2+1*1)/63
  #
  ## for i in range(len(xTlist)):
  ##   if xTlist[i] == [1,5,15,25,17]: print(i)
  ## 2591

elif False: # carrots example from SPRITE paper, https://peerj.com/preprints/26968/

  # actual reported data:
  # vals = [i for i in range(0,56)] # 0...55
  # n = 45; m = 19.4; v = 19.9**2;
  # m_tol = 0.0; v_tol = 0 # no solutions
  # m_tol = 1.0; v_tol = 0 # no solutions
  # m_tol = 0.0; v_tol = 10 # has solutions, rankR = 53, exhaustive search space 46**53 ~= 2**293, infeasible

  # using multiples of 5 for the numbers of carrots:
  vals = [i*5 for i in range(0,12)] # 0, 5, 10, ..., 55
  n = 45; m = 20; v = 20**2; m_tol = 0.0; v_tol = 0
  # has solutions, rankR = 9, exhaustive search space 46**9 ~= 2**50

else: # generated data

  # X = [5, 1, 1]
  X = [1, 2, 3, 1]

  vals = [i+1 for i in range(len(X))];

  n = sum(X);  (m,unusedv) = stats( X, vals);

  print( "n =", n, ", X =", X)

######################################################################
# calculate nm_min, nm_max, p_min, p_max, Arows, Acols
#
nm_min = int(round(n*(m-m_tol)))
nm_max = int(round(n*(m+m_tol)))

nn = (n-1)*n*n

if v_tol == 0:
  vinit = v
  p_min = p_max = 0
else:
  k = n*(m-int(m))
  vinit = 0 if k == 0 or v == 0 else ((n-k)*(m-int(m))**2+k*(int(m)+1-m)**2)/(n-1)
  p_min = int(round((v-v_tol-vinit)*(n-1)/2))
  p_max = int(round((v+v_tol-vinit)*(n-1)/2))

# print( "p_min =", p_min, ", p_max =", p_max)

Arows = 2 if v == 0 else 3

Acols = len(vals)

######################################################################
# solve over the ranges of nm and p
#
xTlist = []
cTlist = []

for nm in range(nm_min,nm_max+1):

  for p in range(p_min,p_max+1):

    t = vinit + p*2/(n-1)

    vnn = int(round(t*(n-1)*n*n))

    print( "nm =", nm, ", m =", nm/n, ", vnn =", vnn, ", v =", t, ", s =", sqrt(t))

    # create A, b
    #
    A = ones(Arows,Acols)

    A[1,:] = Matrix(vals).T

    if Arows == 3:
      A[2,:] = Matrix(1,Acols,lambda i,j: (nm-n*vals[j])**2 )
      b = Matrix([n, nm, vnn])
    else:
      b = Matrix([n, nm])

    # solve equation Ax=b, get x and basis for nullspace of A
    #
    # exec(open("solve.py").read())

    G, H, P, rank = Hermite( Dio, A, b)
    
    xT = -P[rank-1, :-1];  x = xT.T
    
    if A*x != b:
      print("Warning: no solution, A*x != b")
      continue
    
    if rank > Acols:
      print("Warning: no nullspace")
      if nonneg( xT, Acols):
        zTint = [int(val) for val in xT] # for plot, want int, not sym
        xTlist.append(zTint)
        print(zTint)
      continue

    # reduce N(A) basis, HR ~= [ I, xxx...]
    #
    NAT  = P[rank:, :-1];  NA = NAT.T

    HR, PR, rankR = Dio.lllhermite(NAT)
    
    print( "rankR =", rankR)

    # work with rows of transposed matrices for convenience
    
    # make rankR components of x zero
    #
    yT = xT
    for i in range(rankR):
      j = i
      while HR[i,j] == 0:
        ++j
      d = HR[i,j]
      if d != 1:  print( "Warning: HR[",i,",",j,"] =",d," != 1")
      yT -= (yT[j]/d)*HR[i,:]
    top = j+1

    # split off top and bottom parts of y and H
    #
    bot = Acols - top
    yTtop = yT[:,0:top] # might not be all 0's if top != rankR
    yTbot = yT[:,top:]
    HRtop = HR[:,0:top]
    HRbot = HR[:,top:]

    # if top != rankR could have negative values in yTtop
    #
    if not nonneg(yTtop,top):
      print("Warning: yTtop has negative values!")
      print("yT =", yT, ", HR =", HR)
      continue

    # find solutions for x with no negative elements
    #
    cT = zeros(1,rankR)
    
    search_count = 1;  solution_count = 0;  display_count_limit = 20

    zTbot = Matrix(yTbot);  valid = nonneg(zTbot,bot)
    
    while True:
      if valid > 0:
        zT = Matrix.hstack(yTtop+cT*HRtop,zTbot)
        zTint = [int(val) for val in zT] # for plot, want int, not sym
        xTlist.append(zTint)
        cTlist.append([int(val) for val in cT])
        if solution_count < display_count_limit: print(zTint);
        if solution_count == display_count_limit: print("...")
        solution_count += 1
      (count,valid) = inc(cT,rankR,n,zTbot,yTbot,HRbot,bot) # sets zTbot
      search_count += count
      if valid == 0: break
    
    print( "search_count =", search_count, ", solution_count =", solution_count)