conftest.py

from datetime import datetime, timezone
from pathlib import Path
from shutil import copytree
from tempfile import TemporaryDirectory

import h5py
import numpy as np
import pytest
from extra_data import RunDirectory
from extra_data.tests import make_examples
from extra_data.tests.mockdata import write_file

from extra_data.tests.mockdata.motor import Motor
from .mockdata.adq import AdqDigitizer
from .mockdata.detector_motors import (DetectorMotorDataSelector,
                                       get_motor_sources,
                                       write_motor_positions)
from .mockdata.dld import ReconstructedDld
from .mockdata.timepix import Timepix3Receiver, Timepix3Centroids
from .mockdata.timeserver import PulsePatternDecoder, Timeserver
from .mockdata.xgm import XGM, XGMD, XGMReduced, XGMWithData
from .mockdata.mono import MonoMdl
from .mockdata.camera import CameraWithData, GotthardIIWithData


@pytest.fixture(scope='session')
def mock_spb_aux_directory():
    """Mock run directory with SPB auxiliary sources.

    Pulse pattern per train:
        - 0:5, no pulses
        - SA1
            - 10:50, 50 pulses at 1000:1300:6
            - 50:100, 25 pulses at 1000:1300:12
        - SA2
            - 10:100, 62 pulses at 1500:2000:8
        - SA3
            - 10:100, 1 pulse at 200
        - LP_SPB
            - 5:100, 50 pulses at 0:300:6
    """

    sources = [
        Timeserver('SPB_RR_SYS/TSYS/TIMESERVER'),
        PulsePatternDecoder('SPB_RR_SYS/MDL/BUNCH_PATTERN'),
        Timeserver('ODD_TIMESERVER_NAME'),
        PulsePatternDecoder('TRAIN_LESS_DECODER', no_ctrl_data=True),
        Timeserver('TRAIN_LESS_TIMESERVER', no_ctrl_data=True, nsamples=0),
        Timeserver('PULSE_LESS_TIMESERVER', no_pulses=True),
        XGM('SPB_XTD9_XGM/XGM/DOOCS'),
        Motor("MOTOR/MCMOTORYFACE"),
        DetectorMotorDataSelector("SPB_IRU_AGIPD1M/DS", "SPB_IRU_AGIPD1M"),
        DetectorMotorDataSelector("SPB_EXP_AGIPD1M2/DS", "SPB_EXP_AGIPD1M2"),
        Motor('SPB_IRDA_JF4M/MOTOR/X1'),
        Motor('SPB_IRDA_JF4M/MOTOR/X2'),
    ]
    sources += get_motor_sources("SPB_IRU_AGIPD1M")

    with TemporaryDirectory() as td:
        path = Path(td) / 'RAW-R0001-DA01-S00000.h5'
        write_file(path, sources, 100, format_version='1.2')
        with h5py.File(path, 'a') as f:
            motor_ds = f['CONTROL/MOTOR/MCMOTORYFACE/actualPosition/value']
            # Simulate a scan of 10 steps, with intermediate positions for
            # 1 train at each transition between steps.
            motor_ds[:] = np.repeat(np.arange(10), 10)
            motor_ds[10::10] = np.arange(9) + 0.5
            motor_target_ds = f['CONTROL/MOTOR/MCMOTORYFACE/targetPosition/value']
            motor_target_ds[:] = np.repeat(np.arange(10), 10)

            # write agipd quadrand motor positions
            write_motor_positions(f, "SPB_IRU_AGIPD1M")
            # write jf4m halves motor positions
            jfx1 = f['CONTROL/SPB_IRDA_JF4M/MOTOR/X1/actualPosition/value']
            jfx1[:] = 10
            jfx2 = f['CONTROL/SPB_IRDA_JF4M/MOTOR/X2/actualPosition/value']
            jfx2[:10] = 5
            jfx2[10:] = 6

        yield td


def _patch_train_timestamps(run):
    ts = datetime(year=2026, month=4, day=10, hour=11, minute=27) \
        .replace(tzinfo=timezone.utc).timestamp() * 10**9

    for fa in run.files:
        fa.close()
        with h5py.File(fa.filename, 'a') as h5f:
            dset = h5f['INDEX/timestamp']
            dset[:] = [ts + i * 10**8 for i in range(dset.shape[0])]

@pytest.fixture(scope='function')
def mock_spb_aux_run(mock_spb_aux_directory):
    yield RunDirectory(mock_spb_aux_directory)


@pytest.fixture(scope='session')
def mock_agipd1m_directory():
    with TemporaryDirectory() as td:
        make_examples.make_agipd1m_run(td)
        yield td


@pytest.fixture(scope='function')
def mock_agipd1m_run(mock_agipd1m_directory):
    run = RunDirectory(mock_agipd1m_directory)
    _patch_train_timestamps(run)
    yield run


