import numpy as np
import numpy.random as rn
from constants import *

def encounter_map(w,h,*args,**kwargs):
	if kwargs.get('cave'):
		init = rn.randint(2,h-2)	
	else:
		init = rn.randint((h/2)-5,(h/2)+5)

	mx = np.zeros((h,w),dtype=np.uint8)
	mx[init,0] = 1
	jumplist = []
	jumped = 0 
	for col in xrange(w-1):
		for row in xrange(2,h-2):
			try:
				if mx[row,col]:
					rnd = rn.rand()
					mx[row,col+1] = 1 if rnd < 0.8 else 0
					if rn.rand() < 0.05:
						if jumped : 
							jumped = 0
							continue
						nr = rn.randint(-6,7)
						if (h-2) <= (nr + row) <= 2 :
							nr *= -1
						newrow = nr + row
						low,high = (row,newrow) if row < newrow else (newrow,row)
						jumplist.append((col,low,high))
						jumped = 1
						mx[newrow,col+1]= 1
				else:
					if mx[row,col-1]:
						mx[row,col+1] = 1 if rnd < 0.4 else 0
						rnd = rn.rand()
						r = 2 if rnd <= 0.5 else -2
						nr = row + r
						if nr < 2:
							nr += 6
						elif nr > h-2:
							nr -= 6
						low,high = (row,nr) if row < nr else (nr,row)
						jumplist.append((col,low,high))
						mx[nr,col] = 1
			except IndexError, e:
				pass
	strel = np.ones((3,4),dtype=np.uint8)
	for col,low,high in jumplist:
		mx[low:high,col] = 1
	out = dilate(mx,strel)
	strel = np.ones((3,2),dtype=np.uint8)
	out = imopen(flip(out),strel)#dilate(flip(out),strel)

	out[0,0:-1]=1
	if kwargs.get('cave'):
		out[-1,0:-1]=1
		out[0:-1,-1]=1
	else:
		for col in xrange(w):
			for row in xrange(h-2):
				if out[row,col]: 
					out[row,col] = 0
					if not out[row+1,col]:
						break
	return np.rot90(np.rot90(np.rot90(out*2)))

def goal_location(mx,grav_ori):

	si,sj = mx.shape
	itemlist = []
	placed = 0
	for i in xrange(si-2,2,-1):
		if placed: break
		for j in xrange(2,sj-2,1):
			try:
				if mx[i,j]==2: continue
				if grav_ori == ORIENTATION_UP:
					if mx[i,j-1] == 2:
						itemlist.append((i,j))
						placed = 1
						break
				elif grav_ori == ORIENTATION_DOWN:
					if mx[i,j+1] == 2:
						itemlist.append((i,j))
						placed = 1
						break
				elif grav_ori == ORIENTATION_LEFT:
					if mx[i-1,j] == 2:
						itemlist.append((i,j))
						placed = 1
						break
				elif grav_ori == ORIENTATION_RIGHT:
					if mx[i+1,j] == 2:
						itemlist.append((i,j))
						placed = 1
						break
			except IndexError, e:
				print e, i, j

	for _ in xrange(500):
		i=rn.randint(2,si-2)
		j=rn.randint(2,sj-2)
		if mx[i,j]==2: continue

		if grav_ori == ORIENTATION_UP:
			if mx[i,j-1] == 2:
				itemlist.append((i,j))
		elif grav_ori == ORIENTATION_DOWN:
			if mx[i,j+1] == 2:
				itemlist.append((i,j))
		elif grav_ori == ORIENTATION_LEFT:
			if mx[i-1,j] == 2:
				itemlist.append((i,j))
		elif grav_ori == ORIENTATION_RIGHT:
			if mx[i+1,j] == 2:
				itemlist.append((i,j))
	for i,j in itemlist:
		mx[i,j] = 1
		
	rn.shuffle(itemlist)
	return itemlist, mx

