from gardens.models import *

import torch
import os
# TensorFlow and Keras libraries
import tensorflow as tf
from tensorflow import keras
# used to chech device
from tensorflow.python.client import device_lib
# Other libraries
import numpy as np
from PIL import Image

from django.conf import settings

import shutil

leafModel = os.path.join(settings.STATIC_ROOT, 'MachineLearningModels/odLeaves_transferlearning_50epochs.pt')

def scabDetectionModel(photo):
    # model loaded for detecting leaves
    model = torch.hub.load('ultralytics/yolov5', 'custom', path=leafModel)
    # defining the photo  needed to check
    img = photo.image_uri.path

    results = model(img)
    results.crop()
    print(results)

    for folder in os.listdir('runs/detect'):
        directory = f'runs/detect/{folder}/crops/leaf'
        probabilities = []
        for items in os.listdir(directory):
            file_path = os.path.join(directory, items)
            img = Image.open(file_path)

            scabDetectionModel = os.path.join(settings.STATIC_ROOT, 'MachineLearningModels/applescabdls_advanced_tuned_mobilenetv2_adam_20.h5') 

            from tensorflow.keras import backend as K

            def relu6(x):
                return K.relu(x, max_value=6)

            # Model parameters
            image_size = 200
            image_channels = 3
            
            scab = img
            scab.load()

            model = tf.keras.models.load_model(scabDetectionModel, custom_objects={'relu6': relu6} )

            scab = scab.resize((image_size, image_size ))
            data = np.asarray( scab, dtype="int32" )
            data = data / 255.0
            data = np.reshape(data, (1, image_size, image_size, 3))
            results = model.predict( data )
            # probability of each class
            number = results[0][0].item()
            probabilities.append(number)
        hasScab = False
        lowestProbability = 0
        for probability in probabilities:
            if probability < 0.09:
                hasScab = True
                if lowestProbability < probability:
                    lowestProbability = probability
        probability_scab = 1 - lowestProbability
        # if there is scab
        if hasScab == True:
            try:
                state = TreeState.objects.get(tree_id = photo.tree_id, mission_id = photo.mission_id)
                if state.probability_scab < probability_scab:
                    state.probability_scab = probability_scab
                    state.has_scab = True
                    state.save()
            except TreeState.DoesNotExist:
                TreeState.objects.create(tree_id = photo.tree_id, mission_id = photo.mission_id, probability_scab = probability_scab, has_scab = True)
        # if there is no scab
        else:
            try:
                state = TreeState.objects.get(tree_id = photo.tree_id, mission_id = photo.mission_id)
            except TreeState.DoesNotExist:
                TreeState.objects.create(tree_id = photo.tree_id, mission_id = photo.mission_id, probability_scab = 0, has_scab = False)
        shutil.rmtree(f'runs/detect/{folder}/')
        photo.delete()