import os
import cv2
from skimage import io
import sys
# import keras
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import argparse
from natsort import natsorted
# sys.path.append('../')
import segmentation_models_v1 as sm
from segmentation_models_v1 import Unet, Linknet, PSPNet, FPN, AtUnet, ResUnet
sm.set_framework('tf.keras')

from helper_function import plot_history_flu2, save_phase_fl_history, plot_flu_prediction, plot_set_prediction
from helper_function import save_history_for_callback, plot_history_for_callback
from helper_function import precision, recall, f1_score, calculate_psnr, calculate_pearsonr
from sklearn.metrics import confusion_matrix

def str2bool(value):
    return value.lower() == 'true'

def generate_folder(folder_name):
	if not os.path.exists(folder_name):
		os.system('mkdir -p {}'.format(folder_name))

parser = argparse.ArgumentParser()
parser.add_argument("--gpu", type=str, default = '2')
parser.add_argument("--docker", type=str2bool, default = True)
parser.add_argument("--net_type", type=str, default = 'Unet')  #Unet, Linknet, PSPNet, FPN
parser.add_argument("--backbone", type=str, default = 'efficientnetb0')
parser.add_argument("--dataset", type=str, default = 'neuron_float')
parser.add_argument("--subset", type=str, default = 'train')
parser.add_argument("--epoch", type=int, default = 10)
parser.add_argument("--run", type=int, default = 1)
parser.add_argument("--dim", type=int, default = 512)
parser.add_argument("--ch_in", type=int, default = 3)
parser.add_argument("--ch_out", type=int, default = 3)
parser.add_argument("--fl_ch", type=str, default = 'fl12')
parser.add_argument("--rot", type=float, default = 0)
parser.add_argument("--scale", type=float, default = 100)
parser.add_argument("--train", type=int, default = None)
parser.add_argument("--act_fun", type=str, default = 'relu')
parser.add_argument("--loss", type=str, default = 'mse')
parser.add_argument("--batch_size", type=int, default = 6)
parser.add_argument("--lr", type=float, default = 5e-4)
parser.add_argument("--decay", type=float, default = 0.8)
parser.add_argument("--delta", type=float, default = 10)
parser.add_argument("--best_select", type=str2bool, default = True)  ## cancel the selection of best model
parser.add_argument("--pre_train", type=str2bool, default = True)
args = parser.parse_args()
print(args)

model_name = 'Cor-FL1_FL2-net-{}-bone-{}-pre-{}-epoch-{}-batch-{}-lr-{}-dim-{}-train-{}-rot-{}-set-{}-subset-{}-loss-{}-act-{}-scale-{}-decay-{}-delta-{}-chi-{}-cho-{}-chf-{}-bselect-{}-run-{}'.format(args.net_type, args.backbone, args.pre_train,\
		 args.epoch, args.batch_size, args.lr, args.dim, args.train, args.rot, args.dataset, args.subset, args.loss, args.act_fun, args.scale, args.decay, args.delta, args.ch_in, args.ch_out, args.fl_ch, args.best_select, args.run)
print(model_name)

os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu

DATA_DIR = '/data/datasets/{}'.format(args.dataset) if args.docker else './datasets/{}'.format(args.dataset)
train_dim = args.dim

# load the sample names
def read_samples(file_name):
		with open(file_name, 'r+') as f:
				lines = [fn.strip() for fn in f.readlines()]
		return lines

def read_end_points(file_name):
		sample_dict = {}
		with open(file_name, 'r+') as f:
				for line in f.readlines():
						splits = line.strip().split(' ')
						sample_dict[splits[0]] = [int(splits[1]), int(splits[2])]
		return sample_dict

sample_dict = None
if 'neuron' in args.dataset:
		sample_dict = read_end_points(os.path.join(DATA_DIR, 'range.txt'))
train_fns =  read_samples(os.path.join(DATA_DIR, 'train.txt'))
test_fns =   read_samples(os.path.join(DATA_DIR, 'test.txt'))
data_dir = DATA_DIR + '/data'
val_dim = 1760

