from pyxai import Builder, Learning, Explainer, Tools
import shap
import matplotlib.pyplot as plt
import numpy
import os
import json
from sklearn.utils import shuffle

percentages = [1, 2, 4, 8]

main_stopwatch = Tools.Stopwatch()
local_stopwatch = Tools.Stopwatch()

class NpEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, numpy.integer):
            return int(obj)
        if isinstance(obj, numpy.floating):
            return float(obj)
        if isinstance(obj, numpy.ndarray):
            return obj.tolist()
        return super(NpEncoder, self).default(obj)

learner = Learning.Scikitlearn(Tools.Options.dataset, learner_type=Learning.CLASSIFICATION)
models = learner.evaluate(
    method=Learning.LEAVE_ONE_GROUP_OUT,
    output=Learning.RF)

print("here:", Tools.Options.dataset)
dataset_name = Tools.Options.dataset.split(os.sep)[-1].split(".csv")[0]
save_directory = "ijcai2024-anchored"+os.sep+"models"+os.sep+dataset_name+os.sep
print("save_directory:", save_directory)
learner.save(models, save_directory=save_directory)

print("Evaluate and save models time:", main_stopwatch.elapsed_time(), "seconds")

# Save some instance of your choice :)
# accuracies = []
for id, model in enumerate(models):
    local_stopwatch.elapsed_time(reset=True)
    
    instances_anchored_0 = learner.get_instances(
            indexes=Learning.TRAINING,
            model=model,
            labels=[0])
    
    instances_anchored_1 = learner.get_instances(
            indexes=Learning.TRAINING,
            model=model,
            labels=[1])
    instances_anchored_0 = shuffle(instances_anchored_0, random_state=0)
    instances_anchored_1 = shuffle(instances_anchored_1, random_state=0)
    instances_anchored_0 = [tuple(element[0]) for element in instances_anchored_0]
    instances_anchored_1 = [tuple(element[0]) for element in instances_anchored_1]
    print("Number of labeled 0:", len(instances_anchored_0))
    print("Number of labeled 1:", len(instances_anchored_1))
    result_anchored_0 = []
    result_anchored_1 = []
    
    for percentage in percentages:
        n_anchored_0 = int((percentage*len(instances_anchored_0))/100)
        n_anchored_1 = int((percentage*len(instances_anchored_1))/100)
        result_anchored_0 = instances_anchored_0[0:n_anchored_0]
        result_anchored_1 = instances_anchored_0[0:n_anchored_1]
        print("For " + str(percentage)+"%: " + str(len(result_anchored_0)) + " 0-labeled and "  + str(len(result_anchored_1)) + " 1-labeled")
        reference_instances = dict()
        reference_instances[0] = result_anchored_0
        reference_instances[1] = result_anchored_1
        reference_file = save_directory + dataset_name+"."+str(id)+"."+str(percentage)+".references.json"
        with open(reference_file, 'w') as file:
            json.dump(reference_instances, file, cls=NpEncoder)
    
    instances = learner.get_instances(
        indexes=Learning.TEST,
        n=10,
        model=model,
        save_directory=save_directory,
        instances_id=id)
    #accuracies.append(learner.get_details()[id]["metrics"]["accuracy"])

#print("accuracy:", round(numpy.mean(accuracies), 2))    