Source code for coseda.nexus.goniometer

"""Helpers for NeXus goniometer transform chains."""

from __future__ import annotations

from typing import Iterable, Optional, Sequence, Tuple, Union
import re

import h5py
import numpy as np

from coseda.nexus.groups import ensure_nexus_parents
from coseda.nexus.paths import (
    GONIOMETER_GROUP_PATH,
    GONIOMETER_TRANSFORMATIONS_PATH,
    SAMPLE_TRANSFORM_PATH,
)

ORDER_KEY = "goniometer_transform_order"
DEFAULT_ORDER = ("stage_z", "alpha", "beta", "stage_x", "stage_y")
_VALID_KEYS = {"stage_z", "alpha", "beta", "stage_x", "stage_y"}
_TOKEN_MAP = {
    "x": "stage_x",
    "y": "stage_y",
    "z": "stage_z",
    "stagepos_x": "stage_x",
    "stagepos_x_refined": "stage_x",
    "stagepos_y": "stage_y",
    "stagepos_y_refined": "stage_y",
    "stagepos_z": "stage_z",
    "alpha": "alpha",
    "alphatilt": "alpha",
    "alpha_tilt": "alpha",
    "beta": "beta",
    "betatilt": "beta",
    "beta_tilt": "beta",
}


def _pick_dataset(h5file: h5py.File, names: Iterable[str]) -> Optional[h5py.Dataset]:
    for name in names:
        path = f"/entry/data/{name}"
        if path in h5file:
            return h5file[path]
    return None


def _link_transform(
    trans_group: h5py.Group,
    name: str,
    source: h5py.Dataset,
    transform_type: str,
    vector: Tuple[float, float, float],
    units: str,
    depends_on: str,
) -> None:
    if name in trans_group:
        del trans_group[name]
    trans_group[name] = source
    ds = trans_group[name]
    ds.attrs["transformation_type"] = transform_type
    ds.attrs["vector"] = np.array(vector, dtype=np.float32)
    if "units" not in ds.attrs:
        ds.attrs["units"] = units
    ds.attrs["depends_on"] = depends_on


def _normalize_order(order: Optional[Union[Sequence[str], str]]) -> Optional[Tuple[str, ...]]:
    if order is None:
        return None
    if isinstance(order, str):
        tokens = [t for t in re.split(r"[,\s>]+", order) if t]
    else:
        tokens = list(order)
    normalized = []
    for token in tokens:
        key = str(token).strip().lower()
        if not key:
            continue
        key = _TOKEN_MAP.get(key, key)
        if key in _VALID_KEYS and key not in normalized:
            normalized.append(key)
    return tuple(normalized) if normalized else None


def _read_order(goniometer: h5py.Group) -> Optional[Tuple[str, ...]]:
    if "transform_order" in goniometer:
        raw = goniometer["transform_order"][()]
    elif "transform_order" in goniometer.attrs:
        raw = goniometer.attrs["transform_order"]
    else:
        return None

    if isinstance(raw, (bytes, np.bytes_)):
        return _normalize_order(raw.decode("utf-8"))
    if isinstance(raw, np.ndarray):
        items = []
        for item in raw:
            if isinstance(item, (bytes, np.bytes_)):
                items.append(item.decode("utf-8"))
            else:
                items.append(str(item))
        return _normalize_order(items)
    return _normalize_order(str(raw))


def _write_order(goniometer: h5py.Group, order: Sequence[str]) -> None:
    dtype = h5py.string_dtype(encoding="utf-8")
    if "transform_order" in goniometer:
        del goniometer["transform_order"]
    goniometer.create_dataset("transform_order", data=np.array(order, dtype=object), dtype=dtype)
    goniometer.attrs["transform_order"] = ", ".join(order)




def get_goniometer_transform_order(config) -> Optional[Tuple[str, ...]]:
    for section in ("AcquisitionDetails", "Parameters"):
        try:
            if config.has_option(section, ORDER_KEY):
                return _normalize_order(config.get(section, ORDER_KEY))
        except Exception:
            continue
    return None





def ensure_goniometer_transforms(
    h5file: h5py.File,
    transform_order: Optional[Union[Sequence[str], str]] = None,
) -> None:
    """
    Create a goniometer transform chain:
    stage_z -> alpha -> beta -> stage_x -> stage_y (all relative to sample_pos).
    """
    ensure_nexus_parents(h5file)

    goniometer = h5file.require_group(GONIOMETER_GROUP_PATH.lstrip("/"))
    goniometer.attrs["NX_class"] = "NXgoniometer"
    trans = h5file.require_group(GONIOMETER_TRANSFORMATIONS_PATH.lstrip("/"))
    trans.attrs["NX_class"] = "NXtransformations"

    order = _normalize_order(transform_order)
    if order:
        _write_order(goniometer, order)
    else:
        order = _read_order(goniometer) or DEFAULT_ORDER

    chain = []

    for step in order:
        if step == "stage_z":
            stage_z = _pick_dataset(h5file, ("stagepos_z",))
            if stage_z is not None:
                chain.append(("stage_z", stage_z, "translation", (0.0, 0.0, 1.0), "m"))
        elif step == "alpha":
            alpha = _pick_dataset(h5file, ("alphatilt",))
            if alpha is not None:
                chain.append(("alpha", alpha, "rotation", (1.0, 0.0, 0.0), "degree"))
        elif step == "beta":
            beta = _pick_dataset(h5file, ("betatilt",))
            if beta is not None:
                chain.append(("beta", beta, "rotation", (0.0, 1.0, 0.0), "degree"))
        elif step == "stage_x":
            stage_x = _pick_dataset(h5file, ("stagepos_x_refined", "stagepos_x"))
            if stage_x is not None:
                chain.append(("stage_x", stage_x, "translation", (1.0, 0.0, 0.0), "m"))
        elif step == "stage_y":
            stage_y = _pick_dataset(h5file, ("stagepos_y_refined", "stagepos_y"))
            if stage_y is not None:
                chain.append(("stage_y", stage_y, "translation", (0.0, 1.0, 0.0), "m"))

    if not chain:
        return

    depends_on = SAMPLE_TRANSFORM_PATH
    for name, source, ttype, vector, units in chain:
        _link_transform(trans, name, source, ttype, vector, units, depends_on)
        depends_on = f"{GONIOMETER_TRANSFORMATIONS_PATH}/{name}"

    sample = h5file["entry/sample"]
    sample.attrs["depends_on"] = depends_on