"""The one inference path.

`predict_scene()` is the *only* place a height map is produced from an image, and
both the final evaluation (`eval/sliding.py`) and the CLI (`infer/predict.py`)
call it.  That is deliberate: v2 had one code path for validation (centre crop of
a normalised GAMUS tile) and a different one for real images (`predict_image.py`,
which resized any PNG to 512 regardless of its ground sample distance and used
ImageNet normalisation constants the training script also used but the encoder
does not want).  A validation number produced by a path the demo never runs is
not a validation number.

Contract in / out
-----------------
    rgb_u8   (H, W, 3) uint8 at `src_gsd_m` metres per pixel, *unstretched*
    ->  height_m (H, W) float32, nDSM in metres, on the SAME grid as the input

Internally: scene stretch -> resample to `spec.canonical_gsd_m` -> overlapping
`spec.tile_size` windows, batched through the net in parallel -> optional D4 TTA
(all views in the same batch) -> Hann-weighted blend -> resample the height map
back to the input grid.  Heights are metres, so the round trip does not rescale
them.
"""

from __future__ import annotations

import numpy as np
import torch

from config import SEG_IGNORE_INDEX
from dwdata.preprocess import (
    PreprocSpec, gsd_to_shape, hann2d, resize, stretch_scene, tile_origins,
)

# Rough inference footprint of one 512 px tile through the ViT-L encoder, the
# DPT trunk and the full-resolution heads under bf16.  Deliberately generous:
# it only seeds the batch size, and an OOM halves it and retries anyway.
_TILE_BYTES_512 = 640 * 2**20
_MAX_FORWARD_BATCH = 64
# The canvas (source RGB + blend buffers, ~17 B/px) stays on the GPU while it
# takes at most this share of free memory; bigger scenes blend on the host.
_CANVAS_SHARE = 0.35


