#  Copyright (2019-2022) Paul Scherrer Institute
#  SPDX-License-Identifier: GPL-3.0-or-later

import h5py
import sys
import gemmi


def wvl_a_to_ev(val: float) -> float:
    """
    Convert wavelength in A to photon energy in eV
    """
    return 12398.4 / val


def extract_metadata_nxmx(nxmx_file: str) -> dict:
    """
    Create a dictionary of metadata based on HDF5 master file
    """
    output = {}
    with h5py.File(nxmx_file, "r") as f:
        output["filename"] = nxmx_file
        output["photon_energy_eV"] = wvl_a_to_ev(
            f["/entry/instrument/beam/incident_wavelength"][()]
        )
        output["pixel_size_m"] = f["/entry/instrument/detector/x_pixel_size"][()]

        if "/entry/instrument/detector/underload_value" in f:
            output["underload"] = f["/entry/instrument/detector/underload_value"][()]
        output["saturation"] = f["/entry/instrument/detector/saturation_value"][()]

        output["detector_distance_m"] = f[
            "/entry/instrument/detector/detector_distance"
        ][()]
        output["beam_center_x"] = f["/entry/instrument/detector/beam_center_x"][()]
        output["beam_center_y"] = f["/entry/instrument/detector/beam_center_y"][()]

        output["width"] = f["/entry/data/data"].shape[2]
        output["height"] = f["/entry/data/data"].shape[1]

        output["image_time_s"] = f["/entry/instrument/detector/frame_time"][()]

        output["sample_name"] = f["/entry/sample/name"][()]
        if "/entry/sample/unit_cell" in f:
            output["unit_cell"] = f["/entry/sample/unit_cell"][:6]

        if "/entry/sample/space_group" in f:
            output["space_group"] = f["/entry/sample/space_group"]
    return output


def nxmx_to_geom(metadata_nxmx: dict, output_file: str) -> None:
    """
    Write dictionary generated by extract_metadata_nxmx to CrystFEL geometry file
    """
    output = {}

    output["photon_energy"] = "{:f}".format(metadata_nxmx["photon_energy_eV"])
    output["adu_per_eV"] = "{:.10f}".format(1 / metadata_nxmx["photon_energy_eV"])
    output["clen"] = "{:f}".format(metadata_nxmx["detector_distance_m"])
    output["res"] = "{:f}".format(1.0 / metadata_nxmx["pixel_size_m"])

    if "underload" in metadata_nxmx:
        output["flag_lessthen"] = metadata_nxmx["underload"] + 1

    output["rigid_group_0"] = 0
    output["rigid_group_collection_0"] = 0
    output["data"] = "/entry/data/data"
    output["dim0"] = "%"
    output["dim1"] = "ss"
    output["dim2"] = "fs"

    output["mask_file"] = metadata_nxmx["filename"]

    output["mask"] = "/entry/instrument/detector/pixel_mask"
    output["mask_good"] = "0x0"
    output["mask_bad"] = "0xffffffff"

    output["0/min_fs"] = 0
    output["0/min_ss"] = 0
    output["0/max_fs"] = metadata_nxmx["width"]
    output["0/max_ss"] = metadata_nxmx["height"]

    output["0/corner_x"] = "{:f}".format(-1.0 * metadata_nxmx["beam_center_x"])
    output["0/corner_y"] = "{:f}".format(-1.0 * metadata_nxmx["beam_center_y"])
    output["0/fs"] = "x"
    output["0/ss"] = "y"

    with open(output_file, "w") as f:
        for i in output:
            f.write("{} = {}\n".format(i, output[i]))


def nxmx_to_cell(metadata_nxmx: dict, output_file: str):
    """
    Write dictionary generated by extract_metadata_nxmx to CrystFEL cell file (PDB)
    """
    unit_cell_str = "{:9.3f}{:9.3f}{:9.3f}{:7.2f}{:7.2f}{:7.2f}".format(
        *metadata_nxmx["unit_cell"]
    )
    if "space_group" in metadata_nxmx:
        space_group_str = gemmi.find_spacegroup_by_number(
            metadata_nxmx["space_group"]
        ).hm
    else:
        space_group_str = gemmi.find_spacegroup_by_number(1).hm
    with open(output_file, "w") as f:
        f.write("CRYST1{:s} {:11s}{:4d}\n".format(unit_cell_str, space_group_str, 1))


if __name__ == "__main__":
    if (len(sys.argv) != 4) and (len(sys.argv) != 3):
        print(
            "Usage ./jungfraujoch_metadata.py <JUNGFRAU/EIGER master file> <CrystFEL output file> {<cell file>}"
        )
        sys.exit(1)
    else:
        metadata = extract_metadata_nxmx(sys.argv[1])
        nxmx_to_geom(metadata, sys.argv[2])
        if len(sys.argv) == 4:
            nxmx_to_cell(metadata, sys.argv[3])