import gradio as gr
import torch
from torchvision import transforms
from PIL import Image
import joblib

# Load model and metadata
model = torch.jit.load("resnet_grocery_model_scripted.pt", map_location=torch.device("cpu"))
model.eval()

assets = joblib.load("deployment_assets.joblib")
transform = assets['transform']
class_names = assets['class_names']

# Item names and prices
items = ['Bisconni Chocolate Chip Cookies 46.8gm', 'Coca Cola Can 250ml', 'Colgate Maximum Cavity Protection 75gm',
         'Fanta 500ml', 'Fresher Guava Nectar 500ml', 'Fruita Vitals Red Grapes 200ml', 'Islamabad Tea 238gm',
         'Kolson Slanty Jalapeno 18gm', 'Kurkure Chutney Chaska 62gm', 'LU Candi Biscuit 60gm', 'LU Oreo Biscuit 19gm',
         'LU Prince Biscuit 55.2gm', 'Lays Masala 34gm', 'Lays Wavy Mexican Chili 34gm', 'Lifebuoy Total Protect Soap 96gm',
         'Lipton Yellow Label Tea 95gm', 'Meezan Ultra Rich Tea 190gm', 'Peek Freans Sooper Biscuit 13.2gm',
         'Safeguard Bar Soap Pure White 175gm', 'Shezan Apple 250ml', 'Sunsilk Shampoo Soft - Smooth 160ml',
         'Super Crisp BBQ 30gm', 'Supreme Tea 95gm', 'Tapal Danedar 95gm', 'Vaseline Healthy White Lotion 100ml']

prices = [55.20, 31.75, 90.00, 63.50, 50.00, 35.00, 150.00, 15.00, 25.00, 30.00, 10.00, 30.00, 20.00, 20.00, 44.50, 100.00,
          200.00, 10.00, 70.00, 25.00, 120.00, 15.00, 100.00, 100.00, 120.00]

pricelist = dict(zip(items, prices))

# Functions
def add_image(new_img, image_list):
    if new_img:
        image_list.append(new_img)
    return image_list, image_list

def remove_last_image(image_list):
    if image_list:
        image_list.pop()
    return image_list, image_list

def classify_and_track(images, budget):
    results = []
    total_cost = 0
    receipt_lines = []

    receipt_lines.append("        🛒  SCANCART SUPERMARKET  🛒")
    receipt_lines.append("          📍 Online Smart Checkout")
    receipt_lines.append("      ==============================")
    receipt_lines.append("      ITEM                     PRICE")
    receipt_lines.append("      ==============================")

    for img in images:
        if isinstance(img, Image.Image):
            input_tensor = transform(img).unsqueeze(0)
            with torch.no_grad():
                outputs = model(input_tensor)
                predicted_index = torch.argmax(outputs, dim=1).item()
                predicted_class = class_names[predicted_index]
                price = pricelist.get(predicted_class, 0)
                total_cost += price
                results.append((img, f"{predicted_class} - ₱{price:.2f}"))

                item_name = (predicted_class[:22] + '...') if len(predicted_class) > 25 else predicted_class
                receipt_lines.append(f"      {item_name:<25} ₱{price:>6.2f}")
        else:
            results.append((None, "Invalid image"))

    remaining = budget - total_cost
    receipt_lines.append("      ==============================")
    receipt_lines.append(f"      TOTAL:{'':>20}₱{total_cost:>7.2f}")
    receipt_lines.append(f"      BUDGET:{'':>19}₱{budget:>7.2f}")
    receipt_lines.append(f"      REMAINING:{'':>16}₱{remaining:>7.2f}")
    receipt_lines.append("      ==============================")
    if total_cost <= budget:
        receipt_lines.append("      ✅ You're within budget! 🎉")
        receipt_lines.append("      🛍️ Happy Shopping!")
    else:
        receipt_lines.append("      ❌ Over budget! Try removing items 😬")
    receipt_lines.append("      ==============================")
    receipt_lines.append("       THANK YOU FOR SHOPPING WITH US!")
    receipt_lines.append("      Come back for smarter buys! 🧾")

    receipt_text = "\n".join(receipt_lines)
    return results, receipt_text