@pytest.fixture(scope='session')
def mock_legacy_agipd1m_directory():
    # No gain.value or bunchStructure.repetitionRate in /CONTROL/
    with TemporaryDirectory() as td:
        make_examples.make_agipd1m_run(
            td, rep_rate=False,
            gain_setting=False,
            integration_time=False,
            bias_voltage=True
        )
        yield td


@pytest.fixture(scope='function')
def mock_legacy_agipd1m_run(mock_legacy_agipd1m_directory):
    run = RunDirectory(mock_legacy_agipd1m_directory)
    _patch_train_timestamps(run)
    yield run


@pytest.fixture(scope='session')
def mock_agipd500k_directory():
    # No gain.value or bunchStructure.repetitionRate in /CONTROL/
    with TemporaryDirectory() as td:
        make_examples.make_agipd500k_run(td)
        yield td


@pytest.fixture(scope='function')
def mock_agipd500k_run(mock_agipd500k_directory):
    run = RunDirectory(mock_agipd500k_directory)
    _patch_train_timestamps(run)
    yield run


@pytest.fixture(scope="session")
def multi_xgm_directory():
    sources = [
        XGM("SA2_XTD1_XGM/XGM/DOOCS"),
        XGMD("SPB_XTD9_XGM/XGM/DOOCS"),
        XGMReduced("SQS_DIAG1_XGMD/XGM/DOOCS"),
        XGM("NON_ACTUALIZED_XGM/XGM/DOOCS", main_nbunches_property="numberOfBunches"),
        XGM("HOBBIT_XGM/XGM/DOOCS", main_nbunches_property="nummberOfBrunches")
    ]

    with TemporaryDirectory() as td:
        # We need format version 1.1 for the XGM tests, because without the full
        # run metadata we can't get RUN values after a .union() or .select().
        write_file(Path(td) / 'RAW-R0002-DA01-S00000.h5', sources, 100,
                   format_version="1.1", firsttrain=20000)
        yield td


@pytest.fixture(scope="function")
def multi_xgm_run(multi_xgm_directory):
    aliases = {"sa2-xgm": "SA2_XTD1_XGM/XGM/DOOCS"}
    yield RunDirectory(multi_xgm_directory).with_aliases(aliases)


@pytest.fixture(scope='session')
def mock_sqs_remi_directory():
    sources = [
        Timeserver('SQS_RR_UTC/TSYS/TIMESERVER'),
        PulsePatternDecoder('SQS_RR_UTC/TSYS/PP_DECODER'),
        XGM('SA3_XTD10_XGM/XGM/DOOCS'),
        ReconstructedDld('SQS_REMI_DLD6/DET/TOP'),
        ReconstructedDld('SQS_REMI_DLD6/DET/BOTTOM'),
        Motor('SQS_ILH_LAS/MOTOR/DELAY_AX_800'),
        AdqDigitizer('SQS_DIGITIZER_UTC1/ADC/1', channels_per_board=2 * [4]),
        AdqDigitizer('SQS_DIGITIZER_UTC2/ADC/1', channels_per_board=4 * [4],
                     data_channels={(0, 0), (2, 1)})]

    with TemporaryDirectory() as td:
        write_file(Path(td) / 'RAW-R0001-DA01-S00000.h5', sources, 100,
                   format_version='1.2')
        yield td


@pytest.fixture(scope='function')
def mock_sqs_remi_run(mock_sqs_remi_directory):
    yield RunDirectory(mock_sqs_remi_directory)


@pytest.fixture(scope='session')
def mock_sqs_timepix_directory():
    sources = [
        Timeserver('SQS_RR_UTC/TSYS/TIMESERVER'),
        Timepix3Receiver('SQS_EXTRA_TIMEPIX/DET/TIMEPIX3'),
        Timepix3Receiver('SQS_EXP_TIMEPIX/DET/TIMEPIX3'),
        Timepix3Centroids('SQS_EXP_TIMEPIX/CAL/TIMEPIX3')]

    with TemporaryDirectory() as td:
        write_file(Path(td) / 'RAW-R0001-DA01-S00000.h5', sources, 100)
        yield td


@pytest.fixture(scope='function')
def mock_sqs_timepix_run(mock_sqs_timepix_directory):
    yield RunDirectory(mock_sqs_timepix_directory)


@pytest.fixture(scope='function')
def mock_timepix_exceeded_buffer_run(mock_sqs_timepix_directory):
    with TemporaryDirectory() as td:
        copytree(mock_sqs_timepix_directory, td, dirs_exist_ok=True)

        with h5py.File(next(Path(td).glob('*.h5')), 'r+') as f:
            tpx_root = f['INSTRUMENT/SQS_EXP_TIMEPIX/DET/TIMEPIX3'
                         ':daqOutput.chip0']

            size_dset = tpx_root['data/size']
            size_dset[np.argmax(size_dset)] += tpx_root['data/x'].shape[1]

        yield RunDirectory(td).deselect('SQS_EXTRA*')


