import numpy as np
import mpmath as mp
from mpmath import floor, ceil

mp.mp.dps = 40

def nearestInLattice(basis, pt_given, pt, i=0): # basis is a np array v1..n of a LLL basis of R^n, pt_given and pt are np arrays of floats pt1..n representing pt1 v1 + ..., and i is the number of positions gtd to be integers. need to be given pt_given = pt at start
	n = len(pt)
	# print "~~~~~~~~~~~~~~~~", i
	if i == n:
		return pt
	pti = pt[i]
	if floor(pti) == ceil(pti):
		return nearestInLattice(basis, pt_given, pt, i+1)
	ptL = []
	ptR = []
	for j in range(n):
		ptL += [pt[j]]
		ptR += [pt[j]]
	ptL[i] = floor(pti)
	ptR[i] = ceil(pti)
	# print "ptR = ", ptR, "ptL = ", ptL
	snapL = nearestInLattice(basis, pt_given, ptL, i+1)
	dL = np.linalg.norm(np.dot(pt_given-snapL, basis))
	snapR = nearestInLattice(basis, pt_given, ptR, i+1)
	dR = np.linalg.norm(np.dot(pt_given-snapR, basis))
	if dL < dR:
		return snapL
	return snapR

"""
basis = np.array([[1,0],[0.2,1]])
pt = np.array([0.49,0.5])
res = nearestInLattice(basis, pt, pt)
print "RESULT: ", res
"""

def proj(u,v):
	return u*mp.mpf(1)*np.dot(u,v)/np.dot(u,u)

def GramSchmidt(basis): # takes in numpy arrays
	n = len(basis[0])
	newBasis = []
	for i in range(len(basis)):
		v = basis[i]
		s = [0]*n
		for j in range(i):
			s += proj(newBasis[j],v)
		newBasis += [v-s]
	return newBasis

# m = np.array([[3,1],[2,2]])
# print GramSchmidt(m)

def u(b,b2,i,j):
	return mp.mpf(1)*np.dot(b[i],b2[j])/np.dot(b2[j],b2[j])

def hand_coded_det(M):
	n = len(M)
	assert 1 <= n <= 2, "higher dimensions not supported yet (numerical issues with numpy/mpmath)"
	if n == 1:
		return M[0]
	else:
		return M[0][0]*M[1][1] - M[0][1]*M[1][0]

def LLL(basis, delta=mp.mpf(0.75)): # only definitely will work for integer matrices
	assert hand_coded_det(basis) != 0, "not a basis!"
	n = len(basis)-1
	b = []
	for v in basis:
		b += [v]
	b2 = GramSchmidt(b)
	k = 1
	while k <= n:
		for j in range(k-1,-1,-1):
			if abs(u(b,b2,k,j)) > 0.5:
				b[k] = b[k] - round(u(b,b2,k,j))*b[j]
				b2 = GramSchmidt(b)
		if np.dot(b2[k], b2[k]) >= (delta-u(b,b2,k,k-1)**2)*np.dot(b2[k-1], b2[k-1]):
			k += 1
		else:
			bk = b[k]
			bk1 = b[k-1]
			b[k] = bk1
			b[k-1] = bk
			b2 = GramSchmidt(b)
			k = max(k-1,1)
	return np.array(b)

# M = np.array([[1,1,1],[-1,0,2],[3,5,6]])
# print LLL(M)


def LLLfloat(prec, basis, delta=mp.mpf(0.75)): # prec is an integer, will be an exponent
	b = np.vectorize(int)(np.vectorize(round)((mp.mpf(2)**(1.5*prec))*basis))
	return LLL(b,delta)/(mp.mpf(2)**(2*prec))

# M2 = np.array([[1.00003,0.99999,1],[-1,0.00063,2],[2.99998,5.00001,6]])
# print LLLfloat(5,M2)