# Custom Theme
custom_theme = gr.themes.Base(
    primary_hue="rose",
    secondary_hue="amber",
    neutral_hue="stone",
    font=[gr.themes.GoogleFont("Helvetica")]
).set(
    body_background_fill="#ffe5e5"
)

# Gradio Interface
with gr.Blocks(
    theme=custom_theme,
    css="""
    .center-column { 
        display: flex !important; 
        justify-content: center !important; 
        align-items: center !important; 
        width: 100% !important; 
        margin: 0 auto !important; 
    }
    #title { 
        width: 100% !important; 
        margin: 0 auto !important; 
        text-align: center !important; 
        padding: 20px !important;
        border: 1px solid #ccc !important; /* Remove this line after confirming centering */
    }
    #title h1 { 
        font-size: 2.2em !important; 
        margin: 0 auto 0.2em !important; 
        text-align: center !important; 
    }
    #title p { 
        font-size: 1.1em !important; 
        margin: 0 auto !important; 
        text-align: center !important; 
        max-width: 800px !important; 
    }
    .upload-box { 
        background-color: #1a1a1a !important; 
        color: #ffffff !important; 
        text-align: center !important; 
        padding: 20px !important; 
        border-radius: 5px !important; 
        width: 100% !important; 
        height: 150px !important; 
        display: flex !important; 
        flex-direction: column !important; 
        justify-content: center !important; 
    }
    .budget-box { 
        background-color: #1a1a1a !important; 
        color: #ffffff !important; 
        text-align: center !important; 
        padding: 20px !important; 
        border-radius: 5px !important; 
        width: 100% !important; 
        height: 150px !important; 
        display: flex !important; 
        flex-direction: column !important; 
        justify-content: center !important; 
    }
    .button-row button { 
        background-color: #4a4a4a !important; 
        color: #ffffff !important; 
        margin: 0 10px !important; 
        padding: 10px 20px !important; 
        border-radius: 5px !important; 
        border: none !important; 
    }
    """
) as iface:
    with gr.Row():
        with gr.Column(scale=1, min_width=0, elem_classes="center-column"):
            gr.Markdown(
                """
                <div style="width: 100%; margin: 0 auto; text-align: center; padding: 20px;">
                    <h1 style="font-size: 2.2em; margin: 0 auto 0.2em; text-align: center;">ScanCart 🧾</h1>
                    <p style="font-size: 1.1em; margin: 0 auto; text-align: center; max-width: 800px;">Upload images of your grocery items, and ScanCart will identify them, calculate the total cost, and check if it fits within your budget. 🛒📸</p>
                </div>
                """,
                elem_id="title"
            )

    with gr.Row():
        with gr.Column(scale=1):
            with gr.Row():
                image_input = gr.Image(type="pil", label="Upload Grocery Items 🍎", elem_classes="upload-box", height=150)
                budget_input = gr.Number(label="Budget (₱) 🍊", value=500.0, precision=2, elem_classes="budget-box")
            gallery_output = gr.Gallery(label="Cart Preview", columns=3, height="auto", object_fit="contain")
            with gr.Row(elem_classes="button-row"):
                add_btn = gr.Button("Add Items 🚀")
                remove_btn = gr.Button("Reset 🔄")
                classify_btn = gr.Button("Generate Receipt 🧾")

        with gr.Column(scale=2):
            receipt_output = gr.Textbox(label="Receipt", lines=20, interactive=False, show_copy_button=True)

    image_list_state = gr.State([])

    add_btn.click(
        add_image,
        inputs=[image_input, image_list_state],
        outputs=[image_list_state, gallery_output]
    )

    remove_btn.click(
        remove_last_image,
        inputs=image_list_state,
        outputs=[image_list_state, gallery_output]
    )

    classify_btn.click(
        classify_and_track,
        inputs=[image_list_state, budget_input],
        outputs=[gallery_output, receipt_output]
    )

iface.launch()