import os
from NetworkPconvN import Generator
import Utils
from Utils_model import VGG_MSE as Loss
import tensorflow as tf    
config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth=True
sess = tf.compat.v1.Session(config=config)
from keras.optimizers import Adam
import random
from tqdm import tqdm
import numpy as np
import argparse

def get_optimizer():
 
    adam = Adam(lr=1E-4, beta_1=0.9, beta_2=0.999, epsilon=1e-08)
    return adam
np.random.seed(10)
downscale_factor = 2
upsampling = downscale_factor//2

# Validation 
val_dir = ".\\val\\"
val_no_imges = 4 # Number of images to be validated
# Remember to change image shape if you are having different size of images
# image_shape = (384,384,3)
patch_hr = (32,32,3)
patch_lr = (patch_hr[0]//downscale_factor, patch_hr[1]//downscale_factor, patch_hr[2]) # the double // sign is used to get the integer part of this division
strd = patch_lr[0]//2
# strd = 5
if strd > patch_lr[0]:
    strd = patch_lr[0]


# for more info use $python train.py -h
def train(epochs, batch_size, input_dir, output_dir, model_save_dir, number_of_images, train_test_ratio):

    x_train_lr, x_train_hr, x_test_lr, x_test_hr = Utils.load_training_data(input_dir, '.jpg', downscale_factor, (640,640,3), number_of_images, train_test_ratio)
    x_val_hr, x_val_lr = Utils.load_validation_data(val_dir, '.jpg', downscale_factor, (640,640,3), val_no_imges) # by me
    val_patches_hr = Utils.val_patches_generator(x_val_hr, patch_hr[0])
    val_patches_lr = Utils.val_patches_generator(x_val_lr, patch_lr[0])
    print(x_train_lr.shape, x_train_hr.shape,x_test_lr.shape, x_test_hr.shape)
    loss = Loss(patch_hr)  
    best_val = np.empty(x_val_lr.shape[0], dtype="float16")
    best_val[:] = np.inf
    val_mse_best = np.inf 
    
    batch_count = int(x_train_hr.shape[0] / batch_size)
    print(patch_lr)
    generator = Generator(patch_lr).generator(upsampling)
    optimizer = get_optimizer()
    generator.compile(loss=loss.loss, optimizer=optimizer)
    generator.summary()
    with open(output_dir + '\\0Model_summary.txt', 'w') as summ:
        generator.summary(print_fn=lambda x: summ.write(x + '\n'))
        # discriminator.summary(print_fn=lambda x: summ.write(x + '\n'))   >>>>>>>>>>
        # gan.summary(print_fn=lambda x: summ.write(x + '\n'))
        # summ.write(lossd + '\n')
       
    for e in range(1, epochs+1):
        print ('-'*15, 'Epoch %d' % e, '-'*15)
        rand_nums = random.sample(range(0, x_train_hr.shape[0]), x_train_hr.shape[0])
        image_batch_hr = x_train_hr[rand_nums]
        image_batch_lr = x_train_lr[rand_nums]
        batch_ctrl_l = 0 
        batch_ctrl_u = batch_size       
        gan_loss = 0
        s_flag = True
        z = 0 # for calculating epochwise average loss

        for k in tqdm(range(batch_count)):
            # ####->>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>  
            ## train the Generator
            print('\ntraining Generator')
            print(rand_nums[batch_ctrl_l:batch_ctrl_u])
            shift_i = np.random.randint(0,strd,1)[0]
            rows_ = (image_batch_lr.shape[1] - patch_lr[0])//strd + 1
            range_i = np.arange(0, rows_, strd) + shift_i
            range_i = np.mod(range_i, rows_)
            np.random.shuffle(range_i)
            shift_j = np.random.randint(0,strd,1)[0]
            cols_ = (image_batch_lr.shape[2] - patch_lr[1])//strd
            range_j = np.arange(0, cols_, strd) + shift_j
            range_j = np.mod(range_j, cols_)
            np.random.shuffle(range_j)
            for i in range_i:
                for j in range_j:
                    patch_batch_lr = image_batch_lr[batch_ctrl_l:batch_ctrl_u, i:i+patch_lr[0], j:j+patch_lr[1], :]
                    patch_batch_hr = image_batch_hr[batch_ctrl_l:batch_ctrl_u, i*2:i*2+patch_hr[0], j*2:j*2+patch_hr[1], :]                    
                    gan_loss = gan_loss + generator.train_on_batch(patch_batch_lr, patch_batch_hr)
                    z+=len(patch_batch_lr)
            batch_ctrl_l = batch_ctrl_u
            batch_ctrl_u+=batch_size
        # print("gan_loss :%f, k = %d, i = %d, j = %d" %(gan_loss, k,i,j))
        gan_loss = gan_loss/z   
        val_mse = Utils.val_mse(generator, val_patches_hr, val_patches_lr)
        
        psnrg, mseg, idx = Utils.plot_generated_images_psnr_olp3(output_dir, e, generator, x_val_hr, x_val_lr, patch_hr, downscale_factor, True)
        
        gan_loss = str(gan_loss)
        
        loss_file = open(output_dir + '\\0losses_clr.txt' , 'a')
        # loss_file.write('epoch%d : gan_loss = %s ; discriminator_loss = %f; PSNR = %f; MSE = %f\n' %(e, gan_loss, discriminator_loss, psnrg, mseg) )
        loss_file.write('epoch%d : gan_loss = %s ; val_mse = %f; PSNR = %f; MSE = %f; %d\n' %(e, gan_loss, val_mse, psnrg, mseg, idx) )
        loss_file.close()
            
        if e == 1 or e % 5 == 0:
            Utils.plot_generated_images_psnr_olp3(output_dir, e, generator, x_test_hr, x_test_lr, patch_hr, downscale_factor, False)
        if e == 1 or e % 100 == 0:
            s_flag = False # if s_flag == true mean the model may be saved, 
            ###      ##      if flase, mean the model is already saved
            generator.save(model_save_dir + 'gen_model%d.h5' % e)
        if mseg <= best_val[idx] and mseg<1 and s_flag:
            s_flag = False 
            best_val[idx] = mseg
            generator.save(model_save_dir + 'gen_model%d.h5' % e)
        if val_mse < val_mse_best and s_flag:
            val_mse_best = val_mse
            if e > 400:
                generator.save(model_save_dir + 'gen_model%d.h5' % e)

if __name__== "__main__":

    parser = argparse.ArgumentParser()
    
    parser.add_argument('-i', '--input_dir', action='store', dest='input_dir', default='./data/' ,
                    help='Path for input images')
                    
    parser.add_argument('-o', '--output_dir', action='store', dest='output_dir', default='./output/' ,
                    help='Path for Output images')
    
    parser.add_argument('-m', '--model_save_dir', action='store', dest='model_save_dir', default='./model/' ,
                    help='Path for model')

    parser.add_argument('-b', '--batch_size', action='store', dest='batch_size', default=64,
                    help='Batch Size', type=int)
                    
    parser.add_argument('-e', '--epochs', action='store', dest='epochs', default=1000 ,
                    help='number of iteratios for trainig', type=int)
                    
    parser.add_argument('-n', '--number_of_images', action='store', dest='number_of_images', default=1000 ,
                    help='Number of Images', type= int)
                    
    parser.add_argument('-r', '--train_test_ratio', action='store', dest='train_test_ratio', default=0.8 ,
                    help='Ratio of train and test Images', type=float)
    
    values = parser.parse_args()
    # check the presence of output dirtectory 
    if not os.path.exists(values.output_dir):
        os.makedirs(values.output_dir)
    
    train(values.epochs, values.batch_size, values.input_dir, values.output_dir, values.model_save_dir, values.number_of_images, values.train_test_ratio)


