"""
Parser classes for the Soft X-ray (SXR) beamline at the Australian Synchrotron.
"""

from pyNexafs.parsers import parserBase
from pyNexafs.types import parse_fn_ret_type
from pyNexafs.utils.mda import MDAFileReader
from typing import Any
import typing
from numpy.typing import NDArray
import io
import numpy as np
from typing import override


class SXR_NEXAFS(parserBase):
    """
    Australian Synchrotron Soft X-ray (SXR) NEXAFS parser.

    Parses data formats including '.asc' and '.mda' formats from the SXR
    Near Edge X-ray Absorption Fine Structure (NEXAFS) tool.

    Attributes
    ----------
    ALLOWED_EXTENSIONS
    COLUMN_ASSIGNMENTS
    SUMMARY_PARAM_RAW_NAMES
    RELABELS

    Notes
    -----
    Implemented for data as of 2024-Mar.
    """

    ALLOWED_EXTENSIONS = [
        ".asc",  # Ascii files exported from mda2asc tool
        ".mda",  # Binary MDA files from EPICS/Motor scans
        ".txt",  # Files generated by QANT software
    ]  # ascii files exported from the binary .mda files generated.

    COLUMN_ASSIGNMENTS = {
        "x": "SR14ID01PGM_CALC_ENERGY_MONITOR.P",
        "y": [
            # "SR14ID01PGM:LOCAL_SP",
            "SR14ID01IOC68:scaler1.S20",  # Drain Current VF
            "SR14ID01IOC68:scaler1.S17",  # Direct PHD VF
            "SR14ID01IOC68:scaler1.S18",  # I0 VF
            "SR14ID01IOC68:scaler1.S19",  # Ref Foil VF
            "SR14ID01IOC68:scaler1.S21",  # Channeltron Front (PEY)
            "SR14ID01IOC68:scaler1.S22",  # MCP (TFY)
            "SR14ID01IOC68:scaler1.S23",  # Hemispherical Analyser (AEY)
            # No longer included in files...?
            # "SR14ID01IOC68:scaler1.S3", # I0 VF
            # "SR14ID01IOC68:scaler1.S4", # Ref Foil VF
            # "SR14ID01IOC68:scaler1.S8", # Direct PHD VF
        ],
        "y_errs": None,
        "x_errs": None,
    }

    SUMMARY_PARAM_RAW_NAMES = [
        "SR14ID01NEXSCAN:saveData_comment1",
        "SR14ID01NEXSCAN:saveData_comment2",
        "SR14ID01NEX01:R_MTR.RBV",
        "SR14ID01NEX01:RULER_ID",
        "SR14ID01:BRANCH_MODE",
        "SR14ID01NEX01:C_MTR.RBV",
        "SR14ID01NEX01:Z_MTR.RBV",
        "SR14ID01NEX01:X_MTR.RBV",
        "SR14ID01NEX01:Y_MTR.RBV",
    ]

    RELABELS = {
        "SR14ID01PGM:REMOTE_SP": "Photon Energy",
        "SR14ID01PGM:LOCAL_SP": "Local Energy Setpoint",
        "SR14ID01PGM_ENERGY_SP": "Energy Setpoint",
        "SR14ID01PGM_CALC_ENERGY_MONITOR.P": "Encoder Photon Energy",
        (
            "SR14ID01IOC68:scaler1.TPSR14ID01DTACQ06:sample_time",
            "SR14ID01DTACQ07:sample_time",
        ): "Exp Time",
        (
            "SR14ID01IOC68:scaler1.S2",
            "SR14ID01IOC68:scaler1.S20",
            "SR14ID01DTACQ06:ch4_v2f_ave_rbv",
            "SR14ID01DTACQ07:ch3_v2f_ave_rbv",
            "Drain Current VF",
        ): "Drain Current",  # Previously labelled with a VF
        (
            "SR14ID01IOC68:scaler1.S3",
            "SR14ID01IOC68:scaler1.S18",
            "SR14ID01DTACQ06:ch3_v2f_ave_rbv",  # Izero-VF
            "SR14ID01DTACQ07:ch2_v2f_ave_rbv",
            "I0 VF",
        ): "I0",
        (
            "SR14ID01IOC68:scaler1.S4",
            "SR14ID01IOC68:scaler1.S19",
            "SR14ID01DTACQ06:ch2_v2f_ave_rbv",
            "SR14ID01DTACQ07:ch1_v2f_ave_rbv",
            "Ref Foil VF",
        ): "Ref Foil",
        (
            "SR14ID01IOC68:scaler1.S6",
            "SR14ID01IOC68:scaler1.S22",
            "SR14ID01DTACQ06:ch1_counts_rbv",
            "SR14ID01DTACQ07:ch2_counts_rbv",
            "MCP (TFY)",
        ): "Micro Channel Plate (TFY)",
        (
            "SR14ID01IOC68:scaler1.S8",
            "SR14ID01IOC68:scaler1.S17",
            "SR14ID01DTACQ06:ch5_v2f_ave_rbv",
            "SR14ID01DTACQ07:ch4_v2f_ave_rbv",
            "Direct PHD VF",
        ): "Direct PHD",
        (
            "SR14ID01IOC68:scaler1.S9",
            "SR14ID01DTACQ06:ch1_v2f_ave_rbv",
            "BL PHD VF",
        ): "BL PhD",
        (
            "SR14ID01IOC68:scaler1.S10",
            "SR14ID01DTACQ07:ch1_counts_rbv",
        ): "Channeltron Front (PEY)",
        (
            "SR14ID01IOC68:scaler1.S21",
            "SR14ID01DTACQ06:ch2_counts_rbv",
        ): "Channeltron",
        (
            "SR14ID01IOC68:scaler1.S11",
            "TFT PHD VF",
        ): "TFY PHD",
        (
            "SR14ID01IOC68:scaler1.S23",
            "SR14ID01DTACQ07:ch3_counts_rbv",
            "Hemispherical Analyser AEY",
        ): "Hemispherical Analyser (AEY)",
        "SR14ID01AMP01:CURR_MONITOR": "Drain Current - Keithley1",
        "SR14ID01AMP02:CURR_MONITOR": "BL PHD - Keithley2",
        "SR14ID01AMP03:CURR_MONITOR": "I0 - Keithley3",
        "SR14ID01AMP04:CURR_MONITOR": "Ref Foil -Keithley4",
        "SR14ID01AMP05:CURR_MONITOR": "Direct PHD - Keithley5",
        "SR14ID01AMP06:CURR_MONITOR": "Keithley6",
        "SR14ID01AMP07:CURR_MONITOR": "Ref Foil - Keithley7",
        "SR14ID01AMP08:CURR_MONITOR": "Drain Current - Keithley8",
        "SR14ID01AMP09:CURR_MONITOR": "I0 - Keithley9",
        "SR14ID01:BL_GAP_REQUEST": "Undulator Gap Request",
        "SR14ID01:GAP_MONITOR": "Undulator Gap Readback",
        "SR11BCM01:CURRENT_MONITOR": "Ring Current",
        "SR14ID01NEXSCAN:saveData_comment1": "Note 1",
        "SR14ID01NEXSCAN:saveData_comment2": "Note 2",
        "SR14ID01:BRANCH_MODE": "Branch",
        "SR14ID01NEX01:RULER_ID": "Ruler ID",
        "SR14ID01NEX01:R_MTR.RBV": "R",
        "SR14ID01NEX01:C_MTR.RBV": "C",
        "SR14ID01NEX01:Z_MTR.RBV": "Z",
        "SR14ID01NEX01:X_MTR.RBV": "X",
        "SR14ID01NEX01:Y_MTR.RBV": "Y",
        "SR14ID01DTACQ06:ch1_converted_ave_rbv": "BL PHD - SR570-2",
        "SR14ID01DTACQ06:ch2_converted_ave_rbv": "Ref Foil - SR570-4",
        "SR14ID01DTACQ06:ch3_converted_ave_rbv": "I0 - SR570-3",  # Izero-SR570-3"
        "SR14ID01DTACQ06:ch4_converted_ave_rbv": "Drain Current SR570-1",
        "SR14ID01DTACQ06:ch5_converted_ave_rbv": "Direct PHD - SR570-5",
        "SR14ID01DTACQ07:ch1_converted_ave_rbv": "Ref Foil - SR570-8",
        "SR14ID01DTACQ07:ch2_converted_ave_rbv": "I0 - SR570-7",
        "SR14ID01DTACQ07:ch3_converted_ave_rbv": "Drain_Current - SR570-6",
        "SR14ID01DTACQ07:ch4_converted_ave_rbv": "Drain Current - SR570-9",
        "SR14ID01DPP01:mca1.R0": "Dante ROI (PFY)",
        "SR14ID01DPP01:ch1:InputCountRate": "Dante Input Count Rate (TFY)",
        "SR14ID01DPP01:ch1:OutputCountRate": "Dante Output Count Rate (TFY)",
    }

    @classmethod
    def parse_asc_202403(
        cls, file: typing.TextIO, header_only: bool = False
    ) -> parse_fn_ret_type:
        """
        Read Australian Synchrotron SXR NEXAFS `.asc` files.

        Parameters
        ----------
        file : typing.TextIO
            TextIO of the datafile (i.e. open('file.asc', 'r')).
        header_only : bool, optional
            If True, then only the header of the file is read and NDArray is returned as None, by default False.

        Returns
        -------
        tuple[NDArray, list[str], dict[str, Any]]
            Returns a set of data as a numpy array,
            labels as a list of strings,
            units as a list of strings,
            and parameters as a dictionary.

        Raises
        ------
        ValueError
            If the file is not a valid .asc file.
        """

        # Initialise structures
        params = {"filename": file.name}
        labels = []
        units = []

        # Check valid format.
        if not file.name.endswith(".asc"):
            raise ValueError(f"File {file.name} is not a valid .asc file.")

        ### Read file
        # Check header is correct
        assert file.readline() == "## mda2ascii 0.3.2 generated output\n"

        # skip 2 lines
        [file.readline() for i in range(2)]

        ## 1 Initial Parameters including MDA File v, Scan Num, Ovrl scn dim, scn size
        for i in range(4):
            line = file.readline()[2:].strip().split(" = ", 1)
            if len(line) == 2:
                params[line[0]] = line[1]

        # skip 2 lines
        [file.readline() for i in range(2)]
        # Check PV header is correct
        assert file.readline() == "#  Extra PV: name, descr, values (, unit)\n"

        # skip 1 line
        file.readline()

        ## 2 Extra Parameter/Values
        line = file.readline()
        while line != "\n":
            # Discard first 11 characters and linenumber '# Extra PV 1:'.
            line = line.strip().split(":", 1)[1]
            # Separate the name, description, values and units.
            line = line.split(",")
            # Trim values and append to params in a tuple.
            vals = [line[i].strip().replace('"', "") for i in range(1, len(line))]
            # Convert string values to float / int if possible:
            if len(vals) == 4:  # has units column
                try:
                    vals[3] = float(vals[3]) if "." in vals[3] else int(vals[3])
                except ValueError:
                    pass
            # Add vals to params
            params[line[0].strip()] = vals

            # read the newline
            line = file.readline()

        # skip 1 extra line
        file.readline()

        ## 3 Scan Header:
        assert file.readline() == "# 1-D Scan\n"
        for i in range(3):
            line = file.readline()[2:].strip().split(" = ", 1)
            if len(line) == 2:
                params[line[0]] = line[1]

        # skip 1 line
        file.readline()

        ## 4 Scan Column properties:
        # Labels for columns describing data columns.
        # i.e.
        # i.e. #  Positioner: name, descr, step mode, unit, rdbk name, rdbk descr, rdbk unit
        column_types = {}
        line = file.readline()
        while line != "\n":
            line = line.strip().split(": ", 1)
            col_name = line[0][3:]  # ignore "#  " before name
            col_attrs = line[1].split(", ")
            column_types[col_name] = col_attrs
            # read the newline
            line = file.readline()

        # skip 1 line
        file.readline()

        # Column descriptions
        # i.e. #    2  [1-D Positioner 1]  SR14ID01PGM:LOCAL_SP, Mono setpoint, TABLE, eV, SR14ID01PGM:LOCAL_SP, Mono setpoint, eV
        column_descriptions = {}
        column_type_assignments = {}
        line = file.readline()
        while line != "\n":
            # Take index info
            index = int(line[1:6].strip())
            index_line = (
                index == 1
            )  # Boolean to determine if on index description line.
            # Take coltype info
            desc_type = line[8:26]
            assert (
                desc_type[0] == "[" and desc_type[-1] == "]"
            )  # Check parenthesis of the line parameters
            desc_type = desc_type[1:-1].strip()
            # Check if valid col type
            valid = False if not index_line else True
            for coltype in column_types:
                if coltype in desc_type:
                    valid = True
                    break
            if not valid:
                raise ValueError(f"Invalid column type {desc_type} in line {line}")

            # Take info
            desc_info = line[28:].split(", ")
            column_descriptions[index] = desc_info
            column_type_assignments[index] = desc_type
            # Check that the initial parameter begins with the Instrument descriptor.
            if not index_line:
                assert desc_info[0].startswith("SR14ID01")  # code for initial
            # Add to labels and units to lists
            labels += [desc_info[0].strip()] if not index_line else ["Index"]
            if "Positioner" in desc_type:
                pUnit = desc_info[3].strip()  # hardcoded 'unit' position
                units += [pUnit]
            elif "Detector" in desc_type:
                dUnit = desc_info[2].strip()  # hardcoded 'unit' position
                units += [dUnit if dUnit != "" else None]
            elif index_line:
                units += [None]

            # read next line
            line = file.readline()

        # add column data to params
        params["column_types"] = column_types
        params["column_descriptions"] = column_descriptions
        params["column_type_assignments"] = column_type_assignments

        if header_only:
            # Do not process remaining lines
            return None, labels, units, params

        ## 5 Scan Values
        assert file.readline() == "# 1-D Scan Values\n"
        lines = file.readlines()  # read remaining lines efficiently

        # Convert data to numpy array.
        data = np.loadtxt(lines, delimiter="\t")
        # [np.loadtxt(line, delimiter="\t")
        #     for line in lines]
        data = np.array(data)

        return data, labels, units, params

    @classmethod
    def parse_mda_2024_03(
        cls, file: typing.IO | str, header_only: bool = False
    ) -> parse_fn_ret_type:
        """
        Read Australian Synchrotron SXR NEXAFS .mda files.

        Parameters
        ----------
        file : typing.IO | str
            IO of the datafile (i.e. open('file.mda', 'rb')).
        header_only : bool, optional
            If True, only the header of the file is read and NDArray is returned as None, by default False.

        Returns
        -------
        data : NDArray | None
            The data array if header_only is False, else None.
        labels : list[str]
            List of labels for each data column.
        units : list[str]
            List of units for each data column.
        params : dict[str, Any]
            A dictionary of parameters extracted from the file, or an empty dictionary if no parameters are returned.

        Raises
        ------
        ValueError
            If the file is not a valid .mda file.
        """

        # Initialise parameter list
        params = {}
        labels = []

        if isinstance(file, str):
            # Open the file in binary mode
            with open(file, "rb") as f:
                return cls.parse_mda_2024_03(f, header_only=header_only)
        elif isinstance(file, (typing.TextIO, io.TextIOWrapper)):
            return cls.parse_mda_2024_03(file.buffer, header_only=header_only)
        else:
            assert isinstance(file, (typing.BinaryIO, io.BufferedReader, io.BytesIO)), (
                f"File must be a typing.BinaryIO implementation or file path string. Was '{type(file)}'."
            )

        # Check valid format.
        if not file.name.endswith(".mda"):
            raise ValueError(f"File {file.name} is not a valid .mda file.")

        mda = MDAFileReader(file)

        mda_header = mda.read_header_as_dict()
        if mda_header["mda_rank"] != 1:
            raise ValueError("MDA file is not 1D, incompatible for regular NEXAFS.")
        mda_params = mda.read_parameters()
        mda_arrays, mda_scans = mda.read_scans(header_only=header_only)

        # Add values to params dict
        params.update(mda_header)
        params.update(mda_params)
        # Add column types and descriptions to params.
        scan_1d = mda_scans[0]
        column_types = {
            "Positioner": (
                "name",
                "descr",
                "step mode",
                "unit",
                "rdbk name",
                "rdbk descr",
                "rdbk unit",
            ),
            "Detector": (
                "name",
                "descr",
                "unit",
            ),
        }
        column_descriptions = {
            i: [
                p.name,
                p.desc,
                p.step_mode,
                p.unit,
                p.readback_name,
                p.readback_desc,
                p.readback_unit,
            ]
            for i, p in enumerate(scan_1d.positioners)
        }
        column_descriptions.update(
            {
                i + len(scan_1d.positioners): [d.name, d.desc, d.unit]
                for i, d in enumerate(scan_1d.detectors)
            }
        )
        params["column_types"] = column_types
        params["column_descriptions"] = column_descriptions
        # Collect units and labels:
        labels = []
        units = []
        for i, p in enumerate(scan_1d.positioners):
            labels.append(p.name)
            units.append(p.unit)
        for i, d in enumerate(scan_1d.detectors):
            labels.append(d.name)
            units.append(d.unit)

        if header_only:
            return None, labels, units, params
        data: NDArray | tuple[NDArray] | None
        if isinstance(mda_arrays, list):
            if len(mda_arrays) == 1:
                data = mda_arrays[0]
            else:
                data = tuple(*mda_arrays)
        else:
            data = None
        return data, labels, units, params

    @classmethod
    def parse_txt_qant_2025_12(
        cls, file: typing.TextIO | str, header_only: bool = False
    ) -> parse_fn_ret_type:
        """
        A QANT exported txt file reader for SXR NEXAFS data.

        Works for outputs from QANTv1.14.

        Parameters
        ----------
        file : typing.TextIO
            TextIO of the datafile (i.e. open('file.txt', 'r')).
        header_only : bool, optional
            If True, only the header of the file is read and NDArray is returned as None, by default False.

        Returns
        -------
        Data : None or NDArray
            The data array if header_only is False, else None.
        labels : list[str]
            List of labels for each data column.
        units : list[str]
            List of units for each data column.
        tuple[NDArray | None, list[str], list[str], dict[str, Any]]
            A tuple containing the data array, labels as a list of strings,
            units as a list of strings,
            and parameters as a dictionary.

        Raises
        ------
        ValueError
            If the file is not a valid .txt file.
        """
        if isinstance(file, str):
            # Open the file in text mode
            with open(file, "r") as f:
                return cls.parse_txt_qant_2025_12(f, header_only=header_only)

        # Initialise structures
        params = {"filename": file.name}
        labels = []
        units = None

        # Check valid format.
        if not file.name.endswith(".txt"):
            raise ValueError(f"File {file.name} is not a valid .txt file.")

        ### Read file
        # Check header is correct
        assert file.readline().startswith("File Produced by QANT")
        separator_line = "----------------------------------------------------\n"
        assert file.readline() == separator_line
        assert file.readline() == "List of dataseries and their notes:\n"

        # skip 1 lines
        file.readline()

        # Skip noted information about the normalisation steps for each channel
        line = file.readline()
        while line != separator_line:
            if line == "":  # "" is EOF.
                return None, labels, units, params
            line = file.readline()

        # Collect the parameters.
        assert file.readline() == "List of strings or notes regarding this dataset:\n"
        line = file.readline()
        separator_line = (
            "-----------------------------------------------------\n"  # note extra -
        )
        while line != separator_line and line != "":  # "" is EOF.
            if line == "":  # "" is EOF.
                return None, labels, units, params
            # Process the line
            if line != "\n":
                if "=" in line:
                    # Split the line based on the first '=' only.
                    key, val = line.split("=", maxsplit=1)
                    params[key.strip()] = val.strip()
            line = file.readline()

        # Numerical values
        line = file.readline()
        assert line == "List of numeric variables regarding this dataset:\n", (
            f"Unexpected line: {line}"
        )
        line = file.readline()
        while line != separator_line and line != "":  # "" is EOF.
            if line == "":  # "" is EOF.
                return None, labels, units, params
            # Don't process these lines.
            line = file.readline()

        assert file.readline() == "Extra PVs:\n"
        file.readline()  # skip header line

        ## 2 Extra Parameter/Values
        data_line = "Data---------------------------------------------------\n"

        line = file.readline()
        while line != data_line and line != "":  # "" is EOF.
            if line == "":  # "" is EOF.
                return None, labels, units, params
            elif line == "\n" or "\t" not in line:
                # Skip empty lines
                line = file.readline()
                continue

            # Separate the name, description, values and units.
            line = line.split("\t")
            # Trim values and append to params in a tuple.
            vals = [line[i].strip().replace('"', "") for i in range(1, len(line))]

            # Convert string values to float / int if possible:
            if len(vals) == 4:  # has units column
                try:
                    vals[2] = float(vals[2]) if "." in vals[2] else int(vals[2])
                except ValueError:
                    pass
            # Add vals to params
            params[line[0].strip()] = vals

            # read the newline
            line = file.readline()

        # Skip two lines
        [file.readline() for i in range(2)]

        header_line = file.readline()
        assert header_line != "" and "\t" in header_line, (
            "No data header found in QANT .txt file."
        )
        labels = header_line.replace("'", "").strip().split("\t")

        # No units provided.

        if header_only:
            # Do not process remaining lines
            return None, labels, units, params

        ## 5 Scan Values
        lines = file.readlines()  # read remaining lines efficiently
        # Remove empty lines
        lines = [line for line in lines if line.strip() != ""]
        # Convert data to numpy array.
        data = np.loadtxt(lines, delimiter="\t")
        # [np.loadtxt(line, delimiter="\t")
        #     for line in lines]
        data = np.array(data)
        assert data.shape[0] != 0, "No data found in QANT .txt file."
        return data, labels, units, params

    @property
    @override
    def summary_param_values(self) -> list[Any]:
        """
        Return a list of important parameter values of the data file.

        Uses the list element corresponding to 'value' for each file.
        Overrides base method to use the 'value' element of the SXR parameter list.

        Returns
        -------
        list
            List of important parameter values.
        """
        # Get second element which is the parameter number.
        return [self.params[key][1] for key in self.SUMMARY_PARAM_RAW_NAMES]

    @property
    @override
    def summary_param_names_with_units(self) -> list[str]:
        """
        Return a list of important parameter names with units.

        Requires a loaded dataset to return the units of the parameters.
        Not a pre-defined class method.

        Returns
        -------
        list
            List of important parameter names with units.
        """
        pNames = self.SUMMARY_PARAM_RAW_NAMES
        pUnits = [
            (
                self.params[pName][2]
                if (
                    self.params is not None  # Params loaded
                    and pName in self.params  # Parameter listed
                    and hasattr(
                        self.params[pName], "__len__"
                    )  # Parameter has a list of values
                    and len(self.params[pName])
                    == 3  # 3 params for value, description, unit
                    and self.params[pName][2] != ""  # Unit value is not empty.
                )
                else None
            )
            for pName in pNames
        ]
        if not self.relabel:
            return [
                f"{pName} ({unit})" if unit is not None else pName
                for pName, unit in zip(pNames, pUnits)
            ]
        else:
            relabel_names = []
            for pName, unit in zip(pNames, pUnits):
                if unit is not None:
                    relabel_names.append(
                        f"{self.RELABELS[pName]} ({unit})"
                        if pName in self.RELABELS
                        else f"{pName} ({unit})"
                    )
                else:
                    relabel_names.append(
                        self.RELABELS[pName] if pName in self.RELABELS else pName
                    )
            return relabel_names


if __name__ == "__main__":
    # Test each of the parser functions with the test data.
    fname = "sxr129582-QANT_raw.txt"
    fname = "sxr130214-QANT_edge_norm.txt"
    fname = "sxr130214.asc"
    fname = "sxr130214.mda"

    #
    import pkgutil
    import io

    bdata = pkgutil.get_data("pyNexafs", f"../tests/test_data/au/SXR/2024-03/{fname}")
    assert bdata is not None
    bio = io.BytesIO(bdata)
    fdata = io.TextIOWrapper(bio, encoding="utf-8")
    result = SXR_NEXAFS.parse_txt_qant_2025_12(fdata, header_only=False)
    print(result)