# classes for data loading and preprocessing
class Dataset:
    """CamVid Dataset. Read images, apply augmentation and preprocessing transformations.
    
    Args:
        images_dir (str): path to images folder
        masks_dir (str): path to segmentation masks folder
        class_values (list): values of classes to extract from segmentation mask
        augmentation (albumentations.Compose): data transfromation pipeline 
            (e.g. flip, scale, etc.)
        preprocessing (albumentations.Compose): data preprocessing 
            (e.g. noralization, shape manipulation, etc.)
    
    """
    
    def __init__(
            self, 
            data_dir, 
            sample_names,
            end_point_dict,
            fl_ch = None,
            scale = 1.0,
            channels = [3,3],
            augmentation=None, 
            preprocessing=None,
    ):
        self.images_fps = []
        self.masks1_fps = []
        self.masks2_fps = []
        for sn in sample_names:
        		sample_tag = 'T-' + sn.split('_')[3][5:]
        		if end_point_dict:
        				end1, end2 = end_point_dict[sample_tag]
        		else:
        				end1, end2 = 0, np.inf
        		fns = os.listdir(os.path.join(data_dir, sn, 'phase'))
        		for fn in fns:
        				if end1 <= int(fn.split('.')[0].split('-')[-1]) <= end2:
        						self.images_fps.append(os.path.join(data_dir, sn, 'phase', fn))
        						self.masks1_fps.append(os.path.join(data_dir, sn, 'fl1', fn))
        						self.masks2_fps.append(os.path.join(data_dir, sn, 'fl2', fn))
        self.ids = self.images_fps
        print('Load files: image {}, fl1: {}, fl2:{}'.format(len(self.images_fps),len(self.masks1_fps),len(self.masks2_fps)))      
        self.scale = scale
        self.augmentation = augmentation
        self.preprocessing = preprocessing
        self.channels = channels
        self.fl_ch = fl_ch
    
    def __getitem__(self, i):
        
        # load image and fl1 or fl2 or both
        image = np.load(self.images_fps[i]) * 255.
        if self.fl_ch == 'fl1':
            mask  = np.load(self.masks1_fps[i])
            mask = mask * self.scale
        elif self.fl_ch == 'fl2':
            mask  = np.load(self.masks2_fps[i])
            mask = mask * self.scale        
        elif self.fl_ch == 'fl12':
            mask1  = np.load(self.masks1_fps[i])
            mask2  = np.load(self.masks2_fps[i])
            mask = np.stack([mask1[:,:,1], mask2[:,:,1]], axis = -1)
            mask = mask*self.scale
        
        # decide the input and output channels
        if self.channels[0] == 1:
            image[:,:,0], image[:,:,2] = image[:,:,1], image[:,:,1]
        elif self.channels[0] == 2:
            image[:,:,2] = image[:,:,1]		
        
        if self.channels[1] == 1 and not (self.fl_ch=='fl12'):
            mask = mask[:,:,1:2]

        # apply augmentations
        if self.augmentation:
            sample = self.augmentation(image=image, mask=mask)
            image, mask = sample['image'], sample['mask']

        # apply preprocessing
        if self.preprocessing:
            sample = self.preprocessing(image=image, mask=mask)
            image, mask = sample['image'], sample['mask']

        return image, mask
        
    def __len__(self):
        return len(self.ids)
    

class Dataloder(tf.keras.utils.Sequence):
    """Load data from dataset and form batches
    
    Args:
        dataset: instance of Dataset class for image loading and preprocessing.
        batch_size: Integet number of images in batch.
        shuffle: Boolean, if `True` shuffle image indexes each epoch.
    """
    
    def __init__(self, dataset, batch_size=1, shuffle=False):
        self.dataset = dataset
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.indexes = np.arange(len(dataset))

        self.on_epoch_end()

    def __getitem__(self, i):
        
        # collect batch data
        start = i * self.batch_size
        stop = (i + 1) * self.batch_size
        data = []
        for j in range(start, stop):
            data.append(self.dataset[j])
        
        # transpose list of lists
        batch = [np.stack(samples, axis=0) for samples in zip(*data)]
        
        return (batch[0], batch[1])
    
    def __len__(self):
        """Denotes the number of batches per epoch"""
        return len(self.indexes) // self.batch_size
    
    def on_epoch_end(self):
        """Callback function to shuffle indexes each epoch"""
        if self.shuffle:
            self.indexes = np.random.permutation(self.indexes)

