import torch
from model import Modified3DUNet
from torch.utils.data import DataLoader
from kidney_dataloader_1_c_aug_2_classes import KidneyDataset_standardization
import evaluate_tools
# ******************************8
from tqdm import tqdm
import SimpleITK as sitk
import numpy as np
import random
import torch.nn.functional as F
from scipy.ndimage.measurements import label
from scipy.ndimage.morphology import generate_binary_structure
from scipy.ndimage.morphology import distance_transform_cdt
import os
import pdb

def sim_initial_click_dataset_batch(**kwargs):
    '''
    simulate initial clicks for the dataset
    '''
    dataloader = kwargs['dataloader'] 
    model_init = kwargs['model']
    sigma = kwargs['sigma']
    device = kwargs['device']
    model_init.eval()
    clicks_info_init = {}
    for sim_batch in tqdm(dataloader):
        clicks_list = sim_clicks_init_batch(img_ids = sim_batch['kidney_id'], data = sim_batch['data'], mask = sim_batch['mask'],\
                                            mask_meta = sim_batch['mask_meta'], model = model_init, device = device)
        # fill in the dictionary.
        fill_the_dict_init(clicks_info_init,clicks_list,sim_batch['kidney_id'])
    model_init.train()
    return clicks_info_init

def sim_clicks_init_batch(**kwargs):
    '''
    simulate clicks before the training for the batch data
    '''
    img_ids = kwargs['img_ids']
    data = kwargs['data'] 
    mask = kwargs['mask']
    mask_meta = kwargs['mask_meta'] 
    model = kwargs['model']
    device = kwargs['device']
    _, seglayer, = model(torch.unsqueeze(data,1).to(device))
    _, pred_labels = torch.max(F.softmax(seglayer,dim=1), dim=1)
    fp_batch, fn_batch = get_fp_fn_masks_batch(pred_labels, mask.to(device))
    clicks = get_clicks_chamfer_init_without_largets_cc(fp_batch, fn_batch, img_ids, mask_meta)
    return clicks

def get_fp_fn_masks_batch(pred, gt):
    '''
    get fp mask and fn mask for only tumors
    '''
    fp_tensor = ((pred==2)&(gt==0)).type(torch.uint8)
    fn_tensor = ((pred==0)&(gt==2)).type(torch.uint8)
    fp_batch = fp_tensor.cpu().detach().numpy() 
    fn_batch = fn_tensor.cpu().detach().numpy()
    return fp_batch, fn_batch

def get_clicks_chamfer_init_without_largets_cc(fp_batch, fn_batch, img_ids, mask_meta):
    '''
    get one click on chamfer distance for a batch data
    '''
    batch_size = fp_batch.shape[0]
    clicks_list = []
    for i in range(batch_size):        
        cur_fp_arr = fp_batch[i]
        cur_fn_arr = fn_batch[i]
        fp_click, _ = get_chamfer_map_click(cur_fp_arr)
        fn_click, _ = get_chamfer_map_click(cur_fn_arr)
        clicks = np.stack([fp_click,fn_click]).swapaxes(1,0)
        clicks_list.append(clicks)
    return clicks_list

def get_chamfer_map_click(mask):
    '''
    get chamfer distance based on the mask
    '''
    chamfer_map = distance_transform_cdt(mask)  
    points = np.where(chamfer_map==np.max(chamfer_map))
    lower = 0
    upper = points[0].shape[0] - 1
    index = random.randint(lower, upper)
    click = np.array([arr[index] for arr in points]) 
    return click, chamfer_map

def fill_the_dict_init(clicks_info_init, clicks_list, key_list):
    '''
    fill the dict with clicks
    '''
    for i in range(len(key_list)):
        clicks_info_init[key_list[i]]  = [clicks_list[i]]

def sim_clicks_dataset_batch(**kwargs):
    '''
    simulate  clicks during the trianing
    '''
    dataloader = kwargs['dataloader']
    clicks_info_init = kwargs['clicks_info']
    model = kwargs['model']
    num_points = kwargs['num_points']
    sigma = kwargs['sigma']
    device = kwargs['device']
    model.eval()
    for sim_batch in tqdm(dataloader):
        sim_clicks_batch(sim_batch=sim_batch, clicks_info=clicks_info_init, model=model, num_points=num_points, sigma=sigma, device=device)
    model.train()
    return clicks_info_init

def sim_clicks_batch(**kwargs):
    sim_batch = kwargs['sim_batch']
    clicks_info = kwargs['clicks_info']
    model = kwargs['model']
    num_points = kwargs['num_points']
    sigma = kwargs['sigma']
    device = kwargs['device']    
    key_list = sim_batch['kidney_id']
    for k in range(num_points):
        # print(k)
        clicks_list = sim_clicks(sim_batch=sim_batch, model=model, clicks_info=clicks_info, sigma=sigma, device=device, k=k) 
        fill_the_dict(clicks_info, clicks_list, key_list)

