'''
Module for domain decomposition based on the Orthogonal Recursive bisection Algorithm, see [1] and [2] for full details. 
Working : The master processor assigns subdomains determines the min and max bounds from all the bodies (excluding wall, subdomain and box), 
This space is the subdivided into regions based on a splitting axis. The axis of split is along the largest dimension and each split results into 
2 subregions** until a certain 'depth' based on the number of workers is reached.
The maximum depth/number of levels is defined as log2(N_{w}), where N_{w} is the number of workers. The algorithm can be used for any number of workers.

(c) Deepak Kunhappan, deepak.kn1990@gmail.com, deepak.kunhappan@3sr-grenoble.fr 

'''

from math import ceil, log10, log
from yade.wrapper import * 
from minieigen import Vector3
from mpi4py import MPI 
import sys

sys.stderr.write=sys.stdout.write
minVal =  -1e+50
maxVal =  1e+50
maxDim = 3 
_TAG_DATA_=999
_TOL_ = 1e-08


# for coloring bodies. 
import colorsys

class decompBodiesSerial: 


	
	def __init__(self, mpiComm): 
		
		self.comm = mpiComm
		self.numWorkers = self.comm.Get_size() - 1
		# for coloring bodies 
		self.colorscale = []
		for i in range(self.numWorkers): 
			self.colorscale.append(Vector3(colorsys.hsv_to_rgb(i*1./self.numWorkers,1,1)))
		

	def isEqual(self, a,b): 
		'''function to check if 2 floats are equal'''
		
		return abs(a-b) < _TOL_ 
	
	
	def ifPw2(self, numThreads):
		''' Determine if the number of workers is a power of 2 '''
		 
		if numThreads == 0 : return False 
		return numThreads & (numThreads-1) == 0
	
	def numThreadsinPw2(self, numThreads): 
		''' Express the number of workers as sum in poowers of 2, e.g: 11 = 2^{3} + 2^{1} + 2^{0} '''
		parts = [] 
		i = 1 
		while i <= numThreads: 
			if i & numThreads : parts.append(i) 
			i <<= 1 
		return parts 
	      
	def setPartList(self, lst, numThreads): 
		
		''' Used when the number of workers are not a power of 2. Bisectioning divides the space into 2 regions, hence equal divisions are created when the number of workers are in power of 2. If not
		one or more workers end up having more number of bodies compared to other workers, leading to a load imbalance. In-order to overcome this we first express the numW in terms sums of 2^{n}, with each 2^{n} defining 
		a sub group in which bisectioning can be done. 
		For example : we have 11 workers, so : 11 = 2^{0} + 2^{1} + 2^{3} = [1, 2, 8] 
		with 1 worker from 1 level of bisectioning,(a part of the domain is assigned to 1 worker) 
		2 workers from 1 level/depth of bisectioning
		8 workers from 3 levels/depth of bisectioning
		'''
		
		nb = len(lst)//numThreads 
		parts = self.numThreadsinPw2(numThreads) 
		wIds = [[] for x in range(len(parts))]
		for i in range(len(parts)-1): 
			splitAxis = self.domainPreProc(lst)
			lst = sorted(lst, key=lambda x:x[0][splitAxis])
			splitPt = len(lst)-(nb*parts[i])
			wIds[i] = lst[splitPt:]; lst = lst[:splitPt]
		wIds[len(parts)-1] = lst 
		return parts, wIds 
	
	def _partitionDomain(self, bList, numW, strideRank): 
	  
		''' The partitioning algorithm, sorts the list of bodies in increasing order based on a predetermined axis, (split axis)
		for each level. With one half allocated to one worker and the other to another
		at the final level. '''
		
		workerIds = [[] for x in range(numW)]
		workerIds[0] = bList 
		wProcs = [n for n in range(numW)] 
		numLevels = int(log(numW,2)) 
		nlevel = 0; depth = 0
		while (nlevel <= numLevels): 
			lpW2 = 1 << nlevel 
			lpW2_1 = 1 << (nlevel+1)
			sIndxs = [x for x in wProcs if x < lpW2] 
			rIndxs = [x for x in wProcs if x >= lpW2 and x < lpW2_1]
			for s,r in zip(sIndxs, rIndxs): 
				workerIds[s], workerIds[r] = self.partData(workerIds[s], nlevel)
			nlevel += 1 
		
		for i,x in enumerate(workerIds): 
			workerIds[i] = [bId[1] for bId in x]
		
		for i, blst in enumerate(workerIds): 
			for x in blst: 
				O.bodies[x].subdomain = (i+strideRank)+1
				O.bodies[x].shape.color = self.colorscale[(i+strideRank)-1] 

	
	def partData(self, bList, nlevel): 
	
		''' function to sort data and divide into 2 parts. The data is a list containing a tuple of (b.state.pos, b.id) '''
		
		splitAxis = self.domainPreProc(bList) 
		sortD = sorted(bList, key=lambda x:x[0][splitAxis]) 
		splitPt = len(sortD)//2 
		return sortD[:splitPt], sortD[splitPt:]
	
	
	def domainPreProc(self, data): 
	  
		''' determines the max extends of the domain, the splitting axis is based on the axis of the largest extend. Also used to determine the dimensionality of the 
		configuration, eg: 2D or 3D. 
		'''
		
		domainMin = Vector3(maxVal, maxVal, maxVal)
		domainMax = Vector3(minVal, minVal, minVal) 
		ndim = -1; symIndex = -1; domain2D = False 
		splitPlaneIndex = -1; maxDim = 3 
		
		for x in data: 
			pos = x[0] 
			
			domainMin[0] = min(domainMin[0], pos[0]) 
			domainMin[1] = min(domainMin[1], pos[1]) 
			domainMin[2] = min(domainMin[2], pos[2]) 
			
			domainMax[0] = max(domainMax[0], pos[0]) 
			domainMax[1] = max(domainMax[1], pos[1]) 
			domainMax[2] = max(domainMax[2], pos[2]) 
		
		for j in range(maxDim): 
			if (self.isEqual(domainMin[j], domainMax[j])): 
				domain2D = True
				symIndex = j 
				ndim = 2
			else : ndim = 3 
			
		domainAxisL = [abs(domainMax[0]-domainMin[0])] + [abs(domainMax[1]-domainMin[1])] + [abs(domainMax[2]-domainMin[2])]
		splitPlaneIndex = domainAxisL.index(max(domainAxisL))
		
		#TODO: find minimum of intersection areas between subdomains, the idea is to use this as a condition to determine if subdomains have to be reset. 
		interSectionAreas = [domainAxisL[1]*domainAxisL[2]]+[domainAxisL[0]*domainAxisL[2]]+[domainAxisL[0]*domainAxisL[1]]
		areaMin = None 
		if ndim == 3: 
			areaMin = interSectionAreas.index(min(interSectionAreas))
		
		else : 
			if symIndex == 0 : areaMin = min(domainAxisL[1], domainAxisL[2])
			if symIndex == 1 : areaMin = min(domainAxisL[2], domainAxisL[0])
			if symIndex == 2 : areaMin = min(domainAxisL[0], domainAxisL[1]) 
			
		
		return splitPlaneIndex
		
		
	
	def partitionDomain(self):
	  
		''' The main driver function  '''
		
		blst = [(b.state.pos, b.id) for b in O.bodies if not (isinstance(b.shape, Wall) or isinstance(b.shape, Box) or b.isSubdomain)]
		if self.ifPw2(self.numWorkers) : 
			self._partitionDomain(blst, self.numWorkers, strideRank = 0)
		else : 
			parts, wIds = self.setPartList(blst, self.numWorkers)
			nr = self.numWorkers+1
			for i in range(len(parts)): 
				nr = nr - parts[i]
				self._partitionDomain(wIds[i], parts[i], nr-1)


