import random

import matplotlib.pyplot as plt
import pandas as pd
from datasets import load_dataset, ClassLabel, Sequence
import json
import numpy
from transformers import AutoImageProcessor
from torchvision.transforms import RandomResizedCrop, Compose, Normalize, ToTensor
from transformers import DefaultDataCollator
# import evaluate
import numpy as np
from transformers import AutoModelForImageClassification, TrainingArguments, Trainer
from PIL import Image
from matplotlib import cm
from Data_Generation.Shape_Generation_Functions import basic_box, diagonal_box_split, horizontal_vertical_box_split, \
    back_slash_box, forward_slash_box, back_slash_plus_box, forward_slash_plus_box, hot_dog_box, hamburger_box, \
    x_hamburger_box, x_hot_dog_box, x_plus_box

from Data_Generation.Dataset_Generation_Functions import make_boxes



# food = load_dataset("cmudrc/2d-lattices", split="train[:15]")  # Loads the training data samples
food = load_dataset("cmudrc/2d-lattices", split="train+test")  # Loads all of the data, for use after training

# checks to see if the dataset has been assigned a class label
# if type(food.features["label"]) != 'datasets.features.features.ClassLabel': # Cast to ClassLabel
#     food = food.class_encode_column('label')
print(food)
desired_label = 'x_hot_dog_box'
desired_thickness = 1
desired_density = 1

data_frame = pd.DataFrame(food)
# print(data_frame)

shape_rows = data_frame['Shape'] == desired_label
# print(shape_rows)

thickness_rows = data_frame['Thickness'] == desired_thickness
# print(thickness_rows)

density_rows = data_frame['Density'] == desired_density
# print(density_rows)

desired_output = data_frame.loc[shape_rows & thickness_rows & density_rows].iloc[0]['Array']
print(desired_output)
print(type(desired_output))


example_point = numpy.array(json.loads(desired_output))

plt.imshow(example_point)
plt.show()


all_shapes = [basic_box, diagonal_box_split, horizontal_vertical_box_split, back_slash_box, forward_slash_box,
                 back_slash_plus_box, forward_slash_plus_box, hot_dog_box, hamburger_box, x_hamburger_box,
                 x_hot_dog_box, x_plus_box]

base_shapes = [basic_box, back_slash_box, forward_slash_box, hot_dog_box, hamburger_box]
image_size = 11
density = [1]

boxes = make_boxes(image_size, density, all_shapes)


box_arrays, box_shape, box_density, box_thickness,  = list(zip(*boxes))[0], list(zip(*boxes))[1], list(zip(*boxes))[2], list(zip(*boxes))[3]

# indices_1 = [i for i in range(len(boxes)) if boxes[1][i] == str(base_shapes[0]) and boxes[2][i] == density[0] and boxes[3][i] == desired_thickness]
indices_1 = [i for i in range(len(box_arrays)) if box_shape[i] == desired_label and box_density[i] == desired_density and box_thickness[i] == desired_thickness]
print(indices_1)
# indices_1 = random.randint(0, len(box_arrays))


# plt.imshow(box_arrays[indices_1])
plt.imshow(box_arrays[indices_1[0]])
plt.show()


'''trainer.push_to_hub()''' # Need to figure out how to push the model to the hub