@pytest.fixture(scope='session')
def mock_etof_calibration_constants():
    yield (623419.734, 946.026, 11.527)


@pytest.fixture(scope='session')
def mock_etof_mono_energies():
    # 200 trains with 10 energies and 20 trains per energy
    energy = list()
    for e in np.linspace(970.0, 1060.0, 10):
        energy += [e]*20
    energy = np.array(energy)
    yield energy


@pytest.fixture(scope='session')
def mock_sqs_etof_calibration_directory(mock_etof_mono_energies, mock_etof_calibration_constants):
    energy = mock_etof_mono_energies

    # convert energy to time of flight for etofs
    # calibration constants
    c, e0, t0 = mock_etof_calibration_constants
    sigma = 2.0
    A = 1000.0
    Aa = 500.0
    auger = 35.0
    offset = 44.0
    # e = e0+c/(ts-t0)**2
    # ts = t0 + sqrt(c/(e - e0))
    ts = t0 + np.sqrt(c/(energy - e0))
    ts_axis = np.linspace(0.0, int(np.max(ts))+5*sigma+1, int(np.max(ts)+5*sigma)+1+1)
    # create gaussians
    samples = (A*np.exp(-0.5*(ts[:, None] - ts_axis[None, :])**2/(sigma**2))
               + Aa*np.exp(-0.5*(auger - ts_axis[None, :])**2/(sigma**2))
               )
    samples += offset
    # add some samples before data
    samples = -1*np.concatenate((
                                 np.zeros((samples.shape[0], 1000)),            # samples before trigger
                                 samples,                                       # data
                                 np.zeros((samples.shape[0], 3000)),            # samples after
                                 ), axis=-1)
    # Use a fixed seed to make the random data deterministic.
    rng = np.random.default_rng(12345)
    samples += rng.standard_normal(samples.shape)

    sources = [
        Timeserver('SQS_RR_UTC/TSYS/TIMESERVER'),
        XGMWithData('SQS_DIAG1_XGMD/XGM/DOOCS', intensity=(
                rng.standard_normal(energy.shape + (1,)) + 1000
        )),
        MonoMdl('SA3_XTD10_MONO/MDL/PHOTON_ENERGY', energy_data=energy),
        AdqDigitizer('SQS_DIGITIZER_UTC4/ADC/1', channels_per_board=[4],
                     data_channels={(0, 0)},
                     samples=samples)
              ]

    # for tests
    #td = Path("mytest")
    #write_file(Path(td) / 'RAW-R0001-DA01-S00000.h5', sources, 200,
    #           format_version='1.2')
    with TemporaryDirectory() as td:
        write_file(Path(td) / 'RAW-R0001-DA01-S00000.h5', sources, 200,
                   format_version='1.2')
        yield td


@pytest.fixture(scope='session')
def mock_sqs_grating_calibration_directory():
    #energy = mock_grating_mono_energies()
    energy = list()
    for e in np.linspace(992.0, 1008.0, 10):
        energy += [e]*20
    energy = np.array(energy)

    pix = np.linspace(0, 1000, 1000)

    # convert energy to time of flight for etofs
    # calibration constants
    #offset, slope = mock_grating_calibration_constants()
    offset, slope = 990, 20/1000.0
    true_pix_to_e = lambda p: offset + p*slope
    true_e_to_pix = lambda e: (e - offset)/slope
    data = np.exp(-0.5*(pix[None, :] - true_e_to_pix(energy)[:,None])**2)
    # make it 2D
    data2d = np.stack([data]*10, axis=1)
    # add dimension for pulses in Gotthard
    data_gh = np.stack([np.zeros_like(data)]*2+[data]*10, axis=-2)

    sources = [
        Timeserver('SQS_RR_UTC/TSYS/TIMESERVER'),
        MonoMdl('SA3_XTD10_MONO/MDL/PHOTON_ENERGY', energy_data=energy),
        CameraWithData("SQS_DIAG3_BIU/CAM/CAM_6", data=data2d),
        GotthardIIWithData("SQS_EXP_GH2-2/CORR/RECEIVER", data=data_gh,
                           mask=np.zeros_like(data_gh),
                           motor=np.zeros(data_gh.shape[0])),
              ]

    # for tests
    #td = Path("mytest_gr")
    #write_file(Path(td) / 'RAW-R0001-DA01-S00000.h5', sources, 200,
    #           format_version='1.2')
    with TemporaryDirectory() as td:
        write_file(Path(td) / 'RAW-R0001-DA01-S00000.h5', sources, 200,
                   format_version='1.2')
        yield td


@pytest.fixture(scope='function')
def mock_sqs_etof_calibration_run(mock_sqs_etof_calibration_directory):
    yield RunDirectory(mock_sqs_etof_calibration_directory)


@pytest.fixture(scope='function')
def mock_sqs_grating_calibration_run(mock_sqs_grating_calibration_directory):
    yield RunDirectory(mock_sqs_grating_calibration_directory)