import multiprocessing
from os import path
import argparse
import time
from datetime import datetime
import sys 
# sys.path.append('./../code')
# from utils_quintin import * 
from sfransen.utils_quintin import *
# sys.path.append('./../code/DWI_exp')    
# from callbacks import IntermediateImages, dice_coef
# from callbacks import RocCallback
from sfransen.utils_quintin import *
from sfransen.DWI_exp import IntermediateImages, dice_coef
from sfransen.DWI_exp.preprocessing_function import preprocess
import yaml
import numpy as np
from tqdm import tqdm
import tensorflow as tf
tf.compat.v1.disable_eager_execution()

from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, CSVLogger
from tensorflow.keras.optimizers import Adam
from sklearn.model_selection import KFold

from sfransen.DWI_exp.helpers import *
from sfransen.DWI_exp.batchgenerator import BatchGenerator

from sfransen.DWI_exp.unet import build_dual_attention_unet
from focal_loss import BinaryFocalLoss

parser = argparse.ArgumentParser(
    description='Train a U-Net model for segmentation/detection tasks.' + 
                'using cross-validation.')
parser.add_argument('--series', '-s', 
    metavar='[series_name]', required=True, nargs='+',
    help='List of series to include, must correspond with' +
        "path files in ./data/")
parser.add_argument('-experiment',
    help='add experiment title to store the files correctly: test_b50_b400_b800' 
)
args = parser.parse_args()

# Determine the number of input series
num_series = len(args.series)

# Identify this job by the series included in the training
# Output folder will have this name, e.g.: b0_b50_b100
# JOB_NAME = '_'.join(args.series)
JOB_NAME = args.experiment
DATA_DIR = "./../data/Nijmegen paths/"
# DATA_DIR = "./../data/new/"
# PROJECT_DIR = f"/data/pca-rad/sfransen/train_output/{args.experiment}"
PROJECT_DIR = f"/data/pca-rad/sfransen/train_output/{JOB_NAME}"
# 2 x 2mm2 in-plane resolution, 3.6mm slice thickness
TARGET_SPACING = (0.5, 0.5, 3) 
INPUT_SHAPE = (192, 192, 24, num_series) #(64, 64, 20, num_series)
IMAGE_SHAPE = INPUT_SHAPE[:3]
OUTPUT_SHAPE = (192, 192, 24, 1) # One output channel (segmentation)

# Hyperparameters
FOCAL_LOSS_GAMMA = 2
INITIAL_LEARNING_RATE = 1e-4
MAX_EPOCHS = 600
EARLY_STOPPING = 50
# increase batch size
BATCH_SIZE = 12
MODEL_SELECTION_METRIC = 'val_loss'
MODEL_SELECTION_DIRECTION = "min" # Change to 'max' if higher value is better
EARLY_STOPPING_METRIC = 'val_loss'
EARLY_STOPPING_DIRECTION = "min" # Change to 'max' if higher value is better

# Training configuration
# add metric ROC_AUC
TRAINING_METRICS = ["binary_crossentropy", "binary_accuracy", dice_coef]
loss = BinaryFocalLoss(gamma=FOCAL_LOSS_GAMMA)
optimizer = Adam(learning_rate=INITIAL_LEARNING_RATE)

# Create folder structure in the output directory
if path.exists(PROJECT_DIR):
    prepare_project_dir(PROJECT_DIR+'_(2)')
else:
    prepare_project_dir(PROJECT_DIR)

#save params to yaml
params = {
    "focal_loss_gamma": FOCAL_LOSS_GAMMA,
    "initial_learning_rate": INITIAL_LEARNING_RATE,
    "max_epochs": MAX_EPOCHS,
    "MODEL_SELECTION_METRIC": MODEL_SELECTION_METRIC,
    "EARLY_STOPPING_METRIC": EARLY_STOPPING_METRIC,
    "train_output_dir": PROJECT_DIR,
    "batch_size": BATCH_SIZE,
    "optimizer": optimizer,
    "loss": loss,
    "datetime": print(datetime.now().strftime("%Y-%m-%d"))}
dump_dict_to_yaml(params, f"{PROJECT_DIR}", filename=f"params")

