from io import BytesIO
from pathlib import Path

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

from model import SimpleNet, CLASS_NAMES


WEIGHTS_DIR = Path("weights")
EXAMPLES = [
    ["examples/annualcrop_sample.jpg"],
    ["examples/forest_sample.jpg"],
    ["examples/highway_sample.jpg"],
    ["examples/industrial_sample.jpg"],
    ["examples/residential_sample.jpg"],
    ["examples/sealake_sample.jpg"],
]


def load_model():
    model = SimpleNet(num_classes=10)
    weight_bytes = b"".join(
        path.read_bytes() for path in sorted(WEIGHTS_DIR.glob("simple_net_v1.part*"))
    )
    state_dict = torch.load(BytesIO(weight_bytes), map_location="cpu")
    state_dict = {
        name: tensor.float() if torch.is_floating_point(tensor) else tensor
        for name, tensor in state_dict.items()
    }
    model.load_state_dict(state_dict)
    model.eval()
    return model


model = load_model()

preprocess = transforms.Compose([
    transforms.Resize((64, 64)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
])


def predict(image: Image.Image) -> dict[str, float]:
    if image is None:
        return {}

    image = image.convert("RGB")
    tensor = preprocess(image).unsqueeze(0)  # [1, 3, 64, 64]

    with torch.no_grad():
        logits = model(tensor)
        probs = torch.nn.functional.softmax(logits, dim=1)[0]

    return {CLASS_NAMES[i]: float(probs[i]) for i in range(len(CLASS_NAMES))}


demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload a Sentinel-style land image"),
    outputs=gr.Label(num_top_classes=5, label="Top land-use guesses"),
    title="EuroSAT Field Scout",
    description=(
        "A small local-first Gradio classifier for quick land-use triage. "
        "It runs a custom PyTorch CNN trained on EuroSAT and returns the closest scene class."
    ),
    article=(
        "Built for the Build Small Hackathon Backyard AI track. "
        "No cloud inference API, no giant model: the Space loads local weights "
        "and runs CPU inference inside the app."
    ),
    examples=EXAMPLES,
    cache_examples=False,
    allow_flagging="never",
    theme=gr.themes.Soft(),
)

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=2).launch()