#https://youtu.be/V61xy1ZnVTM?si=ZpPwSP5eOnaItPn2
#https://www.kaggle.com/datasets/fadwateimi/food-plates2

import numpy as np
#import pandas as pd
import matplotlib.pyplot as plt
import tensorflow as tf

from tensorflow import keras
from tensorflow.keras import layers

def wykres(epochs_size):
    epochs_range = range(epochs_size)
    plt.figure(figsize=(8,8))
    plt.subplot(1,2,1)
    plt.plot(epochs_range,history.history['accuracy'], label = 'Training Accuracy')
    plt.plot(epochs_range, history.history['val_accuracy'], label='Validation Accuracy')
    plt.title('Accuracy')

    plt.subplot(1, 2, 2)
    plt.plot(epochs_range, history.history['loss'], label='Training Loss')
    plt.plot(epochs_range, history.history['val_loss'], label='Validation Loss')
    plt.title('Loss')

    plt.show()

trainset_path = r"Automatyczny_kelner/trainset"
testset_path = r"Automatyczny_kelner/testset"
useset_path = r"Automatyczny_kelner/validationset"

img_width = 180
img_height = 180

trainset = tf.keras.utils.image_dataset_from_directory(
    trainset_path,
    shuffle = True,
    image_size = (img_width, img_height),
    batch_size = 32,
    validation_split = False
)

data_cat = trainset.class_names

useset = tf.keras.utils.image_dataset_from_directory(
    useset_path,
    shuffle = True,
    image_size = (img_height, img_width),
    batch_size = 32,
    validation_split = False
)

testset = tf.keras.utils.image_dataset_from_directory(
    testset_path,
    shuffle = True,
    image_size = (img_height, img_width),
    batch_size = 32,
    validation_split = False
)

plt.figure(figsize = (10, 10))
for image, labels in trainset.take(1):
    for i in range(4):
        plt.subplot(1, 4, i+1)
        plt.imshow(image[i].numpy().astype('uint8'))
        plt.title(data_cat[labels[i]])

from tensorflow.keras.models import Sequential


model = Sequential([
    layers.Rescaling(1./255),
    layers.Conv2D(16, 3, padding='same', activation='relu'),
    layers.MaxPooling2D(),
    layers.Conv2D(32, 3, padding='same', activation='relu'),
    layers.MaxPooling2D(),
    layers.Conv2D(64, 3, padding='same', activation='relu'),
    layers.MaxPooling2D(),
    layers.Flatten(),
    layers.Dropout(0.2),
    layers.Dense(128),
    layers.Dense(len(data_cat))
])

model.compile(optimizer='adam', loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), metrics=['accuracy'])

epochs_size = 25
history = model.fit(trainset, validation_data = useset, epochs = epochs_size)

model.save("nn.keras")