"""ID table :py:class:`.Element`."""


from __future__ import annotations

from pathlib import Path
from warnings import warn

import numpy as np

from ...constants import clight, e_mass
from .element_object import Element

# 2023jan15 orblancog first release
# 2023jan18 orblancog fix bug with element print
# 2023apr30 orblancog redefinition to function
# 2023jul04 orblancog changing class to save in .mat and .m
# 2025ago15 orblancog include first order kick maps from text files


def _anyarray(value: np.ndarray) -> None:
    # Ensure proper ordering(F) and alignment(A) for "C" access in integrators
    return np.require(value, dtype=np.float64, requirements=["F", "A"])


class InsertionDeviceKickMap(Element):
    """
    Insertion Device Element. Valid for a parallel electron beam.

    This elememt implements tracking through an integrated magnetic
    field map of second order in energy, normalized to an energy
    value that is required to calculate alpha. See Eq. (5) in [#].

    First order maps could be included. See Eq. (3) in [#].
    Note that positive and negative signs are not taken into account in
    this implementation. Input should already include the sign difference.

    Default PassMethod: ``IdTablePass``.

    [#] Pascale ELLEAUME, "A New Approach to the Electron Beam Dynamics in
        Undulators and  Wigglers". EPAC1992 0661.
        European Synchrotron Radiation Facility.
        BP 220, F-38043 Grenoble, France
    """

    _BUILD_ATTRIBUTES = Element._BUILD_ATTRIBUTES + [
        "PassMethod",
        "Filename_in",
        "Normalization_energy",
        "Nslice",
        "Length",
        "xkick",
        "ykick",
        "xkick1",
        "ykick1",
        "xtable",
        "ytable",
    ]

    _conversions = dict(
        Element._conversions,
        Nslice=int,
        xkick=_anyarray,
        ykick=_anyarray,
        xkick1=_anyarray,
        ykick1=_anyarray,
    )

    def __init__(
        self: InsertionDeviceKickMap,
        family_name: str,
        *args: any,
        **kwargs: dict[any, any],
    ) -> None:
        """
        Init IdTable.

        This __init__ takes the input to initialize an InsertionDeviceKickMap
        from an user input with arguments, for example at the moment of the
        element creation, or from all parameters, for example when reading
        a Lattice.

        Args:
            family_name: the family name
            args: postional arguments from user.
            kwargs: dictionary from element.

        """
        _argnames = [
            "PassMethod",
            "Filename_in",
            "Normalization_energy",
            "Nslice",
            "Length",
            "xkick",
            "ykick",
            "xkick1",
            "ykick1",
            "xtable",
            "ytable",
        ]
        if len(args) < 11:
            # get data from user input
            elemargs = self.from_user(*args)
        else:
            # get data from arguments
            elemargs = dict(zip(_argnames, args))
        elemargs.update(kwargs)
        super().__init__(family_name, **elemargs)

    def set_DriftPass(self: InsertionDeviceKickMap) -> None:
        """Set DriftPass tracking pass method."""
        self.PassMethod = "DriftPass"

    def set_IdTablePass(self: InsertionDeviceKickMap) -> str:
        """Set IdTablePass tracking pass method."""
        self.PassMethod = "IdTablePass"

    def get_PassMethod(self: InsertionDeviceKickMap) -> str:
        """Get the current tracking pass method.

        Returns:
            String with the current tracking pass method.
        """
        warn(UserWarning("get_PassMethod is deprecated; do not use"), stacklevel=2)
        return self.PassMethod

    def from_user(
        self: InsertionDeviceKickMap, nslice: int, fname: str, norm_energy: float
    ) -> dict:
        """
        Create an Insertion Device Kick Map from a Radia field map file.

        The following is an example of an Insertion Device element, idelem,
        created from a file 'radiakickmap.txt' with 10 integration steps.
        The tables have been normalized to 3 GeV. The family name is 'IDname'.
        >>> idelem = at.InsertionDeviceKickMap('IDname', 10, 'radiakickmap.txt', 3)

        The input file could be a text file or a dictionary.
        See read_text_radia_field_map for info about the text file format.
        See read_dict_radia_field_map for info about the dict format.

        Family name is part of the base class, and all other arguments are
        parsed here below.

        Arguments:
            nslice: number of slices in integrator.
            fname: input filename. Text of .mat files.
            norm_energy: particle energy in GeV.

        Returns:
            A dict with the file data.
        """

        def sorted_table(
            table_in: np.ndarray, sorted_index: np.ndarray, order_axis: str
        ) -> np.ndarray:
            """Return ordered table.

            Arguments:
                table_in: input table..
                sorted_index: index to sort the table.
                order_axis: sort on 'col' or 'row'.

            Returns:
                Sorted table as fortran array.
            """
            # np.asfortranarray makes a copy of contiguous memory positions
            table_out = np.copy(table_in)
            for i, iis in zip(range(len(sorted_index)), sorted_index):
                if order_axis == "col":
                    table_out[:, i] = table_in[:, iis]
                if order_axis == "row":
                    table_out[i, :] = table_in[iis, :]
            return np.asfortranarray(table_out)

        if isinstance(fname, dict):
            thefields = self.read_dict_radia_field_map(fname)
            fname = ""
        else:
            # assume text file
            thefields = self.read_text_radia_field_map(fname)

        (
            el_length,
            hkickmap2,
            vkickmap2,
            table_colshkick,
            table_rowshkick,
            table_colsvkick,
            table_rowsvkick,
            hkickmap1,
            vkickmap1,
            _,
            _,
        ) = thefields

        # set to float
        table_colshkickarray = np.array(table_colshkick, dtype="float64")
        table_rowshkickarray = np.array(table_rowshkick, dtype="float64")
        table_colsvkickarray = np.array(table_colsvkick, dtype="float64")
        table_rowsvkickarray = np.array(table_rowsvkick, dtype="float64")

        # Reorder table_axes
        cols1sorted_index = np.argsort(table_colshkickarray)
        table_colshkickarray.sort()
        rows1sorted_index = np.argsort(table_rowshkickarray)
        table_rowshkickarray.sort()
        cols2sorted_index = np.argsort(table_colsvkickarray)
        table_colsvkickarray.sort()
        rows2sorted_index = np.argsort(table_rowsvkickarray)
        table_rowsvkickarray.sort()
        # Reorder kickmap2
        hkickmap2_a = sorted_table(hkickmap2, cols1sorted_index, "col")
        hkickmap2 = sorted_table(hkickmap2_a, rows1sorted_index, "row")
        vkickmap2_a = sorted_table(vkickmap2, cols2sorted_index, "col")
        vkickmap2 = sorted_table(vkickmap2_a, rows2sorted_index, "row")
        # Reorder kickmap1
        hkickmap1_a = sorted_table(hkickmap1, cols1sorted_index, "col")
        hkickmap1 = sorted_table(hkickmap1_a, rows1sorted_index, "row")
        vkickmap1_a = sorted_table(vkickmap1, cols2sorted_index, "col")
        vkickmap1 = sorted_table(vkickmap1_a, rows2sorted_index, "row")

        # Field to kick factors
        e_mass_gev = e_mass * 1e-9
        brho = 1e9 * np.sqrt(norm_energy**2 - e_mass_gev**2) / clight
        # kick2 vars
        factor2 = 1.0 / (brho**2)
        xkick = factor2 * hkickmap2
        ykick = factor2 * vkickmap2
        # kick1 vars
        factor1 = 1.0 / (brho)
        xkick1 = factor1 * hkickmap1
        ykick1 = factor1 * vkickmap1
        # axes
        xtable = table_colshkickarray.T
        ytable = table_rowshkickarray.T

        return {
            "PassMethod": "IdTablePass",
            "Filename_in": fname,
            "Normalization_energy": norm_energy,
            "Nslice": np.uint8(nslice),
            "Length": el_length,
            "xkick": xkick,
            "ykick": ykick,
            "xkick1": xkick1,
            "ykick1": ykick1,
            "xtable": xtable,
            "ytable": ytable,
        }

    def read_text_radia_field_map(
        self: InsertionDeviceKickMap, file_in_name: str
    ) -> tuple:
        """
        Read a RadiaField map in text format and return.

        A File, where :
        - comments start with #.
        - the first data line is the length in meters.
        - the second data line is the number of points in the h. plane.
        - the third data line is the number of points in the v. plane.
        - each block is a table with axes.
        - each data block comes after a START.
        - first the horizontal data block, and second the vertical data
        block with the second order kicks.
        There might be two other blocks with the horizontal and
        vertical first order kicks.

        File example (ignore the !SPACE):
        ! #comment in line 1
        ! #comment in line 2
        ! Length_in_m
        ! #comment in line 4
        ! Number of points in horizontal plane :nh
        ! #comment in line 6
        ! Number of points in vertical plane :nv
        ! #comment in line 8
        ! START
        !             pos_point1h pos_point2h ... pos_pointnh
        ! pos_point1v
        ! ...                    horizontal kick_map(nv,nh)
        ! pos_pointnv
        ! START
        !             pos_point1h pos_point2h ... pos_pointnh
        ! pos_point1v
        ! ...                    vertical kick_map(nv,nh)
        ! pos_pointnv
        ! (EOL)

        Arguments:
            file_in_name: the file name.

        Returns:
            Tuple with file tables and axes.

        Raises:
            ValueError: if the number of blocks in less than 2 or equal to 3.
        """
        thepath = Path(file_in_name)
        with thepath.open(encoding="utf-8") as thefile:
            lines = thefile.readlines()
        thefile.close()
        data_lines = 0  # line not starting with '#'
        header_lines = 0  # line starting with '#'
        block_counter = 0  # START of the h.map, START of the v.map
        kick_block_list = []
        kick_haxes_list = []
        kick_vaxes_list = []
        for line in lines:
            sline = line.split()
            if sline[0] == "#":  # line is comment
                header_lines += 1
            else:
                data_lines += 1
                if data_lines == 1:  # get the element length
                    el_length = float(sline[0])
                elif data_lines == 2:  # get the number of hor. points
                    h_points = int(sline[0])
                elif data_lines == 3:  # get the number of ver. points
                    v_points = int(sline[0])
                    # initialize element kicks and table_axes
                    kick_block = np.zeros((v_points, h_points))
                    haxis = np.zeros(h_points)
                    vaxis = np.zeros(v_points)
                else:
                    # read block of data
                    if sline[0] == "START" or sline[0] == "START\n":
                        block_counter += 1
                        block_lines = 0
                    if block_lines == 1:
                        haxis = sline
                    if block_lines > 1:
                        # minus one due to python index starting at 0
                        # and minus another one due
                        # to the column labels in first line
                        vaxis[block_lines - 2] = float(sline[0])
                        kick_block[block_lines - 2][:] = sline[1:]
                    if block_lines > v_points:
                        block_lines = 0
                        kick_block_list.append(kick_block)
                        kick_haxes_list.append(haxis)
                        kick_vaxes_list.append(vaxis)
                    block_lines += 1
        # checking how many kick blocks were added
        lenkick_block_list = len(kick_block_list)
        if lenkick_block_list < 2 or lenkick_block_list == 3:
            _minimumblocknumbererrormsg = (
                "Input file contains only " f"{len(kick_block_list)} block"
            )
            raise ValueError(_minimumblocknumbererrormsg)
        if lenkick_block_list == 2:
            # first order kick not in file
            kick_block_list.append(0.0 * np.copy(kick_block))
            kick_block_list.append(0.0 * np.copy(kick_block))
        elif lenkick_block_list > 4:
            # file contains more blocks that required
            _warn4kickblocks = (
                "Input file contains more than 4 blocks. Additional blocks ignored"
            )
            warn(_warn4kickblocks)

        return (
            el_length,
            kick_block_list[0],
            kick_block_list[1],
            kick_haxes_list[0],
            kick_vaxes_list[0],
            kick_haxes_list[1],
            kick_vaxes_list[1],
            kick_block_list[2],
            kick_block_list[3],
            h_points,
            v_points,
        )

    def read_dict_radia_field_map(
        self: InsertionDeviceKickMap, id_input: dict
    ) -> tuple:
        """Read a dictionary with Radia field map tables.

        The required keys are "Length", "xkick" and "ykick"
        for the second order maps, "xtable" and "ytable" for
        the grid, and "xkick1" and "ykick1" for the first order
        maps.

        Arguments:
            id_input: Radia field map input.

        Returns:
            Tuple with Insertion Device parameters.
        """
        (v_points, h_points) = id_input["xkick"].shape
        return (
            id_input["Length"],
            id_input["xkick"],
            id_input["ykick"],
            id_input["xtable"],
            id_input["ytable"],
            id_input["xtable"],
            id_input["ytable"],
            id_input["xkick1"],
            id_input["ykick1"],
            h_points,
            v_points,
        )


# EOF