def encounter_map_old(w,h,seedcount=60,its=4):
	mx = np.zeros((h,w),dtype=np.uint8)
	strel = np.array([[1,1,1,1,1,1]],dtype=np.uint8)
	strel2 = np.ones((2,8),dtype=np.uint8)
	sti,stj = strel2.shape
	hseeds = rn.random_integers(1,h-(sti/2),seedcount)
	wseeds = rn.random_integers(1,w-(stj/2),seedcount)
	seeds = zip(hseeds,wseeds)
	for i,j in seeds:
		mx[i,j] = 1
	mx = dilate(mx,strel)
	def iter(mx,cnt,cmax):
		out = np.zeros((h,w),dtype=np.uint8)
		for i in xrange(1,h-1):
			for j in xrange(1,w-1):
				p=0
				nh = mx[i-1:i+1,j-1:j+1]
				s = np.sum(nh)
				if mx[i,j] == 0:
					if mx[i-1,j] ==1 and mx[i+1,j]==1 and not cnt == cmax:
						p=1
					elif mx[i,j-1] == 1:
						p=1
					elif mx[i,j-1] == 1 and mx[i,j+1] ==0:
						p=1
					elif s > 4:
						p=1
				else:
					p=1
				out[i,j] = p
		return out
	for n in xrange(its):
		mx = iter(mx,n,its-1)

	strel2 = np.ones((2,rn.randint(4,12)),dtype=np.uint8)
	mx = dilate(mx,strel2)
	strel2 = np.ones((2,rn.randint(4,12)),dtype=np.uint8)
	mx = dilate(mx,strel2)
	strel2 = np.ones((2,rn.randint(4,12)),dtype=np.uint8)
	mx = dilate(mx,strel2)
	strel2 = np.ones((2,2),dtype=np.uint8)
	mx = erode(mx,strel2)
	mx[0,0:-1]=1
	mx[-1,0:-1]=1
	mx[0:-1,0]=1
	mx[0:-1,-1]=1
	return mx

def newmap(w,h, prob):
	mx = np.zeros((h,w),dtype=np.uint8)
	r = [randpick(prob) for _ in xrange((w-2)*(h-2))]
	mx[1:-1,1:-1] = np.reshape(np.array(r),(h-2,w-2))
	return mx

def randpick(prob):
	return 1 if np.random.rand() <= prob else 0

def binmap(mx):
	s=''
	for r in mx:
		s='\n'.join([s, ''.join(['#' if x else ' ' for x in r])])
	return s
def nummap(mx):
	s=''
	for r in mx:
		s='\n'.join([s, ''.join(['[01;3' + str(x) + "m" + str(x) +"[00m" if x else '.' for x in r])])
	return s

# pretty print numerical map
def pp(mx):
	print binmap(mx)
# pretty print numerical map
def pp2(mx):
	print nummap(mx)

def iterate(mx):
	si,sj = mx.shape
	mxt = np.ones((si,sj),dtype=np.uint8)
	for i in xrange(1,si-1):
		for j in xrange(1,sj-1):
			nh = mx[i-1:i+2,j-1:j+2]
			s = np.sum(nh)
			if mx[i,j]==1:
				if s >=6:
					p=1
				elif s == 3:
					p = 0 if mx[i-1,j] == 1 and mx[i+1,j] ==1 else 2
				else:
					p=0
			else:
				if s >= 7:
					p=1
				else:
					p=0
			mxt[i,j]=p
	return mxt

# nicked off http://www.logarithmic.net/pfh/blog/01185880752
def _upscan(f):
    for i, fi in enumerate(f):
        if fi == np.inf: continue
        for j in xrange(1,i+1):
            x = fi+j*j
            if f[i-j] < x: break
            f[i-j] = x

# nicked off http://www.logarithmic.net/pfh/blog/01185880752
def dist_transform(bitmap):
    f = np.where(bitmap, 0.0, np.inf)
    for i in xrange(f.shape[0]):
        _upscan(f[i,:])
        _upscan(f[i,::-1])
    for i in xrange(f.shape[1]):
        _upscan(f[:,i])
        _upscan(f[::-1,i])
    np.sqrt(f,f)
    return f

