#!/usr/bin/env python
"""Check that DIALS places a TILTED detector where Jungfraujoch does.

Run under the DIALS environment:

    source /opt/dials-v3-27-0/dials_env.sh
    dials.python tests/nxmx_geometry_dials_test.py <master.h5>

Why this exists: the NXmx transformation chain the writer produces is only exercised by the rest of
CI at zero tilt, where every plausible encoding of the tilt - right axes or swapped, right
composition order or reversed, right pivot or not - gives exactly the same answer. The errors are
second order and only appear with two angles non-zero at once, or with a non-zero distance being
rotated. A previous encoding placed the detector 40 mm out at a few degrees of tilt while passing
every zero-tilt check.

The reference positions below are the same ones pinned in
tests/DiffractionGeometryTest.cpp ("DiffractionGeometry_Tilted_vs_PyFAI_and_DIALS"), which checks
them against pyFAI as well; that comment explains how to regenerate them. Keep the two in step.

The script patches the geometry into an existing master rather than requiring the writing tool to
take tilt options: the angle VALUES are just scalars, while the axis vectors, the depends_on chain
and the pivot - the parts that were wrong - stay exactly as the writer emitted them, so they remain
under test.
"""

import shutil
import sys

import h5py
import numpy as np
from dxtbx.model.experiment_list import ExperimentListFactory

# Geometry to impose. Mixed signs on purpose, and all three angles non-zero.
BEAM_X, BEAM_Y, DISTANCE_M = 1000.0, 1275.0, 0.150
ROT1, ROT2, ROT3 = 0.05, -0.03, 0.02

# Expected lab position per pixel, in the imgCIF frame DIALS reports (mm).
EXPECTED = {
    (0, 0): (-69.399544982, 98.819979800, -150.623559832),
    (1000, 1275): (7.405508016, 4.642730847, -149.745128473),
    (300, 1800): (-44.232877406, -35.676610671, -153.548927191),
    (1700, 400): (58.519164629, 71.195014535, -145.153947836),
}

# 2 um, matching the C++ test. A wrong axis, order or pivot is a millimetre-scale error.
TOLERANCE_MM = 2e-3


def main(master):
    patched = "geometry_check_master.h5"
    shutil.copy(master, patched)

    with h5py.File(patched, "r+") as h:
        det = h["/entry/instrument/detector"]
        det["beam_center_x"][()] = BEAM_X
        det["beam_center_y"][()] = BEAM_Y
        det["distance"][()] = DISTANCE_M
        transformations = det["transformations"]
        transformations["rot1"][()] = ROT1
        transformations["rot2"][()] = ROT2
        transformations["rot3"][()] = ROT3
        # The writer derives the translation from the beam centre and the distance, so it has to be
        # recomputed here - both the magnitude and the direction it points in.
        pixel_m = float(det["x_pixel_size"][()])
        vector = np.array([BEAM_X * pixel_m, BEAM_Y * pixel_m, DISTANCE_M])
        length = np.linalg.norm(vector)
        transformations["translation"][()] = length
        transformations["translation"].attrs["vector"] = vector / length

    panel = ExperimentListFactory.from_filenames([patched])[0].detector[0]
    origin = np.array(panel.get_origin())
    fast = np.array(panel.get_fast_axis())
    slow = np.array(panel.get_slow_axis())
    pixel_mm = panel.get_pixel_size()[0]

    # Deliberately NOT get_pixel_lab_coord(): it applies a parallax correction from the sensor
    # thickness and material that Jungfraujoch does not model, worth ~0.1 mm at the detector edge.
    failures = []
    for (x, y), expected in sorted(EXPECTED.items()):
        got = origin + x * pixel_mm * fast + y * pixel_mm * slow
        error = np.max(np.abs(got - np.array(expected)))
        status = "ok" if error <= TOLERANCE_MM else "FAILED"
        print(f"  pixel ({x:5d},{y:5d}): max error {error:.2e} mm  {status}")
        if error > TOLERANCE_MM:
            failures.append(f"pixel ({x},{y}): got {got}, expected {expected}")

    if failures:
        print("\nDIALS does not place the tilted detector where Jungfraujoch does:")
        for failure in failures:
            print("  " + failure)
        print("\nThe NXmx transformation chain in writer/HDF5NXmx.cpp is the thing to look at: the "
              "rot1/rot2/rot3 axis vectors, the order of the depends_on chain, and whether "
              "translation sits inside the rotations so the tilt pivots about the sample.")
        return 1

    print(f"\nAll {len(EXPECTED)} pixels agree within {TOLERANCE_MM} mm.")
    return 0


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print(__doc__)
        sys.exit(2)
    sys.exit(main(sys.argv[1]))