def auto_forward_batch(device, tile_size: int, reserve_frac: float = 0.15) -> int:
    """How many tiles one forward pass can take on `device` right now."""
    if device.type != "cuda":
        return 4
    free, _total = torch.cuda.mem_get_info(device)
    per_tile = _TILE_BYTES_512 * (tile_size / 512) ** 2
    return int(max(1, min(_MAX_FORWARD_BATCH, free * (1 - reserve_frac) // per_tile)))


class _Forward:
    """Model call with autocast, memory-sized chunking and OOM back-off.

    The chunk size is shared state: once a batch has hit OOM, later batches
    start at the size that worked instead of rediscovering the limit each time.
    """

    def __init__(self, model, device, amp_dtype, max_batch: int, want_seg: bool, want_std: bool):
        self.model, self.device, self.amp_dtype = model, device, amp_dtype
        self.max_batch, self.want_seg, self.want_std = max(1, int(max_batch)), want_seg, want_std

    def __call__(self, x):
        preds, stds, segs, i = [], [], [], 0
        while i < x.shape[0]:
            n = self.max_batch
            try:
                with _autocast(self.device, self.amp_dtype):
                    out = self.model(x[i:i + n])
            except torch.cuda.OutOfMemoryError:
                if n == 1:
                    raise
                self.max_batch = max(1, n // 2)
                if self.device.type == "cuda":
                    torch.cuda.empty_cache()
                print(f"[engine] OOM at {n} tiles/forward; retrying at {self.max_batch}")
                continue
            preds.append(out["fused"].float())
            if self.want_std:
                stds.append(out["b_std"].float())
            if self.want_seg:
                # the ignore channel is a training target, never a prediction
                segs.append(out["seg"][:, :SEG_IGNORE_INDEX].float().argmax(1))
            del out
            i += n
        return (torch.cat(preds), torch.cat(stds) if self.want_std else None,
                torch.cat(segs) if self.want_seg else None)


@torch.no_grad()
def predict_canonical(
    model, rgb_canon_u8: np.ndarray, spec: PreprocSpec, device,
    *, tta: bool = False, tta_scales=(1.0,), amp_dtype=None,
    overlap: float = 0.25, batch_tiles: int = 0, want_seg: bool = False,
    want_std: bool = False,
    progress=None,
):
    """Sliding-window prediction over an image already at the canonical GSD.

    Every overlapping `tile_size` window is independent, so they go through the
    network as large batches: `batch_tiles=0` sizes the forward batch from free
    GPU memory (a big card takes dozens of tiles per pass), and under TTA all 8
    D4 views of those tiles ride in the same batch rather than 8 passes in a row.

    When the scene fits, the source image and the blend buffers live on the GPU:
    tiles are cut and normalised there and blended in place, so the loop only
    queues kernels and never waits on a device->host copy.  A scene too large
    for that blends on the host, with the next batch cut on a worker thread
    while the GPU runs the current one.
    """
    from models.tta import n_views, tta_predict_parallel

    S = spec.tile_size
    H, W = rgb_canon_u8.shape[:2]
    # Reflect-pad rather than zero-pad: a black border is a strong out-of-
    # distribution edge and v2 trained on ~41% of it, so we never introduce one.
    pad_h, pad_w = max(0, S - H), max(0, S - W)
    if pad_h or pad_w:
        rgb_canon_u8 = np.pad(rgb_canon_u8, ((0, pad_h), (0, pad_w), (0, 0)), mode="reflect")
    Hp, Wp = rgb_canon_u8.shape[:2]

    ys, xs = tile_origins(Hp, S, overlap), tile_origins(Wp, S, overlap)
    coords = [(y, x) for y in ys for x in xs]
    total = len(coords)

    canvas_dev = device
    if device.type == "cuda":
        torch.backends.cudnn.benchmark = True  # every batch is (N, 3, S, S)
        free, _total = torch.cuda.mem_get_info(device)
        if Hp * Wp * (3 + 8 + (6 if want_seg else 0)) > free * _CANVAS_SHARE:
            canvas_dev = torch.device("cpu")
    on_device = canvas_dev.type == device.type

    src = torch.from_numpy(np.ascontiguousarray(rgb_canon_u8))
    if on_device:
        src = src.to(device)  # one upload; every tile is a view of it
    mean = torch.tensor(spec.mean, dtype=torch.float32, device=device).view(1, 3, 1, 1)
    std = torch.tensor(spec.std, dtype=torch.float32, device=device).view(1, 3, 1, 1)
    win = torch.from_numpy(hann2d(S)).to(canvas_dev)
    acc = torch.zeros((Hp, Wp), dtype=torch.float32, device=canvas_dev)
    acc_m2 = torch.zeros_like(acc) if want_std else None
    wsum = torch.zeros((Hp, Wp), dtype=torch.float32, device=canvas_dev)
    # Labels cannot be averaged, so each pixel keeps the label from the tile it
    # sits most centrally in (largest Hann weight).  Last-write-wins put a
    # visible seam at every tile edge, where the head has the least context.
    seg_acc = seg_best = None
    if want_seg:
        seg_acc = torch.zeros((Hp, Wp), dtype=torch.int16, device=canvas_dev)
        seg_best = torch.zeros((Hp, Wp), dtype=torch.float32, device=canvas_dev)

    scales = tuple(float(s) for s in tta_scales)
    if not scales or any(not np.isfinite(s) or s <= 0 for s in scales):
        raise ValueError("inference scales must be finite and positive")
    views = n_views(scales, d4=tta)
    largest_input = max(S, *(max(spec.patch, int(round(S * scale / spec.patch)) * spec.patch)
                             for scale in scales))
    fwd = _Forward(
        model, device, amp_dtype,
        batch_tiles * views if batch_tiles > 0 else auto_forward_batch(device, largest_input),
        want_seg, want_std,
    )
    step = max(1, fwd.max_batch // views)  # tiles per step; TTA multiplies them

    def cut(chunk):
        """(B, S, S, 3) uint8 on the device, ready to normalise."""
        tiles = torch.stack([src[y:y + S, x:x + S] for y, x in chunk])
        if not on_device and device.type == "cuda":
            tiles = tiles.pin_memory()
        return tiles.to(device, non_blocking=True)

    chunks = [coords[b0:b0 + step] for b0 in range(0, total, step)]
    pool = pending = None
    if not on_device and len(chunks) > 1:
        from concurrent.futures import ThreadPoolExecutor

        pool = ThreadPoolExecutor(max_workers=1)
        pending = pool.submit(cut, chunks[0])
    done = 0
    try:
        for ci, chunk in enumerate(chunks):
            if pool is not None:
                tiles = pending.result()
                if ci + 1 < len(chunks):
                    pending = pool.submit(cut, chunks[ci + 1])
            else:
                tiles = cut(chunk)
            t = (tiles.permute(0, 3, 1, 2).float().div_(255.0) - mean) / std
            if tta or scales != (1.0,):
                pred, pred_std, seg = tta_predict_parallel(fwd, t, scales,
                                                          patch=spec.patch, d4=tta)
            else:
                pred, pred_std, seg = fwd(t)
            pred = pred[:, 0].to(canvas_dev)
            if pred_std is not None:
                pred_std = pred_std[:, 0].to(canvas_dev)
            if seg is not None:
                seg = seg.to(device=canvas_dev, dtype=torch.int16)

            for i, (y, x) in enumerate(chunk):
                acc[y:y + S, x:x + S] += pred[i] * win
                if acc_m2 is not None:
                    acc_m2[y:y + S, x:x + S] += (pred_std[i].square() + pred[i].square()) * win
                wsum[y:y + S, x:x + S] += win
                if seg is not None:
                    best = seg_best[y:y + S, x:x + S]
                    better = win > best
                    seg_acc[y:y + S, x:x + S] = torch.where(
                        better, seg[i], seg_acc[y:y + S, x:x + S])
                    seg_best[y:y + S, x:x + S] = torch.where(better, win, best)
            done += len(chunk)
            if progress:
                progress(done, total)
    finally:
        if pool is not None:
            pool.shutdown(wait=True)

    height = (acc / wsum.clamp_min(1e-6))[:H, :W].cpu().numpy()
    std_out = None
    if acc_m2 is not None:
        variance = (acc_m2 / wsum.clamp_min(1e-6) - (acc / wsum.clamp_min(1e-6)).square()).clamp_min(0)
        std_out = variance[:H, :W].sqrt().cpu().numpy().astype(np.float32)
    seg_out = seg_acc[:H, :W].cpu().numpy() if seg_acc is not None else None
    return (height.astype(np.float32), seg_out, std_out) if want_std else (height.astype(np.float32), seg_out)


def _autocast(device, amp_dtype):
    if amp_dtype is None or device.type != "cuda":
        import contextlib

        return contextlib.nullcontext()
    return torch.autocast("cuda", dtype=amp_dtype)


def count_tiles(h: int, w: int, src_gsd_m: float, spec: PreprocSpec,
                overlap: float = 0.25) -> int:
    """Windows `predict_scene` will run for an (h, w) scene — for budgeting."""
    ch, cw = gsd_to_shape(h, w, src_gsd_m, spec.canonical_gsd_m)
    S = spec.tile_size
    return (len(tile_origins(max(ch, S), S, overlap))
            * len(tile_origins(max(cw, S), S, overlap)))


@torch.no_grad()
def predict_scene(
    model, rgb_u8: np.ndarray, src_gsd_m: float, spec: PreprocSpec, device,
    *, tta: bool = False, tta_scales=(1.0,), amp_dtype=None, overlap: float = 0.25,
    batch_tiles: int = 0, want_seg: bool = False, want_std: bool = False, progress=None,
):
    """RGB at any GSD -> nDSM in metres on the input grid (+ optional class map).

    `batch_tiles=0` lets the engine size its batches from free GPU memory.
    """
    H, W = rgb_u8.shape[:2]
    rgb_s = stretch_scene(rgb_u8, spec)
    canon_hw = gsd_to_shape(H, W, src_gsd_m, spec.canonical_gsd_m)
    rgb_c = resize(rgb_s, canon_hw, "bilinear")

    prediction = predict_canonical(
        model, rgb_c, spec, device, tta=tta, tta_scales=tta_scales,
        amp_dtype=amp_dtype, overlap=overlap, batch_tiles=batch_tiles,
        want_seg=want_seg, want_std=want_std, progress=progress,
    )
    height_c, seg_c = prediction[:2]
    height = resize(height_c, (H, W), "bilinear").astype(np.float32)
    seg = resize(seg_c.astype(np.int32), (H, W), "nearest") if seg_c is not None else None
    if want_std:
        std = resize(prediction[2], (H, W), "bilinear").astype(np.float32)
        return height, seg, std
    return height, seg