#!/usr/bin/env python
"""
Created on Tue Sep  2 07:14:28 2014

Author: Oren Freifeld
Email: freifeld@csail.mit.edu
"""

import numpy as np
from of.utils import *
from of.gpu import CpuGpuArray
from pylab import plt 
from . honeycomb import honeycomb

dirname_of_this_file = os.path.dirname(inspect.getfile(inspect.currentframe()))
dirname_precomputed_hex_inits = os.path.join(dirname_of_this_file,'precomputed_hex_inits')
FilesDirs.mkdirs_if_needed(dirname_precomputed_hex_inits)

def create_string(dimx,dimy,nPixels_in_square_side):
   s = 'dimx_{0}_dimy_{1}_n_{2}.npy'.format(dimx,dimy,nPixels_in_square_side) 
   return s


def get_init_seg(dimy,dimx,nPixels_in_square_side,use_hex):
    """
    """
    M=nPixels_in_square_side 

    if use_hex:
        s = create_string(dimx,dimy,nPixels_in_square_side)
#        print (s)
        fname = os.path.join(dirname_precomputed_hex_inits,s)
        try:
            FilesDirs.raise_if_file_does_not_exist(fname)
            print ("Loading",fname)
            seg = np.load(fname)
            
            return seg
        except FileDoesNotExistError:
            pass
        msg = """
        I could not find a precomputed (image-independent)
        honeycomb initilization for this image size and this values of n.     
        So I will compute it from scratch and we will save the result in
        {}
        Next time you will run the code for an image of size
        nRows={}, nCols={}, with n = {},
        it will be faster.
        """.format(fname,dimy,dimx,nPixels_in_square_side)
        print( msg)
        seg = CpuGpuArray.zeros((dimy,dimx),dtype=np.int32)     

        # length of each side   
        a = np.sqrt(M ** 2 / (  1.5 * np.sqrt(3)  ))      
        H =   a
        W = np.sqrt(3)*H
        # XX and YY need to be float
        YY,XX = np.mgrid[0:float(dimy)+0*1.5*H:1.5*H,0:float(dimx)+0*W:W]

        XX[::2]+= float(W)/2
        centers = np.vstack([XX.ravel(),YY.ravel()]).T.copy()
        centers = CpuGpuArray(centers)    

        honeycomb(seg.gpu,centers.gpu,seg.size) 

        seg.gpu2cpu()
        np.save(fname,seg.cpu)
        return seg.cpu

    else:
        seg_cpu = np.zeros((dimy,dimx),dtype=np.int32)  
        yy,xx = np.mgrid[:dimy,:dimx] 
        xx = xx.astype(np.float)
        yy = yy.astype(np.float)
       
        dimx = float(dimx)
        dimy=float(dimy)        
        nTimesInX = np.floor(xx / M).max() + 1
 
        seg_cpu = np.floor(yy / M)  * nTimesInX + np.floor(xx / M)
        seg_cpu = seg_cpu.astype(np.int32)
        return seg_cpu


def random_permute_seg(seg):
    p=np.random.permutation(seg.max()+1)   
    seg2 = np.zeros_like(seg)
    for c in range(seg.max()+1):             
        seg2[seg==c]=p[c]
    return seg2.astype(np.int32)


if __name__ == "__main__":  
    tic = time.clock()
    seg= get_init_seg(500, 500,17,True)      
#    seg= get_init_seg(512, 512,50,False)  
    toc = time.clock()
    print (toc-tic )
    print ('k = ', seg.max()+1)
    plt.figure(1)
    plt.clf()  
    plt.imshow(seg,interpolation="Nearest")
    plt.axis('scaled') 