seed=123#42#321#123
import tensorflow as tf
import coral_ordinal as coral
from sklearn.model_selection import train_test_split
import pandas as pd
import numpy as np
tf.random.set_seed(seed)
np.random.seed(seed)
from scipy import special
import os
os.environ['PYTHONHASHSEED']=str(seed)


# create some toy data
num_classes = 4
samps = 1000
X = np.random.normal(size=(samps,2))
y = np.zeros(samps,dtype=np.int64)

for ind in range(samps):
    x = X[ind,:]
    y[ind] = np.digitize(np.arctan2(x[0],x[1]),
                         np.linspace(-np.pi,np.pi,num=num_classes+1)) - 1


# do a test train split
X_train, X_test, labels_train, labels_test = \
  train_test_split(X, y, test_size = samps//10, random_state = 1)

# build neural net
model = tf.keras.Sequential()
model.add(tf.keras.layers.Input(2))
model.add(tf.keras.layers.Dense(10,activation="relu"))
model.add(tf.keras.layers.Dense(10,activation="relu"))
model.add(coral.CoralOrdinal(num_classes = num_classes)) # Ordinal variable has 4 labels, 0 through 3.
model.compile(loss = coral.OrdinalCrossEntropy(num_classes = num_classes),
              metrics = [coral.MeanAbsoluteErrorLabels(),
                  coral.SparseOrdinalEarthMoversDistance()],
             optimizer=tf.keras.optimizers.Adam())

model.summary()

# fit the model
history = model.fit(x = X_train,
                    y = labels_train,
                    epochs = 100,
                    batch_size = 32,
                    validation_split = 0.2,
                    callbacks = [tf.keras.callbacks.EarlyStopping(patience = 10,
                                                                  min_delta = 0.001,
                                                                  restore_best_weights = True)])
# see how well we did
model.evaluate(X_test, labels_test)
pred_y = model.predict(X_test)

true_y = labels_test
pred_probs = pd.DataFrame(coral.ordinal_softmax(pred_y).numpy())
labels_v1 = pred_probs.idxmax(axis = 1)
print("Accuracy of label version 1:", np.mean(labels_v1 == true_y))
print("Accuracy tolerance 1 of label version 1:", np.mean(np.abs(labels_v1 - true_y) <= 1))


# Compare to logit-based cumulative probs
cum_probs = pd.DataFrame(pred_y).apply(special.expit)
# Calculate the labels using the style of Cao et al.
labels_v2 = cum_probs.apply(lambda x: x > 0.5).sum(axis = 1)
print("Accuracy of label version 2:", np.mean(labels_v2 == true_y))
print("Accuracy1 of label version 2:", np.mean(np.abs(labels_v2 - true_y) <= 1))
print("MAE label version 2:", np.mean(np.abs(labels_v2 - true_y)))
#ordinal_logits = model.predict(X_test)
#tensor_probs = coral.ordinal_softmax(ordinal_logits)
#probs_df = pd.DataFrame(tensor_probs.numpy())
#
## Compute probability on wrong labels and EMD
#wrong = np.zeros(len(labels_test))
#emd = np.zeros(len(labels_test))
#for index, row in probs_df.iterrows():
#    wrong[index] = 1. - row[np.int_(labels_test[index])]
#    emd[index] = np.sum(row*np.abs(labels_test[index]-np.arange(num_classes)))
#
#print("Average percentage incorrect class: ",wrong.mean())
## Average percentage incorrect class:  0.3027844336628914
#print("Average earth movers distance: ",emd.mean())
##Average earth movers distance:  0.34012722706940596
#
## Print the probability in each class
#print(probs_df.mean())
#
## Check bias terms: these should be in descending order.
#last_layer = len(model.layers) - 1
#print(model.layers[last_layer].get_weights()[1])
