import os
import sys
from pathlib import Path

# Keep the official ``src`` package importable after adopting the standard
# conf/model/scripts/weight layout.
PROJECT_ROOT = Path(__file__).resolve().parents[1]
MODEL_ROOT = PROJECT_ROOT / "model"
if str(MODEL_ROOT) not in sys.path:
    sys.path.insert(0, str(MODEL_ROOT))

# The released H5 files use the legacy Keras serialization format.
os.environ.setdefault("TF_USE_LEGACY_KERAS", "1")

import numpy as np
import pandas as pd
import pdb
import tensorflow as tf
# from IPython.display import clear_output
import time
from src.gan.lib import models
from src.gan.lib import utils
import socket
import datetime
from tqdm import tqdm
import random
import matplotlib.pyplot as plt
import argparse

import tensorflow.keras.backend as K
from tensorflow.keras.optimizers import Adam

tf.compat.v1.enable_eager_execution()


parser = argparse.ArgumentParser()
parser.add_argument('-d', type=str, required=False,
                    default=str(PROJECT_ROOT / 'conf' / 'data' / 'utrdb2.csv'))
parser.add_argument('-bs', type=int, required=False ,default=64)
parser.add_argument('-lr', type=int, required=False ,default=5)
parser.add_argument('-mil', type=int, required=False ,default=64)
parser.add_argument('-mxl', type=int, required=False ,default=128)
parser.add_argument('-dim', type=int, required=False ,default=40)
parser.add_argument('-gpu', type=str, required=False ,default='-1')
parser.add_argument('-e', '--epochs', type=int, default=4000,
                    help='Number of complete passes over the filtered UTRdb dataset.')
parser.add_argument('--output-dir', type=str,
                    default=str(PROJECT_ROOT / 'outputs' / 'train'),
                    help='Directory for logs, generated samples, and checkpoints.')
args = parser.parse_args()

MODEL_NAME = 'WGAN-TF2'
RUN_ROOT = Path(args.output_dir).expanduser().resolve()
OUTPUT_PATH = RUN_ROOT / 'outputs' / MODEL_NAME
TRAIN_LOGDIR = RUN_ROOT / 'tensorboard' / MODEL_NAME / 'train_data'
OUTPUT_PATH.mkdir(parents=True, exist_ok=True)
TRAIN_LOGDIR.mkdir(parents=True, exist_ok=True)
file_writer = tf.summary.create_file_writer(str(TRAIN_LOGDIR))

if args.gpu == '-1':
    device = 'cpu'
else:
    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
    device = 'cuda'

def plot(x, y, logdir, name, xlabel=None, ylabel=None, title=None):

    plt.plot(x,y,'-')
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.savefig(logdir+'/'+name+'.png')
    plt.clf()


def plot_valid(x1, y1, x2, y2, logdir, name, xlabel=None, ylabel=None, title=None):

    plt.plot(x2,y2,'-',color='tab:blue')
    plt.plot(x1,y1,'-',color='tab:orange')

    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)

    plt.savefig(logdir+'/'+name+'.png')
    plt.clf()    

def gradient_penalty_loss( y_true, y_pred, discriminator):
  """
  Computes gradient penalty based on prediction and weighted real / fake samples
  """
  alpha = K.random_uniform((DIM, 1, 1))
  averaged_samples = (alpha * y_pred) + ((1 - alpha) * y_true)

  gradients = K.gradients(y_pred, averaged_samples)[0]
  # compute the euclidean norm by squaring ...
  gradients_sqr = K.square(gradients)
  #   ... summing over the rows ...
  gradients_sqr_sum = K.sum(gradients_sqr,
                            axis=np.arange(1, len(gradients_sqr.shape)))
  #   ... and sqrt
  gradient_l2_norm = K.sqrt(gradients_sqr_sum)
  # compute lambda * (1 - ||grad||)^2 still for each single sample
  gradient_penalty = K.square(1 - gradient_l2_norm)
  # return the mean as loss over all the batch samples
  return K.mean(gradient_penalty)

def log(samples_dir=False,suff=None):
    stamp = datetime.date.strftime(datetime.datetime.now(), "%Y.%m.%d-%Hh%Mm%Ss") + "_{}".format(socket.gethostname())
    full_logdir = os.path.join(str(RUN_ROOT / "logs"), stamp)
    if suff:
        full_logdir = full_logdir + suff
    os.makedirs(full_logdir, exist_ok=True)
    if samples_dir: os.makedirs(os.path.join(full_logdir, "samples"), exist_ok=True)
    log_dir = "{}:{}".format(socket.gethostname(), full_logdir)
    
    return full_logdir, 0

data_path = args.d
data_utr = pd.read_csv(data_path)
UTRdf = data_utr['seq'].to_numpy()

seqs = []

UTR_LEN = args.mxl

for i in range(len(UTRdf)):
    if len(UTRdf[i]) < UTR_LEN+1 and len(UTRdf[i]) > int(UTR_LEN/2):
        if UTRdf[i] not in seqs:
            seqs.append(UTRdf[i])

