from flask import Flask, request, jsonify
from flask_cors import CORS
from tensorflow import keras
from keras import layers
import numpy as np
from PIL import Image
from io import BytesIO
import base64
import os

app = Flask(__name__)
CORS(app)
app.config['MAX_CONTENT_LENGTH'] = 16 * 1000 * 1000

cats_and_dogs_model = None

# Load and train the cats and dogs model
def load_cats_and_dogs_model():
    # Define the model architecture
    model = keras.Sequential([
        layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 3)),
        layers.MaxPooling2D((2, 2)),
        layers.Conv2D(64, (3, 3), activation='relu'),
        layers.MaxPooling2D((2, 2)),
        layers.Conv2D(64, (3, 3), activation='relu'),
        layers.Flatten(),
        layers.Dense(64, activation='relu'),
        layers.Dense(1, activation='sigmoid')
    ])

    # Compile the model
    model.compile(optimizer='adam',
                  loss='binary_crossentropy',
                  metrics=['accuracy'])

    # Load or generate your training data (X_train and y_train)
    dataset_directory = 'C:\\Users\\chant\\OneDrive - ZHAW\\ZHAW\\Semester 6\\KI Anwendungen\\Tutorial\\Project2\\dataset'
    train_directory = os.path.join(dataset_directory, 'train')
    test_directory = os.path.join(dataset_directory, 'test')
    train_cats_directory = os.path.join(train_directory, 'cats')
    train_dogs_directory = os.path.join(train_directory, 'dogs')
    test_cats_directory = os.path.join(test_directory, 'cats')
    test_dogs_directory = os.path.join(test_directory, 'dogs')

    train_images = []
    train_labels = []
    test_images = []
    test_labels = []

    for filename in os.listdir(train_cats_directory):
        if filename.endswith('.jpg'):
            img = Image.open(os.path.join(train_cats_directory, filename))
            img = img.resize((28, 28))
            img = img.convert('RGB')
            img = np.array(img)
            train_images.append(img)
            train_labels.append(0)  # Assign label 0 for cats

    for filename in os.listdir(train_dogs_directory):
        if filename.endswith('.jpg'):
            img = Image.open(os.path.join(train_dogs_directory, filename))
            img = img.resize((28, 28))
            img = img.convert('RGB')
            img = np.array(img)
            train_images.append(img)
            train_labels.append(1)  # Assign label 1 for dogs

    for filename in os.listdir(test_cats_directory):
        if filename.endswith('.jpg'):
            img = Image.open(os.path.join(test_cats_directory, filename))
            img = img.resize((28, 28))
            img = img.convert('RGB')
            img = np.array(img)
            test_images.append(img)
            test_labels.append(0)  # Assign label 0 for cats

    for filename in os.listdir(test_dogs_directory):
        if filename.endswith('.jpg'):
            img = Image.open(os.path.join(test_dogs_directory, filename))
            img = img.resize((28, 28))
            img = img.convert('RGB')
            img = np.array(img)
            test_images.append(img)
            test_labels.append(1)  # Assign label 1 for dogs

    X_train = np.array(train_images)
    y_train = np.array(train_labels)
    X_test = np.array(test_images)
    y_test = np.array(test_labels)

    # Preprocess the data
    X_train = X_train.astype('float32') / 255
    X_test = X_test.astype('float32') / 255

    # Train the model
    model.fit(X_train, y_train, epochs=10, batch_size=32)

    # Evaluate the model
    test_loss, test_acc = model.evaluate(X_test, y_test, verbose=2)
    print('Test accuracy:', test_acc)

    # Set the trained model as the global variable
    global cats_and_dogs_model
    cats_and_dogs_model = model

# Define the route for image classification
@app.route('/api/prediction/classify', methods=['POST'])
def classify_image():
    data = request.get_json()
    image_data = data['image']
    image = Image.open(BytesIO(base64.b64decode(image_data)))
    image = image.resize((28, 28))
    image = image.convert('RGB')
    image = np.array(image)
    image = image.astype('float32') / 255
    image = np.expand_dims(image, axis=0)
    result = cats_and_dogs_model.predict(image)[0][0]
    class_name = 'cat' if result < 0.5 else 'dog'
    response = {'class_name': class_name, 'confidence': float(result)}
    return jsonify(response)

# Add this if statement to start the Flask app
if __name__ == "__main__" or __name__ == "app" or __name__ == "flask_app":
    print(("* Loading models and starting the server..."
           "please wait until the server has fully started"))
    load_cats_and_dogs_model()
    app.run()