#!/usr/bin/env python3
"""Unpack the microscope camera's raw frames on a desktop machine.

Recording writes MIPI-packed 12-bit Bayer frames (``SBGGR12_CSI2P`` style: two
pixels in three bytes) straight off the sensor. This tool turns them into files
you can actually look at:

    unpack_raw.py frame_000001.raw                  -> frame_000001.npy   (uint16)
    unpack_raw.py --png frame_000001.raw            -> frame_000001.png   (quick look)
    unpack_raw.py --tiff16 frame_000001.raw         -> frame_000001.tiff  (16-bit mono)
    unpack_raw.py --debayer --rgb16 frame_000001.raw-> frame_000001_rgb.tiff (16-bit RGB)
    unpack_raw.py /mnt/microscope/2026-10-30/21-33-05        (whole session)
    unpack_raw.py --info /mnt/microscope/2026-10-30/21-33-05

Geometry (width, height, stride, Bayer pattern) is picked up automatically from
the session's ``session.json``; if that is missing the tool falls back to the
file size and the command line defaults.
"""

import argparse
import json
import os
import sys

import numpy as np

PATTERNS = {
    # pattern -> (row parity, col parity) of the red sample in the 2x2 tile
    "RGGB": (0, 0),
    "BGGR": (1, 1),
    "GRBG": (0, 1),
    "GBRG": (1, 0),
}


def pattern_from_format(fmt):
    text = (fmt or "SRGGB12").upper()
    if text.startswith("S"):
        text = text[1:]
    head = text[:4]
    return head if head in PATTERNS else "RGGB"


