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
import os
import pandas as pd
import numpy as np
import torch
from ultralytics import YOLO
from PIL import Image, ImageDraw
from django.conf import settings

import shutil

model = os.path.join(settings.STATIC_ROOT, 'MachineLearningModels/FruitDetectionModel.pt')

def divide_picture(image_name, crop_size = 640):
    picture_array = []
    width, height = image_name.size
    crop_width = width // crop_size
    print(f'crop_width: {crop_width}')
    crop_height = height // crop_size
    for i in range(crop_width):
        for j in range(crop_height):
            left = i * crop_size
            top = j * crop_size
            right = (i + 1) * crop_size
            bottom = (j + 1) * crop_size
            picture_array += [image_name.crop((left, top, right, bottom))]
    return picture_array


def objectDetection(photo):
    # model loaded for detecting pears
    model = os.path.join(settings.STATIC_ROOT, 'MachineLearningModels/FruitDetectionModel.pt')
    model = YOLO(model)
    # defining the photo  needed to check
    # img = photo.get('image_uri','0')
    # print(photo)
    original_img = Image.open(photo).convert("RGB")
    img_array = divide_picture(original_img, crop_size=640)
    # total_counts = {
    #     'Pear': {'20': 0,'30': 0,'40': 0,'50': 0,'60': 0,'70': 0, '80': 0, '90': 0},
    #     'PFruitlet': {'20': 0,'30': 0,'40': 0,'50': 0,'60': 0,'70': 0, '80': 0, '90': 0},
    #     'Cherry72': {'20': 0,'30': 0,'40': 0,'50': 0,'60': 0,'70': 0, '80': 0, '90': 0},
    #     'Cherry81': {'20': 0,'30': 0,'40': 0,'50': 0,'60': 0,'70': 0, '80': 0, '90': 0},
    # }
    total_counts = {'Pear640': {'20': 0}, 'PFruitlets640': {'20': 0}, 'CherriesBBCH72': {'20': 0}, 'CherriesBBCH81': {'20': 0}, 'ApplesBBCH76': {'20': 0}, 'ApplesBBCH81': {'20': 0}}
    class_names = {0: 'Pear640', 1: 'PFruitlets640', 2: 'CherriesBBCH72', 3: 'CherriesBBCH81', 4: 'ApplesBBCH76', 5: 'ApplesBBCH81'}

    for cropped_img in img_array:
        results = model.predict(cropped_img, conf=0.2)
            
        for result in results:
            if result.boxes and result.boxes.shape[0] > 0:
                for box in result.boxes.data:
                    confidence_score = box[4] * 100
                    class_index = int(box[5])
                    class_name = class_names.get(class_index, 'Unknown')
                    # if confidence_score >= 90:
                    #     total_counts[class_name]['90'] += 1
                    # if confidence_score >= 80:
                    #     total_counts[class_name]['80'] += 1
                    # if confidence_score >= 70:
                    #     total_counts[class_name]['70'] += 1
                    # if confidence_score >= 60:
                    #     total_counts[class_name]['60'] += 1
                    # if confidence_score >= 50:
                    #     total_counts[class_name]['50'] += 1
                    # if confidence_score >= 40:
                    #     total_counts[class_name]['40'] += 1
                    # if confidence_score >= 30:
                    #     total_counts[class_name]['30'] += 1
                    if confidence_score >= 20:
                        total_counts[class_name]['20'] += 1
    return total_counts


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

    # for folder in os.listdir('runs/detect'):
    #     directory = f'runs/detect/{folder}/crops/apples'
    #     # check if directory exists
    #     if os.path.exists(directory):
    #         amountOfPears = len([entry for entry in os.listdir(directory) if os.path.isfile(os.path.join(directory, entry))])
    #     else:
    #         amountOfPears = 0
    #     try:
    #         state = TreeState.objects.get(tree_id = photo.tree_id, mission_id = photo.mission_id)
    #         state.yields = amountOfPears
    #         state.save()
    #     except TreeState.DoesNotExist:
    #         TreeState.objects.create(tree_id = photo.tree_id, mission_id = photo.mission_id, yields = amountOfPears, has_scab = True)
    #     shutil.rmtree(f'runs/detect/{folder}/')
    #     photo.delete()