import PIL
from torchvision import datasets, transforms, models
import torch
from PIL import Image
import torch.nn as nn 
import pandas as pd
import numpy as np 
import gradio as gr

class_names = ['apple_pie',
 'bibimbap',
 'cannoli',
 'edamame',
 'falafel',
 'french_toast',
 'ramen',
 'sushi',
 'tiramisu']

def pil_loader(path):
        with open(path, 'rb') as f:
            img = Image.open(f)
            return img.convert('RGB')

def predict(img_path):
    # Load and preprocess the image
    # image = pil_loader(img_path)
        # Convert Gradio image input to a NumPy array
    img_array = img_path.astype(np.uint8)

    # # Convert NumPy array to PIL Image
    image = Image.fromarray(img_array)

    test_transforms = 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])
    ])

    # Apply transformations
    image = test_transforms(image)

    inf_model = models.resnet18(pretrained=False)
    num_ftrs = inf_model.fc.in_features
    # Here the size of each output sample is set to 2.
    # Alternatively, it can be generalized to nn.Linear(num_ftrs, len(class_names)).
    inf_model.fc = nn.Linear(num_ftrs, len(class_names))

# model_1 = model_1.to(device)
    inf_model.to(torch.device('cpu'))
    inf_model.load_state_dict(torch.load('./resnet18_tinyfood_classifier.pth', map_location='cpu'))

    # Perform inference
    with torch.no_grad():
        inf_model.eval()
        out = inf_model(image.unsqueeze(0))  # Add batch dimension
    
    # Get the predicted class and confidence
    _, preds = torch.max(out, 1)
    idx = preds.cpu().numpy()[0]
    pred_class = class_names[idx]

    # Assuming `out` is logits, you may need to apply softmax instead of sigmoid
    probabilities = torch.softmax(out, dim=1)  # Apply softmax to get probabilities
    confidence = probabilities[0, idx].item() * 100  # Get confidence for the predicted class

    nutrition_data_path = './food-data.csv'
    # Membaca file CSV
    df = pd.read_csv(nutrition_data_path)

    # Mencocokkan prediksi dengan data CSV
    if pred_class.capitalize() in df["Makanan"].values:
        row = df.loc[df["Makanan"] == pred_class.capitalize()]

        # Mengambil informasi gizi
        calories = row["Kalori"].values[0]
        protein = row["Protein"].values[0]
        fat = row["Lemak"].values[0]
        carbs = row["Karbohidrat"].values[0]
        fiber = row["Serat"].values[0]
        sugar = row["Gula"].values[0]
        price = row["Harga (Rp)"].values[0]

        return pred_class, calories, protein, fat, carbs, fiber, sugar, price
    else:
        nutrition_info = None
        return 'Food not found', 0, 0, 0, 0, 0, 0

    # return pred_class, confidence


# img_path = '/content/drive/MyDrive/Assignment-Citra-SkillacademyAI/bibimbap.jpeg'
# print(predict(img_path))

interface = gr.Interface(
    predict,
    inputs="image",
    title="Selera Cafe App",
    description="This App will provide the information of your food choice in Selera Cafe. The menu includes: Apple Pie, Bibimbap, Cannoli, Edamame, Falafel, French Toast, Ramen, Sushi, Tiramisu. Enjoy your food!",

    outputs=[
      gr.Text(label="Food Label"),
      gr.Number(label="Calories"),
      gr.Number(label="Protein"),
      gr.Number(label="Fat"),
      gr.Number(label="Carbs"),
      gr.Number(label="Fiber"),
      gr.Number(label="Sugar"),
      gr.Number(label="Price")
    ],
    examples = [
        '/content/drive/MyDrive/Assignment-Citra-SkillacademyAI/bibimbap.jpeg',
        '/content/drive/MyDrive/Assignment-Citra-SkillacademyAI/food-101-tiny/apple-pie.jpeg',
        '/content/drive/MyDrive/Assignment-Citra-SkillacademyAI/food-101-tiny/cannoli.jpeg'
    ])
interface.launch()