import tifffile as tiff
import cv2
import numpy as np
from skimage import io

import time
import datetime
import scipy.io as scio
import multiprocessing as mp
import math
import os

from deepwonder.MN.utils import Split_Neuron, listAddtrace
from deepwonder.MN.utils import Neuron_List_Initial, list2mask, listAddtrace, centroid_distance, list_union, initial_mask_list
from deepwonder.MN.utils import list2contours, Joint_Mask_List_Simple, Joint_Mask_List_Mul
from deepwonder.MN.utils import listAddcontours_Laplacian_pytorch, list2contours
from deepwonder.MN.utils import group_mask, calculate_trace, Mining_rest_neuron, clear_neuron

def merge_neuron_SEG_mul(mask_stack, 
                    raw_image, 
                    quit_round_rate = 0.5, 
                    good_round_rate = 0.8,
                    good_round_size_rate = 0.4, 
                    corr_mark=0.9, 
                    max_value=1000,
                    if_nmf = True):
    '''
    save_folder = datetime.datetime.now().strftime("%Y%m%d-%H%M")
    if not os.path.exists(save_folder): 
        os.mkdir(save_folder)
    '''

    threshold = max_value*0.3
    mask_stack[mask_stack>threshold] = max_value
    mask_stack[mask_stack<threshold] = 0

    w_good_neuron_list = []
    w_bad_neuron_list = []
    g_f_contours = np.zeros(mask_stack.shape)
    b_f_contours = np.zeros(mask_stack.shape)

    prev_time = time.time()
    time_start = time.time()

    if_mul=1
    if if_mul:
        num_cores = int(mp.cpu_count())
        pool = mp.Pool(num_cores)
        mask_dict = {}
        for i in range(0, mask_stack.shape[0]):
            mask_dict[str(i)] = mask_stack[i,:,:].squeeze()
        results = [pool.apply_async(initial_mask_list, args=(mask, 0.5, 0.85, 60)) for name, mask in mask_dict.items()]
        for p in results:
            good_neuron_list, bad_neuron_list = p.get()
            w_good_neuron_list = good_neuron_list + w_good_neuron_list
            w_bad_neuron_list = w_bad_neuron_list + bad_neuron_list
            # print('good_neuron_list ---> ',len(good_neuron_list),' bad_neuron_list ---> ',len(bad_neuron_list))
        time_end = time.time()
        time_cost = time_end - time_start
        print('[Neuron Segmentation Mask Initialize Time Cost: %.0d s] \n'%(time_cost), end=' ')
        pool.close()
        pool.join()
    if not if_mul:
        for i in range(0, mask_stack.shape[0]):
            now_mask = mask_stack[i,:,:].squeeze()
            good_neuron_list, bad_neuron_list = initial_mask_list(now_mask, 
                                                            quit_round_rate = 0.5, 
                                                            good_round_rate = 0.85, 
                                                            smallest_neuron_area = 60)

            w_good_neuron_list = good_neuron_list + w_good_neuron_list
            w_bad_neuron_list = w_bad_neuron_list + bad_neuron_list
            ################################################################################################################
            # Determine approximate time left
            batches_done = i+1
            batches_left = 1 * mask_stack.shape[0] - batches_done
            time_left_seconds = int(batches_left * (time.time() - prev_time))
            time_left = datetime.timedelta(seconds=time_left_seconds)
            prev_time = time.time()
            ################################################################################################################
            if i % 1 == 0:
                time_end = time.time()
                time_cost = time_end - time_start  # datetime.timedelta(seconds= (time_end - time_start))
                print(
                    '\r\033[1;33m[MERGE]\033[0m [Patch %d/%d] [Time Cost: %.0d s] [ETA: %s s]     '
                    % ( 
                    i ,
                    mask_stack.shape[0],
                    time_cost,
                    time_left_seconds
                    ), end=' ')
            if (i + 1) %  mask_stack.shape[0] == 0:
                print('\n', end=' ')
            ################################################################################################################

    w_good_neuron_list = listAddtrace(w_good_neuron_list, raw_image)

    if_time = True
    if if_time:
        time_cost = time.time() - time_start
        print('[1 Time Cost: %.0d s] \n'%(time_cost), end=' ')
    
    if_split = True
    if not if_split:
        w_good_neuron_list = Joint_Mask_List_Simple(w_good_neuron_list, [], corr_mark=0.9, area_mark=0.9, active_rate=0, if_coor=True, if_area=True, if_merge=False)
    if if_split:
        w_good_neuron_list = Joint_Mask_List_Mul(w_good_neuron_list, corr_mark=0.9, area_mark=0.9, active_rate=0, if_coor=True, if_area=True, if_merge=False)

    if if_time:
        time_cost = time.time() - time_start
        print('[2 Time Cost: %.0d s] \n'%(time_cost), end=' ')
    # w_good_neuron_list = listAddtrace(w_good_neuron_list, raw_image)
    # print('w_good_neuron_list ---> ',len(w_good_neuron_list))
    w_good_neuron_list = clear_neuron(w_good_neuron_list, [], area_mark=0.2, area_size=100)

    if if_time:
        time_cost = time.time() - time_start
        print('[3 Time Cost: %.0d s] \n'%(time_cost), end=' ')

    '''
    if not if_split:
        w_good_neuron_list = Joint_Mask_List_Simple(w_good_neuron_list, [], corr_mark=0.7, area_mark=0.6, active_rate=0, if_coor=False, if_area=True, if_merge=False)
    if if_split:
        w_good_neuron_list = Joint_Mask_List_Mul(w_good_neuron_list, corr_mark=0.7, area_mark=0.6, active_rate=0, if_coor=False, if_area=True, if_merge=False)
    '''
    # print('w_good_neuron_list ---> ',len(w_good_neuron_list))
    # w_good_neuron_list = Joint_Mask_List_Simple(w_good_neuron_list, [], corr_mark=0.7, area_mark=0.6, active_rate=0, if_coor=False, if_area=True, if_merge=False)
    
    
    if if_time:
        time_cost = time.time() - time_start
        print('[4 Time Cost: %.0d s] \n'%(time_cost), end=' ')
    # ,  save_folder=save_folder
    w_bad_neuron_list = listAddtrace(w_bad_neuron_list, raw_image)
    if not if_split:
        w_bad_neuron_list = Joint_Mask_List_Simple(w_bad_neuron_list, [], corr_mark=0.9, area_mark=0.9, active_rate=0, if_coor=True, if_area=True, if_merge=True)
    if if_split:
        w_bad_neuron_list = Joint_Mask_List_Mul(w_bad_neuron_list, corr_mark=0.9, area_mark=0.9, active_rate=0, if_coor=True, if_area=True, if_merge=True)

    if if_time:
        time_cost = time.time() - time_start
        print('[5 Time Cost: %.0d s] \n'%(time_cost), end=' ')

    add_neuron_list = Mining_rest_neuron(w_good_neuron_list, w_bad_neuron_list, raw_image, quit_round_rate = 0.5, smallest_neuron_area = 200)
    add_neuron_list = listAddtrace(add_neuron_list, raw_image)
    if not if_split:
        add_neuron_list = Joint_Mask_List_Simple(add_neuron_list, [], corr_mark=0.9, area_mark=0.9, active_rate=0, if_coor=True, if_area=True, if_merge=True)
    if if_split:
        add_neuron_list = Joint_Mask_List_Mul(add_neuron_list, corr_mark=0.9, area_mark=0.9, active_rate=0, if_coor=True, if_area=True, if_merge=True)
    
    if if_time:
        time_cost = time.time() - time_start
        print('[6 Time Cost: %.0d s] \n'%(time_cost), end=' ')

    w_g_neuron_list = w_good_neuron_list+add_neuron_list
    # corr_mark, area_mark=0.9, active_rate=0.02, if_coor=True, if_area=True, if_merge=True)

    # w_g_mask_list = listAddcontours_Laplacian_pytorch(w_g_neuron_list, mask_stack.shape[1], mask_stack.shape[2])
    # w_g_f_contours, w_g_w_contours = list2contours(w_g_mask_list, mask_stack.shape[1], mask_stack.shape[2])
    # io.imsave(save_folder+'//'+'good.tif', w_g_f_contours)
    # io.imsave(save_folder+'//'+'good_w.tif', w_g_w_contours)


    # print('w_good_neuron_list -----> ',len(w_good_neuron_list))
    if if_nmf:
        group_neuron_list, arranged_index, cc_mask_j = group_mask(w_g_neuron_list, raw_image)
        '''
        group_neuron_list1 = listAddcontours_Laplacian_pytorch(group_neuron_list, mask_stack.shape[1], mask_stack.shape[2])
        g_n_f_contours, g_n_w_contours = list2contours(group_neuron_list1, mask_stack.shape[1], mask_stack.shape[2])
        io.imsave(save_folder+'//'+'group.tif', g_n_f_contours)
        io.imsave(save_folder+'//'+'cc_mask_j.tif', cc_mask_j)
        '''
        w_g_neuron_list = calculate_trace(raw_image, group_neuron_list, w_g_neuron_list, arranged_index)

    # print('w_good_neuron_list -----> ',len(w_good_neuron_list))
    # return final_mask_list_ct
    '''
    mat_save_name = save_folder+'//'+'neuron_list.mat'
    scio.savemat(mat_save_name, {'w_good_neuron_list':w_good_neuron_list})
    '''
    w_g_neuron_list = Joint_Mask_List_Simple(w_g_neuron_list, [], corr_mark=0.7, area_mark=0.6, active_rate=0, if_coor=False, if_area=True, if_merge=False)

    return w_g_neuron_list