def fill_the_dict(clicks_info, clicks_list, key_list):
    '''
    fill the dictionary
    '''
    for i in range(len(key_list)):
        clicks_info[key_list[i]].append(clicks_list[i])

def sim_clicks(**kwargs):
    '''
    simulate clicks during the training (4 chanel model)
    data : tensor
    mask : tnesor
    clicks_info : list of array
    '''
    batch_data = kwargs['sim_batch']
    model = kwargs['model']
    clicks_info  = kwargs['clicks_info']
    sigma = kwargs['sigma']
    device = kwargs['device']
    click_num = kwargs['k']
    # for key in img_keys:
    train_data_1_c = torch.unsqueeze(batch_data['data'].to(device),1)
    img_ids = batch_data['kidney_id']
    mask_meta = batch_data['mask_meta']
    all_clicks = get_batch_clicks(img_ids,clicks_info) # (3, batch_size *2) numpy arr
    train_maps = torch.tensor(get_prob_maps_sim(all_clicks,batch_data['mask'].shape[1:4],sigma)).to(device)
    train_data = torch.cat((train_data_1_c, train_maps),1)
    _, seglayer = model(train_data)
    _, pred_labels = torch.max(F.softmax(seglayer,dim=1), dim=1)
    # pred = torch.squeeze(pred,1)
    mask = batch_data['mask'].to(device)
    fp_batch, fn_batch = get_fp_fn_masks_batch(pred_labels, mask)
    clicks = get_clicks_chamfer_without_largest_cc(fp_batch=fp_batch, fn_batch=fn_batch, img_ids=img_ids, mask_meta=mask_meta, click_num=click_num)
    return clicks

def get_batch_clicks(img_keys, clicks_info):
    '''
    get click list
    '''
    all_clicks = []
    for img_key in img_keys:
        all_clicks.append(np.concatenate(clicks_info[img_key],axis=1))
    all_clicks = np.stack(all_clicks)
    return all_clicks

def get_clicks_chamfer_without_largest_cc(**kwargs):
    '''
    get one click on chamfer distance for a batch data
    '''
    fp_batch = kwargs['fp_batch']
    fn_batch = kwargs['fn_batch']
    img_ids = kwargs['img_ids']
    mask_meta = kwargs['mask_meta']
    click_num = kwargs['click_num']
    batch_size = fp_batch.shape[0]
    clicks_list = []
    for i in range(batch_size):
        cur_fp_arr = fp_batch[i]
        cur_fn_arr = fn_batch[i]
        fp_click, _ = get_chamfer_map_click(cur_fp_arr)
        fn_click, _ = get_chamfer_map_click(cur_fn_arr)
        clicks = np.stack([fp_click,fn_click]).swapaxes(1,0)
        clicks_list.append(clicks)
    return clicks_list

def get_prob_maps_sim(clicks_arr,shape,sigma):
    '''
    generate probability maps based on click information
    '''
    batch_size = clicks_arr.shape[0]
    case_maps = []
    for i in range(batch_size):
        fp_points = clicks_arr[i][:,::2]
        fn_points = clicks_arr[i][:,1::2]
        M_fp = get_prob_map(fp_points, shape, sigma)
        M_fn = get_prob_map(fn_points, shape, sigma)        
        M = np.stack([M_fp,M_fn])
        case_maps.append(M)
    case_maps = np.stack(case_maps) 
    return case_maps.astype(np.float32)

def get_prob_maps(keys, clicks_info, shape, sigma):
    '''
    get probility map during the trainng
    '''
    prob_maps_fp = []
    prob_maps_fn = []
    for key in keys:
        clicks = clicks_info[key]
        clicks = np.concatenate(clicks,axis=1)
        fp_clicks = clicks[:,::2]
        fn_clicks = clicks[:,1::2] 
        M_fp = get_prob_map(fp_clicks, shape, sigma)
        M_fn = get_prob_map(fn_clicks, shape, sigma)
        prob_maps_fn.append(M_fn)
        prob_maps_fp.append(M_fp)
    prob_maps_fp_arr = np.stack(prob_maps_fp).astype(np.float32)
    prob_maps_fn_arr = np.stack(prob_maps_fn).astype(np.float32)
    probs_M = np.stack([prob_maps_fp_arr,prob_maps_fn_arr],axis=1)
    return  probs_M.astype(np.float32)

def get_prob_map(points, shape, sigma):
    '''
    generate the 3d probability map 
    '''
    z = np.arange(0, shape[2], 1, float)
    y = np.arange(0, shape[1], 1, float)
    y = y[:,np.newaxis]
    x = np.arange(0, shape[0], 1, float)
    x = x[:,np.newaxis,np.newaxis]
    M = np.zeros(shape)

    for i in range(points.shape[1]):
        center = points[:,i]
        x0 = center[0]
        y0 = center[1]
        z0 = center[2]
        temp_probs = np.exp(-4 * np.log(2) * ((x - x0) ** 2 + (y - y0) ** 2 + (z-z0)**2) / sigma ** 2)#.astype(d_type)   
        M = np.maximum(temp_probs, M)
    return M