"""
kitti.py --- loading helpers for the KITTI Road benchmark.

CSCI 3240U (Computer Vision I), Ontario Tech University.

This module does the boring part: finding files, parsing the calibration text
format, and decoding the colour-coded ground truth.  None of that is computer
vision, and every hour you spend debugging it is an hour you are not spending on
the actual lab.

It deliberately does **not** implement anything you are asked to write yourself
--- no edge detection, no region of interest, no scoring, no rectification.

Usage:

    import kitti

    ds = kitti.Dataset("path/to/kitti/data_road")

    print(ds.scenes("um"))                    # ['um_000000', 'um_000001', ...]

    img  = ds.image("um_000032")              # left image,  RGB uint8
    rght = ds.image("um_000032", right=True)  # right image, RGB uint8

    cal = ds.calib("um_000032")
    print(cal.K, cal.fx, cal.baseline, cal.camera_height)

    pos, valid = ds.ground_truth("um_000032")           # road
    pos, valid = ds.ground_truth("um_000032", "lane")   # ego lane (um only)
"""

from __future__ import annotations

import glob
import os

import numpy as np

try:
    import cv2 as cv
except ImportError:                                      # pragma: no cover
    cv = None

CATEGORIES = ("um", "umm", "uu")


# --------------------------------------------------------------------------
# calibration
# --------------------------------------------------------------------------
class Calibration:
    """Parsed contents of a KITTI ``calib/*.txt`` file.

    Attributes
    ----------
    P2, P3 : (3, 4) arrays
        Projection matrices for the left and right colour cameras.
    K : (3, 3) array
        Intrinsics of the left colour camera (the leftmost block of ``P2``).
    fx, fy : float           focal lengths, in pixels
    cx, cy : float           principal point, in pixels
    baseline : float         stereo baseline, in metres
    camera_height : float    height of the camera above the road plane, in metres
    """

    def __init__(self, entries: dict[str, np.ndarray]):
        self._entries = entries
        self.P2 = entries["P2"].reshape(3, 4)
        self.P3 = entries["P3"].reshape(3, 4)
        self.K = self.P2[:, :3].copy()
        self.fx = float(self.K[0, 0])
        self.fy = float(self.K[1, 1])
        self.cx = float(self.K[0, 2])
        self.cy = float(self.K[1, 2])
        # Last column of P is K t, so the horizontal offsets are P[0, 3] / fx.
        self.baseline = float(abs(self.P3[0, 3] - self.P2[0, 3]) / self.fx)
        tr = entries.get("Tr_cam_to_road")
        self.cam_to_road = None if tr is None else tr.reshape(3, 4)
        self.camera_height = (None if self.cam_to_road is None
                              else float(abs(self.cam_to_road[1, 3])))

    def __getitem__(self, key: str) -> np.ndarray:
        """Raw access to any entry, e.g. cal['R0_rect']."""
        return self._entries[key]

    def keys(self):
        return self._entries.keys()

    def __repr__(self):
        height = "?" if self.camera_height is None else f"{self.camera_height:.3f}m"
        return (f"<Calibration fx={self.fx:.1f} cx={self.cx:.1f} cy={self.cy:.1f} "
                f"baseline={self.baseline:.3f}m height={height}>")


def parse_calib(path: str) -> Calibration:
    """Parse a KITTI calibration file into a :class:`Calibration`."""
    entries = {}
    with open(path) as fh:
        for line in fh:
            if ":" not in line:
                continue
            key, values = line.split(":", 1)
            entries[key.strip()] = np.array([float(v) for v in values.split()])
    return Calibration(entries)


