import tensorflow as tf
import numpy as np
import cv2
from pathlib import Path
import scipy.misc
import Network
import os

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # only uses GPU 1


results_dir = './results/'
results_best_dir = './results/best_model/'


save_file_UNET = './results/npo_cubic_e2e_MSE.csv'
save_file_UNET_SSIM = './results/npo_cubic_e2e_SSIM.csv'


# read test image
image_read_dir = './image/'
GT_path = str(Path(image_read_dir + 'test_img.png'))
gt_img = cv2.imread(GT_path, 0)/255
GT = np.tile(gt_img, [21, 1, 1])
GT = np.expand_dims(GT, -1)
GT = tf.convert_to_tensor(GT, dtype=tf.float32)

#save image
image_save_dir = './image/results/'

####### read from the TFRECORD format #################
## for faster reading from Hard disk
def read_tfrecord(TFRECORD_PATH):
    # from tfrecord file to data
    N_w = 326  # size of the images
    N_h = 326
    queue = tf.train.string_input_producer(TFRECORD_PATH, shuffle=True)
    reader = tf.TFRecordReader()

    _, serialized_example = reader.read(queue)

    features = tf.parse_single_example(serialized_example,
                                       features={
                                           'sharp': tf.FixedLenFeature([], tf.string),
                                       })

    RGB_flat = tf.decode_raw(features['sharp'], tf.uint8)
    RGB = tf.reshape(RGB_flat, [N_h, N_w, 1])

    return RGB


########## Preprocess the images #############
##  crop to patches
##  random flip
##  Add uniform noise
############################################  
def data_augment(RGB_batch_float):
    # crop to N_raw x N_raw
    N_raw = 326  # for boundary effect, 256+70, will need cropping after convolution

    data1 = tf.map_fn(lambda img: tf.random_crop(img, [N_raw, N_raw, 1]), RGB_batch_float)

    # flip both images and labels
    data2 = tf.map_fn(lambda img: tf.image.random_flip_up_down(tf.image.random_flip_left_right(img)), data1)

    # only adjust the RGB value of the image
    r1 = tf.random_uniform([]) * 0.3 + 0.8
    RGB_out = data2 * r1

    return RGB_out



############ Put data in batches #############
##  put in batch and shuffle
##  cast to float32
##  call data_augment for image preprocess
## @param{TFRECORD_PATH}: path to the data
## @param{batchsize}: currently 21 for the 21 PSFs
##############################################
def read2batch(TFRECORD_PATH, batchsize):
    # load tfrecord and make them to be usable data
    RGB = read_tfrecord(TFRECORD_PATH)
    #RGB_batch = tf.train.shuffle_batch([RGB], batch_size=batchsize, capacity=200, num_threads=5)
    RGB = tf.expand_dims(RGB, axis=0)
    RGB_batch = tf.tile(RGB, [21,1,1,1])
    RGB_batch_float = tf.image.convert_image_dtype(RGB_batch, tf.float32)

    # padd the target for convolution
    RGB_batch_float = tf.image.resize_image_with_crop_or_pad(RGB_batch_float, 298, 298)

    return RGB_batch_float[:, :, :, 0:1]



def add_gaussian_noise(images, std):
    noise = tf.random_normal(shape=tf.shape(images), mean=0.0, stddev=std, dtype=tf.float32)
    return tf.nn.relu(images + noise)


################  blur the images using PSFs  ##################
## same patch different depths put in a stack
################################################################
def one_wvl_blur(im, PSFs0):
    N_B = PSFs0.shape[1].value
    N_Phi = PSFs0.shape[0].value
    N_im = im.shape[1].value
    N_im_out = N_im - N_B + 1  # the final image size after blurring

    sharp = tf.transpose(tf.reshape(im, [-1, N_Phi, N_im, N_im]),
                         [0, 2, 3, 1])  # reshape to make N_Phi in the last channel
    PSFs = tf.expand_dims(tf.transpose(PSFs0, perm=[1, 2, 0]), -1)
    blurAll = tf.nn.depthwise_conv2d(sharp, PSFs, strides=[1, 1, 1, 1], padding='VALID')
    blurStack = tf.transpose(
        tf.reshape(tf.transpose(blurAll, perm=[0, 3, 1, 2]), [-1, 1, N_im_out, N_im_out]),
        perm=[0, 2, 3, 1])  # stack all N_Phi images to the first dimension

    return blurStack


def blurImage_diffPatch_diffDepth(RGB, PSFs):
    blur = one_wvl_blur(RGB[:, :, :, 0], PSFs[:, :, :, 0])

    return blur


