"""Foreground ground-truth masks from Guillaumin et al. 2014.

Loads `data/gtsegs_ijcv.mat` (HDF5 v7.3) and exposes per-image binary masks
indexed by ImageNet id (e.g. "n01322343_1025") or row index.

The .mat file structure (HDF5):
    /value/n              shape (1, 1) — total count (4276)
    /value/id[i, 0]       reference -> uint16 array, ASCII "nXXXXXXXX_NNNN"
    /value/gt[i, 0]       reference -> group with refs -> mask (H, W) uint8 {0, 1}
    /value/img[i, 0]      reference -> image (3, H, W) uint8
    /value/target[i, 0]   reference -> class target

Usage:
    from utils.foreground import GTMaskLoader

    loader = GTMaskLoader()                 # uses default data/gtsegs_ijcv.mat
    mask = loader.get_mask("n01322343_1025")  # (H, W) uint8 binary
"""

from __future__ import annotations

import os
from pathlib import Path
from typing import Dict, Optional

import numpy as np


def _resolve_default_mat_path() -> Path:
    """Resolve the default GT masks file path with a fallback chain.

    Priority:
        1. Env var VITVIZ_GT_MASKS_PATH (explicit override)
        2. <cwd>/data/gtsegs_ijcv.mat (matches SLURM script `cd $WORKDIR` setup)
        3. <__file__>/../../data/gtsegs_ijcv.mat (matches local dev cwd inside repo)

    The bundle layout on Apuana has utils/foreground.py inside the bundle dir
    (~/vitviz_jobs/<stamp>/) but data/ is in the project dir (~/ViTViz/).
    The SLURM script `cd`s to the project dir before invoking the runner,
    so cwd-based resolution works. Falling back to __file__-relative keeps
    local invocations from arbitrary directories working.
    """
    env_path = os.environ.get("VITVIZ_GT_MASKS_PATH")
    if env_path:
        return Path(env_path)
    cwd_path = Path.cwd() / "data" / "gtsegs_ijcv.mat"
    if cwd_path.exists():
        return cwd_path
    return Path(__file__).resolve().parent.parent / "data" / "gtsegs_ijcv.mat"


_DEFAULT_MAT_PATH = _resolve_default_mat_path()


def _decode_id(uint16_array: np.ndarray) -> str:
    """Decode a MATLAB-encoded char array (uint16) into a Python string."""
    return "".join(chr(int(c)) for c in np.asarray(uint16_array).flatten())


def normalize_image_id(name: str) -> str:
    """Strip extension and directory components from a filename.

    Handles common ImageNet-style identifiers like "n01322343_1025.JPEG"
    -> "n01322343_1025". Pass-through for already-normalized ids.
    """
    stem = Path(name).stem
    return stem


class GTMaskLoader:
    """Lazy loader for Guillaumin 2014 foreground masks.

    The HDF5 file is opened on first access and the id index is built lazily.
    Safe to instantiate even when the .mat file is missing — `available`
    returns False and `get_mask` returns None in that case.
    """

    def __init__(self, mat_path: Optional[Path] = None):
        self.mat_path = Path(mat_path) if mat_path is not None else _DEFAULT_MAT_PATH
        self._h5 = None
        self._index: Optional[Dict[str, int]] = None

    @property
    def available(self) -> bool:
        return self.mat_path.exists()

    def _ensure_open(self) -> None:
        if self._h5 is not None:
            return
        if not self.available:
            raise FileNotFoundError(
                f"GT masks file not found: {self.mat_path}. "
                "Download `gtsegs_ijcv.mat` (Guillaumin et al. 2014) into data/."
            )
        import h5py
        self._h5 = h5py.File(str(self.mat_path), "r")

    def _ensure_index(self) -> None:
        if self._index is not None:
            return
        self._ensure_open()
        val = self._h5["value"]
        n = int(np.asarray(val["n"]).squeeze())
        index: Dict[str, int] = {}
        for i in range(n):
            ref = val["id"][i, 0]
            raw = self._h5[ref][()]
            index[_decode_id(raw)] = i
        self._index = index

    def __len__(self) -> int:
        self._ensure_index()
        return len(self._index or {})

    def has(self, image_id: str) -> bool:
        if not self.available:
            return False
        self._ensure_index()
        return normalize_image_id(image_id) in (self._index or {})

    def get_mask(self, image_id: str) -> Optional[np.ndarray]:
        """Return binary mask (H, W) uint8 for `image_id`, or None if absent."""
        if not self.available:
            return None
        self._ensure_index()
        idx = (self._index or {}).get(normalize_image_id(image_id))
        if idx is None:
            return None
        return self._mask_at(idx)

    def get_image(self, image_id: str) -> Optional[np.ndarray]:
        """Return RGB image (H, W, 3) uint8 for `image_id`, or None if absent."""
        if not self.available:
            return None
        self._ensure_index()
        idx = (self._index or {}).get(normalize_image_id(image_id))
        if idx is None:
            return None
        return self._image_at(idx)

    def _mask_at(self, row_idx: int) -> np.ndarray:
        self._ensure_open()
        val = self._h5["value"]
        ref_outer = val["gt"][row_idx, 0]
        inner = self._h5[ref_outer]
        ref_inner = inner[0, 0]
        mask = self._h5[ref_inner][()]
        return np.asarray(mask, dtype=np.uint8)

    def _image_at(self, row_idx: int) -> np.ndarray:
        self._ensure_open()
        val = self._h5["value"]
        ref = val["img"][row_idx, 0]
        img = self._h5[ref][()]
        arr = np.asarray(img)
        if arr.ndim == 3 and arr.shape[0] == 3:
            arr = np.transpose(arr, (1, 2, 0))
        return arr.astype(np.uint8, copy=False)

    def close(self) -> None:
        if self._h5 is not None:
            self._h5.close()
            self._h5 = None

    def __del__(self):
        try:
            self.close()
        except Exception:
            pass