import albumentations as A

def round_clip_0_1(x, **kwargs):
    return x.round().clip(0, 1)

# define heavy augmentations
def get_training_augmentation(dim, rot = 0):
    train_transform = [
        A.HorizontalFlip(p=0.5),
        A.PadIfNeeded(min_height=dim, min_width=dim, always_apply=True, border_mode=0),
        A.RandomCrop(height=dim, width=dim, always_apply=True),]
    return A.Compose(train_transform)


def get_validation_augmentation(dim = 992):
    """Add paddings to make image shape divisible by 32"""
    test_transform = [
        A.PadIfNeeded(dim, dim)
    ]
    return A.Compose(test_transform)

def get_preprocessing(preprocessing_fn):
    """Construct preprocessing transform
    
    Args:
        preprocessing_fn (callbale): data normalization function 
            (can be specific for each pretrained neural network)
    Return:
        transform: albumentations.Compose
    
    """
    _transform = [
        A.Lambda(image=preprocessing_fn),
    ]
    return A.Compose(_transform)

## create models
BACKBONE = args.backbone
BATCH_SIZE = args.batch_size
LR = args.lr
EPOCHS = args.epoch

# processing configuration
preprocess_input = sm.get_preprocessing(BACKBONE)

# define network parameters
n_classes = args.ch_out if args.fl_ch == 'fl1' or args.fl_ch == 'fl2' else 2
activation = '{}'.format(args.act_fun)

#create model
net_func = globals()[args.net_type]

encoder_weights='imagenet' if args.pre_train else None

model = net_func(BACKBONE, encoder_weights=encoder_weights, classes=n_classes, activation=activation)

# define optomizer
optim = tf.keras.optimizers.Adam(LR)

if args.loss == 'mse':
	loss = tf.keras.losses.MSE
elif args.loss == 'mae':
	loss = tf.keras.losses.MAE
elif args.loss == 'huber':
	loss = tf.keras.losses.Huber(reduction=tf.keras.losses.Reduction.NONE)

from tensorflow.keras import backend as K
def pearson(y_true, y_pred):
    x = y_true
    y = y_pred
    mx = K.mean(x)
    my = K.mean(y)
    xm, ym = x-mx, y-my
    r_num = K.sum(tf.multiply(xm,ym))
    r_den = K.sqrt(tf.multiply(K.sum(K.square(xm)), K.sum(K.square(ym))))
    r = r_num / r_den
    return r

metrics = [sm.metrics.PSNR(max_val=args.scale), pearson]

# compile keras model with defined optimozer, loss and metrics
model.compile(optim, loss, metrics)

# Dataset for train images
train_dataset = Dataset(
		data_dir = data_dir,
		sample_names = train_fns,
		end_point_dict = sample_dict,
    fl_ch = args.fl_ch,
    channels = [args.ch_in, args.ch_out],
    scale = args.scale,
    augmentation=get_training_augmentation(train_dim, args.rot),
    preprocessing=get_preprocessing(preprocess_input),
)

# Dataset for validation images
valid_dataset = Dataset(
		data_dir = data_dir,
		sample_names = test_fns,
		end_point_dict = sample_dict,
    fl_ch = args.fl_ch,
    scale = args.scale,
    channels = [args.ch_in, args.ch_out],
    augmentation=get_validation_augmentation(val_dim),
    preprocessing=get_preprocessing(preprocess_input),
)