def merge_neuron_SEG(mask_stack, 
                    raw_image, 
                    quit_round_rate = 0.5, 
                    good_round_rate = 0.8,
                    good_round_size_rate = 0.4, 
                    corr_mark=0.9, 
                    max_value=1000,
                    if_nmf = True):

    save_folder = datetime.datetime.now().strftime("%Y%m%d-%H%M")
    if not os.path.exists(save_folder): 
        os.mkdir(save_folder)


    threshold = max_value*0.3
    mask_stack[mask_stack>threshold] = max_value
    mask_stack[mask_stack<threshold] = 0

    w_good_neuron_list = []
    w_bad_neuron_list = []
    g_f_contours = np.zeros(mask_stack.shape)
    b_f_contours = np.zeros(mask_stack.shape)

    prev_time = time.time()
    time_start = time.time()
    for i in range(0, mask_stack.shape[0]):
        now_mask = mask_stack[i,:,:].squeeze()
        good_neuron_list, bad_neuron_list = initial_mask_list(now_mask, 
                                                        quit_round_rate = 0.5, 
                                                        good_round_rate = 0.85, 
                                                        smallest_neuron_area = 60)

        # good_mask_list = listAddcontours_Laplacian_pytorch(good_neuron_list, mask_stack.shape[1], mask_stack.shape[2])
        # good_final_contours, good_whole_contours = list2contours(good_mask_list, mask_stack.shape[1], mask_stack.shape[2])
        # g_f_contours[i,:,:] = good_final_contours
        # io.imsave(save_folder+'//'+str(i)+'_good.tif', good_final_contours)

        # bad_mask_list = listAddcontours_Laplacian_pytorch(bad_neuron_list, mask_stack.shape[1], mask_stack.shape[2])
        # bad_final_contours, bad_whole_contours = list2contours(bad_mask_list, mask_stack.shape[1], mask_stack.shape[2])
        # b_f_contours[i,:,:] = bad_final_contours
        # io.imsave(save_folder+'//'+str(i)+'_bad.tif', bad_final_contours)

        w_good_neuron_list = good_neuron_list + w_good_neuron_list
        w_bad_neuron_list = w_bad_neuron_list + bad_neuron_list
        ################################################################################################################
        # Determine approximate time left
        batches_done = i+1
        batches_left = 1 * mask_stack.shape[0] - batches_done
        time_left_seconds = int(batches_left * (time.time() - prev_time))
        time_left = datetime.timedelta(seconds=time_left_seconds)
        prev_time = time.time()
        ################################################################################################################
        if i % 1 == 0:
            time_end = time.time()
            time_cost = time_end - time_start  # datetime.timedelta(seconds= (time_end - time_start))
            print(
                '\r\033[1;33m[MERGE]\033[0m [Patch %d/%d] [Time Cost: %.0d s] [ETA: %s s]     '
                % ( 
                i ,
                mask_stack.shape[0],
                time_cost,
                time_left_seconds
                ), end=' ')
        if (i + 1) %  mask_stack.shape[0] == 0:
            print('\n', end=' ')
        ################################################################################################################

    # io.imsave(save_folder+'//'+'g_f_contours.tif', g_f_contours)
    # io.imsave(save_folder+'//'+'b_f_contours.tif', b_f_contours)

    w_good_neuron_list = listAddtrace(w_good_neuron_list, raw_image)
    w_good_neuron_list = Joint_Mask_List_Simple(w_good_neuron_list, [], corr_mark=0.9, area_mark=0.9, active_rate=0, if_coor=True, if_area=True, if_merge=False)
    # w_good_neuron_list = listAddtrace(w_good_neuron_list, raw_image)
    print('w_good_neuron_list ---> ',len(w_good_neuron_list))
    w_good_neuron_list = clear_neuron(w_good_neuron_list, [], area_mark=0.2, area_size=100)
    w_good_neuron_list = Joint_Mask_List_Simple(w_good_neuron_list, [], corr_mark=0.7, area_mark=0.7, active_rate=0, if_coor=False, if_area=True, if_merge=False)
    print('w_good_neuron_list ---> ',len(w_good_neuron_list))
    # w_good_neuron_list = Joint_Mask_List_Simple(w_good_neuron_list, [], corr_mark=0.7, area_mark=0.6, active_rate=0, if_coor=False, if_area=True, if_merge=False)

    # ,  save_folder=save_folder
    w_bad_neuron_list = listAddtrace(w_bad_neuron_list, raw_image)
    w_bad_neuron_list = Joint_Mask_List_Simple(w_bad_neuron_list, [], corr_mark=0.9, area_mark=0.9, active_rate=0, if_coor=True, if_area=True, if_merge=True)
    add_neuron_list = Mining_rest_neuron(w_good_neuron_list, w_bad_neuron_list, raw_image, quit_round_rate = 0.5, smallest_neuron_area = 200)
    add_neuron_list = listAddtrace(add_neuron_list, raw_image)
    add_neuron_list = Joint_Mask_List_Simple(add_neuron_list, [], corr_mark=0.9, area_mark=0.9, active_rate=0, if_coor=True, if_area=True, if_merge=True)
    
    w_g_neuron_list = w_good_neuron_list+add_neuron_list
    # corr_mark, area_mark=0.9, active_rate=0.02, if_coor=True, if_area=True, if_merge=True)

    w_g_mask_list = listAddcontours_Laplacian_pytorch(w_g_neuron_list, mask_stack.shape[1], mask_stack.shape[2])
    w_g_f_contours, w_g_w_contours = list2contours(w_g_mask_list, mask_stack.shape[1], mask_stack.shape[2])
    io.imsave(save_folder+'//'+'good.tif', w_g_f_contours)
    io.imsave(save_folder+'//'+'good_w.tif', w_g_w_contours)


    # print('w_good_neuron_list -----> ',len(w_good_neuron_list))
    if if_nmf:
        group_neuron_list, arranged_index, cc_mask_j = group_mask(w_good_neuron_list, raw_image)
        '''
        group_neuron_list1 = listAddcontours_Laplacian_pytorch(group_neuron_list, mask_stack.shape[1], mask_stack.shape[2])
        g_n_f_contours, g_n_w_contours = list2contours(group_neuron_list1, mask_stack.shape[1], mask_stack.shape[2])
        io.imsave(save_folder+'//'+'group.tif', g_n_f_contours)
        io.imsave(save_folder+'//'+'cc_mask_j.tif', cc_mask_j)
        '''
        w_good_neuron_list = calculate_trace(raw_image, group_neuron_list, w_good_neuron_list, arranged_index)

    # print('w_good_neuron_list -----> ',len(w_good_neuron_list))
    # return final_mask_list_ct
    '''
    mat_save_name = save_folder+'//'+'neuron_list.mat'
    scio.savemat(mat_save_name, {'w_good_neuron_list':w_good_neuron_list})
    '''
    return w_good_neuron_list