class PNGMaskLoader:
    """Loader for pre-rendered binary foreground masks stored as PNG files.

    Used for datasets whose ground-truth foreground does not come from Guillaumin
    2014 — e.g. GTSRB, where the "foreground" is the sign ROI box (Roi.X1..Y2)
    rasterized to a filled rectangle. One PNG per image, named `<stem>.png`,
    single-channel, non-zero = foreground, at the SAME pixel dims as the image
    file (so the patch-grid pooling matches, exactly like GTMaskLoader). Same
    public interface (`available`, `has`, `get_mask`) so `run_single` can use
    either loader transparently.
    """

    def __init__(self, mask_dir):
        self.mask_dir = Path(mask_dir)

    @property
    def available(self) -> bool:
        return self.mask_dir.is_dir()

    def _path(self, image_id: str) -> Path:
        return self.mask_dir / f"{normalize_image_id(image_id)}.png"

    def has(self, image_id: str) -> bool:
        return self.available and self._path(image_id).exists()

    def get_mask(self, image_id: str) -> Optional[np.ndarray]:
        """Return binary mask (H, W) uint8 for `image_id`, or None if absent."""
        if not self.available:
            return None
        p = self._path(image_id)
        if not p.exists():
            return None
        from PIL import Image
        arr = np.asarray(Image.open(p).convert("L"))
        return (arr > 0).astype(np.uint8)


_DEFAULT_LOADER: Optional[GTMaskLoader] = None
_ACTIVE_LOADER = None


def get_default_loader() -> GTMaskLoader:
    """Module-level singleton for the default mask file in data/."""
    global _DEFAULT_LOADER
    if _DEFAULT_LOADER is None:
        _DEFAULT_LOADER = GTMaskLoader()
    return _DEFAULT_LOADER


def set_active_loader(loader) -> None:
    """Override the mask loader used by the sweep (e.g. PNGMaskLoader for GTSRB).

    Called once per process in `run_attack_sweep.main()` when config selects a
    non-Guillaumin mask source. Leaves the Guillaumin default in place otherwise.
    """
    global _ACTIVE_LOADER
    _ACTIVE_LOADER = loader


def get_active_loader():
    """Loader used by `run_single`: the override if one was set, else Guillaumin."""
    return _ACTIVE_LOADER if _ACTIVE_LOADER is not None else get_default_loader()


def compute_phi_per_patch(
    mask: np.ndarray,
    grid_h: int,
    grid_w: int,
) -> np.ndarray:
    """Foreground proportion per patch.

    Pools a pixel-level binary mask (H, W) onto a (grid_h, grid_w) patch grid
    by averaging within each patch. The result φ ∈ [0, 1]^(grid_h, grid_w)
    is the per-patch fraction of foreground pixels — boundary patches receive
    fractional values; no thresholding is applied.

    Args:
        mask: Binary mask (H, W). Non-zero entries are treated as foreground.
        grid_h: Number of patch rows on the attention grid.
        grid_w: Number of patch columns on the attention grid.

    Returns:
        np.ndarray of shape (grid_h, grid_w), dtype float64, values in [0, 1].
    """
    if grid_h <= 0 or grid_w <= 0:
        raise ValueError(f"grid dims must be positive, got {grid_h}x{grid_w}")

    arr = np.asarray(mask)
    if arr.ndim != 2:
        raise ValueError(f"mask must be 2D (H, W), got shape {arr.shape}")

    binary = (arr > 0).astype(np.float64)

    # Resize via block-mean: divide image into grid_h x grid_w cells of equal size.
    # Use bilinear-ish pooling: for each cell, take the mean of the corresponding pixel block.
    H, W = binary.shape
    if H == grid_h and W == grid_w:
        return binary

    # Build float-valued row/col edges for non-uniform block sizes when H/W don't divide evenly.
    row_edges = np.linspace(0, H, grid_h + 1)
    col_edges = np.linspace(0, W, grid_w + 1)

    out = np.empty((grid_h, grid_w), dtype=np.float64)
    for i in range(grid_h):
        r0, r1 = int(np.floor(row_edges[i])), int(np.ceil(row_edges[i + 1]))
        r1 = max(r1, r0 + 1)
        for j in range(grid_w):
            c0, c1 = int(np.floor(col_edges[j])), int(np.ceil(col_edges[j + 1]))
            c1 = max(c1, c0 + 1)
            block = binary[r0:r1, c0:c1]
            out[i, j] = float(block.mean()) if block.size else 0.0
    return out