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

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import cv2.cv2
import numpy as np
import tensorflow as tf
from by3d import unet3d
import SimpleITK as sitk
from util import imgutil, fileutil
from tensorflow.python.keras.optimizer_v2.adam import Adam


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)


def predict_one(model, ct_nii_path, dose_plan_nii_path, ptv_nii_path, in_shape, twice_pred=False):
    ct = imgutil.load_nii(ct_nii_path, dtype=np.float16, zfirst=True)
    dose_plan = imgutil.load_nii(dose_plan_nii_path, dtype=np.float16, zfirst=True)
    ptv = imgutil.load_nii(ptv_nii_path, dtype=np.float16, zfirst=True)
    z, x, y = ct.shape

    out_img = np.zeros(shape=ct.shape, dtype=np.float32)
    for b in range(0, z, in_shape[0]):
        imgs = np.zeros(shape=(1, in_shape[0], in_shape[1], in_shape[2], 3), dtype=np.float16)

        begin = z - in_shape[0] if b + in_shape[0] >= z else b
        end = z if b + in_shape[0] >= z else b + in_shape[0]
        imgs[0, :, :, :, 0] = ct[begin:end, :, :]
        imgs[0, :, :, :, 1] = dose_plan[begin:end, :, :]
        imgs[0, :, :, :, 2] = ptv[begin:end, :, :]
        out_batch_img = model.predict(imgs, steps=1, verbose=0)
        out_img[begin:end, :, :] = out_batch_img[0, :, :, :, 0]

    if twice_pred:
        for b in range(in_shape[0] // 2, z, in_shape[0]):
            imgs = np.zeros(shape=(1, in_shape[0], in_shape[1], in_shape[2], 3), dtype=np.float16)

            begin = z - in_shape[0] if b + in_shape[0] >= z else b
            end = z if b + in_shape[0] >= z else b + in_shape[0]
            imgs[0, :, :, :, 0] = ct[begin:end, :, :]
            imgs[0, :, :, :, 1] = dose_plan[begin:end, :, :]
            imgs[0, :, :, :, 2] = ptv[begin:end, :, :]
            out_batch_img = model.predict(imgs, steps=1, verbose=0)
            quarter = in_shape[0] // 4
            out_img[begin + quarter:end - quarter, :, :] = out_batch_img[0, quarter:in_shape[0] - quarter, :, :, 0]

    # out_img = out_img + 0.5
    out_img[out_img < 0] = 0
    return out_img


if __name__ == '__main__':
    in_dir = '/home/allen/Temp/PlanQA/test'
    model_path = f'../model/model_8_3_mse_relu_v2.hdf5'
    activation = 'relu'
    in_shape = [8, 512, 512, 3]
    twice_pred = True

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

    model = unet3d.load_unet3d(input_shape=in_shape, activation=activation)
    model.compile(optimizer='adam', loss='mean_squared_error', metrics=[ssim, 'accuracy'])
    model.load_weights(model_path)

    for d in fileutil.list_dir(in_dir):
        print(f'predicting {os.path.basename(d)} ...')
        # if os.path.basename(d) != 'RT192924': continue
        ct_nii_path = f'{d}/ct.nii.gz'
        dose_plan_nii_path = f'{d}/dose_plan.nii.gz'
        ptv_nii_path = f'{d}/rs_ptv.nii.gz'
        out_path = f'{d}/dose_pred_8_3_mse_relu_v1.nii.gz'
        out_img = predict_one(model, ct_nii_path, dose_plan_nii_path, ptv_nii_path, in_shape, twice_pred=twice_pred)
        nii = sitk.ReadImage(dose_plan_nii_path)
        imgutil.write_nii(out_img, out_path, spacing=nii.GetSpacing(), origin=nii.GetOrigin())
