#
# Copyright The NOMAD Authors.
#
# This file is part of NOMAD. See https://nomad-lab.eu for further info.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""Reader for FAU/Erlangen's HDF5-based formats introduced with the pyccapt library."""

# pylint: disable=fixme

import os
from typing import Any

import h5py
import numpy as np
import pandas as pd
from ase.data import atomic_numbers

# from ifes_apt_tc_data_modeling.utils.combinatorics import apply_combinatorics
from ifes_apt_tc_data_modeling.utils.custom_logging import logger
from ifes_apt_tc_data_modeling.utils.molecular_ions import get_chemical_symbols
from ifes_apt_tc_data_modeling.utils.nx_ion import NxIon
from ifes_apt_tc_data_modeling.utils.pint_custom_unit_registry import ureg
from ifes_apt_tc_data_modeling.utils.utils import (
    MAX_NUMBER_OF_ATOMS_PER_ION,
    isotope_to_hash,
    nuclide_hash_to_nuclide_list,
)

# this implementation focuses on the following state of the pyccapt repository
# https://github.com/mmonajem/pyccapt/commit/e955beb4f2627befb8b4d26f2e74e4c52e00394e

# during the course of an atom probe measurement and analysis with FAU/Erlangen's Oxcart instrument
# several HDF5 files are generated with essentially two software tools. One is pyccapt which has a
# a control module plus a calibration module (whereby the voltage/bowl calibration and reconstruction
# is performed), the other one is a module/functionality to document ranging i.e. made ion type
# definitions. These results can be processed further with the FAU/Erlangen Matlab Atom Probe Toolbox;
# instructed as a set of Matlab live scripts this toolbox offers data analysis functionalities.
# Results obtained with this Atom Probe Toolbox are stored via an HDF5 file but currently not parsed
# by this pyccapt parser

# specific comments
# pyccapt/control
# an HDF5 file keeping relevant quantities

# pyccapt/calibration
# unfortunately the generated HDF5 file has internally no provenance information
# with which pyccapt version it was generated. Therefore, developers of pyccapt should rather
# write the content of the HDF5 file explicitly dataset by dataset e.g. using h5py instead
# of the pandas HDF5 dump convenience functionality
# of course pandas stores its own version but that is not conclusive enough to infer with
# which pyccapt version and most importantly from which other context the file was generated
# this is an aspect of the FAIR RDM principles which the pyccapt approach currently ignores
SUPPORTED_PYCCAPT_VERSION = "e955beb4f2627befb8b4d26f2e74e4c52e00394e"


def get_nuclide_hash_from_fau_list(elements, complexs, isotopes) -> np.ndarray:
    """Compute nuclide_hash from specific representation used at FAU/Erlangen."""
    # TODO:: add raise ValueError checks
    ivec = np.zeros((MAX_NUMBER_OF_ATOMS_PER_ION,), np.uint16)
    hash_vector: list = []
    for idx in np.arange(0, len(elements)):
        symbol = elements[idx]
        if symbol in get_chemical_symbols():
            proton_number = atomic_numbers[symbol]
            neutron_number = isotopes[idx] - proton_number
            hash_vector.extend(
                [isotope_to_hash(proton_number, neutron_number)] * complexs[idx]
            )
    ivec[0 : len(hash_vector)] = np.sort(
        np.asarray(hash_vector, np.uint16), kind="stable"
    )[::-1]
    return ivec


class ReadPyccaptControlFileFormat:
    """Read FAU/Erlangen pyccapt (control module) HDF5 file format."""

    def __init__(self, file_path: str, verbose: bool = False):
        self.supported = 0  # voting-based
        if not file_path.lower().endswith((".h5", ".hdf5")):
            logger.warning(
                f"{file_path} is likely not an HDF5 file as generated by pyccapt control"
            )
            return
        self.file_path = file_path
        self.verbose = verbose
        self.file_size = os.path.getsize(self.file_path)
        self.number_of_events = None
        self.version = SUPPORTED_PYCCAPT_VERSION

        # check that the formatting matches that of an pyccapt control module output HDF5 file
        with h5py.File(self.file_path, "r") as h5r:
            required_group_names = ["apt", "dld", "tdc"]
            for required_group_name in required_group_names:
                if required_group_name in h5r.keys():
                    self.supported += 1
            if self.supported == 3:
                logger.debug(
                    f"{self.file_path} is a supported pyccapt/control HDF5 file."
                )
            else:
                logger.warning(
                    f"{self.file_path} is not a supported pyccapt/control HDF5 file."
                )
                return
        # parse out relevant pieces of information


