"""
Load an Inscoper/Micro-Manager RAW_DATA tile set into napari, correctly placed in
stage coordinates, ready for the napari-stitcher "Stitcher" widget.

Usage
-----
    python stitch_napari.py                       # auto-find RAW_DATA below cwd
    python stitch_napari.py <dir>                 # RAW_DATA folder, or any folder above it
    python stitch_napari.py <dir> --source index  # place tiles from grid indices
    python stitch_napari.py <dir> --t 0 --z 5     # load a single plane per tile
    python stitch_napari.py <dir> --c 0           # multi-channel: pick one channel

Nothing in this script is tied to a particular microscope: the stage axes, the
pixel size, and the stage-to-image calibration are all read from the acquisition
metadata.

Multi-channel data: napari-stitcher works on one channel at a time, and it names
the remaining axes from their count alone, so a channel axis left in place is
mistaken for z (or for t). Select a single channel with `--c`, which drops the
axis entirely, and stitch the channels one after another.

The tiles are added as individual napari image layers carrying `scale` (um/px)
and `translate` (um). napari-stitcher reads exactly those two attributes when you
press "Load layers", so the metadata positions become its starting guess and
phase correlation only has to refine them.

Registering and fusing are the plugin's job -- this script only loads and places.
"""

import argparse
import json
from pathlib import Path

import numpy as np
import tifffile
import xarray as xr

# ---------------------------------------------------------------------------
# metadata
# ---------------------------------------------------------------------------

# The stage->image transform is not hard-coded: every plane record carries
#   "axis"              -> which device/subdevice plays the X and Y stage axis
#   "MosaicCalibration" -> {StageCalibAngle (rad), StageCalibStaticScale,
#                           StageCalibMirror}
# which together give  image_xy = R(angle) @ (swap if mirror) @ stage_xy.
# Fitting that matrix against the TilingXIndex / TilingYIndex grid reproduces the
# recorded calibration. Residual off-diagonal terms are real: they encode the
# small angle between the stage axes and the camera axes.
# StageCalibStaticScale is NOT part of this transform; see stage_to_image_matrix.
# Stage values are in nanometres.
STAGE_UNIT_UM = 1e-3             # nm -> um


def stage_keys(entry):
    """Resolve the metadata keys that hold the X and Y stage position.

    "axis" says which subdevice plays each role, and the position itself lives
    under "<DeviceId>-<SubDeviceId>" among the flat plane keys:

        "axis": {"ExtraAxis": [],
                 "Focus": {"DeviceId": "NikonTi2", "SubDeviceId": "Focus"},
                 "X":     {"DeviceId": "NikonTi2", "SubDeviceId": "xAxis"},
                 "Y":     {"DeviceId": "NikonTi2", "SubDeviceId": "yAxis"}},
        ...
        "NikonTi2-xAxis": -13188708.64,
        "NikonTi2-yAxis": 1207816.82,

    so "X" resolves to "NikonTi2-xAxis". Roles other than X and Y (Focus, and
    the ExtraAxis list) are not tiling axes and are ignored.
    """
    axis = entry.get("axis") or {}
    keys = []
    for role in ("X", "Y"):
        a = axis.get(role)
        if not isinstance(a, dict) or "DeviceId" not in a or "SubDeviceId" not in a:
            raise SystemExit(
                "metadata 'axis' does not name a %s stage axis (got %r)" % (role, a))
        key = "%s-%s" % (a["DeviceId"], a["SubDeviceId"])
        if key not in entry:
            raise SystemExit(
                "'axis' names %s as the %s stage axis, but '%s' is not a key of the "
                "plane record" % (a, role, key))
        keys.append(key)
    return keys


def stage_value(entry, key, file_name):
    """Stage coordinate of one tile, in the metadata's own unit (nm)."""
    if key not in entry:
        raise SystemExit("no '%s' in the record of %s" % (key, file_name))
    return float(entry[key])


def stage_to_image_matrix(entry):
    """2x2 matrix mapping (stage_x, stage_y) um onto (image_x, image_y) um."""
    cal = entry.get("MosaicCalibration")
    if not cal:
        print("  ! no MosaicCalibration in metadata -- stage axes used as-is")
        return np.eye(2)
    a = float(cal["StageCalibAngle"])
    rot = np.array([[np.cos(a), -np.sin(a)], [np.sin(a), np.cos(a)]])
    swap = np.array([[0.0, 1.0], [1.0, 0.0]])
    mirror = swap if cal.get("StageCalibMirror") else np.eye(2)
    # StageCalibStaticScale is deliberately NOT applied. The stage coordinates are
    # already true micrometres, so applying it stretches every step, and the error
    # accumulates across the mosaic. Measured against phase-correlated neighbour
    # pairs, rotation + mirror alone places tiles about twice as accurately as
    # raw stage axes, and about four times as accurately as including the scale.
    return rot @ mirror


