"""Batched inference functions for depth (DA3), edges (Sobel), and segmentation.

All functions accept explicit parameters — no global config imports.
"""

from __future__ import annotations

import logging
from typing import Any

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image

logger = logging.getLogger(__name__)

_IMGNET_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
_IMGNET_STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
_cached_device: torch.device | None = None
_cached_mean: torch.Tensor | None = None
_cached_std: torch.Tensor | None = None


def _get_imgnet_stats(
    device: torch.device, dtype: torch.dtype,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return cached ImageNet mean/std on device."""
    global _cached_device, _cached_mean, _cached_std
    if _cached_device != device:
        _cached_mean = _IMGNET_MEAN.to(device, dtype=dtype)
        _cached_std = _IMGNET_STD.to(device, dtype=dtype)
        _cached_device = device
    return _cached_mean, _cached_std  # type: ignore[return-value]


# ---------------------------------------------------------------------------
# Depth
# ---------------------------------------------------------------------------

@torch.inference_mode()
def infer_depth_batch(
    model: nn.Module,
    images_raw: torch.Tensor,
    device: torch.device,
) -> torch.Tensor:
    """Run depth estimation on a batch.

    Args:
        images_raw: [B, 3, H, W] float32 [0, 1].

    Returns:
        depth: [B, 1, H, W] float32 [0, 1] (per-frame normalized).
    """
    B, _, H, W = images_raw.shape

    # DA3 API: model.inference([PIL/np images]) → Prediction with .depth [N, H, W].
    if hasattr(model, "inference"):
        img_list = [
            Image.fromarray(
                (images_raw[i].permute(1, 2, 0).numpy() * 255).astype(np.uint8)
            )
            for i in range(B)
        ]
        prediction = model.inference(img_list, process_res=H)
        depth_np = np.asarray(prediction.depth)  # [N, H', W']
        depths_list = []
        for i in range(B):
            d = torch.from_numpy(depth_np[i]).float()
            if d.ndim == 2:
                d = d.unsqueeze(0)
            if d.shape[-2:] != (H, W):
                d = F.interpolate(
                    d.unsqueeze(0), size=(H, W), mode="bilinear", align_corners=False,
                ).squeeze(0)
            d_min, d_max = d.min(), d.max()
            d = (d - d_min) / (d_max - d_min + 1e-8)
            depths_list.append(d)
        return torch.stack(depths_list)

    # DA V2 fallback (transformers API).
    mean, std = _get_imgnet_stats(device, images_raw.dtype)
    x = (images_raw.to(device) - mean) / std
    if next(model.parameters()).dtype == torch.float16:
        x = x.half()

    pred = model(pixel_values=x).predicted_depth
    depth = F.interpolate(
        pred.unsqueeze(1).float(), size=(H, W), mode="bilinear", align_corners=False,
    )
    for i in range(B):
        d = depth[i]
        d_min, d_max = d.min(), d.max()
        depth[i] = (d - d_min) / (d_max - d_min + 1e-8)
    return depth.cpu()


# ---------------------------------------------------------------------------
# CHMv2 (DINOv3 cross-view height map)
# ---------------------------------------------------------------------------

@torch.inference_mode()
def infer_chmv2_batch(
    model: nn.Module,
    processor: Any,
    images_raw: torch.Tensor,
    device: torch.device,
) -> torch.Tensor:
    """Run CHMv2 depth estimation on a batch.

    Args:
        model: CHMv2ForDepthEstimation.
        processor: CHMv2ImageProcessor (used for post-processing).
        images_raw: [B, 3, H, W] float32 [0, 1].

    Returns:
        depth: [B, 1, H, W] float32 [0, 1] (per-frame normalized).
    """
    B, _, H, W = images_raw.shape

    # Convert to PIL for processor
    pil_images = [
        Image.fromarray(
            (images_raw[i].permute(1, 2, 0).numpy() * 255).astype(np.uint8)
        )
        for i in range(B)
    ]
    inputs = processor(images=pil_images, return_tensors="pt")
    # CHMv2 must run in FP32 (NaN in FP16).
    pixel_values = inputs["pixel_values"].to(device, dtype=torch.float32)

    outputs = model(pixel_values=pixel_values)

    # Post-process to get depth maps at original resolution
    target_sizes = [(H, W)] * B
    results = processor.post_process_depth_estimation(
        outputs, target_sizes=target_sizes,
    )

    depths_list = []
    for r in results:
        d = r["predicted_depth"].float()
        if d.ndim == 2:
            d = d.unsqueeze(0)
        d_min, d_max = d.min(), d.max()
        d = (d - d_min) / (d_max - d_min + 1e-8)
        depths_list.append(d)
    return torch.stack(depths_list).cpu()


# ---------------------------------------------------------------------------
# Edges
# ---------------------------------------------------------------------------

_SOBEL_X = torch.tensor(
    [[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32,
).view(1, 1, 3, 3) / 8.0
_SOBEL_Y = torch.tensor(
    [[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32,
).view(1, 1, 3, 3) / 8.0


def compute_edges_from_depth(depth: torch.Tensor) -> torch.Tensor:
    """Compute structural edges from depth via Sobel filters.

    Args:
        depth: [B, 1, H, W] float32 [0, 1].

    Returns:
        edges: [B, 1, H, W] float32 [0, 1] (edge magnitude).
    """
    depth_padded = F.pad(depth, (1, 1, 1, 1), mode="replicate")
    dz_dx = F.conv2d(depth_padded, _SOBEL_X)
    dz_dy = F.conv2d(depth_padded, _SOBEL_Y)
    edges = torch.sqrt(dz_dx ** 2 + dz_dy ** 2)

    for i in range(edges.shape[0]):
        e = edges[i]
        e_max = e.max()
        if e_max > 0:
            edges[i] = e / e_max
    return edges


# ---------------------------------------------------------------------------
# Segmentation
# ---------------------------------------------------------------------------

@torch.inference_mode()
def infer_segmentation_batch(
    model: Any,
    seg_config: dict[str, Any],
    images_raw: torch.Tensor,
    device: torch.device,
) -> torch.Tensor:
    """Run semantic segmentation on a batch.

    Args:
        model: SegEarth-OV3 pipeline or SegFormer nn.Module.
        seg_config: Must contain 'type' and optionally 'prompts', 'processor'.
        images_raw: [B, 3, H, W] float32 [0, 1].

    Returns:
        seg_ids: [B, 1, H, W] uint8.
    """
    if seg_config.get("type") == "segearth-ov3":
        prompts = seg_config.get("prompts", [])
        return _infer_segearth_ov3(model, images_raw, prompts)
    else:
        processor = seg_config["processor"]
        proc_mean = torch.tensor(processor.image_mean).view(1, 3, 1, 1)
        proc_std = torch.tensor(processor.image_std).view(1, 3, 1, 1)
        return _infer_segformer(model, proc_mean, proc_std, images_raw, device)


def _infer_segearth_ov3(
    model: Any,
    images_raw: torch.Tensor,
    prompts: list[str],
) -> torch.Tensor:
    """Run SegEarth-OV3 on a batch of images.

    Uses model.predict_pil_batch() for batched backbone inference when available,
    falls back to per-image predict_pil().
    """
    B, _, H, W = images_raw.shape

    # Convert all tensors to PIL up front
    pil_images = []
    for i in range(B):
        img_np = (images_raw[i].permute(1, 2, 0).numpy() * 255).astype(np.uint8)
        pil_images.append(Image.fromarray(img_np))

    ori_shape = (H, W)

    # Batched backbone path — chunk into sub-batches of up to 16
    _MAX_SEG_BATCH = 16
    if hasattr(model, "predict_pil_batch"):
        try:
            seg_list = []
            for start in range(0, B, _MAX_SEG_BATCH):
                chunk = pil_images[start:start + _MAX_SEG_BATCH]
                chunk_results = model.predict_pil_batch(
                    chunk, ori_shapes=[ori_shape] * len(chunk),
                )
                for seg_pred, _ in chunk_results:
                    t = seg_pred.cpu().to(torch.uint8)
                    if t.ndim == 2:
                        t = t.unsqueeze(0)
                    seg_list.append(t)
            return torch.stack(seg_list)
        except Exception as exc:
            logger.warning("⚠️  predict_pil_batch failed, falling back to per-image: %s", exc)

    # Fallback: per-image inference
    seg_list = []
    for i, pil_img in enumerate(pil_images):
        try:
            if hasattr(model, "predict_pil"):
                seg_pred, _ = model.predict_pil(pil_img, ori_shape=ori_shape)
                t = seg_pred.cpu().to(torch.uint8)
            elif hasattr(model, "predict"):
                seg_map = model.predict(pil_img, text_prompts=prompts)
                seg_np = np.asarray(seg_map).squeeze()
                t = torch.from_numpy(seg_np.astype(np.uint8))
            else:
                raise AttributeError(f"Unknown SegEarth API: {type(model)}")

            if t.ndim == 2 and t.shape != (H, W):
                t = F.interpolate(
                    t.float().unsqueeze(0).unsqueeze(0), size=(H, W), mode="nearest",
                ).squeeze(0).squeeze(0).to(torch.uint8)
            if t.ndim == 2:
                t = t.unsqueeze(0)
            seg_list.append(t)
        except Exception as exc:
            logger.warning("⚠️  SegEarth-OV3 failed on image %d: %s", i, exc)
            seg_list.append(torch.zeros(1, H, W, dtype=torch.uint8))
    return torch.stack(seg_list)


@torch.inference_mode()
def _infer_segformer(
    model: nn.Module,
    proc_mean: torch.Tensor,
    proc_std: torch.Tensor,
    images_raw: torch.Tensor,
    device: torch.device,
) -> torch.Tensor:
    """Run SegFormer-B5 on a batch (fallback)."""
    _, _, H, W = images_raw.shape
    mean = proc_mean.to(device)
    std = proc_std.to(device)
    x = (images_raw.to(device) - mean) / std
    if next(model.parameters()).dtype == torch.float16:
        x = x.half()
    logits = model(pixel_values=x).logits
    upsampled = F.interpolate(
        logits.float(), size=(H, W), mode="bilinear", align_corners=False,
    )
    return upsampled.argmax(dim=1, keepdim=True).cpu().to(torch.uint8)


# ---------------------------------------------------------------------------
# Post-processing heuristics
# ---------------------------------------------------------------------------

def postprocess_segmentation(
    seg_ids: torch.Tensor,
    images_raw: torch.Tensor,
    water_class: int = 4,
    dark_water_mean_thr: float = 0.24,
    dark_water_std_thr: float = 0.18,
    fix_uniform_water: bool = True,
    uniform_water_bg_thr: float = 0.80,
    uniform_water_std_thr: float = 0.10,
    uniform_water_ch_std_thr: float = 0.05,
    uniform_water_mean_thr: float = 0.60,
    reclassify_wetland: bool = False,
    wetland_class: int = 14,
    vegetation_class: int = 3,
    bare_soil_class: int = 11,
    green_thr: float = 0.35,
) -> torch.Tensor:
    """Fix known segmentation failure modes with simple heuristics.

    Applied per-image after model inference.

    Rules:
        1a. Dark water: if background pixels have mean_rgb < dark_water_mean_thr
            and std < dark_water_std_thr, reclassify all background pixels as water.
        1b. Uniform water: if background ratio > uniform_water_bg_thr AND
            (whole-image std < uniform_water_std_thr OR
             mean-per-channel std < uniform_water_ch_std_thr AND
             mean brightness < uniform_water_mean_thr),
            reclassify background as water.
            First condition catches muddy/grey river water (UAV_VisLoc).
            Second condition catches blue ocean where cross-channel variance is
            high but each channel is spatially uniform (GTA-UAV satellite).
        2.  Wetland reclassification (optional, for GTA-UAV): reclassify wetland pixels
            by local color — green-dominant → vegetation, else → bare soil.

    Args:
        seg_ids: [B, 1, H, W] uint8 class IDs.
        images_raw: [B, 3, H, W] float32 [0, 1] original RGB.
        water_class: class ID for water.
        dark_water_mean_thr: mean RGB threshold (0-1) below which bg → water.
        dark_water_std_thr: std threshold below which bg → water.
        fix_uniform_water: if True, apply Rule 1b for uniform unclassified water.
        uniform_water_bg_thr: background ratio above which Rule 1b triggers.
        uniform_water_std_thr: whole-image std below which Rule 1b triggers.
        uniform_water_ch_std_thr: mean per-channel std below which Rule 1b triggers
            (second condition, for blue ocean with high cross-channel variance).
        uniform_water_mean_thr: mean brightness below which the ch_std condition
            applies (filters out bright terrain like sand/concrete).
        reclassify_wetland: if True, split wetland into vegetation/bare_soil.
        wetland_class: class ID for muddy/wetland.
        vegetation_class: class ID for vegetation.
        bare_soil_class: class ID for bare soil.
        green_thr: green-channel ratio threshold for vegetation vs bare soil.

    Returns:
        seg_ids: [B, 1, H, W] uint8, corrected in-place.
    """
    B = seg_ids.shape[0]
    seg = seg_ids.clone()

    for i in range(B):
        s = seg[i, 0]        # [H, W] uint8
        rgb = images_raw[i]  # [3, H, W] float32

        # Rule 1a: dark uniform images → background becomes water
        bg_mask = s == 0
        if bg_mask.any():
            bg_pixels = rgb[:, bg_mask]  # [3, N]
            mean_val = bg_pixels.mean().item()
            std_val = bg_pixels.std().item()
            if mean_val < dark_water_mean_thr and std_val < dark_water_std_thr:
                s[bg_mask] = water_class
            # Rule 1b: uniform unclassified water
            elif fix_uniform_water:
                bg_ratio = bg_mask.sum().item() / bg_mask.numel()
                if bg_ratio > uniform_water_bg_thr:
                    img_std = rgb.std().item()
                    # Condition A: low overall std (muddy/grey river water)
                    is_uniform = img_std < uniform_water_std_thr
                    # Condition B: low per-channel std + not bright
                    # (blue ocean — high cross-channel variance but
                    # each channel is spatially uniform)
                    if not is_uniform:
                        ch_std = (
                            rgb[0].std().item()
                            + rgb[1].std().item()
                            + rgb[2].std().item()
                        ) / 3.0
                        mean_val = rgb.mean().item()
                        is_uniform = (
                            ch_std < uniform_water_ch_std_thr
                            and mean_val < uniform_water_mean_thr
                        )
                    if is_uniform:
                        s[bg_mask] = water_class

        # Rule 2: reclassify wetland by color
        if reclassify_wetland:
            wet_mask = s == wetland_class
            if wet_mask.any():
                g = rgb[1, wet_mask]  # green channel
                r = rgb[0, wet_mask]
                # Green-dominant → vegetation, else → bare soil
                is_green = g > (r + green_thr * 0.5)
                reclassed = torch.where(is_green, vegetation_class, bare_soil_class)
                s[wet_mask] = reclassed.to(torch.uint8)

        seg[i, 0] = s

    return seg
