import os
import numpy as np
import cv2
import matplotlib.pyplot as plt
import random
from sklearn.model_selection import train_test_split

cnn_image_shape = (128, 128)
"""
Animes will be class 0 and cartoons class 1.

Import anime data.
"""
anime_path = '../input/anime-and-cartoon-image-classification/Training Data/Anime/'

X = []
y = []

for folder in os.scandir(anime_path):
    for file in os.scandir(anime_path + folder.name):
        img = cv2.imread(anime_path + folder.name + '/' + file.name, cv2.IMREAD_COLOR)
        img = cv2.resize(img, cnn_image_shape)
        img = np.array(img, dtype='float32')
        X.append(img)
        y.append(0)
"""
Import cartoon data.
"""
cartoon_path = '../input/anime-and-cartoon-image-classification/Training Data/Cartoon/'

for folder in os.scandir(cartoon_path):
    for file in os.scandir(cartoon_path + folder.name):
        img = cv2.imread(cartoon_path + folder.name + '/' + file.name, cv2.IMREAD_COLOR)
        img = cv2.resize(img, cnn_image_shape)
        img = np.array(img, dtype='float32')
        X.append(img)
        y.append(1)
"""
convert arrays to numpy
"""
X = np.array(X)
y = np.array(y)
"""
Vamos ver quantas imagens temos de cada um dos 
"""
print(f'Number of anime images: {np.count_nonzero(y == 0)}')

print(f'Number of cartoons images: {np.count_nonzero(y == 1)}')
"""
Plot some images
"""
labels = {0:'Anime', 1:'Cartoon'}

for i in range(5):
    n = random.randint(0, X.shape[0])
    plt.figure()
    plt.title(labels[y[n]])
    plt.imshow(X[n]/255)
"""
split the database into 3split the database into 3
* 80% for training
* 10% for validation
* 10% for test
"""
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, stratify=y, shuffle=True, random_state=42)
X_val, X_test, y_val, y_test = train_test_split(X_val, y_val, test_size=0.5, stratify=y_val, random_state=42)
"""
Create the model
"""
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Input, Lambda, GlobalAveragePooling2D, Dropout, Dense
from tensorflow.keras.callbacks import EarlyStopping
from tensorflow.keras.optimizers import Adam


before_mobilenet = Sequential([Input((128, 128, 3)),
                               Lambda(preprocess_input)])

mobilenet = MobileNetV2(input_shape=(128, 128, 3), include_top=False)

after_mobilenet = Sequential([GlobalAveragePooling2D(),
                              Dropout(0.3),
                              Dense(2, activation='softmax')])

model = Sequential([before_mobilenet, mobilenet, after_mobilenet])

opt = Adam(learning_rate=0.0001)

model.build(((None, 128, 128, 3)))

model.compile(optimizer=opt, loss='sparse_categorical_crossentropy', metrics=['accuracy'])

model.summary()
early_stopping = EarlyStopping(
    patience=10,
    min_delta=0.001,
    restore_best_weights=True,
)

history = model.fit(
    X_train, 
    y_train, 
    epochs=300, 
    validation_data=(X_test, y_test), 
    callbacks=[early_stopping]
)
plt.plot(history.history['loss'])
plt.plot(history.history['val_loss'])
plt.title('model loss')
plt.ylabel('loss')
plt.xlabel('epoch')
plt.legend(['train', 'validation'], loc='upper right')
plt.show()
plt.plot(history.history['accuracy'])
plt.plot(history.history['val_accuracy'])
plt.title('model accuracy')
plt.ylabel('accuracy')
plt.xlabel('epoch')
plt.legend(['train', 'validation'], loc='upper left')
plt.show()
from sklearn.metrics import classification_report

predictions = model.predict(X_test)

predict = []
for i in predictions:
    predict.append(np.argmax(i))
    
print(classification_report(y_test, predict))
"""
Test with outside image
"""
"""
![](https://a-static.mlcdn.com.br/574x431/painel-de-festa-cavaleiros-do-zodiaco-cdz-02-colormyhome/colormyhome/3455/2b09a906f973deadbf918d1008084f28.jpg)
"""
!wget https://a-static.mlcdn.com.br/574x431/painel-de-festa-cavaleiros-do-zodiaco-cdz-02-colormyhome/colormyhome/3455/2b09a906f973deadbf918d1008084f28.jpg -O cdz.jpg
img = cv2.imread('./cdz.jpg', cv2.IMREAD_COLOR)
img = cv2.resize(img, cnn_image_shape)
img = np.array(img, dtype='float32')

img = img.reshape((1, 128, 128, 3))

img.shape
prediction = model.predict(img)
prediction_probabilities = np.array(prediction)
argmaxs = np.argmax(prediction_probabilities, axis=1)
print(argmaxs)