def unpack_mipi12(packed, width, height, stride):
    """Unpack MIPI CSI-2 RAW12 (3 bytes -> 2 pixels).

    Byte order matters: byte0 holds the low 8 bits of pixel 0, byte1 holds the
    high 4 bits of pixel 0 in its HIGH nibble and the high 4 bits of pixel 1 in
    its LOW nibble.  Treating byte1's low nibble as pixel 0's high bits (and
    its high nibble as pixel 1's) transposes the 12-bit words within each pair,
    which scrambles the Bayer mosaic and produces colour-fringed output.
    """
    rows = packed[: stride * height].reshape(height, stride)
    useful = rows[:, : (width * 3) // 2]
    triplets = useful.reshape(height, width // 2, 3).astype(np.uint16)
    b0 = triplets[:, :, 0]
    b1 = triplets[:, :, 1]
    b2 = triplets[:, :, 2]
    even = (b0 << 4) | (b1 >> 4)
    odd = ((b1 & 0x0F) << 8) | b2
    out = np.empty((height, width), dtype=np.uint16)
    out[:, 0::2] = even
    out[:, 1::2] = odd
    return out


def demosaic(raw, pattern="RGGB"):
    rp, cp = PATTERNS.get(pattern, (0, 0))
    h, w = raw.shape
    tiles = raw.reshape(h // 2, 2, w // 2, 2).transpose(0, 2, 1, 3)
    red = tiles[..., rp, cp].astype(np.uint32)
    blue = tiles[..., 1 - rp, 1 - cp].astype(np.uint32)
    green_r = tiles[..., rp, 1 - cp]
    green_b = tiles[..., 1 - rp, cp]

    def upsample(plane):
        p = np.pad(plane, 1, mode="edge")
        center = p[1:-1, 1:-1]
        left, right = p[1:-1, :-2], p[1:-1, 2:]
        up, down = p[:-2, 1:-1], p[2:, 1:-1]
        col_n = right if cp == 0 else left
        row_n = down if rp == 0 else up
        corner = p[2:, 2:] if (cp == 0 and rp == 0) else \
            p[2:, :-2] if (cp == 0 and rp == 1) else \
            p[:-2, 2:] if (cp == 1 and rp == 0) else p[:-2, :-2]
        full = np.empty((h, w), dtype=np.uint16)
        full[rp::2, cp::2] = center
        full[rp::2, (1 - cp)::2] = (center + col_n) // 2
        full[(1 - rp)::2, cp::2] = (center + row_n) // 2
        full[(1 - rp)::2, (1 - cp)::2] = (center + col_n + row_n + corner) // 4
        return full

    red_full, blue_full = upsample(red), upsample(blue)
    g1 = np.pad(green_r.astype(np.uint32), 1, mode="edge")
    g2 = np.pad(green_b.astype(np.uint32), 1, mode="edge")
    green_full = np.empty((h, w), dtype=np.uint16)
    green_full[rp::2, (1 - cp)::2] = green_r
    green_full[(1 - rp)::2, cp::2] = green_b
    green_full[rp::2, cp::2] = (g1[1:-1, 1:-1] + g1[1:-1, :-2]) // 2
    green_full[(1 - rp)::2, (1 - cp)::2] = (g2[1:-1, 1:-1] + g2[1:-1, 2:]) // 2

    rgb = np.empty((h, w, 3), dtype=np.uint16)
    rgb[..., 1] = green_full
    return rgb, red_full, blue_full


def apply_rotation(image, rotation):
    """Rotate an array clockwise by 0/90/180/270 degrees."""
    rotation = int(rotation or 0) % 360
    if rotation == 0:
        return image
    if rotation == 90:
        return np.ascontiguousarray(np.rot90(image, k=3))
    if rotation == 180:
        return np.ascontiguousarray(np.rot90(image, k=2))
    if rotation == 270:
        return np.ascontiguousarray(np.rot90(image, k=1))
    return image


_PATTERN_TILE = {
    "RGGB": (("R", "G"), ("G", "B")),
    "BGGR": (("B", "G"), ("G", "R")),
    "GRBG": (("G", "R"), ("B", "G")),
    "GBRG": (("G", "B"), ("R", "G")),
}


def rotate_pattern(pattern, rotation):
    """CFA pattern after rotating the Bayer plane clockwise.

    Rotating a mosaic without moving its pattern would put red samples where
    blue belong, so a rotated .npy/.tiff must record the rotated pattern.
    """
    steps = (int(rotation or 0) % 360) // 90
    if not steps:
        return pattern
    tile = [list(row) for row in _PATTERN_TILE.get(pattern, _PATTERN_TILE["RGGB"])]
    for _ in range(steps):
        tile = [[tile[1][0], tile[0][0]], [tile[1][1], tile[0][1]]]  # 90 deg clockwise
    text = "".join(tile[0] + tile[1])
    for name, known in _PATTERN_TILE.items():
        if "".join(known[0] + known[1]) == text:
            return name
    return pattern


def find_session(path):
    """Walk up from *path* looking for session.json / photo.json."""
    current = os.path.abspath(path if os.path.isdir(path) else os.path.dirname(path))
    for _ in range(4):
        for name in ("session.json", "photo.json"):
            candidate = os.path.join(current, name)
            if os.path.isfile(candidate):
                return current, candidate
        parent = os.path.dirname(current)
        if parent == current:
            break
        current = parent
    return None, None


def geometry_for(path, args):
    width, height, stride, pattern = args.width, args.height, args.stride, args.pattern
    session_dir, meta_file = find_session(path)
    meta = {}
    if meta_file:
        try:
            with open(meta_file, "r", encoding="utf-8") as handle:
                meta = json.load(handle)
        except (OSError, ValueError):
            meta = {}
    raw = meta.get("raw") or {}
    if raw:
        width = width or raw.get("width")
        height = height or raw.get("height")
        stride = stride or raw.get("stride")
        pattern = args.pattern or pattern_from_format(raw.get("format"))
    path_size = os.path.getsize(path)
    height = height or 1520
    if not width:
        # packed 12-bit: bytes per row = width * 3/2
        if stride:
            width = int(stride * 2 / 3)
        else:
            width = 2028
    if not stride:
        stride = path_size // height
    expected = width * height if args.unpacked else stride * height
    if path_size < expected:
        raise SystemExit(
            "%s is %d bytes but %dx%d needs %d — pass --width/--height/--stride"
            % (os.path.basename(path), path_size, width, height, expected)
        )
    rotation = args.rotate if args.rotate is not None else int(meta.get("rotation") or 0)
    return width, height, stride, pattern, session_dir, meta, rotation


def contrast_stretch(image):
    lo, hi = np.percentile(image, (0.5, 99.5))
    if hi <= lo:
        hi = lo + 1
    scaled = np.clip((image.astype(np.float32) - lo) * (255.0 / (hi - lo)), 0, 255)
    return scaled.astype(np.uint8)


def srgb_encode(rgb16):
    """Linear 16-bit RGB -> display-ready 8-bit sRGB.

    Without this the frame looks flat and washed out, because linear sensor data
    held in a 16-bit container has all its mid-tones crammed into the dark end.
    """
    x = np.clip(rgb16.astype(np.float32) / 65535.0, 0.0, 1.0)
    x = np.where(x <= 0.0031308, x * 12.92,
                 1.055 * np.power(x, 1.0 / 2.4) - 0.055)
    return np.clip(x * 255.0 + 0.5, 0, 255).astype(np.uint8)


def white_balance(rgb, red_full, blue_full, gains):
    """Apply camera white balance. Only called when --white-balance is given:
    gains above 1.0 clip at 65535, which loses highlight data."""
    if not gains:
        gains = (1.0, 1.0)
    out = rgb
    out[..., 0] = np.clip(red_full.astype(np.float32) * float(gains[0]), 0, 65535).astype(np.uint16)
    out[..., 2] = np.clip(blue_full.astype(np.float32) * float(gains[1]), 0, 65535).astype(np.uint16)
    return out


def process(path, args):
    width, height, stride, pattern, session_dir, meta, rotation = geometry_for(path, args)
    stem = os.path.splitext(os.path.basename(path))[0]
    outdir = args.outdir or os.path.dirname(os.path.abspath(path))
    os.makedirs(outdir, exist_ok=True)
    payload = np.fromfile(path, dtype=np.uint8)

    if args.unpacked:
        image = payload[: width * height * 2].view(np.uint16).reshape(height, width)
    else:
        image = unpack_mipi12(payload, width, height, stride)

    print("%-28s %dx%d  %s  stride=%d  rotation=%d°  max=%d  mean=%.1f"
          % (os.path.basename(path), width, height, pattern, stride, rotation,
             int(image.max()), float(image.mean())))

    # The session's rotation is the operator's display orientation, so honour it
    # in everything this tool writes. Bayer exports are rotated *with* their CFA
    # pattern; the RGB path is demosaiced first and then turned, which is
    # equivalent and avoids re-deriving the pattern.
    bayer = apply_rotation(image, rotation)
    bayer_pattern = rotate_pattern(pattern, rotation)

    made = []
    if args.npy or not (args.png or args.tiff16 or args.rgb16):
        out = os.path.join(outdir, stem + ".npy")
        np.save(out, bayer)
        made.append(out)
    if args.png:
        from PIL import Image
        out = os.path.join(outdir, stem + ".png")
        Image.fromarray(contrast_stretch(bayer)).save(out)
        made.append(out)
    if args.tiff16:
        import tifffile
        out = os.path.join(outdir, stem + ".tiff")
        tifffile.imwrite(out, bayer,
                         description="16-bit unpacked Bayer %s (rotated %d deg clockwise)"
                                     % (bayer_pattern, rotation))
        made.append(out)
    if args.rgb16:
        import tifffile
        gains = meta.get("awb")
        # Demosaic in sensor orientation, fill all three channels, then rotate:
        # rotating an RGB triple is equivalent to rotating the mosaic and
        # re-deriving its CFA pattern, without the pattern bookkeeping.
        rgb, red_full, blue_full = demosaic(image, pattern)
        rgb[..., 0] = np.clip(red_full.astype(np.float32) * 16.0, 0, 65535).astype(np.uint16)
        rgb[..., 2] = np.clip(blue_full.astype(np.float32) * 16.0, 0, 65535).astype(np.uint16)
        rgb[..., 1] = np.clip(rgb[..., 1].astype(np.float32) * 16.0, 0, 65535).astype(np.uint16)
        del red_full, blue_full, image
        if rotation:
            rgb = apply_rotation(rgb, rotation)
        applied = False
        if args.white_balance:
            rgb = white_balance(rgb, rgb[..., 0], rgb[..., 2], gains)
            applied = True
        out = os.path.join(outdir, stem + "_rgb.tiff")
        tifffile.imwrite(out, rgb, photometric="rgb",
                         description="linear 16-bit RGB, %s, rotation=%d deg clockwise, "
                                     "white_balance_applied=%s camera_white_balance=%s"
                                     % (pattern, rotation, applied, gains))
        made.append(out)
    for item in made:
        print("  -> %s (%.1f MB)" % (item, os.path.getsize(item) / 1e6))
    return made


def show_info(target):
    session_dir, meta_file = find_session(target)
    if not meta_file:
        raise SystemExit("no session.json/photo.json found near %s" % target)
    with open(meta_file, "r", encoding="utf-8") as handle:
        meta = json.load(handle)
    print(json.dumps(meta, indent=2))


def collect_inputs(paths, recursive):
    files = []
    for path in paths:
        if os.path.isdir(path):
            raw_dir = os.path.join(path, "raw")
            scan = raw_dir if os.path.isdir(raw_dir) else path
            for name in sorted(os.listdir(scan)):
                if name.endswith(".raw"):
                    files.append(os.path.join(scan, name))
        else:
            files.append(path)
    return files


def main():
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("paths", nargs="+", help="raw files, session dirs, or date dirs")
    parser.add_argument("--outdir", help="where to write results (default: next to input)")
    parser.add_argument("--npy", action="store_true", help="write 16-bit .npy (default)")
    parser.add_argument("--png", action="store_true", help="write 8-bit contrast-stretched PNG")
    parser.add_argument("--tiff16", action="store_true", help="write 16-bit Bayer TIFF")
    parser.add_argument("--rgb16", action="store_true", help="write 16-bit RGB TIFF (demosaiced)")
    parser.add_argument("--debayer", action="store_true", help="alias for --rgb16")
    parser.add_argument("--rotate", type=int, choices=(0, 90, 180, 270),
                        help="override the session's display rotation (degrees clockwise)")
    parser.add_argument("--no-rotate", dest="rotate", action="store_const", const=0,
                        help="ignore the rotation recorded in the session metadata")
    parser.add_argument("--white-balance", action="store_true",
                        help="apply the camera's colour gains to the RGB output "
                             "(clips highlights; off by default, gains are recorded anyway)")
    parser.add_argument("--info", action="store_true", help="print session metadata and exit")
    parser.add_argument("--pattern", choices=sorted(PATTERNS), help="Bayer pattern override")
    parser.add_argument("--width", type=int, help="frame width override")
    parser.add_argument("--height", type=int, help="frame height override")
    parser.add_argument("--stride", type=int, help="row stride in bytes override")
    parser.add_argument("--unpacked", action="store_true",
                        help="input is already unpacked 16-bit little-endian")
    args = parser.parse_args()

    if args.debayer:
        args.rgb16 = True
    if args.info:
        for path in args.paths:
            show_info(path)
        return 0

    files = collect_inputs(args.paths, recursive=True)
    if not files:
        raise SystemExit("no .raw files found")
    total = 0
    for path in files:
        total += len(process(path, args))
    print("done: %d file(s) from %d input(s)" % (total, len(files)))
    return 0


if __name__ == "__main__":
    sys.exit(main())