def find_raw_data(start):
    """Accept a RAW_DATA dir, its parent, or any ancestor; return the RAW_DATA dir."""
    start = Path(start).resolve()
    if (start / "metadata.json").is_file():
        return start
    for cand in sorted(start.rglob("RAW_DATA")):
        if (cand / "metadata.json").is_file():
            return cand
    raise SystemExit("No RAW_DATA/metadata.json found under %s" % start)


def read_metadata(raw_dir):
    """Collapse the per-plane metadata list into one record per position."""
    with open(raw_dir / "metadata.json") as fh:
        entries = json.load(fh)

    per_file = {}
    for e in entries:
        per_file.setdefault(e["FileName"], e)  # first plane carries the position

    e_any = next(iter(per_file.values()))
    key_x, key_y = stage_keys(e_any)

    tiles = []
    for name, e in per_file.items():
        path = raw_dir / name
        if not path.is_file():
            print("  ! listed in metadata but missing on disk, skipped: %s" % name)
            continue
        tiles.append(
            dict(
                name=name,
                path=path,
                col=e["TilingXIndex"],
                row=e["TilingYIndex"],
                pos=e["PositionIndex"],
                stage_x=stage_value(e, key_x, name) * STAGE_UNIT_UM,
                stage_y=stage_value(e, key_y, name) * STAGE_UNIT_UM,
            )
        )
    if not tiles:
        raise SystemExit("No tile files found in %s" % raw_dir)

    # One file per stage position is required. If I.S. was set to "One TIFF per
    # image", every plane is its own file and several files share a position,
    # which would otherwise load as a pile of duplicate single-plane tiles.
    by_pos = {}
    for t in tiles:
        by_pos.setdefault(t["pos"], []).append(t["name"])
    split = {p: n for p, n in by_pos.items() if len(n) > 1}
    if split:
        p, names = sorted(split.items())[0]
        raise SystemExit(
            "%d files share stage position %s (e.g. %s, %s).\n"
            "This acquisition was saved as one TIFF per image. The loader needs "
            "one file per position,\nwith the planes inside it: re-save with the "
            "'All in one file' image format in I.S." % (
                len(names), p, names[0], names[1]))

    tiles.sort(key=lambda t: t["pos"])

    e0 = per_file[tiles[0]["name"]]
    info = dict(
        pixel_um=float(e0["PixelSizeUm"]),
        z_step_um=float(e0.get("ZStack Step", 0) or 0) * 1e-3,  # nm -> um
        overlap_x=e0.get("Tiling Overlap X (%)", 0),
        overlap_y=e0.get("Tiling Overlap Y (%)", 0),
        mode=e0.get("TilingMode", "?"),
        stage_axes=(key_x, key_y),
        stage_to_image=stage_to_image_matrix(e0),
    )
    return tiles, info


# ---------------------------------------------------------------------------
# pixels
# ---------------------------------------------------------------------------

def tiff_axes_shape(path):
    """Axis labels (lower-case) and shape of the TIFF series."""
    with tifffile.TiffFile(path) as tif:
        s = tif.series[0]
        return tuple(s.axes.lower()), tuple(s.shape)


def lazy_tile(path):
    """One dask chunk per TIFF page -- a TZYX stack is never read whole."""
    from multiview_stitcher.tif_utils import tif_to_dask_plane_chunks

    return tif_to_dask_plane_chunks(path)


def tile_positions(tiles, info, source):
    """Return {name: (y_um, x_um)} with the top-left tile at the origin."""
    if source == "stage":
        stage = np.array([[t["stage_x"], t["stage_y"]] for t in tiles])
        image = stage @ info["stage_to_image"].T
        xs, ys = image[:, 0], image[:, 1]
    else:  # "index": ideal grid rebuilt from the requested overlap
        step_x = info["tile_w_um"] * (1 - info["overlap_x"] / 100.0)
        step_y = info["tile_h_um"] * (1 - info["overlap_y"] / 100.0)
        xs = np.array([t["col"] * step_x for t in tiles], dtype=float)
        ys = np.array([t["row"] * step_y for t in tiles], dtype=float)
    xs = xs - xs.min()
    ys = ys - ys.min()
    return {t["name"]: (float(y), float(x)) for t, y, x in zip(tiles, ys, xs)}


# ---------------------------------------------------------------------------
# building layers
# ---------------------------------------------------------------------------

def parse_selector(spec):
    """'3' -> 3 (drops the axis); '0:4' -> slice(0, 4); None -> slice(None)."""
    if spec is None:
        return slice(None)
    if ":" in spec:
        lo, _, hi = spec.partition(":")
        return slice(int(lo) if lo else None, int(hi) if hi else None)
    return int(spec)


