import os

import gradio.helpers
from PIL import Image
import torch
import gradio as gr
from torchvision import transforms
from transformers import AutoModelForImageClassification, AutoConfig


# Define the mapping of label indices to human-readable names
label_names = {
    0: 'Apple___Apple_scab', 1: 'Apple___Black_rot', 2: 'Apple___Cedar_apple_rust', 3: 'Apple___healthy',
    4: 'Background_without_leaves', 5: 'Blueberry___healthy', 6: 'Cherry___Powdery_mildew', 7: 'Cherry___healthy',
    8: 'Corn___Cercospora_leaf_spot Gray_leaf_spot', 9: 'Corn___Common_rust', 10: 'Corn___Northern_Leaf_Blight',
    11: 'Corn___healthy', 12: 'Grape___Black_rot', 13: 'Grape___Esca_(Black_Measles)',
    14: 'Grape___Leaf_blight_(Isariopsis_Leaf_Spot)', 15: 'Grape___healthy',
    16: 'Orange___Haunglongbing_(Citrus_greening)', 17: 'Peach___Bacterial_spot', 18: 'Peach___healthy',
    19: 'Pepper,_bell___Bacterial_spot', 20: 'Pepper,_bell___healthy', 21: 'Potato___Early_blight',
    22: 'Potato___Late_blight', 23: 'Potato___healthy', 24: 'Raspberry___healthy', 25: 'Soybean___healthy',
    26: 'Squash___Powdery_mildew', 27: 'Strawberry___Leaf_scorch', 28: 'Strawberry___healthy',
    29: 'Tomato___Bacterial_spot', 30: 'Tomato___Early_blight', 31: 'Tomato___Late_blight',
    32: 'Tomato___Leaf_Mold', 33: 'Tomato___Septoria_leaf_spot',
    34: 'Tomato___Spider_mites Two-spotted_spider_mite', 35: 'Tomato___Target_Spot',
    36: 'Tomato___Tomato_Yellow_Leaf_Curl_Virus', 37: 'Tomato___Tomato_mosaic_virus', 38: 'Tomato___healthy'
}

baslik="Bitki Hastalık Tespitinde Yapay Zeka ile Sürdürülebilir Tarımın İnşası"
aciklama="""
Bu projenin ana amacı, bitki hastalıklarının erken teşhisini sağlayan bir yapay zeka sistemini geliştirmek ve çiftçilerin 
tarımsal ürün kayıplarını önlemelerine yardımcı olmaktır. Projemiz aynı zamanda Birleşmiş Milletler Sürdürülebilir 
Kalkınma Hedefleri'nden biri olan \"Açlığı Sıfıra İndirme\" (SDG 2) hedefine katkıda bulunmayı amaçlamaktadır.
Projenin yaygın etkisi tarım topluluklarına ve çiftçilere yöneliktir. Erken bitki hastalığı tespiti, çiftçilerin mahsullerini
korumalarına ve verimliliği artırmalarına yardımcı olur. Bu da gıda güvenliğinin artırılmasına katkıda bulunur.\n
Çiftçiler, bitki resimlerini yükleyebilir ve hastalıkların tespitini gerçekleştirebilirler. 
Şayet yüklenen resimdeki bitki hastalıklı ise sistem tarafından bir öneri mekanizması devreye girecek ve çiftçiye, bitkinin hastalığına karşı neler yapabileceği
hakkında özet bir bilgi verecektir.

"""
images = [
    ["examples/Apple_scab(24).JPG", "Apple Scab"],
    ["examples/Background_without_leaves(194).jpg", "A photo"],
    ["examples/Background_without_leaves(290).jpg", "A photo"],
    ["examples/Corn_leaf_spot(108).JPG", "Corn leaf spot"],
    ["examples/Orange_Haunglongbing(92).JPG", "Orange Haunglongbing"],
    ["examples/Peach_Bacterial_spot(9).JPG", "Peach bacterial spot"],
    ["examples/Potato_early_blight(2).JPG", "Potato early blight"],
    ["examples/Strawberry_healthy(94).JPG", "Strawberry healthy"],
    ["examples/Tomato_bacterial_spot(6).JPG", "Tomato bacterial spot"],
    ["examples/Tomato_mosaic_virus(42).JPG", "Tomato mosaic virus"],
]

YOUR_NUMBER_OF_CLASSES = 39

# Load a configuration from a known model but don't load pretrained weights
config = AutoConfig.from_pretrained("facebook/convnext-tiny-224")
config.num_labels = YOUR_NUMBER_OF_CLASSES  # Adjust number of classes

model_path = "model/model-fine-tune.pth"
state_dict = torch.load(model_path, map_location=torch.device('cpu'))  # Ensure to load to the right device

model = AutoModelForImageClassification.from_config(config)
model.load_state_dict(state_dict)

# Set the model to evaluation mode (now it's a proper PyTorch model)
model.eval()

# Define the image transformation
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

def treatment(label):
    md_file = open("treatments/" + label + ".md", "r", encoding="utf-8")
    return md_file.read()
def predict(image):
    tensor_image = transform(image).unsqueeze(0)
    with torch.no_grad():
        outputs = model(tensor_image)
        _, predicted_idx = torch.max(outputs.logits, 1)
        predicted_label = predicted_idx.item()
        predicted_label_name = label_names.get(predicted_label, "Background_without_leaves")

    return treatment(predicted_label_name)

demo = gr.Interface(fn=predict,
                    inputs=gr.Image(type="pil", image_mode="RGB",label="Fotoğraf"),
                    outputs=gr.Markdown(value="Çözümü burada görüntülemek için lütfen ilgili alana fotoğraf yükleyiniz."),
                    title=baslik,
                    description=aciklama,
                    allow_flagging="never",
                    submit_btn="Gönder",
                    clear_btn="Temizle",
                    theme='sudeepshouche/minimalist',
                    examples=images,
                    examples_per_page=5
                    )


if __name__ == "__main__":
    demo.launch()