"""Overlap-tiled inference helpers for WildFIRE-FM probability maps."""

from __future__ import annotations

from typing import Iterable, Tuple

import torch
import torch.nn.functional as F


def _starts(length: int, tile_size: int, stride: int) -> list[int]:
    if length <= tile_size:
        return [0]
    out = list(range(0, max(length - tile_size + 1, 1), stride))
    last = length - tile_size
    if out[-1] != last:
        out.append(last)
    return out


def _crop_slices(top: int, left: int, tile_size: int, height: int, width: int, halo: int) -> Tuple[slice, slice, slice, slice]:
    y0 = 0 if top == 0 else halo
    x0 = 0 if left == 0 else halo
    y1 = tile_size if top + tile_size >= height else tile_size - halo
    x1 = tile_size if left + tile_size >= width else tile_size - halo
    return slice(y0, y1), slice(x0, x1), slice(top + y0, top + y1), slice(left + x0, left + x1)


def predict_probability_tiled(
    model: torch.nn.Module,
    x: torch.Tensor,
    tile_size: int = 32,
    stride: int = 16,
    halo: int = 8,
    device: torch.device | str | None = None,
    batch_size: int = 16,
) -> torch.Tensor:
    """Predict a full probability map from an input tensor using overlap tiles.

    Parameters
    ----------
    model:
        WildFIRE-FM model returning logits or ``(logits, aux_logits)``.
    x:
        Input tensor in ``[C, H, W]`` or ``[1, C, H, W]`` order.
    tile_size:
        Spatial tile size used for model calls.
    stride:
        Distance between tile origins. Use a value smaller than ``tile_size``
        for overlap.
    halo:
        Number of pixels cropped away from interior tile borders before
        stitching. Border tiles keep the image edge.
    device:
        Device for inference. Defaults to the model parameter device.
    batch_size:
        Number of tiles evaluated per model call.

    Returns
    -------
    torch.Tensor
        Probability map in ``[H, W]`` order on CPU.
    """
    if x.ndim == 3:
        x = x.unsqueeze(0)
    if x.ndim != 4 or x.shape[0] != 1:
        raise ValueError("x must have shape [C, H, W] or [1, C, H, W].")
    if tile_size <= 0 or stride <= 0:
        raise ValueError("tile_size and stride must be positive.")
    if halo < 0 or halo * 2 >= tile_size:
        raise ValueError("halo must be non-negative and smaller than tile_size / 2.")

    if device is None:
        try:
            device = next(model.parameters()).device
        except StopIteration:
            device = torch.device("cpu")
    device = torch.device(device)
    model.eval()

    _, channels, height, width = x.shape
    pad_h = max(tile_size - height, 0)
    pad_w = max(tile_size - width, 0)
    if pad_h or pad_w:
        x_work = F.pad(x, (0, pad_w, 0, pad_h), mode="replicate")
    else:
        x_work = x
    _, _, work_h, work_w = x_work.shape

    output = torch.zeros((work_h, work_w), dtype=torch.float32)
    weight = torch.zeros((work_h, work_w), dtype=torch.float32)
    coords = [(top, left) for top in _starts(work_h, tile_size, stride) for left in _starts(work_w, tile_size, stride)]

    with torch.no_grad():
        for start in range(0, len(coords), batch_size):
            batch_coords = coords[start : start + batch_size]
            tiles = torch.cat(
                [x_work[:, :, top : top + tile_size, left : left + tile_size] for top, left in batch_coords],
                dim=0,
            ).to(device)
            pred = model(tiles)
            logits = pred[0] if isinstance(pred, tuple) else pred
            probs = torch.sigmoid(logits.float()).detach().cpu()[:, 0]
            for prob, (top, left) in zip(probs, batch_coords):
                sy, sx, dy, dx = _crop_slices(top, left, tile_size, work_h, work_w, halo)
                output[dy, dx] += prob[sy, sx]
                weight[dy, dx] += 1.0
    output = output / weight.clamp_min(1.0)
    return output[:height, :width].contiguous()


__all__ = ["predict_probability_tiled"]