def build(raw_dir, source="stage", tsel=None, zsel=None, csel=None):
    tiles, info = read_metadata(raw_dir)

    file_axes, shape = tiff_axes_shape(tiles[0]["path"])
    if "y" not in file_axes or "x" not in file_axes:
        raise SystemExit("Unexpected TIFF axes %r" % (file_axes,))
    info["tile_h_um"] = shape[file_axes.index("y")] * info["pixel_um"]
    info["tile_w_um"] = shape[file_axes.index("x")] * info["pixel_um"]

    print("  tiles       : %d  (%s)" % (len(tiles), info["mode"]))
    print("  axes/shape  : %s %s" % ("".join(file_axes).upper(), shape))
    print("  pixel size  : %.4f um" % info["pixel_um"])
    print("  tile size   : %.1f x %.1f um" % (info["tile_w_um"], info["tile_h_um"]))
    print("  overlap     : %s%% x / %s%% y" % (info["overlap_x"], info["overlap_y"]))
    print("  stage axes  : %s (x) / %s (y)" % info["stage_axes"])
    flat = " ".join(np.array2string(
        info["stage_to_image"], precision=4, suppress_small=True).split())
    print("  stage->image: %s" % flat)
    if "z" in file_axes:
        print("  z step      : %.3f um" % info["z_step_um"])

    # optional t/z/c subsetting -- an integer selector drops the axis entirely,
    # a range keeps it, so `--t 0 --z 5` turns a TZYX set into plain 2D tiles
    picks = {"t": parse_selector(tsel), "z": parse_selector(zsel),
             "c": parse_selector(csel)}
    index = tuple(picks.get(ax, slice(None)) for ax in file_axes)
    axes = tuple(ax for ax, sel in zip(file_axes, index) if not isinstance(sel, int))
    chosen = [(a, s) for a, s in zip(file_axes, index)
              if a in picks and s != slice(None)]
    if chosen:
        print("  selection   : %s -> %s" % (
            " ".join("%s=%s" % (a, s) for a, s in chosen), "".join(axes).upper()))

    if "c" in axes:
        print("  ! the channel axis is still present (%d channels). napari-stitcher"
              % shape[file_axes.index("c")])
        print("    names axes from their count alone, so it will be mistaken for")
        print("    another axis. Pass a single index to --c to pick one channel.")

    pos = tile_positions(tiles, info, source)

    # napari scale/translate follow the layer's own dimension order
    scale = []
    for ax in axes:
        if ax in ("y", "x"):
            scale.append(info["pixel_um"])
        elif ax == "z":
            scale.append(info["z_step_um"] or 1.0)
        else:  # t, and anything else non-spatial
            scale.append(1.0)
    iy, ix = axes.index("y"), axes.index("x")

    # napari-stitcher infers dimensions as ['t','z','y','x'][-ndim:] unless the
    # layer data is an xarray, so a TYX stack would be misread as ZYX. Wrapping in
    # an xarray fixes that -- but only do it when there IS a t axis: for data
    # without one the plugin does `['t'] + dims`, which raises TypeError on the
    # tuple that xarray hands back (napari-stitcher 0.1.4).
    wrap = "t" in axes

    layers = []
    for t in tiles:
        data = lazy_tile(t["path"])[index]
        if wrap:
            data = xr.DataArray(data, dims=axes)
        tr = [0.0] * len(axes)
        tr[iy], tr[ix] = pos[t["name"]]
        layers.append(
            dict(
                data=data,
                name="Pos%02d_c%dr%d" % (t["pos"], t["col"], t["row"]),
                scale=tuple(scale),
                translate=tuple(tr),
                blending="additive",
            )
        )
    return layers, tiles, info, axes


# ---------------------------------------------------------------------------
# entry points
# ---------------------------------------------------------------------------

def open_napari(layers):
    import napari

    viewer = napari.Viewer()
    for spec in layers:
        viewer.add_image(**spec)

    viewer.dims.ndisplay = 2
    viewer.reset_view()

    # dock napari-stitcher; then press Load layers -> Register -> Fuse
    try:
        viewer.window.add_plugin_dock_widget("napari-stitcher", "Stitcher")
    except Exception as exc:  # plugin missing or renamed
        print("  ! could not auto-dock the Stitcher widget (%s)" % exc)
        print("    open it manually: Plugins > napari-stitcher > Stitcher")

    napari.run()


def main():
    ap = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    ap.add_argument("path", nargs="?", default=".",
                    help="RAW_DATA folder, or any folder above it")
    ap.add_argument("--source", choices=["stage", "index"], default="stage",
                    help="tile placement from stage coordinates (default) or grid indices")
    ap.add_argument("--t", metavar="N|A:B",
                    help="keep only these timepoints; a bare index drops the t axis")
    ap.add_argument("--z", metavar="N|A:B",
                    help="keep only these z slices; a bare index drops the z axis")
    ap.add_argument("--c", metavar="N|A:B",
                    help="keep only these channels; a bare index drops the c axis. "
                         "napari-stitcher stitches one channel at a time, so pass a "
                         "single index on multi-channel data")
    args = ap.parse_args()

    raw_dir = find_raw_data(args.path)
    print("RAW_DATA: %s" % raw_dir)
    layers, _tiles, _info, _axes = build(
        raw_dir, args.source, tsel=args.t, zsel=args.z, csel=args.c)
    open_napari(layers)


if __name__ == "__main__":
    main()
