# -*- coding: utf-8 -*-
# Author: allen
# Copyright (c) 2022 RuiCare All rights reserved.
import os

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ["SM_FRAMEWORK"] = "tf.keras"
import tensorflow as tf
from by3d import unet3d
from generator3d_in_ct_dose import Generator
# from generator3d_in_ct_dose_ptv import Generator
from tensorflow.python.keras.optimizer_v2.adam import Adam
from keras.callbacks import ModelCheckpoint, TensorBoard


def ssim(y_truth, y_pred):
    return tf.reduce_mean(tf.image.ssim(y_truth, y_pred, 60))


def ssim_loss(y_truth, y_pred):
    return 1 - ssim(y_truth, y_pred)


if __name__ == '__main__':

    gen = Generator(train_dir='/home/allen/Temp/PlanQA/train-normalized',
                    batch_size=1,
                    shape=[8, 512, 512, 2])

    save_model_path = f'../model/model_8_mse_sigmoid_v2.hdf5'
    pre_model_path = f'../model/model_8_mse_sigmoid_v1.hdf5'

    config = tf.compat.v1.ConfigProto()
    config.gpu_options.allow_growth = True
    session = tf.compat.v1.Session(config=config)

    model = unet3d.load_unet3d(input_shape=gen.shape, activation='sigmoid')
    model.compile(optimizer=Adam(lr=0.0001), loss='mean_squared_error', metrics=[ssim, 'accuracy'])

    model._get_distribution_strategy = lambda: None
    if pre_model_path and os.path.exists(pre_model_path):  model.load_weights(pre_model_path)
    checkpoint = ModelCheckpoint(save_model_path, monitor='loss', verbose=1, save_best_only=True)
    board = TensorBoard(log_dir='../log', histogram_freq=1, update_freq='batch')

    model.fit(gen, None, batch_size=gen.batch_size,
              shuffle=True, epochs=100, steps_per_epoch=gen.steps_per_epoch,
              use_multiprocessing=False, workers=10, max_queue_size=10, callbacks=[checkpoint, board])
