import gradio as gr
import torch
import torchvision.transforms as transforms
from torchvision import models
from PIL import Image

model = models.resnet18()
model.fc = torch.nn.Linear(model.fc.in_features, 2)
model.load_state_dict(torch.load("papillon_model.pth", map_location="cpu"))
model.eval()

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

def classify_dog(image):
    if image is None:
        return "## drop a dog pic first! 🐾"
    image_tensor = transform(image).unsqueeze(0)
    with torch.no_grad():
        output = model(image_tensor)
        _, predicted = torch.max(output, 1)
        confidence = torch.nn.functional.softmax(output, dim=1)[0]
        score = confidence[predicted.item()].item() * 100
    if predicted.item() == 1:
        return f"## yep! that's a papillon 🦋\n**{score:.0f}% sure**"
    else:
        return f"## nope! not a papillon 🐶\n**{score:.0f}% sure**"

css = """
@import url('https://fonts.googleapis.com/css2?family=Fredoka+One&family=Nunito:wght@400;600;800&display=swap');

/* === WOOF PALETTE === */
:root {
    --sage: #8fbfb0;
    --yellow: #e8d44d;
    --rust: #c85a2a;
    --cream: #f5f0e0;
    --dark: #1e3d3a;
    --white: #faf8f2;
}

/* === GLOBAL === */
body, .gradio-container {
    background-color: var(--sage) !important;
    font-family: 'Nunito', sans-serif !important;
    min-height: 100vh !important;
}

.gradio-container {
    max-width: 520px !important;
    margin: 0 auto !important;
    padding: 0 16px 40px !important;
}

/* === HEADER === */
.woof-header {
    text-align: center;
    padding: 28px 0 8px;
    position: relative;
}

.woof-title {
    font-family: 'Fredoka One', cursive !important;
    font-size: 72px !important;
    color: var(--dark) !important;
    letter-spacing: -2px;
    line-height: 1;
    margin: 0;
    /* sketchy text shadow trick */
    text-shadow:
        3px 3px 0 var(--rust),
        -1px -1px 0 var(--dark);
}

.woof-subtitle {
    font-family: 'Fredoka One', cursive;
    font-size: 18px;
    color: var(--dark);
    margin: 4px 0 20px;
    opacity: 0.85;
    letter-spacing: 0.3px;
}

/* === UPLOAD ZONE === */
.upload-zone {
    background: var(--cream) !important;
    border: 3.5px solid var(--dark) !important;
    border-radius: 24px !important;
    /* hand-drawn uneven border using box-shadow */
    box-shadow:
        4px 4px 0 var(--dark),
        inset 0 0 0 2px rgba(255,255,255,0.4) !important;
    overflow: hidden;
    margin-bottom: 16px;
}

.upload-zone:hover {
    transform: translate(-1px, -1px) !important;
    box-shadow: 6px 6px 0 var(--dark) !important;
    transition: all 0.15s ease !important;
}

/* target the Gradio upload area */
.upload-zone .wrap {
    background: transparent !important;
    border: none !important;
}

.upload-zone .upload-container {
    background: transparent !important;
}

/* the drag & drop text inside */
.upload-zone .icon-wrap,
.upload-zone .wrap span {
    color: var(--dark) !important;
    font-family: 'Fredoka One', cursive !important;
    font-size: 16px !important;
}

/* === CLASSIFY BUTTON === */
button#classify-btn,
.classify-btn button,
button.primary-btn {
    background: var(--yellow) !important;
    color: var(--dark) !important;
    font-family: 'Fredoka One', cursive !important;
    font-size: 22px !important;
    letter-spacing: 0.5px;
    border: 3.5px solid var(--dark) !important;
    border-radius: 50px !important;
    box-shadow: 4px 4px 0 var(--dark) !important;
    padding: 14px 0 !important;
    width: 100% !important;
    cursor: pointer;
    transition: all 0.12s ease !important;
    margin-bottom: 16px;
}

button.primary-btn:hover,
button#classify-btn:hover {
    transform: translate(-2px, -2px) !important;
    box-shadow: 6px 6px 0 var(--dark) !important;
}

button.primary-btn:active,
button#classify-btn:active {
    transform: translate(2px, 2px) !important;
    box-shadow: 2px 2px 0 var(--dark) !important;
}

/* gradio's actual submit button */
.gradio-container button.primary {
    background: var(--yellow) !important;
    color: var(--dark) !important;
    font-family: 'Fredoka One', cursive !important;
    font-size: 20px !important;
    border: 3.5px solid var(--dark) !important;
    border-radius: 50px !important;
    box-shadow: 4px 4px 0 var(--dark) !important;
    padding: 14px 0 !important;
    transition: all 0.12s ease !important;
}

.gradio-container button.primary:hover {
    transform: translate(-2px,-2px) !important;
    box-shadow: 6px 6px 0 var(--dark) !important;
}

/* === RESULT BOX === */
.result-box {
    background: var(--rust) !important;
    border: 3.5px solid var(--dark) !important;
    border-radius: 20px !important;
    box-shadow: 4px 4px 0 var(--dark) !important;
    padding: 0 !important;
    overflow: hidden;
}

.result-box textarea,
.result-box .prose,
.result-box p,
.result-box * {
    background: transparent !important;
    color: var(--cream) !important;
    font-family: 'Fredoka One', cursive !important;
    font-size: 20px !important;
    line-height: 1.4 !important;
    text-align: center !important;
    border: none !important;
    padding: 20px !important;
}

.result-box label span {
    color: var(--dark) !important;
    font-family: 'Fredoka One', cursive !important;
    font-size: 13px !important;
    background: var(--yellow) !important;
    padding: 4px 12px !important;
    border-radius: 0 0 10px 0 !important;
    display: inline-block !important;
}

/* === ALL LABELS === */
.gradio-container label span,
.gradio-container .label-wrap span {
    font-family: 'Fredoka One', cursive !important;
    font-size: 15px !important;
    color: var(--dark) !important;
}

/* === FOOTER === */
.woof-footer {
    text-align: center;
    margin-top: 20px;
    font-family: 'Fredoka One', cursive;
    font-size: 13px;
    color: var(--dark);
    opacity: 0.7;
}

/* === decorative dog doodle strip === */
.doodle-strip {
    display: flex;
    justify-content: center;
    gap: 12px;
    margin: 10px 0 18px;
    font-size: 28px;
    filter: drop-shadow(1px 1px 0 rgba(0,0,0,0.2));
}

/* === MOBILE camera hint === */
.camera-hint {
    background: var(--dark);
    color: var(--cream);
    font-family: 'Fredoka One', cursive;
    font-size: 13px;
    text-align: center;
    border-radius: 12px;
    padding: 8px 14px;
    margin-bottom: 14px;
    border: 2px solid var(--dark);
    box-shadow: 2px 2px 0 rgba(0,0,0,0.3);
}

/* remove gradio default footer */
footer { display: none !important; }
.gradio-container > .footer { display: none !important; }

/* image preview styling */
.upload-zone img {
    border-radius: 16px !important;
}
"""

with gr.Blocks(css=css, title="Papillon Checker") as app:

    gr.HTML("""
    <div class='woof-header'>
        <div class='doodle-strip'>🐾 🦋 🐾</div>
        <div class='woof-title'>woof.</div>
        <div class='woof-subtitle'>is that a papillon?</div>
    </div>
    """)

    gr.HTML("""
    <div class='camera-hint'>
        📱 on iphone? tap the upload box → choose photo library or take a photo!
    </div>
    """)

    with gr.Column():
        image_input = gr.Image(
            type="pil",
            label="drop your dog here 🐶",
            elem_classes=["upload-zone"],
            sources=["upload", "webcam", "clipboard"],
        )

        submit_btn = gr.Button(
            "is it a papillon? 🔍",
            variant="primary",
        )

        result_output = gr.Markdown(
            value="*upload a dog photo to find out!*",
            label="verdict",
            elem_classes=["result-box"],
        )

    submit_btn.click(fn=classify_dog, inputs=image_input, outputs=result_output)

    gr.HTML("""
    <div class='woof-footer'>
        built by amelie romano · powered by pytorch 🦋
    </div>
    """)

app.launch(share=True)