def erode(mx,strel=None):
	if strel==None:
		strel = np.array([[0,1,0],[1,1,1],[0,1,0]])
	try:
		si,sj = mx.shape
		mxt = np.ones((si,sj),dtype=np.uint8)
		sti,stj = strel.shape
		for i in xrange(sti/2, si-(sti/2)):
			for j in xrange(stj/2, sj-(stj/2)):
		#		if i in [0,si] or j in [0,sj]:
		#			p=1
		#			continue
				iadd = 1 if sti % 2 else 0
				jadd = 1 if stj % 2 else 0
				nh = mx[i-(sti/2):i+(sti/2)+iadd,j-(stj/2):j+(stj/2)+jadd]	
				if np.sum(nh*strel) == np.sum(strel):
					p = 1
				else:
					p = 0
				mxt[i,j]=p
		return mxt
	except IndexError, e:
		print e

def dilate(mx,strel=None):
	if strel==None:
		strel = np.array([[0,1,0],[1,1,1],[0,1,0]])
	try:
		si,sj = mx.shape
		mxt = np.zeros((si,sj),dtype=np.uint8)
		sti,stj = strel.shape
		for i in xrange(sti/2, si-(sti/2)):
			for j in xrange(stj/2, sj-(stj/2)):
		#		if i in [0,si] or j in [0,sj]:
		#			p=1
		#			continue
				iadd = 1 if sti % 2 else 0
				jadd = 1 if stj % 2 else 0
				nh = mxt[i-(sti/2):i+(sti/2)+iadd,j-(stj/2):j+(stj/2)+jadd]	
				if mx[i,j] == 1:
					nh += strel
		return np.array((mxt >= 1), dtype=np.uint8)
	except IndexError, e:
		print e

def flip(mx):
	si,sj=mx.shape
	return np.array([0 if x else 1 for x in (mx==1).reshape(1,si*sj)[0]],dtype=np.uint8).reshape((si,sj))

def imopen(mx,strel=None):
	return dilate(erode(mx,strel),strel)
def imclose(mx,strel=None):
	return erode(dilate(mx,strel),strel)

def nm(w,h,p,n):
	m=newmap(w,h,p)
#	print 'initial'
#	pp(m
	#m = np.array((iterate(np.array((m==0),dtype=np.uint8),ERODE)==0),dtype=np.uint8)
	#m=imopen(m)
	#m=dilate(m)

	for x in xrange(n):
		m = iterate(m)
#		print 'iter %s' % x
#		pp(m)

	return dilate(m)

# segment binary map into numerical areas (beach, plain, mountain, whatever..)
def segment_map(map):
	n = dist_transform(flip(map))
	beach = ((n >=1) & (n <2))*2
	plain = ((n >=2)) * 4 # & (n <=4))*4
	mtn = np.array((n>4),dtype=np.uint8)
	mtn = dilate(mtn, np.array([[1,1,1],[1,1,1]])) * 2
	output = beach + plain + mtn
	return output

def valid_terrain(ter_type):
	return ter_type in [0,2,4,6]

def gen_encounters(map, enc_list):
	# enc_list is of type (ter_type<int>, enc_id<int>), where ter_type is a terrain type that exists on the map, enc_id is the encounter that will be placed there. returns (enc_id, (x,y))
	# and yes, this is ridiculously ugly
	xmax,ymax = map.shape
	out = []
	map_enum = [x for x in np.ndenumerate(map)]
	np.random.shuffle(map_enum)
	for ter_type, enc_id in enc_list:
		if not valid_terrain(ter_type):
			continue
		# fugly
		for ix,((i,j), ter) in enumerate(map_enum):
			if ter == ter_type:
				out.append((enc_id, (i,j)))
				del map_enum[ix]	
				break
	return out

def gen(i=150,j=60,prob=0.67,its=4):
  m = nm(i,j,prob,its)
  map = segment_map(m)
  pp2(map)
  return m,map
  #dummy_encounters = [(rn.randint(0,9), x) for x in range(15)]
  #encounters = gen_encounters(map, dummy_encounters)
  #print encounters

if __name__=='__main__':
	#m,map=gen(prob=0.72)
	map = encounter_map(200,40,cave=1)
	items, mx = goal_location(map,3)
	print items
	pp2(mx)
