# -*- coding: utf-8 -*-
"""
Created on Sat Jan 24 16:57:16 2015

@author: Yin Lin

Generate monte-carlo toy box without clustering bias
"""
import scipy.io
import numpy as np

def generate_monte_box(boxName, pixDim, phyDim, ionizedFrac, saveNameList, 
                       saveNameBox):
    """
    Generate monte-carlo toybox given a list of bubble distribution until the
    desired ionization fraction is reached. Each bubble is allowed to overlap 
    with each other. A bubble is removed from the final distribution list if
    it is completely overlap by other bubbles.
    
    Keyword arguments:
    boxName --- input box in mat format with pixelList entry (presumably 
               generated by adaptive_montecarlo.m)    
    pixDim --- pixel dimension of box
    phyDim --- physical dimension of box in Mpc
    ionizedFrac --- desired ionization fraction of the output toybox
    saveNameList    --- output name for the final bubble size distribution
                        in mat format with pixelList entry
    saveNameBox     --- output name for the toy box in binary format (same as
                        the output format in 21cmFAST)
    """

   
    ratio = float(pixDim) / float(phyDim)
    radiusList = scipy.io.loadmat(boxName)['pixelList'][0] #in Mpc
    
    
    #convert to pixel unit and shuffle the order for random shuffling
    randPixList = (np.array(radiusList) * ratio).astype('float32') #in no. of pixels
    np.random.shuffle(randPixList)
    
    
    numBubbles = np.shape(randPixList)[0]
    
    #padded the box so that we won't have index that exceed the dim of the box
    maxR  = int(np.max(randPixList))
    boxDim = pixDim + int(2*maxR)
    toyBox = np.zeros([boxDim,boxDim,boxDim]).astype('int64')
    
    #generate the coordinates for each bubbles, padded with maxR to avoid the 
    #boundary of box
    randX = np.add((np.random.random(numBubbles) * pixDim).astype(np.int),maxR)
    randY = np.add((np.random.random(numBubbles) * pixDim).astype(np.int),maxR)
    randZ = np.add((np.random.random(numBubbles) * pixDim).astype(np.int),maxR)
    
    
    pixelList = [];
    current_lindex = 0
    count = 0
    while True:
        pixRadius = randPixList[count]
        pixelList.append(pixRadius)
    
    
        #create a sphere that is to be overlaped over the array
        x,y,z = np.ogrid[-pixRadius:pixRadius+1,
                         -pixRadius:pixRadius+1,
                         -pixRadius:pixRadius+1]
        sphere = (x**2 + y**2 + z**2 <= pixRadius**2)
    

        #loop over different random coordinate until we can place one without
        #completely overlapped by other bubbles
        while True:
            centerY,centerX, centerZ = (randX[current_lindex], 
                                        randY[current_lindex], 
                                        randZ[current_lindex])
            #count the pixels already ionized 
            temp = toyBox[centerX-pixRadius:centerX+pixRadius+1,
                          centerY-pixRadius:centerY+pixRadius+1,
                          centerZ-pixRadius:centerZ+pixRadius+1]
            temp = temp > 0
            ion_before = np.sum(temp)
            
            # if all the cells where we try to place a new bubble is occupied,
            # we continue with next coordinate
            #(we assume the entire box is ionized, but in fact the bubble we
            #put down in sphereical, so it STILL can be completely overlapped
            #if the condition is not met, but the chance is slim)
            if ion_before == (2*pixRadius + 1)**3:
                current_lindex += 1
                continue
            
            #update the toy box
            (toyBox[centerX-pixRadius:centerX+pixRadius+1,
                   centerY-pixRadius:centerY+pixRadius+1,
                   centerZ-pixRadius:centerZ+pixRadius+1])[sphere] = count + 1   
            break
        count += 1
        current_lindex += 1
    #remove padding 
    toyBox = toyBox[maxR:maxR+pixDim,maxR:maxR+pixDim,maxR:maxR+pixDim]
    
    
    print 'saving...' 
    #eliminate any entried they are completely overlaped
    binn = np.bincount(toyBox.flatten(),minlength=count+1)
    binn = np.delete(binn,0)
    pixelList = np.array(pixelList)
    pixelList = pixelList[np.where(binn>0)]
    scipy.io.savemat(saveNameList,{'pixelList':pixelList})
    
    toyBox = (toyBox == 0)    #to make it the same convection as 21cmFAST
    toyBox=toyBox.astype('float32')
    toyBox.tofile(saveNameBox)
