# -*- coding: utf-8 -*-

# Note:
# This file is a slightly modified version of the original application_script.py for use with MorphoGraphX/MorphoDynamX


import os
import sys
import numpy as np
import torch
from argparse import ArgumentParser
from skimage import io
from skimage.segmentation import watershed
from skimage.feature import peak_local_max
from scipy.ndimage import distance_transform_edt, label

from models.UNet3D import UNet3D
from dataloader.h5_dataloader import MeristemH5Tiler as Tiler
from torch.autograd import Variable

SEED = 1337
torch.manual_seed(SEED)
np.random.seed(SEED)


def main(hparams):
    
    
    """
    Main testing routine specific for this project
    :param hparams:
    """

    # ------------------------
    # 1 INIT LIGHTNING MODEL
    # ------------------------
    model = UNet3D(hparams)
    model = model.load_from_checkpoint(hparams.ckpt_path)

    #example = torch.rand(1, 1, 32, 128, 128)
    #traced_script_module = torch.jit.trace(model, example)
    #traced_script_module.save("tracedModel.pt")

    model = model.cuda()
    
    # ------------------------
    # 2 INIT DATA TILER
    # ------------------------
    tiler = Tiler(hparams.test_list, hparams.data_root, patch_size=hparams.patch_size, overlap=hparams.overlap,\
                  image_group=hparams.image_group, mask_groups=None)
    os.makedirs(os.path.abspath(hparams.output_path), exist_ok=True)
    
    # ------------------------
    # 3 PROCESS EACH IMAGE
    # ------------------------
    for image_idx in range(len(tiler.data_list)):
        tiler.set_data_idx(image_idx)
                
        application_size = tuple(np.maximum(hparams.patch_size, tiler.data_shape))
        
        # Create a map indicating how many patches overlap at each position
        overlap_map = np.zeros((hparams.out_channels,)+application_size, dtype=np.float32)
        for patch_start in tiler.locations:
            patch_end = [start+patch_dim for start, patch_dim in zip(patch_start, hparams.patch_size)]  
            slicing = tuple(map(slice, (0,)+tuple(patch_start), (hparams.out_channels,)+tuple(patch_end)))
            overlap_map[slicing] = overlap_map[slicing]+np.ones((hparams.out_channels,)+hparams.patch_size)
            
        # Initialize maps            
        predicted_img = np.zeros((hparams.out_channels,)+application_size, dtype=np.float32)        
        norm_map = np.zeros((hparams.out_channels,)+application_size, dtype=np.float32)
        
        fading_map_default = np.zeros((hparams.out_channels,)+hparams.patch_size, dtype=np.float32)
        fading_map_default[:,1:-1,1:-1,1:-1] = 1
        fading_map_default = distance_transform_edt(fading_map_default).astype(np.float32)
        fading_map_default /= fading_map_default.max()
        fading_map_default = np.clip(fading_map_default, 1e-4, 1)
        
        for patch_idx in range(tiler.__len__()):
            
            print('Processing patch {0}/{1} for image {2}/{3}...'.format(patch_idx+1, tiler.__len__(), image_idx+1, len(tiler.data_list)))
            
            # Get the mask
            sample = tiler.__getitem__(patch_idx)
            data = Variable(torch.from_numpy(sample['image'][np.newaxis,...]).cuda())
            data = data.float()
            
            # Predict the image
            pred_patch = model(data)
            pred_patch = pred_patch.cpu().data.numpy()
            pred_patch = np.squeeze(pred_patch)
            pred_patch = np.clip(pred_patch, 0, 1)
            
            # Get the current slice position
            slicing = tuple(map(slice, (0,)+tuple(tiler.patch_start), (hparams.out_channels,)+tuple(tiler.patch_end)))
            
            # Compute the fading map for the current position
            fading_map = fading_map_default.copy()
            fading_map[overlap_map[slicing]==overlap_map[slicing].min()] = 1 
            
            # Add predicted patch and fading weights to the corresponding maps
            predicted_img[slicing] = predicted_img[slicing]+pred_patch*fading_map
            norm_map[slicing] = norm_map[slicing]+fading_map
            
        # Normalize and save the predicted image
        predicted_img = predicted_img / norm_map          
        
        # Crop image to the original size
        slicing = tuple(map(slice, (0,0,0,0), (hparams.out_channels,)+tuple(tiler.data_shape)))
        predicted_img = predicted_img[slicing]
        
        # Save the image
#        save_path = os.path.abspath(os.path.join(hparams.output_path, 'pred_'+os.path.split(tiler.data_list[image_idx][0])[-1][:-3]+'.tif'))
#        io.imsave(save_path, np.transpose(predicted_img, (1,2,3,0)).astype(np.float16))
        save_path = os.path.abspath(os.path.join(hparams.data_root, os.path.split(tiler.data_list[image_idx][0])[-1][:-3]+'Predict.tif'))
        print("Saving wall predictions to:" + save_path)
        io.imsave(save_path, (predicted_img[2]*65536).astype(np.uint16))
        
        # Generate the instance mask using seeded watershed
        if hparams.segment == 'Yes':
            print("Segmenting image stack")
            predicted_img[0,...] = ~(predicted_img[0,...]>hparams.bg_thresh)
            predicted_img[1,...] = label(peak_local_max(predicted_img[1,...], threshold_rel=hparams.seed_thresh, min_distance=10, indices=False))[0]
            predicted_img[2,...] = distance_transform_edt(~(predicted_img[2,...]>hparams.boundary_thresh))
            predicted_instances = watershed(-predicted_img[2,...], markers=predicted_img[1,...], mask=predicted_img[0,...])
        
            # Save the instance mask
            save_path = os.path.abspath(os.path.join(hparams.data_root, os.path.split(tiler.data_list[image_idx][0])[-1][:-3]+'Segment.tif'))
            print("Saving segmented cells to:" + save_path)
            io.imsave(save_path, predicted_instances.astype(np.uint16))

        #io.imsave(os.path.abspath(os.path.join(hparams.output_path, 'instances_'+os.path.split(tiler.data_list[image_idx][0])[-1][:-3]+'.tif')), predicted_instances.astype(np.uint16))

if __name__ == '__main__':
    # ------------------------
    # TRAINING ARGUMENTS
    # ------------------------
    # these are project-wide arguments

    parent_parser = ArgumentParser(add_help=False)

    parent_parser.add_argument(
        '--output_path',
        type=str,
        default='results/UNet3D_pyTorch',
        help='output path for test results'
    )
    
    parent_parser.add_argument(
        '--ckpt_path',
        type=str,
        default='results/UNet3D_pyTorch/latest.ckpt',
        help='output path for test results'
    )
    
    parent_parser.add_argument(
        '--overlap',
        type=int,
        default=20,
        help='overlap of adjacent patches'
    )
    
    parent_parser.add_argument(
        '--bg_thresh',
        type=float,
        default=0.8,
        help='background threshold'
    )
    
    parent_parser.add_argument(
        '--seed_thresh',
        type=float,
        default=0.2,
        help='seed threshold'
    )
    
    parent_parser.add_argument(
        '--boundary_thresh',
        type=float,
        default=0.5,
        help='boundary threshold'
    )

    parent_parser.add_argument(
        '--segment',
        type=str,
        default='No',
        help='run segmentation after prediction'
    ) 
    # each LightningModule defines arguments relevant to it
    parser = UNet3D.add_model_specific_args(parent_parser)
    hyperparams = parser.parse_args()

    # ---------------------
    # RUN TRAINING
    # ---------------------
    main(hyperparams)
