import sys
import os
import pickle
import numpy as np
from sklearn.model_selection import train_test_split
from scipy.stats import entropy, wasserstein_distance
import tensorflow as tf
from tensorflow.keras.optimizers import Adam
import tensorflow_probability as tfp
tfd = tfp.distributions

import src.helper as h
from src.model import dirichlet_nll_loss_with_regularization, predictor

exp_seed_str = sys.argv[1]
exp_seed = int(exp_seed_str)
predictor_order = sys.argv[2]
annotator_num_str = sys.argv[3]
annotator_num = int(annotator_num_str)


#in order not to use GPU
os.environ["CUDA_VISIBLE_DEVICES"]="-1"  
tf.keras.utils.set_random_seed(exp_seed)
tf.config.experimental.enable_op_determinism()

with open(os.path.join("data", "synthetic", "1500_points_3_classes.pkl"), 'rb') as f:
    data = pickle.load(f)

simplex = np.load(os.path.join("data", "simplex", "0.001.npy"))

model = predictor(order=predictor_order, feature_dim=data["X"].shape[1])
if predictor_order == "first": 
    loss_fn = tf.keras.losses.CategoricalCrossentropy()
else: 
    loss_fn = dirichlet_nll_loss_with_regularization
model.compile(optimizer=Adam(learning_rate=1e-4), loss=loss_fn)
X_train, X_test, lambda_train, lambda_test, lambda_tilde_train, lambda_tilde_test = train_test_split(data["X"], data["W"],  data["W_annotated"][annotator_num], test_size=100, random_state=2024+exp_seed)
X_proper_train, X_calib, lambda_proper_train, lambda_calib, lambda_tilde_proper_train, lambda_tilde_calib = train_test_split(X_train, lambda_train, lambda_tilde_train, test_size=100, random_state=2024+exp_seed)
model.fit(X_proper_train, lambda_tilde_proper_train, epochs=100, batch_size=8)


if predictor_order == "first": 
    lambda_hat_calib = model.predict(X_calib)
    lambda_hat_test = model.predict(X_test)
    distance_functions = ["tv", "kl", "ws", "inner"]
    calib_scores = {k: [] for k in distance_functions} #calib scores with lambda_tilde
    calib_scores_ideal = {k: [] for k in distance_functions} #calib scores with lambda (just for later comparisons)
    for i in range(lambda_hat_calib.shape[0]):
        calib_scores["tv"].append(h.tv(lambda_hat_calib[i], lambda_tilde_calib[i]))
        calib_scores["kl"].append(entropy(lambda_tilde_calib[i], lambda_hat_calib[i], base=2))
        calib_scores["ws"].append(wasserstein_distance([0,1,2], [0,1,2], lambda_hat_calib[i], lambda_tilde_calib[i]))
        calib_scores["inner"].append(1-np.inner(lambda_hat_calib[i], lambda_tilde_calib[i]))
        calib_scores_ideal["tv"].append(h.tv(lambda_hat_calib[i], lambda_calib[i]))
        calib_scores_ideal["kl"].append(entropy(lambda_calib[i], lambda_hat_calib[i], base=2))
        calib_scores_ideal["ws"].append(wasserstein_distance([0,1,2], [0,1,2], lambda_hat_calib[i], lambda_calib[i]))
        calib_scores_ideal["inner"].append(1-np.inner(lambda_hat_calib[i], lambda_calib[i]))
    data_to_save = {"calib_score":calib_scores, "calib_scores_ideal":calib_scores_ideal, "lambda_hat_test":lambda_hat_test, "lambda_test":lambda_test, "lambda_tilde_test": lambda_tilde_test}


else:
    alpha_calib = model.predict(X_calib) + 1 
    alpha_test = model.predict(X_test) + 1
    dist_lambda_hat_calib = tfd.Dirichlet(alpha_calib)
    calib_scores = []
    calib_scores_ideal = []
    for k in range(X_calib.shape[0]):
        probs = (dist_lambda_hat_calib[k].prob(simplex)).numpy()
        p_lambda = dist_lambda_hat_calib[k].prob(lambda_calib[k]).numpy()
        p_lambda_tilde = dist_lambda_hat_calib[k].prob(lambda_tilde_calib[k]).numpy()
        calib_scores_ideal.append(1-p_lambda/(np.max(probs)))
        calib_scores.append(1-p_lambda_tilde/(np.max(probs)))
    data_to_save = {"calib_score":calib_scores, "calib_scores_ideal":calib_scores_ideal, "alpha_test":alpha_test, "lambda_test":lambda_test, "lambda_tilde_test": lambda_tilde_test}


os.makedirs(os.path.join("results", "synthetic", "annotator_num"+annotator_num_str, exp_seed_str), exist_ok=True)      
with open(os.path.join("results", "synthetic", "annotator_num"+annotator_num_str, exp_seed_str, predictor_order+"_order_data.pkl"), 'wb') as f:
    pickle.dump(data_to_save, f)


