import os
import torch
from planet import PlaNetModel, decode_tile_to_latlon
from offset_predictor import get_offset
from map_overlay import generate_map
from PIL import Image
import torchvision.transforms as T
import numpy as np
from geopy.geocoders import Nominatim
import gradio as gr

model = PlaNetModel(num_cells=100000)
checkpoint = torch.load("./planet_weights.pth", map_location="cpu", weights_only=False)
model.load_state_dict(checkpoint["model_state"])
model.eval()

tfm = T.Compose([T.ToTensor()])
geolocator = Nominatim(user_agent="geo")

def predict(image):
    image = image.convert("RGB").resize((224, 224))
    inp = tfm(image).unsqueeze(0)
    with torch.no_grad():
        probs = model(inp).cpu().numpy().squeeze()
    cell = int(np.argmax(probs))
    lat, lon = decode_tile_to_latlon(cell)
    lat_off, lon_off = get_offset(lat, lon)
    full_lat, full_lon = lat + lat_off, lon + lon_off

    try:
        location = geolocator.reverse((full_lat, full_lon), exactly_one=True, addressdetails=True, language='en')
        address = location.address if location else "Unknown"
        addr = location.raw.get("address", {}) if location else {}
        city = addr.get("city") or addr.get("town") or addr.get("village") or "Unknown"
        state = addr.get("state", "Unknown")
        country = addr.get("country", "Unknown")
        summary = f"{city}, {state}, {country}"
    except Exception as e:
        address = "Address lookup failed"
        summary = "Unknown, Unknown, Unknown"

    return full_lat, full_lon, summary, address, generate_map(full_lat, full_lon)

gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Number(label="Latitude"),
        gr.Number(label="Longitude"),
        gr.Textbox(label="City, State, Country"),
        gr.Textbox(label="Full Address"),
        gr.Image(label="Map View")
    ],
    title="Geolocation AI",
    description="Upload a photo. This AI predicts your latitude, longitude, city, state, country, full address, and shows it on a map."
).launch()