# Build the U-Net model
detection_model = build_dual_attention_unet(INPUT_SHAPE)
detection_model.summary(line_length=120)

# Load all numpy images into RAM
images, image_paths = {s: [] for s in args.series}, {}
segmentations = []
print_(f"> Loading images into RAM...")

# Read the image paths from the data directory.
# Texts files are expected to have the name "[series_name].txt"
for s in args.series:
    with open(path.join(DATA_DIR, f"{s}.txt"), 'r') as f:
        image_paths[s] = [l.strip() for l in f.readlines()]
with open(path.join(DATA_DIR, f"seg.txt"), 'r') as f:
    seg_paths = [l.strip() for l in f.readlines()]
num_images = len(seg_paths)

# Read and preprocess each of the paths for each series, and the segmentations.
for img_idx in tqdm(range(num_images)):  # [:40]):for less images
    img_s = {s: sitk.ReadImage(image_paths[s][img_idx], sitk.sitkFloat32) 
        for s in args.series}
    seg_s = sitk.ReadImage(seg_paths[img_idx], sitk.sitkFloat32)
    img_n, seg_n = preprocess(img_s, seg_s, 
        shape=IMAGE_SHAPE, spacing=TARGET_SPACING)
    for seq in img_n:
        images[seq].append(img_n[seq])
    segmentations.append(seg_n)

# Split train and validation 
# We use KFold to split the data, but we don't actually do cross validation, we 
# just use it to split the data 1:9.
# kfold = KFold(10, shuffle=True, random_state=123)
# train_idxs, valid_idxs = list(kfold.split(segmentations))[0]
# train_idxs = list(train_idxs)
# valid_idxs = list(valid_idxs)

yml_paths = read_yaml_to_dict('./../data/Nijmegen paths/train_val_test_idxs.yml')
train_idxs = yml_paths['train_set0']
valid_idxs = yml_paths['val_set0']


detection_model.compile(
    optimizer=optimizer, 
    loss=loss,
    metrics=TRAINING_METRICS
    )

train_generator = BatchGenerator(images, segmentations,
    sequences=args.series,
    shape=IMAGE_SHAPE,
    indexes=train_idxs, 
    batch_size=BATCH_SIZE,
    shuffle=True, 
    augmentation_function=augment
)

valid_generator = get_generator(images, segmentations,
    sequences=args.series,
    shape=IMAGE_SHAPE,
    indexes=valid_idxs, 
    batch_size=None,
    shuffle=False, 
    augmentation=None
)
valid_data = next(valid_generator)
print_(f"The shape of valid_data input = {np.shape(valid_data[0])}")
print_(f"The shape of valid_data label = {np.shape(valid_data[1])}")

callbacks = [
    EarlyStopping(
        monitor=EARLY_STOPPING_METRIC,
        mode=EARLY_STOPPING_DIRECTION,
        patience=EARLY_STOPPING,
        verbose=1),
    ModelCheckpoint(
        filepath=path.join(PROJECT_DIR, "models", JOB_NAME + ".h5"), 
        monitor=MODEL_SELECTION_METRIC,
        mode=MODEL_SELECTION_DIRECTION,
        verbose=1,
        save_best_only=True),
    # ModelCheckpoint(
    #     filepath=path.join(PROJECT_DIR, "models_dice", JOB_NAME + ".h5"), 
    #     monitor='val_dice_coef',
    #     mode='max',
    #     verbose=0,
    #     save_best_only=True),
    CSVLogger(
        filename=path.join(PROJECT_DIR, "logs", f"{JOB_NAME}.csv")),
    IntermediateImages(
        validation_set=valid_data,
        sequences=args.series,
        prefix=path.join(PROJECT_DIR, "output", JOB_NAME),
        num_images=25)
    # RocCallback(
    #      validation_set=valid_data,
    #      num_images=25)
]

detection_model.fit(train_generator,
    validation_data    = valid_data,
    steps_per_epoch    = len(train_idxs) // BATCH_SIZE, 
    epochs             = MAX_EPOCHS,
    callbacks          = callbacks,
    verbose            = 1,
    )