"""Gradio HF Space for PhysFlow-Earth.

The public Space uses synthetic coarse fields and a CPU bilinear surrogate to
show the interface, resolution change, and physics-violation dashboard without
private checkpoints or external Earth-observation downloads.
"""
from __future__ import annotations

import gradio as gr

VARIABLES = ["Sentinel-2 RGBN (4x SR)", "ERA5 precipitation (5x downscale)", "ERA5 wind (5x downscale)"]
SCENARIOS = ["Historical", "SSP2-4.5", "SSP5-8.5"]


def downscale(
    variable: str,
    aoi_geojson: str | None,
    year: int,
    scenario: str,
    enforce_physics: bool,
    progress=gr.Progress(track_tqdm=True),
):
    progress(0.1, desc="Loading pipeline")
    pipeline = _load_pipeline(variable)
    progress(0.4, desc="Fetching coarse data")
    x_lr = _fetch_coarse(variable, aoi_geojson, year, scenario)
    progress(0.6, desc="Running flow")
    sr = pipeline(x_lr)
    progress(0.85, desc="Computing physics violation")
    metrics = _violation_metrics(variable, x_lr, sr)
    return _to_image(x_lr), _to_image(sr), metrics


def _load_pipeline(variable):
    """Return a CPU-safe demo pipeline for the public Space.

    The trainable `PhysFlowPipeline.from_pretrained(...)` path stays in the
    library, but the hosted demo should not depend on private checkpoints.
    """
    import torch.nn.functional as F

    scale = 4 if "Sentinel" in variable else 5

    def pipeline(x_lr):
        sr = F.interpolate(x_lr, scale_factor=scale, mode="bilinear", align_corners=False)
        return sr.clamp(-1, 1)

    return pipeline


def _fetch_coarse(variable, aoi_geojson, year, scenario):
    import torch

    channels = 4 if "Sentinel" in variable else 1
    grid_x = torch.linspace(-1, 1, 64).view(1, 1, 1, 64)
    grid_y = torch.linspace(-1, 1, 64).view(1, 1, 64, 1)
    phase = ((year - 1990) % 37) / 37.0
    base = torch.sin(3.14 * (grid_x + phase)) * torch.cos(3.14 * (grid_y - phase))
    if scenario != "Historical":
        base = base + (0.08 if scenario == "SSP2-4.5" else 0.16)
    if aoi_geojson:
        base = base + min(len(aoi_geojson), 100) / 1000.0
    return base.repeat(1, channels, 1, 1).clamp(-1, 1)


def _violation_metrics(variable, x_lr, sr) -> str:
    residual = float((sr.mean() - x_lr.mean()).abs())
    band = float(sr.std().clamp(max=1.0))
    return (
        f"Variable: {variable}\n"
        f"Mass conservation residual: {residual:.3f}\n"
        f"Band-ratio violation proxy: {band:.3f}\n"
        "Divergence residual: n/a\n"
    )


def _to_image(t):
    import numpy as np
    arr = t[0].clamp(-1, 1).add(1).div(2).cpu().numpy()
    arr = (arr.transpose(1, 2, 0) * 255).clip(0, 255).astype("uint8")
    if arr.shape[-1] == 1:
        arr = np.concatenate([arr] * 3, axis=-1)
    elif arr.shape[-1] == 4:
        arr = arr[..., :3]
    return arr


def build_ui():
    with gr.Blocks(title="PhysFlow-Earth") as demo:
        gr.Markdown(
            "# PhysFlow-Earth\n"
            "CPU-safe synthetic-field demo for physics-informed Earth observation super-resolution."
        )
        with gr.Row():
            var = gr.Dropdown(VARIABLES, value=VARIABLES[0], label="Variable")
            scenario = gr.Dropdown(SCENARIOS, value="Historical", label="Scenario (climate only)")
            year = gr.Slider(1990, 2100, value=2030, step=1, label="Year")
        aoi = gr.Textbox(label="AOI bbox (GeoJSON or 'lon_min,lat_min,lon_max,lat_max')")
        enforce = gr.Checkbox(value=True, label="Enforce physics constraints at sampling time")
        with gr.Row():
            coarse = gr.Image(label="Coarse input")
            sr = gr.Image(label="PhysFlow output (HR, physics-consistent)")
        violation = gr.Textbox(label="Physics violation dashboard", lines=4, interactive=False)
        gr.Button("Downscale").click(
            downscale, [var, aoi, year, scenario, enforce], [coarse, sr, violation]
        )
    return demo


if __name__ == "__main__":
    build_ui().launch(server_name="0.0.0.0")