import tensorflow as tf

import os
import time
import matplotlib.pyplot as plt

BUFFER_SIZE = 50
BATCH_SIZE = 1
IMG_WIDTH = 256
IMG_HEIGHT = 256

def load(image_file):
  image = tf.io.read_file(image_file)
  image = tf.image.decode_jpeg(image)

  w = tf.shape(image)[1]

  w = w // 2
  real_image = image[:, :w, :]
  input_image = image[:, w:, :]

  input_image = tf.cast(input_image, tf.float32)
  real_image = tf.cast(real_image, tf.float32)

  return input_image, real_image

def resize(input_image, real_image, height, width):
  input_image = tf.image.resize(input_image, [height, width],
                                method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
  real_image = tf.image.resize(real_image, [height, width],
                               method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)

  return input_image, real_image

def random_crop(input_image, real_image):
  stacked_image = tf.stack([input_image, real_image], axis=0)
  cropped_image = tf.image.random_crop(
      stacked_image, size=[2, IMG_HEIGHT, IMG_WIDTH, 3])

  return cropped_image[0], cropped_image[1]

def normalize(input_image, real_image):
  input_image = (input_image / 127.5) - 1
  real_image = (real_image / 127.5) - 1

  return input_image, real_image

def random_jitter(input_image, real_image):
  # resizing to 286 x 286 x 3
  input_image, real_image = resize(input_image, real_image, 286, 286)

  # randomly cropping to 256 x 256 x 3
  input_image, real_image = random_crop(input_image, real_image)

  #if tf.random.uniform(()) > 0.5:
    # random mirroring
   # input_image = tf.image.flip_left_right(input_image)
   # real_image = tf.image.flip_left_right(real_image)

  return input_image, real_image


def load_image_train(image_file):
  input_image, real_image = load(image_file)
  input_image, real_image = random_jitter(input_image, real_image)
  input_image, real_image = normalize(input_image, real_image)

  return input_image, real_image

def load_image_test(image_file):
  input_image, real_image = load(image_file)
  input_image, real_image = resize(input_image, real_image,
                                   IMG_HEIGHT, IMG_WIDTH)
  input_image, real_image = normalize(input_image, real_image)

  return input_image, real_image



OUTPUT_CHANNELS = 3

def downsample(filters, size, apply_batchnorm=True):
  initializer = tf.random_normal_initializer(0., 0.02)

  result = tf.keras.Sequential()
  result.add(
      tf.keras.layers.Conv2D(filters, size, strides=2, padding='same',
                             kernel_initializer=initializer, use_bias=False))

  if apply_batchnorm:
    result.add(tf.keras.layers.BatchNormalization())

  result.add(tf.keras.layers.LeakyReLU())

  return result

def upsample(filters, size, apply_dropout=False):
  initializer = tf.random_normal_initializer(0., 0.02)

  result = tf.keras.Sequential()
  result.add(
    tf.keras.layers.Conv2DTranspose(filters, size, strides=2,
                                    padding='same',
                                    kernel_initializer=initializer,
                                    use_bias=False))

  result.add(tf.keras.layers.BatchNormalization())

  if apply_dropout:
      result.add(tf.keras.layers.Dropout(0.5))

  result.add(tf.keras.layers.ReLU())

  return result

def Generator():
  down_stack = [
    downsample(64, 4, apply_batchnorm=False), # (bs, 128, 128, 64)
    downsample(128, 4), # (bs, 64, 64, 128)
    downsample(256, 4), # (bs, 32, 32, 256)
    downsample(512, 4), # (bs, 16, 16, 512)
    downsample(512, 4), # (bs, 8, 8, 512)
    downsample(512, 4), # (bs, 4, 4, 512)
    downsample(512, 4), # (bs, 2, 2, 512)
    downsample(512, 4), # (bs, 1, 1, 512)
  ]

  up_stack = [
    upsample(512, 4, apply_dropout=True), # (bs, 2, 2, 1024)
    upsample(512, 4, apply_dropout=True), # (bs, 4, 4, 1024)
    upsample(512, 4, apply_dropout=True), # (bs, 8, 8, 1024)
    upsample(512, 4), # (bs, 16, 16, 1024)
    upsample(256, 4), # (bs, 32, 32, 512)
    upsample(128, 4), # (bs, 64, 64, 256)
    upsample(64, 4), # (bs, 128, 128, 128)
  ]

  initializer = tf.random_normal_initializer(0., 0.02)
  last = tf.keras.layers.Conv2DTranspose(OUTPUT_CHANNELS, 4,
                                         strides=2,
                                         padding='same',
                                         kernel_initializer=initializer,
                                         activation='tanh') # (bs, 256, 256, 3)

  concat = tf.keras.layers.Concatenate()

  inputs = tf.keras.layers.Input(shape=[None,None,3])
  x = inputs

  # Downsampling through the model
  skips = []
  for down in down_stack:
    x = down(x)
    skips.append(x)

  skips = reversed(skips[:-1])

  # Upsampling and establishing the skip connections
  for up, skip in zip(up_stack, skips):
    x = up(x)
    x = concat([x, skip])

  x = last(x)

  return tf.keras.Model(inputs=inputs, outputs=x)

generator = Generator()

def Discriminator():
  initializer = tf.random_normal_initializer(0., 0.02)

  inp = tf.keras.layers.Input(shape=[None, None, 3], name='input_image')
  tar = tf.keras.layers.Input(shape=[None, None, 3], name='target_image')

  x = tf.keras.layers.concatenate([inp, tar]) # (bs, 256, 256, channels*2)

  down1 = downsample(64, 4, False)(x) # (bs, 128, 128, 64)
  down2 = downsample(128, 4)(down1) # (bs, 64, 64, 128)
  down3 = downsample(256, 4)(down2) # (bs, 32, 32, 256)

  zero_pad1 = tf.keras.layers.ZeroPadding2D()(down3) # (bs, 34, 34, 256)
  conv = tf.keras.layers.Conv2D(512, 4, strides=1,
                                kernel_initializer=initializer,
                                use_bias=False)(zero_pad1) # (bs, 31, 31, 512)

  batchnorm1 = tf.keras.layers.BatchNormalization()(conv)

  leaky_relu = tf.keras.layers.LeakyReLU()(batchnorm1)

  zero_pad2 = tf.keras.layers.ZeroPadding2D()(leaky_relu) # (bs, 33, 33, 512)

  last = tf.keras.layers.Conv2D(1, 4, strides=1,
                                kernel_initializer=initializer)(zero_pad2) # (bs, 30, 30, 1)

  return tf.keras.Model(inputs=[inp, tar], outputs=last)


def Discriminator2():
  initializer = tf.random_normal_initializer(0., 0.02)

  #inp = tf.keras.layers.Input(shape=[None, None, 3], name='input_image')
  #tar = tf.keras.layers.Input(shape=[None, None, 3], name='target_image')

  #x = tf.keras.layers.concatenate([inp, tar]) # (bs, 256, 256, channels*2)
  x = tf.keras.layers.Input(shape=[None, None, 3], name='target_image')
  down1 = downsample(64, 4, False)(x) # (bs, 128, 128, 64)
  down2 = downsample(128, 4)(down1) # (bs, 64, 64, 128)
  down3 = downsample(256, 4)(down2) # (bs, 32, 32, 256)

  zero_pad1 = tf.keras.layers.ZeroPadding2D()(down3) # (bs, 34, 34, 256)
  conv = tf.keras.layers.Conv2D(512, 4, strides=1,
                                kernel_initializer=initializer,
                                use_bias=False)(zero_pad1) # (bs, 31, 31, 512)

  batchnorm1 = tf.keras.layers.BatchNormalization()(conv)

  leaky_relu = tf.keras.layers.LeakyReLU()(batchnorm1)

  zero_pad2 = tf.keras.layers.ZeroPadding2D()(leaky_relu) # (bs, 33, 33, 512)

  last = tf.keras.layers.Conv2D(1, 4, strides=1,
                                kernel_initializer=initializer)(zero_pad2) # (bs, 30, 30, 1)

  return tf.keras.Model(inputs=[x], outputs=last)

discriminator = Discriminator()

LAMBDA = 100

def generate_images(model, test_input, tar):
  # the training=True is intentional here since
  # we want the batch statistics while running the model
  # on the test dataset. If we use training=False, we will get
  # the accumulated statistics learned from the training dataset
  # (which we don't want)
  prediction = model(test_input, training=True)
  #dt = datetime.now()
  #print(dt.microsecond)
  #plt.figure(figsize=(15,15))

  #display_list = [test_input[0], tar[0], prediction[0]]
  #title = ['Input Image', 'Ground Truth', 'Predicted Image']

  #for i in range(3):
  #  plt.subplot(1, 3, i+1)
  #  plt.title(title[i])
    # getting the pixel values between [0, 1] to plot it.
  #  plt.imshow(display_list[i] * 0.5 + 0.5)
    #plt.savefig('books_read.png')
  #  plt.axis('off')
  #plt.show()
  return tar[0], prediction[0]

def load_test(image_file_1,image_file_2):
  input_image = tf.io.read_file(image_file_1)
  input_image = tf.image.decode_jpeg(input_image)
  
  real_image = tf.io.read_file(image_file_2)
  real_image = tf.image.decode_jpeg(real_image)
  
  input_image = tf.cast(input_image, tf.float32)
  real_image = tf.cast(real_image, tf.float32)
  
  #input_image, real_image = load_img(image_file)
  input_image, real_image = resize(input_image, real_image,
                                   IMG_HEIGHT, IMG_WIDTH)
  input_image, real_image = normalize(input_image, real_image)
  
  return input_image, real_image

import numpy as np
def mse(imageA, imageB):
	# the 'Mean Squared Error' between the two images is the
	# sum of the squared difference between the two images;
	# NOTE: the two images must have the same dimension
	err = np.sum((imageA - imageB) ** 2)
	err /= float(imageA.shape[0] * imageA.shape[1])
	
	# return the MSE, the lower the error, the more "similar"
	# the two images are
	return err

generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)

