import sys
import itertools
from keras.layers import Input, Dense, Reshape, Flatten
from keras import layers, initializers
from keras.models import Model, load_model
import keras.backend as K
import tensorflow as tf
import numpy as np
from seqtools import SequenceTools as ST
from gfp_gp import SequenceGP
from util import AA, AA_IDX
from util import build_vae
from sklearn.model_selection import train_test_split, ShuffleSplit
from keras.callbacks import EarlyStopping
import matplotlib.pyplot as plt
import pandas as pd
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import RBF, ConstantKernel as C
import scipy.stats
from scipy.stats import norm
from scipy.optimize import minimize
from keras.utils.generic_utils import get_custom_objects
from util import one_hot_encode_aa, partition_data, get_balaji_predictions, get_samples, get_argmax
from util import convert_idx_array_to_aas, build_pred_vae_model, get_experimental_X_y
from util import get_gfp_X_y_aa
from losses import neg_log_likelihood
import json

import isolearn.io as isoio
import isolearn.keras as isol

import warnings
warnings.filterwarnings("ignore")

from keras.backend.tensorflow_backend import set_session

def contain_tf_gpu_mem_usage() :
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    sess = tf.Session(config=config)
    set_session(sess)

contain_tf_gpu_mem_usage()

#Evaluate ELBO distribution on GFP dataset, decoder epsilon = 1e-6

n_z_samples = 128

for it in range(3) :

    TRAIN_SIZE = 5000
    train_size_str = "%ik" % (TRAIN_SIZE/1000)
    num_models = [1, 5, 20][it]
    RANDOM_STATE = it + 1

    X_train, y_train, gt_train  = get_experimental_X_y(random_state=RANDOM_STATE, train_size=TRAIN_SIZE)

    L = X_train.shape[1]

    vae_suffix = '_%s_%i' % (train_size_str, RANDOM_STATE)
    
    print(vae_suffix)

    vae_0 = build_vae(latent_dim=20, n_tokens=20, seq_length=L, enc1_units=50)

    vae_0.encoder_.load_weights("models/vae_0_encoder_weights%s.h5" % vae_suffix)
    vae_0.decoder_.load_weights("models/vae_0_decoder_weights%s.h5"% vae_suffix)
    vae_0.vae_.load_weights("models/vae_0_vae_weights%s.h5"% vae_suffix)

    #Compute multi-sample ELBO on test set
    log_mean_p_vae_test, mean_log_p_vae_test, log_p_vae_test = evaluate_elbo(vae_0.encoder_, vae_0.decoder_, X_train, n_samples=n_z_samples)

    #Log Likelihood Plot
    plot_min_val = None
    plot_max_val = None

    f = plt.figure(figsize=(6, 4))

    log_p_vae_test_hist, log_p_vae_test_edges = np.histogram(log_mean_p_vae_test, bins=50, density=True)
    bin_width_test = log_p_vae_test_edges[1] - log_p_vae_test_edges[0]

    mean_log_p_vae_test = np.mean(log_mean_p_vae_test)
    mode_log_p_vae_test = log_p_vae_test_edges[np.argmax(log_p_vae_test_hist)] + bin_width_test / 2.

    print("mean log(likelihood) = " + str(round(mean_log_p_vae_test, 4)))
    print("mode log(likelihood) = " + str(round(mode_log_p_vae_test, 4)))


    plt.bar(log_p_vae_test_edges[1:] - bin_width_test/2., log_p_vae_test_hist, width=bin_width_test, linewidth=2, edgecolor='black', color='orange')

    plt.xticks(fontsize=14)
    plt.yticks(fontsize=14)

    if plot_min_val is not None and plot_max_val is not None :
        plt.xlim(plot_min_val, plot_max_val)

    plt.xlabel("VAE Log Likelihood", fontsize=14)
    plt.ylabel("Data Density", fontsize=14)

    plt.axvline(x=mean_log_p_vae_test, linewidth=2, color='red', linestyle="--")
    plt.axvline(x=mode_log_p_vae_test, linewidth=2, color='purple', linestyle="--")

    plt.tight_layout()
    plt.show()