# --------------------------------------------------------------------------
# ground truth
# --------------------------------------------------------------------------
def decode_ground_truth(gt_rgb: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Decode a KITTI ground-truth image.

    KITTI colour-codes the labels: magenta is the positive region, red is the
    negative region, and black marks pixels the annotator declined to label.

    The image is anti-aliased at region edges, so exact colour equality does not
    work --- the channels are thresholded instead.

    Parameters
    ----------
    gt_rgb : (H, W, 3) uint8 array, in **RGB** order.

    Returns
    -------
    positive : (H, W) bool   the labelled region (road, or ego lane)
    valid    : (H, W) bool   pixels that count; exclude everything else from
                             any score you compute
    """
    r = gt_rgb[..., 0].astype(np.int16)
    g = gt_rgb[..., 1].astype(np.int16)
    b = gt_rgb[..., 2].astype(np.int16)
    positive = (r > 200) & (g < 60) & (b > 200)      # magenta
    void = (r < 60) & (g < 60) & (b < 60)            # black
    return positive, ~void


# --------------------------------------------------------------------------
# dataset
# --------------------------------------------------------------------------
class Dataset:
    """Convenience wrapper around an extracted ``data_road`` directory.

    Parameters
    ----------
    root : str
        Path to the extracted ``data_road`` folder (the one containing
        ``training/`` and ``testing/``).
    split : {'training', 'testing'}
    """

    def __init__(self, root: str, split: str = "training"):
        if os.path.isdir(os.path.join(root, "data_road")):
            root = os.path.join(root, "data_road")       # tolerate the parent
        self.root = root
        self.split = split
        base = os.path.join(root, split)
        if not os.path.isdir(base):
            raise FileNotFoundError(
                f"No '{split}' folder under {root!r}. Point this at the "
                f"extracted data_road directory --- see the 'Getting the KITTI "
                f"data' handout.")
        self.base = base

    # -- listing ----------------------------------------------------------
    def scenes(self, category: str | None = None) -> list[str]:
        """Scene ids, e.g. ``'um_000032'``. Optionally filtered by category."""
        pattern = f"{category}_*.png" if category else "*.png"
        found = glob.glob(os.path.join(self.base, "image_2", pattern))
        return sorted(os.path.splitext(os.path.basename(p))[0] for p in found)

    # -- images -----------------------------------------------------------
    def image(self, scene: str, right: bool = False) -> np.ndarray:
        """Load an image as an **RGB** uint8 array."""
        folder = "image_3" if right else "image_2"
        path = os.path.join(self.base, folder, scene + ".png")
        if not os.path.exists(path):
            extra = ("  The right-camera images are a separate download --- "
                     "re-run fetch-kitti.sh without --left-only."
                     if right else "")
            raise FileNotFoundError(f"{path} not found.{extra}")
        if cv is None:
            raise ImportError("OpenCV is required to load images.")
        bgr = cv.imread(path, cv.IMREAD_COLOR)
        return cv.cvtColor(bgr, cv.COLOR_BGR2RGB)

    # -- calibration ------------------------------------------------------
    def calib(self, scene: str) -> Calibration:
        return parse_calib(os.path.join(self.base, "calib", scene + ".txt"))

    # -- ground truth -----------------------------------------------------
    def ground_truth(self, scene: str, kind: str = "road"):
        """Return ``(positive, valid)`` boolean masks.

        ``kind`` is ``'road'`` (whole road surface, all training scenes) or
        ``'lane'`` (ego lane only, available for the ``um`` category).
        """
        if self.split != "training":
            raise ValueError("Ground truth exists for the training split only.")
        category, index = scene.rsplit("_", 1)
        path = os.path.join(self.base, "gt_image_2",
                            f"{category}_{kind}_{index}.png")
        if not os.path.exists(path):
            if kind == "lane":
                raise FileNotFoundError(
                    f"No lane ground truth for {scene!r}. Lane labels exist "
                    f"only for the 'um' category; use kind='road' instead.")
            raise FileNotFoundError(f"{path} not found.")
        if cv is None:
            raise ImportError("OpenCV is required to load images.")
        gt = cv.cvtColor(cv.imread(path, cv.IMREAD_COLOR), cv.COLOR_BGR2RGB)
        return decode_ground_truth(gt)

    def __repr__(self):
        return f"<kitti.Dataset {self.base!r}: {len(self.scenes())} scenes>"