####################### system ##########################
## @param{PSFs}: the PSFs
## @param{RGB_batch_float}: patches
## @param{phase_BN}: batch normalization, True only during training
########################################################
def system(PSFs, RGB_batch_float, phase_BN=False):
    with tf.variable_scope("system", reuse=tf.AUTO_REUSE):
        blur = blurImage_diffPatch_diffDepth(RGB_batch_float, PSFs)  # size [batch_size * N_Phi, Nx, Ny, 3]

        # noise
        sigma = 0.01
        blur_noisy = add_gaussian_noise(blur, sigma)

        RGB_hat = Network.UNet(blur_noisy, phase_BN)

        return blur_noisy, RGB_hat


######################  RMS cost #############################
## @param{GT}: ground truth
## @param{hat}: reconstruction
##############################################################
def cost_rms(GT, hat):
    cost = tf.sqrt(tf.reduce_mean(tf.reduce_mean((tf.square(GT - hat)),1),1))
    return cost

######################  SSIM cost #############################
## @param{GT}: ground truth
## @param{hat}: reconstruction
##############################################################
def cost_ssim(GT, hat):
    cost = tf.image.ssim(GT, hat, 1.0) # assume img intensity ranges from 0 to 1
    cost = tf.expand_dims(cost, axis = 1)
    return cost

##########  compare the reconstruction reblured with U-net input?  ############
## important for EDOF to utilize the PSF information
## @param{RGB_hat}: Unet reconstructed image
## @param{PSFs}: PSF used
## @param{blur}: all-in-focus image conv PSF
## @param{N_B}: size of blur kernel
## @return{reblur}: reconstruction blurred
## @return{cost}: l2 norm between blur_GT and reblur
##############################################################################
def cost_reblur(RGB_hat, PSFs, blur, N_B):
    reblur = blurImage_diffPatch_diffDepth(RGB_hat, PSFs)
    blur_GT = blur[:, int((N_B - 1) / 2):-int((N_B - 1) / 2), int((N_B - 1) / 2):-int((N_B - 1) / 2),
              :]  # crop the patch to 256x256

    cost = tf.sqrt(tf.reduce_mean(tf.square(blur_GT - reblur)))

    return reblur, cost


########################################################  PARAMETER ####################################################
N_B = 71

N_Phi = 21
batch_size = N_Phi
Phi_list = np.linspace(-10, 10, N_Phi, np.float32)  # defocus

PSFs = np.load(results_dir + 'PSFs.npy')
PSFs = tf.convert_to_tensor(PSFs, dtype=tf.float32)


####################################################### architecture ###################################################
RGB_batch_float_test = GT

[blur_test, RGB_hat_test] = system(PSFs, RGB_batch_float_test)

# cost function
with tf.name_scope("cost"):
    RGB_GT_test = RGB_batch_float_test[:, int((N_B - 1) / 2):-int((N_B - 1) / 2),
                  int((N_B - 1) / 2):-int((N_B - 1) / 2), :]  # crop the all-in-focus to be


    cost_rms_test = cost_rms(RGB_GT_test, RGB_hat_test)
    cost_ssim_test = cost_ssim(RGB_GT_test, RGB_hat_test)



###################################################  reload model ##################################################
saver_best = tf.train.Saver()

with tf.Session() as sess:

    # threading for parallel
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess=sess, coord=coord)

    sess.run(tf.global_variables_initializer())
    sess.run(tf.local_variables_initializer())

    model_path = tf.train.latest_checkpoint(results_dir)
    saver_best.restore(sess, model_path)

    for i in range(1):
        [loss_blur_test, loss_estimate_test, loss_rms_test, loss_ssim_test, GT_test] = sess.run(
        	[blur_test, RGB_hat_test, cost_rms_test, cost_ssim_test, RGB_GT_test])

        sharp_crop = GT_test[0, :, :, 0]
        gt_min = np.amin(np.ndarray.flatten(sharp_crop))
        gt_max = np.amax(np.ndarray.flatten(sharp_crop))
        scipy.misc.toimage(sharp_crop, cmin=gt_min, cmax=gt_max).save(image_save_dir + 'sharp_crop.png')

        np.savetxt(save_file_UNET, loss_rms_test, delimiter=',', newline='\n')
        np.savetxt(save_file_UNET_SSIM, loss_ssim_test, delimiter=',', newline='\n')

    np.save('blur.npy', loss_blur_test)
    np.save('estimate.npy', loss_estimate_test)


    coord.request_stop()
    coord.join(threads)

print('Now saving the images')

mask_blur = np.load('blur.npy')
estimate = np.load('estimate.npy')


def npy_to_images(npy_stack, save_name):
    for i in range(21):
        img_cur = npy_stack[i, :, :, 0]
        img_cur = 1 - img_cur
        img_min = np.amin(np.ndarray.flatten(img_cur))
        img_max = np.amax(np.ndarray.flatten(img_cur))
        save_name_cur = '00_'+ str(i) + '_' + save_name
        scipy.misc.toimage(img_cur, cmin=img_min, cmax=img_max).save(image_save_dir + save_name_cur)


npy_to_images(mask_blur, 'deepDOF_blur.png')
npy_to_images(estimate, 'deepDOF_hat.png')