"""Stage 1: DL inference. RGB (+ optional DEM) -> probability raster (ONNX).

Ported from FraXtex's `fraxtex_tools/inference.py`, which is canonical for this
pipeline (see CLAUDE.md). The logic below is deliberately kept line-comparable with
the plugin: same contract reading, same Hann-blended overlapping tiles, same DEM
handling, same output convention (high value = fracture).

Two deliberate differences, both because the Space knows things the plugin cannot:

* No input-mode enum. The plugin lets the user load any .onnx, so it must offer
  auto/rgb/rgb_dem to resolve what the file needs. Here the model comes from a fixed
  registry, so the channel count from the graph is always authoritative.
* Input size limits (`limits.py`), because this container is shared.
"""

import numpy as np
import onnxruntime as ort
import rasterio

from .limits import check_size


def load_session(model_path):
    """Open an ONNX model and read the input contract out of it.

    Returns `(session, input_name, contract)` where contract carries the channel
    count, patch size and whether the caller must scale to [0, 1].
    """
    session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
    input_name = session.get_inputs()[0].name
    shape = session.get_inputs()[0].shape

    # Symbolic dims (e.g. 'channels') are strings, so model_ch stays None.
    model_ch = shape[1] if len(shape) == 4 and isinstance(shape[1], int) else None

    # Tile size: the graph when it pins one (512 for SAM 2, 256 for the RGB+DEM
    # pair), else the historical default for dynamic exports.
    dim_h, dim_w = (shape[2], shape[3]) if len(shape) == 4 else (None, None)
    if isinstance(dim_h, int) and isinstance(dim_w, int):
        ps, ps_pinned = dim_h, True
        square = dim_h == dim_w
    else:
        ps, ps_pinned, square = 256, False, True

    meta = session.get_modelmeta().custom_metadata_map or {}

    return session, input_name, {
        "channels": model_ch,
        "patch_size": ps,
        "patch_pinned": ps_pinned,
        "square": square,
        "meta": meta,
    }


def _resolve_input_range(contract, use_dem, feedback):
    """Decide whether the caller scales to [0, 1], from what the model declares.

    Preprocessing is not visible in the graph, so the model declares it: some models
    normalize internally and want raw 0-255. The channel count is the fallback for
    exports made before the keys existed.
    """
    meta = contract["meta"]
    declared = (meta.get("input.range") or "").strip()
    if not declared and meta.get("preprocessing.in_graph", "").strip().lower() == "true":
        declared = "0-255"

    if declared in ("0-1", "0-255"):
        feedback.pushInfo(f"Model declares input range {declared}.")
        return declared == "0-1", False

    if declared:
        feedback.pushWarning(
            f"Ignoring unrecognised declared input range '{declared}'; "
            "expected '0-1' or '0-255'."
        )
    scale_255 = use_dem
    feedback.pushInfo(
        "Model declares no input range; assuming "
        + ("0-1, scaling by /255" if scale_255 else "0-255, raw values")
        + f" from its {4 if use_dem else 3}-channel input."
    )
    return scale_255, True


def _read_rgb(rgb_path):
    """Read an orthophoto to HWC uint8. rasterio handles both PNG and GeoTIFF."""
    with rasterio.open(rgb_path) as src:
        img = src.read()
        if img.shape[0] >= 3:
            img = img[:3].transpose(1, 2, 0)
        elif img.shape[0] == 2:
            img = img[0:1].repeat(3, axis=0).transpose(1, 2, 0)
        else:
            img = img.repeat(3, axis=0).transpose(1, 2, 0)
        if img.dtype != np.uint8:
            img = (img * 255).astype(np.uint8) if img.max() <= 1.0 else img.astype(np.uint8)
        return img, src.profile.copy(), src.crs


