import os
import gradio as gr
import torch
import numpy as np
import cv2
from PIL import Image
import rasterio
from rasterio.transform import xy
from model_arch import get_model

# ---------- Load Model ----------
MODEL_PATH = os.path.join(os.path.dirname(__file__), "marida_final_model.pth")
model = get_model(MODEL_PATH)

# ---------- Prediction Function ----------
def predict(file):

    # Read GeoTIFF
    with rasterio.open(file.name) as src:
        transform = src.transform
        img = src.read()  # (bands, H, W)

    # Use first 11 bands
    img = img[:11].astype(np.float32)

    # Normalize
    img = img / (np.max(img) + 1e-8)

    # Convert to tensor
    img_tensor = torch.tensor(img).unsqueeze(0)

    # ---- PAD to multiple of 32 ----
    _, c, h, w = img_tensor.shape
    pad_h = (32 - h % 32) % 32
    pad_w = (32 - w % 32) % 32
    img_tensor = torch.nn.functional.pad(img_tensor, (0, pad_w, 0, pad_h))

    # Model inference
    with torch.no_grad():
        y = model(img_tensor)

    # Remove padding from output
    mask = y.squeeze().numpy()
    mask = mask[:h, :w]

    # Threshold
    mask = (mask > 0.5).astype(np.uint8) * 255

    # Find regions
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    geo_boxes = []

    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        lon1, lat1 = xy(transform, y, x)
        lon2, lat2 = xy(transform, y + h, x + w)

        geo_boxes.append({
            "top_left": (lat1, lon1),
            "bottom_right": (lat2, lon2)
        })

    mask_img = Image.fromarray(mask)

    return mask_img, str(geo_boxes)


# ---------- Gradio Interface ----------
demo = gr.Interface(
    fn=predict,
    inputs=gr.File(label="Upload Satellite GeoTIFF"),
    outputs=["image", "text"],
    title="Satellite Plastic Detection System",
    description="Upload a multi-band GeoTIFF to detect plastic and get map coordinates"
)

demo.launch()