class ReadPyccaptCalibrationFileFormat:
    """Read FAU/Erlangen pyccapt (calibration module) HDF5 file format."""

    def __init__(self, file_path: str, verbose: bool = False):
        self.supported = 0  # voting-based
        if not file_path.lower().endswith((".h5", ".hdf5")):
            logger.warning(
                f"{file_path} is likely not a HDF5 as generated by pyccapt calibration"
            )
            return
        self.file_path = file_path
        self.verbose = verbose
        self.file_size = os.path.getsize(self.file_path)
        self.number_of_events = None
        self.version = SUPPORTED_PYCCAPT_VERSION

        with h5py.File(self.file_path, "r") as h5r:
            required_entries = [
                "df",
                "df/axis0",
                "df/axis1",
                "df/block0_items",
                "df/block0_values",
                "df/block1_items",
                "df/block1_values",
            ]
            for entry in required_entries:
                if entry in h5r.keys():
                    self.supported += 1
            if self.supported == 7:
                logger.debug(
                    f"{self.file_path} is a supported pyccapt/calibration HDF5 file."
                )
            else:
                logger.warning(
                    f"{self.file_path} is not a supported pyccapt/calibration HDF5 file."
                )
                return

        self.df = pd.read_hdf(self.file_path)
        self.number_of_events = len(self.df)

    def get_named_quantities(self, term: str):
        """Get named quantities from dataframe."""
        # see https://pyccapt.readthedocs.io/en/latest/Calibration_DATA_STRUCTURE.html
        # for the semantics, typically float64 are returned by pyccapt !!
        if term in self.df.keys():
            return self.df[term]
        return None

    def get_reconstructed_positions(self):
        """Read xyz columns."""
        values = np.zeros((self.number_of_events, 3), np.float32)
        all_values = True
        for dim, prefix in enumerate(["x", "y", "z"]):
            found = False
            for typo in [
                f"{prefix.capitalize()} (nm)",
                f"{prefix.capitalize()}(nm)",
                f"{prefix} (nm)",
                f"{prefix}(nm)",
                f"{prefix.capitalize()} (cm)",
                f"{prefix.capitalize()}(cm)",
                f"{prefix} (cm)",
                f"{prefix}(cm)",
            ]:  # several typos across versions, pick the first found
                if typo in self.df:
                    data = self.get_named_quantities(typo)
                    if data is not None:
                        np.copyto(values[:, dim], data, casting="unsafe")
                        found = True
                    else:
                        all_values = False
                        logger.warning(
                            f"Unable to get_reconstructed_positions dim {dim}"
                        )
                    break
            if not found:
                all_values = False
                logger.warning(f"Unable to get_reconstructed_positions dim {dim}")
        if all_values:
            return ureg.Quantity(values, ureg.nanometer)
        logger.warning("Unable to get_reconstructed_positions")

    def get_mass_to_charge_state_ratio(self):
        """Read (calibrated) mass-to-charge-state-ratio column."""
        values = np.zeros((self.number_of_events,), np.float32)
        for typo in ["mc (Da)", "mc(Da)"]:
            if typo in self.df:
                data = self.get_named_quantities(typo)
                if data is not None:
                    np.copyto(values[:], data, casting="unsafe")
                    return ureg.Quantity(values, ureg.dalton)
                break
        logger.warning("Unable to get_mass_to_charge_state_ratio")

    def get_standing_voltage(self):
        """Read high voltage mapping it to the standing voltage."""
        values = np.zeros((self.number_of_events,), np.float32)
        for typo in ["high-voltage (V)", "high_voltage (V)"]:
            if typo in self.df:
                data = self.get_named_quantities(typo)
                if data is not None:
                    np.copyto(values[:], data, casting="unsafe")
                    return ureg.Quantity(values, ureg.volt)
                break
        logger.warning("Unable to get_standing_voltage")

    def get_pulse_voltage(self):
        """Read pulse mapping it to pulse voltage."""
        # !! the concept can either be a voltage or a laser energy !!
        # in the past FAU-Erlangen-Nürnberg OXCART instrument had no laser
        # TODO how to know whether voltage or laser run?
        values = np.zeros((self.number_of_events,), np.float32)
        data = self.get_named_quantities("pulse")
        if data is not None:
            np.copyto(values[:], data, casting="unsafe")
            return ureg.Quantity(values, ureg.volt)
        logger.warning("Unable to get_pulse_voltage")

    def get_raw_time_of_flight(self):
        """Read uncalibrated time of flight."""
        values = np.zeros((self.number_of_events,), np.float32)
        for typo in ["t(ns)", "t (ns)"]:
            if typo in self.df:
                data = self.get_named_quantities(typo)
                if data is not None:
                    np.copyto(values[:], data, casting="unsafe")
                    return ureg.Quantity(values, ureg.nanosecond)
                break
        logger.warning("Unable to get_raw_time_of_flight")

    def get_calibrated_time_of_flight(self):
        """Read bowl and voltage calibrated time of flight."""
        values = np.zeros((self.number_of_events,), np.float32)
        for typo in ["t_c(ns)", "t_c (ns)"]:
            if typo in self.df:
                data = self.get_named_quantities(typo)
                if data is not None:
                    np.copyto(values[:], data, casting="unsafe")
                    return ureg.Quantity(values, ureg.nanosecond)
                break
        logger.warning("Unable to get_calibrated_time_of_flight")

    def get_detector_hit_positions(self):
        """Read (calibrated) hit positions on the detector."""
        values = np.zeros((self.number_of_events, 2), np.float32)
        all_values = True
        for dim, prefix in enumerate(["x", "y"]):
            found = False
            for typo in [f"{prefix}_det (cm)", f"{prefix}_det(cm)"]:
                if typo in self.df:
                    data = self.get_named_quantities(typo)
                    if data is not None:
                        np.copyto(values[:, dim], data, casting="unsafe")
                        found = True
                    else:
                        all_values = False
                    break
            if not found:
                all_values = False
                logger.warning(f"Unable to get_detector_hit_positions dim {dim}")
        if all_values:
            return ureg.Quantity(values, ureg.centimeter).to(ureg.millimeter)
        logger.warning("Unable to get_detector_hit_positions")