def run_inference(rgb_path, model_path, feedback, dem_path=None, overlap=0.25):
    """Run tiled ONNX inference and return `(pred_uint8, profile, info)`.

    `pred_uint8` follows the plugin's convention: **high values mean fracture**.
    """
    session, input_name, contract = load_session(model_path)
    model_ch = contract["channels"]
    ps = contract["patch_size"]

    if model_ch is not None and model_ch not in (3, 4):
        feedback.reportError(
            f"This model declares {model_ch} input channels. Only 3 (RGB) and "
            "4 (RGB+DEM) are supported.",
            fatalError=True,
        )
    if not contract["square"]:
        feedback.reportError(
            "This model expects non-square inputs. Only square patches are supported.",
            fatalError=True,
        )

    has_dem = dem_path is not None
    # The registry always yields a declared channel count; the None branch exists so
    # a third-party model dropped into the registry still behaves like the plugin.
    use_dem = has_dem if model_ch is None else model_ch == 4
    if use_dem and not has_dem:
        feedback.reportError(
            "This model expects 4 input channels (RGB+DEM) but no DEM was provided. "
            "Provide a DEM, or choose an RGB-only model.",
            fatalError=True,
        )
    if has_dem and not use_dem:
        feedback.pushWarning(
            "Running RGB-only: the supplied DEM is not fed to the model. It still "
            "provides the output georeferencing."
        )

    scale_255, range_guessed = _resolve_input_range(contract, use_dem, feedback)

    img, rgb_profile, rgb_crs = _read_rgb(rgb_path)
    check_size(img.shape[0], img.shape[1], feedback)

    # The DEM donates the output profile whenever it is supplied, since the RGB may
    # be a plain PNG with no CRS.
    dem_arr = None
    if has_dem:
        with rasterio.open(dem_path) as src:
            profile = src.profile.copy()
            if (src.height, src.width) != img.shape[:2]:
                feedback.reportError(
                    f"RGB and DEM must have the same dimensions "
                    f"(RGB {img.shape[1]}x{img.shape[0]}, DEM {src.width}x{src.height}).",
                    fatalError=True,
                )
            if use_dem:
                dem_arr = src.read(1).astype(np.float32)
        if use_dem:
            dem_lo, dem_hi = float(dem_arr.min()), float(dem_arr.max())
            feedback.pushInfo(f"DEM value range: [{dem_lo:.3f}, {dem_hi:.3f}]")
            if dem_lo < 0.0 or dem_hi > 255.0:
                feedback.pushWarning(
                    "DEM values fall outside 0-255. The DEM band is scaled exactly like "
                    "the RGB bands, so a DEM in raw metres enters the model far outside "
                    "the range it was trained on. Byte-scale the DEM the way it was "
                    "scaled for training."
                )
    else:
        # A PNG source profile cannot be written back, so pin the driver.
        profile = rgb_profile
        profile.update(driver="GTiff")
        if rgb_crs is None:
            feedback.pushWarning(
                "The RGB orthophoto has no CRS, so the output probability raster is "
                "not georeferenced and vectorization works in pixel coordinates."
            )

    # Overlapping patches, Hann-blended: no tile seams.
    if use_dem:
        combined = np.concatenate((img[:, :, :3], np.expand_dims(dem_arr, 2)), axis=2)
        del img, dem_arr  # combined owns its data; the sources are dead weight now
    else:
        combined = img[:, :, :3]  # a view, so img stays alive either way
    h, w = combined.shape[:2]
    n_ch = combined.shape[2]

    # Origins clamp to the far edge, so only sub-patch images need padding.
    pad_h = max(0, ps - h)
    pad_w = max(0, ps - w)
    if pad_h or pad_w:
        padded = np.pad(combined, ((0, pad_h), (0, pad_w), (0, 0)), mode="symmetric")
    else:
        padded = combined
    ph, pw = padded.shape[:2]

    stride = max(1, int(round(ps * (1.0 - overlap))))

    def origins(size):
        offs = list(range(0, size - ps + 1, stride))
        if offs[-1] != size - ps:
            offs.append(size - ps)  # flush against the far edge
        return offs

    ys, xs = origins(ph), origins(pw)

    # Taper: ~0 at the patch border, 1 at its centre. Min weight ~1.5e-4, so
    # dividing a single-coverage pixel by it is safe in float32.
    w1d = np.hanning(ps + 2)[1:-1].astype(np.float32)
    win = np.outer(w1d, w1d)

    # win is separable over a full ys x xs grid, so the total weight per pixel is
    # outer(row_cov, col_cov), so no second full-size accumulator.
    def coverage(offs, size):
        cov = np.zeros(size, dtype=np.float32)
        for o in offs:
            cov[o:o + ps] += w1d
        return cov

    row_cov = coverage(ys, ph)
    col_cov = coverage(xs, pw)

    pred_sum = np.zeros((ph, pw), dtype=np.float32)

    total = len(ys) * len(xs)
    done = 0
    feedback.pushInfo(
        f"Tiling: {len(xs)} x {len(ys)} = {total} patches "
        f"({ps} px, {int(overlap * 100)}% overlap, stride {stride}, {n_ch} channels)"
    )

    for y in ys:
        for x in xs:
            patch = padded[y:y + ps, x:x + ps, :]
            patch_t = patch.transpose(2, 0, 1)[np.newaxis, ...].astype(np.float32)
            if scale_255:
                patch_t /= 255.0
            if done == 0:
                try:
                    out = session.run(None, {input_name: patch_t})
                except Exception as exc:
                    hint = (
                        f"The model declares {model_ch} input channels and was given "
                        "that many, so the channel count is not the problem."
                        + (
                            f" The {ps}x{ps} patch size came from the model itself."
                            if contract["patch_pinned"] else
                            f" It declares no patch size, so {ps}x{ps} was assumed."
                        )
                    )
                    feedback.reportError(
                        f"The model rejected a {n_ch}-channel {ps}x{ps} input: {exc}\n{hint}",
                        fatalError=True,
                    )
            else:
                out = session.run(None, {input_name: patch_t})
            pred_sum[y:y + ps, x:x + ps] += out[0][0, 0] * win
            done += 1
            feedback.setProgress(int(done / total * 100))

    # In place; both profiles are > 0 since every pixel is covered.
    pred_sum /= row_cov[:, np.newaxis]
    pred_sum /= col_cov[np.newaxis, :]
    pred = pred_sum[:h, :w]

    # Logits would wrap modulo 256 in the cast below.
    #
    # Deviation from the canonical plugin: it tests `lo < 0.0 or hi > 1.0` exactly.
    # A sigmoid that saturates to 1.0 (normal, and what unet-rgb does on the example
    # tile) picks up ~1.2e-07 of float32 error from the Hann blend and the coverage
    # division, so the exact test fires on every run of a perfectly good model. The
    # tolerance is far above that error and far below any real logit. Reported
    # upstream; revert this once the plugin carries the same tolerance.
    RANGE_TOL = 1e-5
    lo, hi = float(pred.min()), float(pred.max())
    if lo < -RANGE_TOL or hi > 1.0 + RANGE_TOL:
        feedback.pushWarning(
            f"Model output range [{lo:.3f}, {hi:.3f}] falls outside 0-1 and was clipped. "
            "This expects a model that emits probabilities (e.g. with a final sigmoid), "
            "not raw logits."
        )
    # Clipped unconditionally, not just in the warning branch. Real logits are what
    # this guards: 8.0 * 255 = 2040 wraps to 248 in the uint8 cast below. The
    # sub-tolerance overshoot is harmless by comparison (255.00003 truncates to 255)
    # but clipping it too keeps the output provably in range.
    pred = np.clip(pred, 0.0, 1.0)

    # A near-constant map is the fingerprint of feeding the wrong input range.
    # Measured on the four published models: wrong range spans 0.000-0.016, right
    # range spans 0.79-1.00, so 0.05 sits in the gap between them.
    if range_guessed and hi - lo < 0.05:
        feedback.pushWarning(
            f"Output is nearly constant (range [{lo:.4f}, {hi:.4f}]) and this model "
            "declares no input range, so it was inferred from the channel count. The "
            "model may expect the other range."
        )

    pred_uint8 = (pred * 255).astype(np.uint8)
    profile.update(count=1, dtype=rasterio.uint8, nodata=None, compress="lzw")

    info = {
        "crs": profile.get("crs"),
        "transform": profile.get("transform"),
        "patches": total,
        "patch_size": ps,
        "channels": n_ch,
        "used_dem": use_dem,
        "output_range": (lo, hi),
    }
    feedback.pushInfo("Inference complete.")
    return pred_uint8, profile, info