train_dataloader = Dataloder(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
valid_dataloader = Dataloder(valid_dataset, batch_size=1, shuffle=False)

print(train_dataloader[0][0].shape)
print(train_dataloader[0][1].shape)
print(train_dataloader[0][1].min(), train_dataloader[0][1].max())
# check shapes for errors
assert train_dataloader[0][0].shape == (BATCH_SIZE, train_dim, train_dim, 3)
assert train_dataloader[0][1].shape == (BATCH_SIZE, train_dim, train_dim, n_classes)

model_folder = '/data/2d_models/{}/{}'.format(args.dataset, model_name) if args.docker else './2d_models/{}/{}'.format(args.dataset, model_name)
generate_folder(model_folder)

def concat_tile(im_list_2d):
    return cv2.vconcat([cv2.hconcat(im_list_h) for im_list_h in im_list_2d])

def save_images(file_name, vols):
		vols = vols[:,:,:,1] if vols.shape[-1] >= 2 else vols[:,:,:,0]
		shp = vols.shape
		ls, lx, ly = shp
		sx, sy = int(lx/128), int(ly/128)
		vols = vols[:,::sx,::sy]
		slice_list, rows = [], []
		for si in range(vols.shape[0]):
				slice = vols[si,:,:]
				slice[0, :] = 255
				slice[:, 0] = 255
				slice[:, -1] = 255
				slice[-1, :] = 255
				rows.append(slice)
				if si%8 == 7 and not si == vols.shape[0]-1:
						slice_list.append(rows)
						rows = []
		save_img = concat_tile(slice_list)		
		cv2.imwrite(file_name, save_img)

class HistoryPrintCallback(tf.keras.callbacks.Callback):
		def __init__(self):
				super(HistoryPrintCallback, self).__init__()
				self.history = {}

		def on_epoch_end(self, epoch, logs=None):
				if logs:
						for key in logs.keys():
								if epoch == 0:
										self.history[key] = []
								self.history[key].append(logs[key])
				if epoch%5 == 0:
						plot_history_for_callback(model_folder+'/train_history.png', self.history)
						save_history_for_callback(model_folder, self.history)
						img_vols, gt_vols, pr_vols = [],[],[]
						for i in range(0, len(valid_dataset),int(len(valid_dataset)/64)):
								img_vols.append(np.load(valid_dataloader.dataset.images_fps[i]))
								gt_vols.append(valid_dataloader[i][1])
								pr_vols.append(self.model.predict(valid_dataloader[i]))
						img_vols = np.stack(img_vols, axis = 0)
						gt_vols = np.concatenate(gt_vols, axis = 0)
						pr_vols = np.concatenate(pr_vols, axis = 0)
						save_images(model_folder+'/epoch-{}-img.png'.format(epoch), np.uint8(img_vols))
						save_images(model_folder+'/epoch-{}-gt.png'.format(epoch), gt_vols/args.scale*255)
						save_images(model_folder+'/epoch-{}-pr.png'.format(epoch), pr_vols/args.scale*255)


if not args.best_select:
		callbacks = [
				tf.keras.callbacks.ModelCheckpoint(model_folder+'/weights_{epoch:02d}.h5', save_weights_only=True, save_best_only=False, period=5),
				tf.keras.callbacks.ReduceLROnPlateau(factor=args.decay),
				HistoryPrintCallback(),
		]
else:
		callbacks = [
				tf.keras.callbacks.ModelCheckpoint(model_folder+'/best_model-{epoch:03d}.h5', monitor='val_pearson', save_weights_only=True, save_best_only=True, mode='max'),
				tf.keras.callbacks.ReduceLROnPlateau(factor=args.decay),
				HistoryPrintCallback(),
		]


# train model
history = model.fit_generator(
    train_dataloader, 
    steps_per_epoch=len(train_dataloader), 
    epochs=EPOCHS, 
    callbacks=callbacks, 
    validation_data=valid_dataloader, 
    validation_steps=len(valid_dataloader),
)

# evaluate model
test_dataset = Dataset(
    x_test_dir, 
    y1_test_dir,
    y2_test_dir,
    fl_ch = args.fl_ch,
    channels = [args.ch_in, args.ch_out],
    scale = args.scale,
    augmentation=get_validation_augmentation(val_dim),
    preprocessing=get_preprocessing(preprocess_input),
)
tensorflow.keras.callbacks.ModelCheckpoint tensorflow.multiply numpy.uint8 tensorflow.keras.callbacks.ReduceLROnPlateau numpy.stack numpy.concatenate tensorflow.keras.losses.Huber tensorflow.keras.backend.square tensorflow.keras.optimizers.Adam numpy.random.permutation tensorflow.keras.backend.mean numpy.load
[(14, 'segmentation_models_v1.set_framework', 'sm.set_framework', (['"""tf.keras"""'], {}), True, 'import segmentation_models_v1 as sm\n'), (28, 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {}), False, 'import argparse\n'), (259, 'segmentation_models_v1.get_preprocessing', 'sm.get_preprocessing', (['BACKBONE'], {}), True, 'import segmentation_models_v1 as sm\n'), (273, 'tensorflow.keras.optimizers.Adam', 'tf.keras.optimizers.Adam', (['LR'], {}), True, 'import tensorflow as tf\n'), (81, 'os.path.join', 'os.path.join', (['DATA_DIR', '"""train.txt"""'], {}), False, 'import os\n'), (82, 'os.path.join', 'os.path.join', (['DATA_DIR', '"""test.txt"""'], {}), False, 'import os\n'), (227, 'albumentations.Compose', 'A.Compose', (['train_transform'], {}), True, 'import albumentations as A\n'), (235, 'albumentations.Compose', 'A.Compose', (['test_transform'], {}), True, 'import albumentations as A\n'), (250, 'albumentations.Compose', 'A.Compose', (['_transform'], {}), True, 'import albumentations as A\n'), (286, 'tensorflow.keras.backend.mean', 'K.mean', (['x'], {}), True, 'from tensorflow.keras import backend as K\n'), (287, 'tensorflow.keras.backend.mean', 'K.mean', (['y'], {}), True, 'from tensorflow.keras import backend as K\n'), (294, 'segmentation_models_v1.metrics.PSNR', 'sm.metrics.PSNR', ([], {'max_val': 'args.scale'}), True, 'import segmentation_models_v1 as sm\n'), (357, 'cv2.imwrite', 'cv2.imwrite', (['file_name', 'save_img'], {}), False, 'import cv2\n'), (25, 'os.path.exists', 'os.path.exists', (['folder_name'], {}), False, 'import os\n'), (80, 'os.path.join', 'os.path.join', (['DATA_DIR', '"""range.txt"""'], {}), False, 'import os\n'), (224, 'albumentations.HorizontalFlip', 'A.HorizontalFlip', ([], {'p': '(0.5)'}), True, 'import albumentations as A\n'), (225, 'albumentations.PadIfNeeded', 'A.PadIfNeeded', ([], {'min_height': 'dim', 'min_width': 'dim', 'always_apply': '(True)', 'border_mode': '(0)'}), True, 'import albumentations as A\n'), (226, 'albumentations.RandomCrop', 'A.RandomCrop', ([], {'height': 'dim', 'width': 'dim', 'always_apply': '(True)'}), True, 'import albumentations as A\n'), (233, 'albumentations.PadIfNeeded', 'A.PadIfNeeded', (['dim', 'dim'], {}), True, 'import albumentations as A\n'), (248, 'albumentations.Lambda', 'A.Lambda', ([], {'image': 'preprocessing_fn'}), True, 'import albumentations as A\n'), (289, 'tensorflow.multiply', 'tf.multiply', (['xm', 'ym'], {}), True, 'import tensorflow as tf\n'), (388, 'tensorflow.keras.callbacks.ModelCheckpoint', 'tf.keras.callbacks.ModelCheckpoint', (["(model_folder + '/weights_{epoch:02d}.h5')"], {'save_weights_only': '(True)', 'save_best_only': '(False)', 'period': '(5)'}), True, 'import tensorflow as tf\n'), (389, 'tensorflow.keras.callbacks.ReduceLROnPlateau', 'tf.keras.callbacks.ReduceLROnPlateau', ([], {'factor': 'args.decay'}), True, 'import tensorflow as tf\n'), (394, 'tensorflow.keras.callbacks.ModelCheckpoint', 'tf.keras.callbacks.ModelCheckpoint', (["(model_folder + '/best_model-{epoch:03d}.h5')"], {'monitor': '"""val_pearson"""', 'save_weights_only': '(True)', 'save_best_only': '(True)', 'mode': '"""max"""'}), True, 'import tensorflow as tf\n'), (395, 'tensorflow.keras.callbacks.ReduceLROnPlateau', 'tf.keras.callbacks.ReduceLROnPlateau', ([], {'factor': 'args.decay'}), True, 'import tensorflow as tf\n'), (138, 'numpy.load', 'np.load', (['self.images_fps[i]'], {}), True, 'import numpy as np\n'), (140, 'numpy.load', 'np.load', (['self.masks1_fps[i]'], {}), True, 'import numpy as np\n'), (203, 'numpy.stack', 'np.stack', (['samples'], {'axis': '(0)'}), True, 'import numpy as np\n'), (214, 'numpy.random.permutation', 'np.random.permutation', (['self.indexes'], {}), True, 'import numpy as np\n'), (280, 'tensorflow.keras.losses.Huber', 'tf.keras.losses.Huber', ([], {'reduction': 'tf.keras.losses.Reduction.NONE'}), True, 'import tensorflow as tf\n'), (337, 'cv2.hconcat', 'cv2.hconcat', (['im_list_h'], {}), False, 'import cv2\n'), (371, 'helper_function.plot_history_for_callback', 'plot_history_for_callback', (["(model_folder + '/train_history.png')", 'self.history'], {}), False, 'from helper_function import save_history_for_callback, plot_history_for_callback\n'), (372, 'helper_function.save_history_for_callback', 'save_history_for_callback', (['model_folder', 'self.history'], {}), False, 'from helper_function import save_history_for_callback, plot_history_for_callback\n'), (378, 'numpy.stack', 'np.stack', (['img_vols'], {'axis': '(0)'}), True, 'import numpy as np\n'), (379, 'numpy.concatenate', 'np.concatenate', (['gt_vols'], {'axis': '(0)'}), True, 'import numpy as np\n'), (380, 'numpy.concatenate', 'np.concatenate', (['pr_vols'], {'axis': '(0)'}), True, 'import numpy as np\n'), (121, 'os.path.join', 'os.path.join', (['data_dir', 'sn', '"""phase"""'], {}), False, 'import os\n'), (143, 'numpy.load', 'np.load', (['self.masks2_fps[i]'], {}), True, 'import numpy as np\n'), (290, 'tensorflow.keras.backend.square', 'K.square', (['xm'], {}), True, 'from tensorflow.keras import backend as K\n'), (290, 'tensorflow.keras.backend.square', 'K.square', (['ym'], {}), True, 'from tensorflow.keras import backend as K\n'), (381, 'numpy.uint8', 'np.uint8', (['img_vols'], {}), True, 'import numpy as np\n'), (146, 'numpy.load', 'np.load', (['self.masks1_fps[i]'], {}), True, 'import numpy as np\n'), (147, 'numpy.load', 'np.load', (['self.masks2_fps[i]'], {}), True, 'import numpy as np\n'), (148, 'numpy.stack', 'np.stack', (['[mask1[:, :, (1)], mask2[:, :, (1)]]'], {'axis': '(-1)'}), True, 'import numpy as np\n'), (375, 'numpy.load', 'np.load', (['valid_dataloader.dataset.images_fps[i]'], {}), True, 'import numpy as np\n'), (124, 'os.path.join', 'os.path.join', (['data_dir', 'sn', '"""phase"""', 'fn'], {}), False, 'import os\n'), (125, 'os.path.join', 'os.path.join', (['data_dir', 'sn', '"""fl1"""', 'fn'], {}), False, 'import os\n'), (126, 'os.path.join', 'os.path.join', (['data_dir', 'sn', '"""fl2"""', 'fn'], {}), False, 'import os\n')]