import gradio as gr
import requests
from PIL import Image
import torch
from torchvision import models
from huggingface_hub import hf_hub_download
from torchvision import transforms

# Mengunduh dan mempersiapkan model
model_path = hf_hub_download(repo_id="ahmadalfian/fruits_vegetables_classifier", filename="resnet50_finetuned.pth")
model = models.resnet50(pretrained=False)
num_classes = 36
model.fc = torch.nn.Linear(in_features=2048, out_features=num_classes)
model.load_state_dict(torch.load(model_path))
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()

# Transformasi untuk pra-pemrosesan gambar
preprocess = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# Fungsi untuk memprediksi kelas
def predict(image):
    image = image.convert("RGB")
    image_tensor = preprocess(image).unsqueeze(0).to(device)  # Tambahkan dimensi batch dan pindahkan ke device
    with torch.no_grad():
        outputs = model(image_tensor)
    predictions = outputs.argmax(dim=1)
    return predictions.item()

# Fungsi untuk mengambil informasi nutrisi
def get_nutritional_info(food):
    api_key = "3pm2NGZzYongVN1gRjnroVLUpsHC8rKWJFyx5moq"
    url = "https://api.nal.usda.gov/fdc/v1/foods/search"
    
    params = {
        "query": food,
        "pageSize": 5,
        "api_key": api_key
    }
    
    response = requests.get(url, params=params)
    data = response.json()

    if "foods" in data and len(data["foods"]) > 0:
        nutrients_totals = {
            "Energy": 0,
            "Carbohydrate, by difference": 0,
            "Fiber, total dietary": 0,
            "Vitamin C, total ascorbic acid": 0
        }
        item_count = len(data["foods"])

        for food in data["foods"]:
            for nutrient in food['foodNutrients']:
                nutrient_name = nutrient['nutrientName']
                nutrient_value = nutrient['value']

                if nutrient_name in nutrients_totals:
                    nutrients_totals[nutrient_name] += nutrient_value

        average_nutrients = {name: total / item_count for name, total in nutrients_totals.items()}

        return average_nutrients
    else:
        return None

# Fungsi utama Gradio
def classify_and_get_nutrition(image):
    predicted_class_idx = predict(image)
    class_labels = [
        'apple', 'banana', 'beetroot', 'bell pepper', 'cabbage', 'capsicum',
        'carrot', 'cauliflower', 'chilli pepper', 'corn', 'cucumber', 
        'eggplant', 'garlic', 'ginger', 'grapes', 'jalepeno', 'kiwi', 
        'lemon', 'lettuce', 'mango', 'onion', 'orange', 'paprika', 
        'pear', 'peas', 'pineapple', 'pomegranate', 'potato', 'raddish', 
        'soy beans', 'spinach', 'sweetcorn', 'sweetpotato', 'tomato', 
        'turnip', 'watermelon'
    ]
    predicted_label = class_labels[predicted_class_idx]
    
    nutrisi = get_nutritional_info(predicted_label)
    
    if nutrisi:
        return {
            "Predicted Class": predicted_label,
            "Energy (kcal)": nutrisi["Energy"],
            "Carbohydrates (g)": nutrisi["Carbohydrate, by difference"],
            "Fiber (g)": nutrisi["Fiber, total dietary"],
            "Vitamin C (mg)": nutrisi["Vitamin C, total ascorbic acid"]
        }
    else:
        return {
            "Predicted Class": predicted_label,
            "Nutritional Information": "Not Found"
        }

# Antarmuka Gradio
iface = gr.Interface(
    fn=classify_and_get_nutrition,
    inputs=gr.Image(type="pil"),
    outputs=gr.JSON(),
    title="Fruits and Vegetables Classifier",
    description="Upload an image of a fruit or vegetable to classify and get its nutritional information."
)

iface.launch()