sequences = np.array(seqs)

sequences = [x.upper() for x in sequences]

rna_vocab = {"A":0,
             "C":1,
             "G":2,
             "U":3,
             "*":4}

rev_rna_vocab = {v:k for k,v in rna_vocab.items()}

def one_hot_encode(seq, SEQ_LEN=UTR_LEN):
    mapping = dict(zip("ACGT*", range(5)))    
    seq2 = [mapping[i] for i in seq]
    if len(seq2) < SEQ_LEN:
        extra = [np.eye(5)[4]] * (SEQ_LEN - len(seq2))
        return np.vstack([np.eye(5)[seq2] , extra])
    return np.eye(5)[seq2]

def one_hot_encode_2(seq, SEQ_LEN=UTR_LEN):
    mapping = dict(zip("ACGT", range(4)))    
    seq2 = [mapping[i] for i in seq]
    return np.eye(4)[seq2]    

ohe_sequences = np.asarray([one_hot_encode(x) for x in sequences])


BATCH_SIZE = args.bs # Batch size
ITERS = args.epochs # How many complete dataset iterations to train for
SEQ_LEN = UTR_LEN # Sequence length in characters
DIM = args.dim # Model dimensionality.
CRITIC_ITERS = 5 # How many critic iterations per generator iteration. 
LAMBDA = 10 # Gradient penalty lambda hyperparameter.
# README defines ``-lr 5`` as 1e-5.  The upstream ``e**args.lr`` expression
# is inconsistent with that documented interface.
LR = np.power(10.0, -args.lr)


LAMBDA = 10 # For gradient penalty

CURRENT_EPOCH = 1 # Epoch start from
SAVE_EVERY_N_EPOCH = 50 # Save checkpoint at every n epoch

MIN_LR = 0.000001 # Minimum value of learning rate
DECAY_FACTOR=1.00004 # learning rate decay factor
'''
Set seed for reproducibility
'''
seed = 35
np.random.seed(seed)
# tf.set_random_seed(seed)
tf.random.set_seed(seed)

logdir, checkpoint_baseline = log(samples_dir=True)
logdir2 = ''

'''
Build GAN
'''
model_type = "resnet"
data_enc_dim = 5
data_size = SEQ_LEN * data_enc_dim
# data_size = 256
gen_layers = 3
disc_layers = 3
lmbda = 10. #lipschitz penalty hyperparameter.

SAMPLE_SIZE = 128

N_CHANNELS = DIM

G = models.resnet_g2(DIM,N_CHANNELS,SEQ_LEN,5,res_layers=gen_layers)
D = models.resnet_d2(N_CHANNELS,SEQ_LEN,5,res_layers=disc_layers)

G.summary()
D.summary()

D_optimizer = Adam(learning_rate=LR, beta_1=0.5, beta_2=0.99)
G_optimizer = Adam(learning_rate=LR, beta_1=0.5, beta_2=0.99)

EPOCHs = ITERS

@tf.function
def WGAN_GP_train_d_step(real_sample, batch_size, step):

    noise = tf.random.normal([batch_size, DIM])
    epsilon = tf.random.uniform(shape=[batch_size, 1, 1], minval=0, maxval=1)
    ###################################
    # Train D
    ###################################
    with tf.GradientTape(persistent=True) as d_tape:
        with tf.GradientTape() as gp_tape:
            fake_sample = G([noise], training=True)
            fake_sample_mixed = epsilon * tf.dtypes.cast(real_sample, tf.float32) + ((1 - epsilon) * fake_sample)
            fake_mixed_pred = D([fake_sample_mixed], training=True)
            
        # Compute gradient penalty
        grads = gp_tape.gradient(fake_mixed_pred, fake_sample_mixed)
        grad_norms = tf.sqrt(tf.reduce_sum(tf.square(grads), axis=[1,2])) # Originally axis=[1,2]
        # grad_norms = tf.norm(grads, axis=[1,2])
        gradient_penalty = tf.reduce_mean(tf.square(grad_norms - 1.))
        
        fake_pred = D([fake_sample], training=True)
        real_pred = D([real_sample], training=True)
        
        D_loss = tf.reduce_mean(fake_pred) - tf.reduce_mean(real_pred) + LAMBDA * gradient_penalty
    # Calculate the gradients for discriminator
    D_gradients = d_tape.gradient(D_loss,D.trainable_variables)
    # Apply the gradients to the optimizer
    D_optimizer.apply_gradients(zip(D_gradients,D.trainable_variables))
    # Write loss values to tensorboard
    if step % 10 == 0:
        with file_writer.as_default():
            tf.summary.scalar('D_loss', tf.reduce_mean(D_loss), step=step)

    return D_loss, gradient_penalty

