"""
Crop Boundary Segmentation — Gradio demo.
Upload a Sentinel-2 NetCDF file to detect field boundaries.
"""

import sys
import tempfile
from pathlib import Path

import gradio as gr
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import netCDF4 as nc
import numpy as np
import torch
from huggingface_hub import hf_hub_download

sys.path.insert(0, str(Path(__file__).parent))
from model.segformer import CropBoundarySegFormer


# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
MODEL_REPO = "NandiniLokeshReddy/imerit-segformer-cropboundary"
VARIABLES  = ["B2", "B3", "B4", "B8", "NDVI"]
N_TIME     = 6
_NORM = {
    "B2": (500.0, 400.0), "B3": (700.0, 450.0),
    "B4": (600.0, 450.0), "B8": (2500.0, 1500.0), "NDVI": (0.37, 0.22),
}
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


# ---------------------------------------------------------------------------
# Model (loaded once at startup)
# ---------------------------------------------------------------------------
def _load_model():
    print("Downloading model weights from HF Hub...")
    ckpt_path = hf_hub_download(repo_id=MODEL_REPO, filename="best.pt")
    m = CropBoundarySegFormer(
        backbone="nvidia/mit-b2", num_labels=2, in_channels=30, pretrained=False
    ).to(device)
    ckpt = torch.load(ckpt_path, map_location=device)
    m.load_state_dict(ckpt["model"])
    m.eval()
    print(f"Model ready (epoch {ckpt['epoch']}, best mIoU {ckpt['best_miou']:.4f})")
    return m

model = _load_model()


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def read_nc(path: str) -> np.ndarray:
    with nc.Dataset(path) as ds:
        bands = []
        for var in VARIABLES:
            arr = ds.variables[var][:]
            if hasattr(arr, "filled"):
                arr = arr.filled(0.0)
            arr = np.nan_to_num(arr.astype(np.float32), nan=0.0)
            bands.append(arr)
    stacked = np.stack(bands, axis=0)
    return stacked.reshape(stacked.shape[0] * stacked.shape[1], *stacked.shape[2:])


def normalize(img: np.ndarray) -> np.ndarray:
    out = img.copy()
    for i, var in enumerate(VARIABLES):
        m, s = _NORM[var]
        out[i*N_TIME:(i+1)*N_TIME] = (out[i*N_TIME:(i+1)*N_TIME] - m) / (s + 1e-8)
    return out


def stretch(arr, p_low=2, p_high=98):
    lo, hi = np.percentile(arr, p_low), np.percentile(arr, p_high)
    return np.clip((arr - lo) / (hi - lo + 1e-8), 0, 1)


def make_rgb(img):
    return np.stack([stretch(img[12:18].mean(0)),
                     stretch(img[6:12].mean(0)),
                     stretch(img[0:6].mean(0))], axis=-1)


def make_fci(img):
    return np.stack([stretch(img[18:24].mean(0)),
                     stretch(img[12:18].mean(0)),
                     stretch(img[6:12].mean(0))], axis=-1)


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------
def run_inference(nc_file):
    if nc_file is None:
        return None, "Please upload a .nc file."

    try:
        img_raw  = read_nc(nc_file)
        img_norm = normalize(img_raw)
    except Exception as e:
        return None, f"Failed to read file: {e}"

    x = torch.from_numpy(img_norm).unsqueeze(0).to(device)
    with torch.no_grad():
        pred = model(x)["logits"].argmax(1).squeeze(0).cpu().numpy()  # (H, W)

    # Boundary IoU (no GT available, just show prediction stats)
    boundary_px = int(pred.sum())
    total_px    = pred.size
    pct         = 100.0 * boundary_px / total_px

    rgb = make_rgb(img_raw)
    fci = make_fci(img_raw)

    # Overlay prediction on RGB
    overlay = rgb.copy()
    mask = pred.astype(bool)
    for c, val in enumerate([0.1, 0.9, 0.1]):   # green boundary
        ch = overlay[:, :, c]
        ch[mask] = ch[mask] * 0.3 + val * 0.7

    fig, axes = plt.subplots(1, 4, figsize=(20, 5))
    fig.patch.set_facecolor("#1a1a2e")

    panels = [
        (rgb,     "True Colour RGB\n(B4/B3/B2)",              {}),
        (fci,     "False Colour IR\n(B8/B4/B3) — veg=red",    {}),
        (pred,    "Predicted Boundary Mask",                   {"cmap": "hot", "vmin": 0, "vmax": 1}),
        (overlay, "Prediction Overlay\n(green = boundary)",    {}),
    ]
    for ax, (data, title, kwargs) in zip(axes, panels):
        ax.set_facecolor("#0f0f1a")
        ax.imshow(data, interpolation="nearest", **kwargs)
        ax.set_title(title, color="white", fontsize=9)
        ax.axis("off")

    plt.tight_layout()
    out_path = tempfile.mktemp(suffix=".png")
    plt.savefig(out_path, dpi=130, bbox_inches="tight", facecolor=fig.get_facecolor())
    plt.close()

    info = f"Boundary pixels: {boundary_px:,} / {total_px:,} ({pct:.1f}%)"
    return out_path, info


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
with gr.Blocks(title="Crop Boundary Segmentation") as demo:
    gr.Markdown("""
    # Crop Boundary Segmentation
    **Model**: SegFormer-B2 fine-tuned on [AI4Boundaries](https://github.com/waldnerf/ai4boundaries) dataset
    **Input**: Sentinel-2 NetCDF — 5 variables (B2, B3, B4, B8, NDVI) × 6 monthly time steps → 30 channels, 256×256 px
    **Output**: Binary field boundary mask
    """)

    with gr.Row():
        nc_input = gr.File(label="Upload Sentinel-2 .nc file", file_types=[".nc"])
        run_btn  = gr.Button("Run Inference", variant="primary")

    with gr.Row():
        out_img  = gr.Image(label="Result", type="filepath")
        out_info = gr.Textbox(label="Stats", lines=2)

    run_btn.click(fn=run_inference, inputs=nc_input, outputs=[out_img, out_info])

if __name__ == "__main__":
    demo.launch()