class ReadPyccaptRangingFileFormat:
    """Read FAU/Erlangen pyccapt (ranging module) HDF5 file format."""

    def __init__(self, file_path: str, verbose: bool = False):
        self.supported = 0  # voting-based
        if not file_path.lower().endswith((".h5", ".hdf5")):
            logger.warning(
                f"{file_path} is likely not a HDF5 as generated by pyccapt ranging"
            )
            return
        self.file_path = file_path
        self.verbose = verbose
        self.file_size = os.path.getsize(self.file_path)
        self.number_of_events = None
        self.version = SUPPORTED_PYCCAPT_VERSION
        self.df = None

        with h5py.File(self.file_path, "r") as h5r:
            required_entries = [
                "df",
                "df/axis0",
                "df/axis1",
                "df/block0_items",
                "df/block0_values",
                "df/block1_items",
                "df/block1_values",
                "df/block2_items",
                "df/block2_values",
            ]
            for entry in required_entries:
                if entry in h5r.keys():
                    self.supported += 1
            if self.supported == 9:
                logger.debug(
                    f"{self.file_path} is a supported pyccapt/ranging HDF5 file."
                )
            else:
                logger.warning(
                    f"{self.file_path} is not a supported pyccapt/ranging HDF5 file."
                )
                return

        self.df = pd.read_hdf(self.file_path)
        self.rng: dict[str, Any] = {}
        self.rng["molecular_ions"] = []
        logger.debug(np.shape(self.df)[0])
        for idx in map(int, np.arange(0, np.shape(self.df)[0])):
            if isinstance(self.df.iat[idx, 6], str):  # type: ignore[index]
                if self.df.iat[idx, 6] == "unranged":  # type: ignore[index]
                    continue

            ivec = get_nuclide_hash_from_fau_list(
                elements=self.df.iat[idx, 6],  # type: ignore[index]
                complexs=self.df.iat[idx, 7],  # type: ignore[index]
                isotopes=self.df.iat[idx, 8],  # type: ignore[index]
            )
            m_ion = NxIon()
            m_ion.nuclide_hash = ivec
            m_ion.nuclide_list = nuclide_hash_to_nuclide_list(ivec)
            m_ion.charge_state = np.int8(self.df.iat[idx, 9])  # type: ignore[arg-type,index]
            m_ion.add_range(self.df.iat[idx, 3], self.df.iat[idx, 4])  # type: ignore[arg-type,index]
            m_ion.apply_combinatorics()
            # m_ion.report()
            self.rng["molecular_ions"].append(m_ion)
        logger.info(f"{self.file_path} parsed successfully.")