@tf.function
def WGAN_GP_train_g_step(real_sample, batch_size, step):

    noise = tf.random.normal([batch_size, DIM])
    ###################################
    # Train G
    ###################################
    with tf.GradientTape() as g_tape:
        fake_sample = G([noise], training=True)
        fake_pred = D([fake_sample], training=True)
        G_loss = -tf.reduce_mean(fake_pred)

    G_gradients = g_tape.gradient(G_loss,
                                            G.trainable_variables)
    # Apply the gradients to the optimizer
    G_optimizer.apply_gradients(zip(G_gradients,
                                                G.trainable_variables))
    # Write loss values to tensorboard
    if step % 10 == 0:
        with file_writer.as_default():
            tf.summary.scalar('G_loss', G_loss, step=step)

    return G_loss, noise

checkpoint_path = os.path.join(str(RUN_ROOT), "checkpoints", "tensorflow", MODEL_NAME)

ckpt = tf.train.Checkpoint(generator=G,
                        discriminator=D,
                        G_optimizer=G_optimizer,
                        D_optimizer=D_optimizer)

ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=40)

def generate_and_save_samples(model, epoch, test_input, figure_size=(12,6), subplot=(3,6), save=True, is_flatten=False):
    '''
        Generate samples and plot it.
    '''
    predictions = model.predict(test_input)
    utils.save_samples(logdir, predictions, epoch, rev_rna_vocab, annotated=False)


'''
load data
'''

Train = True
validate = True

data = ohe_sequences


if validate:
    split = len(data) // 10
    # print(split)
    valid_data = data[:split]
    train_data = data[split:]
    if len(train_data) == 1: train_data = train_data[0]
    if len(valid_data) == 1: valid_data = valid_data[0]
else:
    train_data = data


train_data = train_data.astype('float32')
valid_data = valid_data.astype('float32')

train_seqs = tf.data.Dataset.from_tensor_slices(train_data).shuffle(len(train_data)).batch(BATCH_SIZE)
valid_seqs = tf.data.Dataset.from_tensor_slices(valid_data).shuffle(len(valid_data)).batch(split)


plot_iter = 10

current_learning_rate = LR
trace = True
n_critic_count = 0

d_losses = []
g_losses = []
gradient_penalties = []
iterations = 0
iteration_numbers = []
iteration_numbers_valid = []
d_validation_losses = []

random_name = time.strftime("%Y%m%d-%H%M%S")

gen_iters = 0

if Train:
    sample_noise = tf.random.normal([BATCH_SIZE, DIM])
    generate_and_save_samples(G, 0, [sample_noise], figure_size=(12,6), subplot=(3,6), save=False, is_flatten=False)

    pbar = tqdm(range(EPOCHs))
    for epoch in pbar:
        start = time.time()

        tdataset = train_seqs.enumerate()

        for step, tdata in tdataset.as_numpy_iterator():
            current_batch_size = tdata.shape[0]

            d_loss, gp = WGAN_GP_train_d_step(tdata, batch_size=tf.constant(current_batch_size, dtype=tf.int64), step=tf.constant(step, dtype=tf.int64))
            n_critic_count += 1
            if n_critic_count >= CRITIC_ITERS: 
                g_loss, noise = WGAN_GP_train_g_step(tdata, batch_size= tf.constant(current_batch_size, dtype=tf.int64), step=tf.constant(step, dtype=tf.int64))
                gen_iters += 1
                n_critic_count = 0
            
            if step % 10 == 0:
                print ('.', end='')

        
        if epoch % SAVE_EVERY_N_EPOCH == 0 and epoch != 0:
            ckpt_save_path = ckpt_manager.save()
            utils.save_checkpoints(logdir,G,epoch)

        if epoch % 50 == 0:    
            
            generate_and_save_samples(G, epoch, [sample_noise], figure_size=(12,6), subplot=(3,6), save=True, is_flatten=False)

        os.system('clear')

        iteration_numbers.append(iterations)
        g_losses.append(-g_loss)
        d_losses.append(-d_loss)
        gradient_penalties.append(gp)
        # plot(iteration_numbers, d_losses, logdir, 'discriminator_loss', xlabel="Iteration", ylabel="Discriminator Cost")
        # plot(iteration_numbers, g_losses, logdir, 'generator_loss', xlabel="Iteration", ylabel="Generator Cost")
        # plot(iteration_numbers, gradient_penalties, logdir, 'gradient_penalty', xlabel="Iteration", ylabel="Gradient Penalty")

        iterations+=1

        iteration_numbers_valid.append(iterations)
        fake_sample_valid = G([noise], training=True)
        fake_pred_valid = D([fake_sample_valid], training=True)
        real_pred_valid = D([tf.convert_to_tensor(list(valid_seqs.as_numpy_iterator())[0])], training=True)
        
        D_loss_valid = tf.reduce_mean(fake_pred_valid) - tf.reduce_mean(real_pred_valid)
        d_validation_losses.append(-D_loss_valid)

        # plot_valid(iteration_numbers_valid, d_validation_losses, iteration_numbers, d_losses, logdir, 'validation_loss', xlabel="Iteration", ylabel="D Validation Loss")

    ckpt_save_path = ckpt_manager.save()
    print ('Saving checkpoint for epoch {} at {}'.format(EPOCHs,
                                                            ckpt_save_path))


print("####################################################")
print(f"############### Gen Iterations : {gen_iters} #####")
print("####################################################")