#!/usr/bin/env python3 
# -*- coding: utf-8 -*- 
""" 
""" 
from operator import mod
from tomo_encoders.misc.voxel_processing import modified_autocontrast, TimerGPU
from tomo_encoders.reconstruction.recon import recon_patches_3d
import cupy as cp
import numpy as np
from skimage.filters import threshold_otsu
from tomo_encoders import Grid



def get_values_cyl_mask(vol, mask_fac):
    
    vol_shape = vol.shape
    assert vol_shape[1] == vol_shape[2], "must be a tomographic volume where shape y = shape x"
    
    shape_yx = vol_shape[1]
    shape_z = vol_shape[0]
    rad = int(mask_fac*shape_yx/2)
    
    pts = cp.arange(-int(shape_yx//2), int(cp.ceil(shape_yx//2)))
    yy, xx = cp.meshgrid(pts, pts, indexing = 'ij')
    circ = (cp.sqrt(yy**2 + xx**2) < rad).astype(cp.uint8) # inside is positive
    circ = circ[cp.newaxis, ...]
    cyl = cp.repeat(circ, shape_z, axis = 0)
    return vol[cyl > 0]


def cylindrical_mask(out_vol, mask_fac, mask_val = 0):
    
    vol_shape = out_vol.shape
    assert vol_shape[1] == vol_shape[2], "must be a tomographic volume where shape y = shape x"
    
    shape_yx = vol_shape[1]
    shape_z = vol_shape[0]
    rad = int(mask_fac*shape_yx/2)
    
    pts = cp.arange(-int(shape_yx//2), int(cp.ceil(shape_yx//2)))
    yy, xx = cp.meshgrid(pts, pts, indexing = 'ij')
    circ = (cp.sqrt(yy**2 + xx**2) < rad).astype(cp.uint8) # inside is positive
    circ = circ[cp.newaxis, ...]
    cyl = cp.repeat(circ, shape_z, axis = 0)
    out_vol[cyl == 0] = mask_val
    
    return


def segment_otsu(vol, s = 0.05):
    '''segment volume with otsu'''
    timer = TimerGPU()
    timer.tic()
    tmp_values = vol[::4,::4,::4].get()
    # rec_min_max = modified_autocontrast(tmp_values, s = s, normalize_sampling_factor=1)
    thresh = cp.float32(threshold_otsu(tmp_values.reshape(-1)))
    vol = (vol < thresh).astype(cp.uint8)
    timer.toc("otsu thresholding")
    return vol

def edge_map(Y):

    '''
    this algorithm was inspired by: https://github.com/tomochallenge/tomochallenge_utils/blob/master/foam_phantom_utils.py
    '''
    msk = cp.zeros_like(Y)
    tmp = Y[:-1]!=Y[1:]
    msk[:-1][tmp] = 1
    msk[1:][tmp] = 1
    tmp = Y[:,:-1]!=Y[:,1:]
    msk[:,:-1][tmp] = 1
    msk[:,1:][tmp] = 1
    tmp = Y[:,:,:-1]!=Y[:,:,1:]
    msk[:,:,:-1][tmp] = 1
    msk[:,:,1:][tmp] = 1
    return msk > 0

def guess_surface(V_bin, b, wd):
    
    # find patches on surface
    wdb = int(wd//b)
    p3d = Grid(V_bin.shape, width = wdb)
    
    x = p3d.extract(V_bin)
    is_surf = (np.std(x, axis = (1,2,3)) > 0.0)
    is_ones = (np.sum(x, axis = (1,2,3))/(wdb**3) == 1)
    is_zeros = (np.sum(x, axis = (1,2,3))/(wdb**3) == 0)
    
    p3d = p3d.rescale(b)
    p3d_surf = p3d.filter_by_condition(is_surf)
    p3d_ones = p3d.filter_by_condition(is_ones)
    p3d_zeros = p3d.filter_by_condition(is_zeros)
    eff = len(p3d_surf)*(wd**3)/np.prod(p3d_surf.vol_shape)
    print(f"\tSTAT: r value: {eff*100.0:.2f}")        
    return p3d_surf, p3d_ones, p3d_zeros

def process_patches(projs, theta, center, fe, p_surf, min_max, TIMEIT = False):

    # SCHEME 1: integrate reconstruction and segmention (segments data on gpu itself)
    # st_proc = cp.cuda.Event(); end_proc = cp.cuda.Event(); st_proc.record()
    # x_surf, p_surf = recon_patches_3d(projs, theta, center, p_surf, \
    #                                   apply_fbp = True, segmenter = fe, \
    #                                   segmenter_batch_size = 256)
    # end_proc.record(); end_proc.synchronize(); t_surf = cp.cuda.get_elapsed_time(st_proc,end_proc)
    
    
    # SCHEME 2: reconstruct and segment separately (copies rec data from gpu to cpu)
    st_rec = cp.cuda.Event(); end_rec = cp.cuda.Event(); st_rec.record()
    x_surf, p_surf = recon_patches_3d(projs, theta, center, p_surf, \
                                      apply_fbp =True)
    end_rec.record(); end_rec.synchronize(); t_rec = cp.cuda.get_elapsed_time(st_rec,end_rec)
    st_seg = cp.cuda.Event(); end_seg = cp.cuda.Event(); st_seg.record()
    
    x_surf = np.clip(x_surf, *min_max)
    x_surf = fe.predict_patches("segmenter", x_surf[...,np.newaxis], 256, None, min_max = min_max)[...,0]
    end_seg.record(); end_seg.synchronize(); t_seg = cp.cuda.get_elapsed_time(st_seg,end_seg)
    
    print(f'\tTIME: local reconstruction - {t_rec/1000.0:.2f} secs')    
    print(f'\tTIME: local segmentation - {t_seg/1000.0:.2f} secs')
    print(f'\tSTAT: total patches in neighborhood: {len(p_surf)}')    
    if TIMEIT:
        return x_surf, p_surf, t_rec, t_seg
    else:
        return x_surf, p_surf
    
    
    


[((25, 13, 25, 51), 'cupy.meshgrid', 'cp.meshgrid', (), '', True, 'import cupy as cp\n'), ((28, 10, 28, 44), 'cupy.repeat', 'cp.repeat', (), '', True, 'import cupy as cp\n'), ((42, 13, 42, 51), 'cupy.meshgrid', 'cp.meshgrid', (), '', True, 'import cupy as cp\n'), ((45, 10, 45, 44), 'cupy.repeat', 'cp.repeat', (), '', True, 'import cupy as cp\n'), ((53, 12, 53, 22), 'tomo_encoders.misc.voxel_processing.TimerGPU', 'TimerGPU', ({}, {}), '()', False, 'from tomo_encoders.misc.voxel_processing import modified_autocontrast, TimerGPU\n'), ((67, 10, 67, 26), 'cupy.zeros_like', 'cp.zeros_like', ({(67, 24, 67, 25): 'Y'}, {}), '(Y)', True, 'import cupy as cp\n'), ((83, 10, 83, 40), 'tomo_encoders.Grid', 'Grid', (), '', False, 'from tomo_encoders import Grid\n'), ((109, 13, 109, 28), 'cupy.cuda.Event', 'cp.cuda.Event', ({}, {}), '()', True, 'import cupy as cp\n'), ((109, 40, 109, 55), 'cupy.cuda.Event', 'cp.cuda.Event', ({}, {}), '()', True, 'import cupy as cp\n'), ((110, 21, 111, 54), 'tomo_encoders.reconstruction.recon.recon_patches_3d', 'recon_patches_3d', (), '', False, 'from tomo_encoders.reconstruction.recon import recon_patches_3d\n'), ((112, 53, 112, 93), 'cupy.cuda.get_elapsed_time', 'cp.cuda.get_elapsed_time', ({(112, 78, 112, 84): 'st_rec', (112, 85, 112, 92): 'end_rec'}, {}), '(st_rec, end_rec)', True, 'import cupy as cp\n'), ((113, 13, 113, 28), 'cupy.cuda.Event', 'cp.cuda.Event', ({}, {}), '()', True, 'import cupy as cp\n'), ((113, 40, 113, 55), 'cupy.cuda.Event', 'cp.cuda.Event', ({}, {}), '()', True, 'import cupy as cp\n'), ((115, 13, 115, 38), 'numpy.clip', 'np.clip', ({(115, 21, 115, 27): 'x_surf', (115, 29, 115, 37): '*min_max'}, {}), '(x_surf, *min_max)', True, 'import numpy as np\n'), ((117, 53, 117, 93), 'cupy.cuda.get_elapsed_time', 'cp.cuda.get_elapsed_time', ({(117, 78, 117, 84): 'st_seg', (117, 85, 117, 92): 'end_seg'}, {}), '(st_seg, end_seg)', True, 'import cupy as cp\n'), ((86, 15, 86, 40), 'numpy.std', 'np.std', (), '', True, 'import numpy as np\n'), ((94, 32, 94, 59), 'numpy.prod', 'np.prod', ({(94, 40, 94, 58): 'p3d_surf.vol_shape'}, {}), '(p3d_surf.vol_shape)', True, 'import numpy as np\n'), ((24, 43, 24, 63), 'cupy.ceil', 'cp.ceil', ({(24, 51, 24, 62): 'shape_yx // 2'}, {}), '(shape_yx // 2)', True, 'import cupy as cp\n'), ((41, 43, 41, 63), 'cupy.ceil', 'cp.ceil', ({(41, 51, 41, 62): 'shape_yx // 2'}, {}), '(shape_yx // 2)', True, 'import cupy as cp\n'), ((87, 15, 87, 40), 'numpy.sum', 'np.sum', (), '', True, 'import numpy as np\n'), ((88, 16, 88, 41), 'numpy.sum', 'np.sum', (), '', True, 'import numpy as np\n'), ((26, 12, 26, 34), 'cupy.sqrt', 'cp.sqrt', ({(26, 20, 26, 33): 'yy ** 2 + xx ** 2'}, {}), '(yy ** 2 + xx ** 2)', True, 'import cupy as cp\n'), ((43, 12, 43, 34), 'cupy.sqrt', 'cp.sqrt', ({(43, 20, 43, 33): 'yy ** 2 + xx ** 2'}, {}), '(yy ** 2 + xx ** 2)', True, 'import cupy as cp\n')]