import io
 
import spaces
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from PIL import Image
import torch
import gradio as gr
from huggingface_hub import hf_hub_download
 
from river_model import load_trained_model, build_processor, WATER_CLASS_ID
 
 
 
# Configuration
 
REPO_ID       = "beaunix/river-segmentation"
CKPT_FILENAME = "best_model.pt"
 
# Cyberpunk HUD palette - navy + turquoise edition
BG        = "#000000"
PANEL     = "#050912"
NAVY      = "#0A1C40"
TURQUOISE = "#00E5D4"
LIGHTBLUE = "#5FD8FF"
SILVER    = "#C0C0C0"
GRID      = "#10182a"
 
WATER_COLOR_RGB = (95, 216, 255)   # lightblue, for the mask overlay
BG_COLOR_RGB    = (0, 0, 0)        # black fallback for all other classes
 
 
 
# Model loading (CPU, once at startup)
 
print("[INIT] Loading River-Segmentation model...")
_ckpt_path = hf_hub_download(repo_id=REPO_ID, filename=CKPT_FILENAME)
MODEL = load_trained_model(_ckpt_path)
PROCESSOR = build_processor()
print("[INIT] Model ready (CPU).")
 
 
 
# Inference (single GPU allocation)
 
@spaces.GPU(duration=60)
def run_inference_gpu(image: Image.Image):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    MODEL.to(device)
    MODEL.eval()
 
    orig_w, orig_h = image.size
    inputs = PROCESSOR(images=image, return_tensors="pt")
    pixel_values = inputs["pixel_values"].to(device)
 
    with torch.no_grad():
        if device == "cuda":
            with torch.autocast(device_type="cuda", dtype=torch.float16):
                outputs = MODEL(pixel_values=pixel_values)
        else:
            outputs = MODEL(pixel_values=pixel_values)
        logits = outputs.logits
 
    logits_up = torch.nn.functional.interpolate(
        logits, size=(orig_h, orig_w), mode="bilinear", align_corners=False,
    )
    pred_mask = logits_up.argmax(dim=1).squeeze(0).cpu().numpy().astype(np.uint8)
 
    MODEL.to("cpu")
    return pred_mask
 
 
 
# Post-processing: collapse to binary water mask
 
def binary_water_mask(pred_mask: np.ndarray) -> np.ndarray:
    return (pred_mask == WATER_CLASS_ID).astype(np.uint8)
 
 
def mask_to_rgb(binary_mask: np.ndarray) -> np.ndarray:
    h, w = binary_mask.shape
    rgb = np.zeros((h, w, 3), dtype=np.uint8)
    rgb[binary_mask == 1] = WATER_COLOR_RGB
    rgb[binary_mask == 0] = BG_COLOR_RGB
    return rgb
 
 
def build_overlay(orig_image: Image.Image, rgb_mask: np.ndarray) -> np.ndarray:
    orig_np = np.array(orig_image.convert("RGB"))
    overlay = (orig_np * 0.55 + rgb_mask * 0.45).astype(np.uint8)
    return overlay
 
 
def build_hud_figure(orig_image, rgb_mask, overlay, water_pct):
    plt.rcParams.update({
        "font.family": "monospace",
        "text.color": SILVER,
        "axes.edgecolor": TURQUOISE,
    })
 
    fig, axes = plt.subplots(1, 3, figsize=(15, 5.2))
    fig.patch.set_facecolor(BG)
 
    for ax in axes:
        ax.set_facecolor(PANEL)
        ax.axis("off")
        for s in ax.spines.values():
            s.set_color(TURQUOISE)
 
    axes[0].imshow(np.array(orig_image.convert("RGB")))
    axes[0].set_title("ORIGINAL", color=LIGHTBLUE, fontsize=11)
 
    axes[1].imshow(rgb_mask)
    axes[1].set_title("WATER MASK (binary)", color=LIGHTBLUE, fontsize=11)
 
    axes[2].imshow(overlay)
    axes[2].set_title("OVERLAY", color=LIGHTBLUE, fontsize=11)
 
    fig.suptitle(
        f"AEGIS-RIVER-SEGMENTATION // WATER BODY DETECTION   "
        f"water_ratio={water_pct:.2f}%",
        color=TURQUOISE, fontsize=13, y=1.02
    )
 
    fig.tight_layout()
    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=140, bbox_inches="tight", facecolor=BG)
    plt.close(fig)
    buf.seek(0)
    return Image.open(buf)
 
 
 
# Main handler
 
def analyze(image: Image.Image):
    empty_df = pd.DataFrame()
    if image is None:
        return "Please upload an image.", None, empty_df
 
    pred_mask   = run_inference_gpu(image)
    water_mask  = binary_water_mask(pred_mask)
    rgb_mask    = mask_to_rgb(water_mask)
    overlay     = build_overlay(image, rgb_mask)
 
    total_px    = water_mask.size
    water_px    = int(water_mask.sum())
    water_pct   = round((water_px / total_px) * 100, 2)
 
    hud_img = build_hud_figure(image, rgb_mask, overlay, water_pct)
 
    status = f"Analysis complete  //  Water coverage: {water_pct:.2f}%"
 
    summary_df = pd.DataFrame({
        "Metric": ["Image size", "Water pixels", "Total pixels", "Water ratio (%)"],
        "Value": [
            f"{image.size[0]} x {image.size[1]}",
            f"{water_px:,}", f"{total_px:,}", f"{water_pct:.2f}",
        ],
    })
 
    return status, hud_img, summary_df
 
 
 
# Gradio UI - Navy + Turquoise Cyberpunk HUD
 
CSS = """
.gradio-container { background: #000000 !important; }
h1, h2, h3, p, span, label { color: #00E5D4 !important; font-family: monospace !important; }
.block, .form { border: 1px solid #0A1C40 !important; border-radius: 6px !important;
                background: #050912 !important; }
.gr-button { border: 1px solid #00E5D4 !important; color: #00E5D4 !important;
             background: #050912 !important; font-family: monospace !important; }
"""
 
with gr.Blocks(css=CSS, title="Aegis-River-Segmentation") as demo:
    gr.Markdown("# AEGIS-RIVER-SEGMENTATION // WATER BODY MASK")
    gr.Markdown(
        "SegFormer-B2 semantic segmentation, collapsed to a binary water mask. "
        "Upload a river or waterway image (drag and drop) to extract the "
        "water region and its coverage ratio."
    )
 
    with gr.Row():
        with gr.Column(scale=1):
            image_in = gr.Image(label="Input image", type="pil")
            run_btn  = gr.Button("RUN ANALYSIS", variant="primary")
            status   = gr.Markdown()
 
    hud_out     = gr.Image(label="Water segmentation HUD", type="pil")
    summary_out = gr.Dataframe(label="Summary metrics", interactive=False)
 
    run_btn.click(
        fn=analyze,
        inputs=[image_in],
        outputs=[status, hud_out, summary_out],
    )
 
 
if __name__ == "__main__":
    
    demo.queue().launch(show_api=False, ssr_mode=False)