seed=42#321#123
import tensorflow as tf
import condor_tensorflow as condor
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)
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)

# Create neural network model
model = tf.keras.models.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(tf.keras.layers.Dense(num_classes,activation="softmax"))

model.summary()

# compile and fit the model
model.compile(loss = tf.keras.losses.SparseCategoricalCrossentropy(),
              metrics = [tf.keras.metrics.SparseCategoricalAccuracy()],
              optimizer = tf.keras.optimizers.Adam())

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)

y_pred = model.predict(X_test)
probs_df = pd.DataFrame(y_pred)

true_y = labels_test
pred_probs = pd.DataFrame(y_pred)
x = y_pred
binProbs = np.cumsum(x[:,::-1],axis=1)[:,::-1][:,1:]
labels_v1 = pred_probs.idxmax(axis = 1)
print("Accuracy: ", np.mean(labels_v1 == true_y))

print("Accuracy tol1: ", np.mean(np.abs(labels_v1 - true_y)<=1))
print("MAE: ", np.mean(np.abs(labels_v1 - true_y)))
print("BCE: ",condor.SparseCondorOrdinalCrossEntropy()(true_y,binProbs))
print("EMD: ",condor.SparseOrdinalEarthMoversDistance()(true_y,binProbs))

# save results
pred_probs['true_label'] = true_y
pred_probs.to_csv('toyPredictedProbsCategorical.csv')