import numpy as np
import cv2

from sklearn.metrics import precision_recall_curve
import matplotlib.pyplot as plt
from inspect import signature
import numpy as np
from sklearn.metrics import roc_curve
from sklearn.metrics import roc_auc_score


#checkpoint_dir = '/media/313727ac-f148-4b50-9224-3f3986e68217/Lakshmi/anomaly_detection/DA/model_dual/trainingAno_0p1_checkpoints/'
checkpoint_dir = '/media/313727ac-f148-4b50-9224-3f3986e68217/Lakshmi/anomaly_detection/DA/model/trainingAno_dualEpo150_checkpoints/'
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer,
                                 generator=generator)
                                 
checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))


PATH='/media/313727ac-f148-4b50-9224-3f3986e68217/Lakshmi/anomaly_detection/DA/test/anomaly_test/run_test/'
#test_dataset = tf.data.Dataset.list_files(PATH+'test/*.jpg')
# shuffling so that for every epoch a different image is generated
# to predict and display the progress of our model.

#test_dataset = test_dataset.map(load_image_test)
#test_dataset = test_dataset.batch(1)
from datetime import datetime
#N = 400 # DB
N = 390 # DA run
#N = 144 #DA bend
#N = 121 #DA bend_1
#N = 420 #bag
#N = 183 # good bag
#N = 240 #fight
i = 0
errCurr = np.zeros((N,1))
for i in range(0,N):
    image_file = PATH+str(i)+'.jpg'
    inp, tar = load_image_test(image_file)
    #for inp, tar in test_dataset.take(i):
    #if i>0:
    dt = datetime.now()
    print(dt.microsecond)
    gt, pred = generate_images(generator, tf.expand_dims(inp,0), tf.expand_dims(tar,0))
    errCurr[i] = mse(gt, pred)
    print(i)
    print(errCurr[i])

#for i in range(N):
#    if errCurr[i]>0.25:
#       errCurr[i] = 0
plt.plot(errCurr)
plt.show()   

from array import array
output_file = open('DA_bend_cond', 'wb')
y_test = array('d', errCurr)
y_test.tofile(output_file)

y_score = np.zeros((N,1))

#DA run
for i in range(310,370):
    y_score[i] = 1
    
#DA bend
#for i in range(35,65):
#    y_score[i] = 1

#for i in range(90,138):
#    y_score[i] = 1

#DA bend_1
#for i in range(35,60):
#    y_score[i] = 1

#DA bag
#for i in range(200,320):
#    y_score[i] = 1

#for i in range(334,390):
#    y_score[i] = 1

#DA bag_1
#for i in range(77,96):
#    y_score[i] = 1

#for i in range(130,180):
#    y_score[i] = 1

#DA fight
#for i in range(75,140):
#    y_score[i] = 1


ns_auc = roc_auc_score(y_score, errCurr)
fpr, tpr, thresholds = roc_curve(y_score, errCurr)
print(ns_auc, np.sum(fpr), np.sum(tpr))
