"""
Module of various types of time_evolution
"""

from copy import deepcopy
from numba import float64, int64, jit, void
from numpy import arange, array, cos, cross, pi, sin, sqrt, zeros
from scipy.linalg import norm

class Integrator:
    """
    Class used to assign integrator type.

    Attributes
    ----------
    dt : float
        Timestep.

    kB : float
        Boltzmann constant.

    magnetized : bool
        Magnetized simulation flag.

    species_num : numpy.ndarray
        Number of particles of each species.

    species_plasma_frequencies : numpy.ndarray
        Plasma frequency of each species.

    box_lengths : numpy.ndarray
        Length of each box side.

    pbox_lengths : numpy.ndarray
        Initial particle box sides' lengths.

    verbose : bool
        Verbose output flag.

    type : str
        Integrator type.

    """

    dt: float = None
    kB: float = None

    # attributes
    type: str = None
    supported_integrators = {}
    equilibration_type: str = "verlet"
    magnetization_type: str = "magnetic_verlet"
    production_type: str = "verlet"

    species_num = None
    species_plasma_frequencies = None

    # Thermostat attributes
    thermalization: bool = True
    thermostat_type: str = "berendsen"
    thermalization_rate: float = 2.0
    thermalization_timestep: int = 0
    berendsen_tau: float = None
    thermostat_temperatures = None
    thermostat_temperatures_eV = None

    # Magnetic attributes
    magnetized: bool = False
    magnetic_field_uvector = None
    magnetic_field = None
    omega_c = None
    species_cyclotron_frequencies = None
    ccodt = None
    cdt = None
    ssodt = None
    sdt = None
    v_B = None
    v_F = None

    # Langevin attributes
    c1 = None
    c2 = None
    sigma = None
    box_lengths = None
    pbox_lengths = None

    boundary_conditions = None

    supported_boundary_conditions = {}

    verbose: bool = False

    # def __repr__(self):
    #     sortedDict = dict(sorted(self.__dict__.items(), key=lambda x: x[0].lower()))
    #     disp = 'Integrator( \n'
    #     for key, value in sortedDict.items():
    #         disp += "\t{} : {}\n".format(key, value)
    #     disp += ')'
    #     return disp

    def __copy__(self):
        """Make a shallow copy of the object using copy by creating a new instance of the object and copying its __dict__."""
        # Create a new object
        _copy = type(self)()
        # copy the dictionary
        _copy.from_dict(input_dict=self.__dict__)
        return _copy

    def __deepcopy__(self, memodict={}):
        """Make a deepcopy of the object.

        Parameters
        ----------
        memodict: dict
            Dictionary of id's to copies

        Returns
        -------
        _copy: :class:`sarkas.time_evolution.integrators.Integrator`
            A new Integrator class.
        """
        id_self = id(self)  # memorization avoids unnecessary recursion
        _copy = memodict.get(id_self)
        if _copy is None:
            _copy = type(self)()
            # Make a deepcopy of the mutable arrays using numpy copy function
            for k, v in self.__dict__.items():
                if k != "thread_ls":
                    _copy.__dict__[k] = deepcopy(v, memodict)

        return _copy

    def from_dict(self, input_dict: dict):
        """
        Update attributes from input dictionary.

        Parameters
        ----------
        input_dict: dict
            Dictionary to be copied.

        """
        self.__dict__.update(input_dict)

    def copy_params(self, params):
        """
        Copy necessary parameters.

        Parameters
        ----------
        params: :class:`sarkas.core.Parameters`
            Simulation's parameters.

        """
        self.box_lengths = params.box_lengths
        self.pbox_lengths = params.pbox_lengths
        self.dimensions = params.dimensions
        self.kB = params.kB
        self.eV2K = params.eV2K
        self.total_num_ptcls = params.total_num_ptcls
        self.species_num = params.species_num.copy()
        self.species_plasma_frequencies = params.species_plasma_frequencies.copy()
        self.species_masses = params.species_masses.copy()
        self.species_temperatures = params.species_temperatures.copy()
        self.verbose = params.verbose
        # Enforce consistency
        if not self.boundary_conditions:
            self.boundary_conditions = params.boundary_conditions.lower()

        # Check whether you input temperatures in eV or K
        if self.thermalization and self.thermostat_temperatures:
            self.thermostat_temperatures_eV = self.thermostat_temperatures.copy() / self.eV2K
        elif self.thermalization and self.thermostat_temperatures_eV:
            self.thermostat_temperatures = self.thermostat_temperatures_eV.copy() * self.eV2K
        elif self.thermalization and not self.thermostat_temperatures:
            self.thermostate_temperatures = params.species_temperatures.copy()
            self.thermostate_temperatures_eV = params.species_temperatures_eV.copy()

        # Backwards compatibility
        if hasattr(self, "equilibration_steps"):
            params.equilibration_steps = self.equilibration_steps

        if hasattr(self, "magnetization_steps"):
            params.equilibration_steps = self.magnetization_steps

        if hasattr(self, "production_steps"):
            params.production_steps = self.production_steps

        if hasattr(self, "eq_dump_step"):
            params.eq_dump_step = self.eq_dump_step

        if hasattr(self, "mag_dump_step"):
            params.mag_dump_step = self.mag_dump_step

        if hasattr(self, "eq_dump_step"):
            params.prod_dump_step = self.prod_dump_step

        if not self.boundary_conditions:
            self.boundary_conditions = params.boundary_conditions

        if not hasattr(params, "boundary_conditions"):
            params.boundary_conditions = self.boundary_conditions

        if params.magnetized:
            self.magnetized = True
            self.magnetic_field = params.magnetic_field.copy()
            self.species_cyclotron_frequencies = params.species_cyclotron_frequencies.copy()

    def setup(self, params, potential):
        """
        Assign attributes from simulation's parameters and classes.

        Parameters
        ----------
        params : :class:`sarkas.core.Parameters`
            Parameters class.

        potential : :class:`sarkas.potentials.core.Potential`
            Potential class.

        """
        if self.dt is None:
            raise ValueError("integrator.dt is None. Please define Integrator.dt")

        self.copy_params(params)

        if self.magnetized:
            self.magnetic_setup()

        if self.type:
            self.type = self.type.lower()
            self.equilibration_type = self.type
            self.production_type = self.type

        self.boundary_condition_setup()

        if self.thermalization:
            self.thermostat_setup()

        self.pot_acc_setup(potential)

    def pot_acc_setup(self, potential):
        """
        Link the :meth:`.update_accelerations` method depending on the potential algorithm.

        Parameters
        ----------
        potential : :class:`sarkas.potentials.core.Potential`
            Potential class.

        """

        self.potential_type = potential.type
        if potential.method != "fmm":
            if potential.pppm_on:
                self.update_accelerations = potential.update_pppm
            else:
                if potential.linked_list_on:
                    self.update_accelerations = potential.update_linked_list
                else:
                    self.update_accelerations = potential.update_brute
        else:
            self.update_accelerations = (
                potential.update_fmm_coulomb if potential.type == "coulomb" else potential.update_fmm_yukawa
            )

    def boundary_condition_setup(self):

        self.supported_boundary_conditions = {
            "periodic": self.periodic_bc,
            "absorbing": self.absorbing_bc,
            "reflective": self.reflecting_bc,
            "open": self.open_bc,
        }
        msg = (
            f"Unsupported boundary conditions. "
            f"Choose one of the supported boundary conditions\n{self.supported_boundary_conditions.keys()}",
        )
        # Assign integrator.enforce_bc to the correct method
        self.enforce_bc = self.supported_boundary_conditions.get(self.boundary_conditions, ValueError(msg))

    def thermostat_setup(self):
        """
        Assign attributes from simulation's parameters.

        Raises
        ------
        ValueError
            If a thermostat different than Berendsen is chosen.

        """

        self.thermostat_type = self.thermostat_type.lower()

        if self.thermostat_type != "berendsen":
            raise ValueError("Only Berendsen thermostat is supported.")

        if self.berendsen_tau:
            self.thermalization_rate = 1.0 / self.berendsen_tau
        else:
            self.berendsen_tau = 1.0 / self.thermalization_rate

    def type_setup(self, int_type):
        """

        Parameters
        ----------
        int_type: str
            Integrator type to use.

        Raises
        ------
        : ValueError
            If `int_type` is not a supported integrator.

        """

        # if int_type not in self.supported_integrators:
        #     raise ValueError(
        #         "Integrator not supported. " "Please choose one of the supported integrators \n",
        #         self.supported_integrators,
        #     )

        # Assign integrator.update to the correct method

        if int_type == "langevin":

            self.sigma = sqrt(2.0 * self.langevin_gamma * self.kB * self.species_temperatures / self.species_masses)
            self.c1 = 1.0 - 0.5 * self.langevin_gamma * self.dt
            self.c2 = 1.0 / (1.0 + 0.5 * self.langevin_gamma * self.dt)

        elif int_type == "magnetic_verlet":

            # Calculate functions for magnetic integrator
            # This could be used when the generalization to Forest-Ruth and MacLachlan algorithms will be implemented
            # In a magnetic Velocity-Verlet the coefficient is 1/2, see eq.~(78) in :cite:`Chin2008`
            self.magnetic_helpers(0.5)

            if self.magnetic_field_uvector @ array([0.0, 0.0, 1.0]) == 1.0:  # dot product
                int_type = "magnetic_verlet_zdir"

        elif int_type == "magnetic_pos_verlet":
            # Calculate functions for magnetic integrator
            # This could be used when the generalization to Forest-Ruth and MacLachlan algorithms will be implemented
            # In a magnetic Velocity-Verlet the coefficient is 1/2, see eq.~(78) in :cite:`Chin2008`
            self.magnetic_helpers(1.0)

            if self.magnetic_field_uvector @ array([0.0, 0.0, 1.0]) == 1.0:  # dot product
                int_type = "magnetic_pos_verlet_zdir"

        elif int_type == "magnetic_boris":

            # In a leapfrog-type algorithm the coefficient is different for the acceleration and magnetic rotation
            # see eq.~(79) in :cite:`Chin2008`
            self.magnetic_helpers(1.0)

            if self.magnetic_field_uvector @ array([0.0, 0.0, 1.0]) == 1.0:  # dot product
                int_type = "magnetic_boris_zdir"

        elif int_type == "cyclotronic":
            # Calculate functions for magnetic integrator
            # This could be used when the generalization to Forest-Ruth and MacLachlan algorithms will be implemented
            # In a magnetic Velocity-Verlet the coefficient is 1/2, see eq.~(78) in :cite:`Chin2008`
            self.magnetic_helpers(0.5)

            if self.magnetic_field_uvector @ array([0.0, 0.0, 1.0]) == 1.0:  # dot product
                int_type = "cyclotronic_zdir"

        self.supported_integrators = {
            "verlet": self.verlet,
            "langevin": self.langevin,
            "magnetic_verlet": self.magnetic_verlet,
            "magnetic_verlet_zdir": self.magnetic_verlet_zdir,
            "magnetic_pos_verlet": self.magnetic_pos_verlet,
            "magnetic_pos_verlet_zdir": self.magnetic_pos_verlet_zdir,
            "magnetic_boris": self.magnetic_boris,
            "magnetic_boris_zdir": self.magnetic_boris_zdir,
            "cyclotronic": self.cyclotronic,
            "cyclotronic_zdir": self.cyclotronic_zdir,
        }

        msg = f"Integrator not supported. Please choose one of the supported integrators \n{self.supported_integrators.keys()}"

        return self.supported_integrators.get(int_type, ValueError(msg))

    def magnetic_setup(self):
        # Create the unit vector of the magnetic field
        self.magnetic_field_uvector = self.magnetic_field / norm(self.magnetic_field)
        self.omega_c = zeros((self.total_num_ptcls, 3))

        sp_start = 0
        sp_end = 0
        for ic, sp_np in enumerate(self.species_num):
            sp_end += sp_np
            self.omega_c[sp_start:sp_end, :] = self.species_cyclotron_frequencies[ic]
            sp_start += sp_np

        # array to temporary store velocities
        # Luciano: I have the vague doubt that allocating memory for these arrays is faster than calculating them
        # each time step
        self.v_B = zeros((self.total_num_ptcls, 3))
        self.v_F = zeros((self.total_num_ptcls, 3))

    def langevin(self, ptcls):
        """
        Update particles class using the velocity verlet algorithm and Langevin damping.

        Parameters
        ----------
        ptcls: :class:`sarkas.particles.Particles`
            Particles data.

        """

        beta = ptcls.gaussian(0.0, 1.0, (self.total_num_ptcls, self.dimensions))
        sp_start = 0  # start index for species loop
        sp_end = 0
        for ic, num in enumerate(self.species_num):
            sp_end += num
            ptcls.pos[sp_start:sp_end, : self.dimensions] += (
                self.c1 * self.dt * ptcls.vel[sp_start:sp_end, : self.dimensions]
                + 0.5 * self.dt**2 * ptcls.acc[sp_start:sp_end, : self.dimensions]
                + 0.5 * self.sigma[ic] * self.dt**1.5 * beta
            )
            sp_start += num

        # Enforce boundary condition
        self.enforce_bc(ptcls)

        acc_old = ptcls.acc.copy()
        self.update_accelerations(ptcls)

        sp_start = 0
        sp_end = 0
        for ic, num in enumerate(self.species_num):
            sp_end += num

            ptcls.vel[sp_start:sp_end, : self.dimensions] = (
                self.c1 * self.c2 * ptcls.vel[sp_start:sp_end, : self.dimensions]
                + 0.5
                * self.c2
                * self.dt
                * (ptcls.acc[sp_start:sp_end, : self.dimensions] + acc_old[sp_start:sp_end, : self.dimensions])
                + self.c2 * self.sigma[ic] * sqrt(self.dt) * beta
            )
            sp_start += num

    def verlet(self, ptcls):
        """
        Update particles' class based on velocity verlet algorithm.
        More information can be found here: https://en.wikipedia.org/wiki/Verlet_integration
        or on the Sarkas website.

        Parameters
        ----------
        ptcls: :class:`sarkas.particles.Particles`
            Particles data.

        """
        # First half step velocity update
        ptcls.vel += 0.5 * ptcls.acc * self.dt
        # Full step position update
        ptcls.pos += ptcls.vel * self.dt
        # Enforce boundary condition
        self.enforce_bc(ptcls)
        # Compute total potential energy and acceleration for second half step velocity update
        self.update_accelerations(ptcls)
        # Second half step velocity update
        ptcls.vel += 0.5 * ptcls.acc * self.dt

    def magnetic_helpers(self, coefficient):
        """Calculate the trigonometric functions of the magnetic integrators.

        Parameters
        ----------
        coefficient: float
            Timestep coefficient.

        Notes
        -----
        This is useful for the Leapfrog magnetic algorithm and future Forest-Ruth and MacLachlan algorithms.

        """
        theta = self.omega_c * self.dt * coefficient
        self.sdt = sin(theta)
        self.cdt = cos(theta)
        self.ccodt = 1.0 - self.cdt
        self.ssodt = 1.0 - self.sdt / theta

    def magnetic_verlet_zdir(self, ptcls):
        """
        Update particles' class based on velocity verlet method in the case of a
        constant magnetic field along the :math:`z` axis. For more info see eq. (78) of Ref. :cite:`Chin2008`

        Parameters
        ----------
        ptcls: :class:`sarkas.particles.Particles`
            Particles data.

        Returns
        -------
        potential_energy : float
             Total potential energy.

        Notes
        -----
        This integrator is faster than `magnetic_verlet` but valid only for a magnetic field in the :math:`z`-direction.
        This is the preferred choice in this case.
        """

        # First half step of velocity update
        # # Magnetic rotation x - velocity
        # (B x v)_x  = -v_y, (B x B x v)_x = -v_x
        self.v_B[:, 0] = ptcls.vel[:, 1] * self.sdt[:, 0] + ptcls.vel[:, 0] * self.cdt[:, 0]
        # Magnetic rotation y - velocity
        # (B x v)_y  = v_x, (B x B x v)_y = -v_y
        self.v_B[:, 1] = -ptcls.vel[:, 0] * self.sdt[:, 0] + ptcls.vel[:, 1] * self.cdt[:, 1]

        # Magnetic + Const force field x - velocity
        # (B x a)_x  = -a_y, (B x B x a)_x = -a_x
        self.v_F[:, 0] = (
            self.ccodt[:, 1] / self.omega_c[:, 1] * ptcls.acc[:, 1]
            + self.sdt[:, 0] / self.omega_c[:, 0] * ptcls.acc[:, 0]
        )
        # Magnetic + Const force field y - velocity
        # (B x a)_y  = a_x, (B x B x a)_y = -a_y
        self.v_F[:, 1] = (
            -self.ccodt[:, 0] / self.omega_c[:, 0] * ptcls.acc[:, 0]
            + self.sdt[:, 1] / self.omega_c[:, 1] * ptcls.acc[:, 1]
        )

        ptcls.vel[:, 0] = self.v_B[:, 0] + self.v_F[:, 0]
        ptcls.vel[:, 1] = self.v_B[:, 1] + self.v_F[:, 1]
        ptcls.vel[:, 2] += 0.5 * self.dt * ptcls.acc[:, 2]

        # Position update
        ptcls.pos += ptcls.vel * self.dt

        # Enforce boundary condition
        self.enforce_bc(ptcls)

        # Compute total potential energy and acceleration for second half step velocity update
        potential_energy = self.update_accelerations(ptcls)

        # # Magnetic rotation x - velocity
        # (B x v)_x  = -v_y, (B x B x v)_x = -v_x
        self.v_B[:, 0] = ptcls.vel[:, 1] * self.sdt[:, 0] + ptcls.vel[:, 0] * self.cdt[:, 0]
        # Magnetic rotation y - velocity
        # (B x v)_y  = v_x, (B x B x v)_y = -v_y
        self.v_B[:, 1] = -ptcls.vel[:, 0] * self.sdt[:, 0] + ptcls.vel[:, 1] * self.cdt[:, 1]

        # Magnetic + Const force field x - velocity
        # (B x a)_x  = -a_y, (B x B x a)_x = -a_x
        self.v_F[:, 0] = (
            self.ccodt[:, 1] / self.omega_c[:, 1] * ptcls.acc[:, 1]
            + self.sdt[:, 0] / self.omega_c[:, 0] * ptcls.acc[:, 0]
        )
        # Magnetic + Const force field y - velocity
        # (B x a)_y  = a_x, (B x B x a)_y = -a_y
        self.v_F[:, 1] = (
            -self.ccodt[:, 0] / self.omega_c[:, 0] * ptcls.acc[:, 0]
            + self.sdt[:, 1] / self.omega_c[:, 1] * ptcls.acc[:, 1]
        )

        ptcls.vel[:, 0] = self.v_B[:, 0] + self.v_F[:, 0]
        ptcls.vel[:, 1] = self.v_B[:, 1] + self.v_F[:, 1]
        ptcls.vel[:, 2] += 0.5 * self.dt * ptcls.acc[:, 2]

        return potential_energy

    def magnetic_verlet(self, ptcls):
        """
        Update particles' class based on velocity verlet method in the case of an arbitrary direction of the
        constant magnetic field. For more info see eq. (78) of Ref. :cite:`Chin2008`

        Parameters
        ----------
        ptcls: :class:`sarkas.particles.Particles`
            Particles data.

        Returns
        -------
        potential_energy : float
             Total potential energy.

        Notes
        -----
        :cite:`Chin2008` equations are written for a negative charge. This allows him to write
        :math:`\\dot{\\mathbf v} = \\omega_c \\hat{B} \\times \\mathbf v`. In the case of positive charges we will have
        :math:`\\dot{\\mathbf v} = - \\omega_c \\hat{B} \\times \\mathbf v`.
        Hence the reason of the different signs in the formulas below compared to Chin's.

        Warnings
        --------
        This integrator is valid for a magnetic field in an arbitrary direction. However, while the integrator works for
        an arbitrary direction, methods in :mod:`sarkas.tools.observables` work only for a magnetic field in the
        :math:`z` - direction. Hence, if you choose to use this integrator remember to change your physical observables.

        """
        # Calculate the cross products
        b_cross_v = cross(self.magnetic_field_uvector, ptcls.vel)
        b_cross_b_cross_v = cross(self.magnetic_field_uvector, b_cross_v)
        b_cross_a = cross(self.magnetic_field_uvector, ptcls.acc)
        b_cross_b_cross_a = cross(self.magnetic_field_uvector, b_cross_a)

        # First half step of velocity update
        ptcls.vel += -self.sdt * b_cross_v + self.ccodt * b_cross_b_cross_v

        ptcls.vel += (
            0.5 * ptcls.acc * self.dt
            - self.ccodt / self.omega_c * b_cross_a
            + 0.5 * self.dt * self.ssodt * b_cross_b_cross_a
        )

        # Position update
        ptcls.pos += ptcls.vel * self.dt

        # Enforce boundary condition
        self.enforce_bc(ptcls)

        # Compute total potential energy and acceleration for second half step velocity update
        potential_energy = self.update_accelerations(ptcls)

        # Re-calculate the cross products
        b_cross_v = cross(self.magnetic_field_uvector, ptcls.vel)
        b_cross_b_cross_v = cross(self.magnetic_field_uvector, b_cross_v)
        b_cross_a = cross(self.magnetic_field_uvector, ptcls.acc)
        b_cross_b_cross_a = cross(self.magnetic_field_uvector, b_cross_a)

        # Second half step velocity update
        ptcls.vel += -self.sdt * b_cross_v + self.ccodt * b_cross_b_cross_v

        ptcls.vel += (
            0.5 * ptcls.acc * self.dt
            - self.ccodt / self.omega_c * b_cross_a
            + 0.5 * self.dt * self.ssodt * b_cross_b_cross_a
        )

        return potential_energy

    def magnetic_boris_zdir(self, ptcls):
        """
        Update particles' class using the Boris algorithm in the case of a
        constant magnetic field along the :math:`z` axis. For more info see eqs. (80) - (81) of Ref. :cite:`Chin2008`

        Parameters
        ----------
        ptcls: :class:`sarkas.particles.Particles`
            Particles data.

        Returns
        -------
        potential_energy : float
             Total potential energy.

        """
        # First half step of velocity update: Apply exp(dt * V_F / 2)
        ptcls.vel += 0.5 * ptcls.acc * self.dt

        # Rotate: Apply exp( dt * V)
        # B cross v
        self.v_B[:, 0] = -self.sdt[:, 1] * ptcls.vel[:, 1]
        self.v_B[:, 1] = self.sdt[:, 0] * ptcls.vel[:, 0]

        # B cross B cross v
        self.v_B[:, 0] -= self.ccodt[:, 0] * ptcls.vel[:, 0]
        self.v_B[:, 1] -= self.ccodt[:, 1] * ptcls.vel[:, 1]
        # Update velocities
        ptcls.vel[:, :2] += self.v_B[:, :2]

        # Second Acceleration half step: Apply exp(dt * V_F / 2)
        ptcls.vel += 0.5 * ptcls.acc * self.dt

        # Full step position update
        ptcls.pos += ptcls.vel * self.dt

        # Enforce boundary condition
        self.enforce_bc(ptcls)

        # Compute total potential energy and acceleration for second half step velocity update
        potential_energy = self.update_accelerations(ptcls)

        return potential_energy

    def magnetic_boris(self, ptcls):
        """
        Update particles' class using the Boris algorithm in the case of a
        constant magnetic field along the :math:`z` axis. For more info see eqs. (80) - (81) of Ref. :cite:`Chin2008`

        Parameters
        ----------
        ptcls: :class:`sarkas.particles.Particles`
            Particles data.

        Returns
        -------
        potential_energy : float
             Total potential energy.

        """

        # First half step of velocity update: Apply exp(eV_F/2)
        ptcls.vel += 0.5 * ptcls.acc * self.dt

        # Rotate: Apply exp( dt * V)
        # B cross v
        b_cross_v = cross(self.magnetic_field_uvector, ptcls.vel)
        # B cross B cross v
        b_cross_b_cross_v = cross(self.magnetic_field_uvector, b_cross_v)
        ptcls.vel += self.sdt * b_cross_v + self.ccodt * b_cross_b_cross_v

        # Second Acceleration half step: Apply exp(dt * V_F / 2)
        ptcls.vel += 0.5 * ptcls.acc * self.dt

        # Full step position update
        ptcls.pos += ptcls.vel * self.dt

        # Periodic boundary condition
        enforce_pbc(ptcls.pos, ptcls.pbc_cntr, self.box_lengths)

        # Compute total potential energy and acceleration for second half step velocity update
        potential_energy = self.update_accelerations(ptcls)

        return potential_energy

    def magnetic_pos_verlet_zdir(self, ptcls):
        """
        Update particles' class based on position verlet method in the case of a
        constant magnetic field along the :math:`z` axis. For more info see eq. (79) of Ref. :cite:`Chin2008`

        Parameters
        ----------
        ptcls: :class:`sarkas.particles.Particles`
            Particles data.

        Returns
        -------
        potential_energy : float
             Total potential energy.

        Notes
        -----
        This integrator is faster than `magnetic_verlet` but valid only for a magnetic field in the :math:`z`-direction.
        This is the preferred choice in this case.
        """

        # Position update
        ptcls.pos += 0.5 * ptcls.vel * self.dt

        # Enforce boundary condition
        self.enforce_bc(ptcls)

        # Compute total potential energy and acceleration for second half step velocity update
        potential_energy = self.update_accelerations(ptcls)

        # First half step of velocity update
        # # Magnetic rotation x - velocity
        # (B x v)_x  = -v_y, (B x B x v)_x = -v_x
        self.v_B[:, 0] = ptcls.vel[:, 1] * self.sdt[:, 0] + ptcls.vel[:, 0] * self.cdt[:, 0]
        # Magnetic rotation y - velocity
        # (B x v)_y  = v_x, (B x B x v)_y = -v_y
        self.v_B[:, 1] = -ptcls.vel[:, 0] * self.sdt[:, 0] + ptcls.vel[:, 1] * self.cdt[:, 1]

        # Magnetic + Const force field x - velocity
        # (B x a)_x  = -a_y, (B x B x a)_x = -a_x
        self.v_F[:, 0] = (
            self.ccodt[:, 1] / self.omega_c[:, 1] * ptcls.acc[:, 1]
            + self.sdt[:, 0] / self.omega_c[:, 0] * ptcls.acc[:, 0]
        )
        # Magnetic + Const force field y - velocity
        # (B x a)_y  = a_x, (B x B x a)_y = -a_y
        self.v_F[:, 1] = (
            -self.ccodt[:, 0] / self.omega_c[:, 0] * ptcls.acc[:, 0]
            + self.sdt[:, 1] / self.omega_c[:, 1] * ptcls.acc[:, 1]
        )

        ptcls.vel[:, 0] = self.v_B[:, 0] + self.v_F[:, 0]
        ptcls.vel[:, 1] = self.v_B[:, 1] + self.v_F[:, 1]
        ptcls.vel[:, 2] += self.dt * ptcls.acc[:, 2]

        # Position update
        ptcls.pos += 0.5 * ptcls.vel * self.dt

        # Enforce boundary condition
        self.enforce_bc(ptcls)

        return potential_energy

    def magnetic_pos_verlet(self, ptcls):
        """
        Update particles' class based on position verlet method in the case of an arbitrary direction of the
        constant magnetic field. For more info see eq. (79) of Ref. :cite:`Chin2008`

        Parameters
        ----------
        ptcls: :class:`sarkas.particles.Particles`
            Particles data.

        Returns
        -------
        potential_energy : float
             Total potential energy.

        Notes
        -----
        :cite:`Chin2008` equations are written for a negative charge. This allows him to write
        :math:`\\dot{\\mathbf v} = \\omega_c \\hat{B} \\times \\mathbf v`. In the case of positive charges we will have
        :math:`\\dot{\\mathbf v} = - \\omega_c \\hat{B} \\times \\mathbf v`.
        Hence the reason of the different signs in the formulas below compared to Chin's.

        Warnings
        --------
        This integrator is valid for a magnetic field in an arbitrary direction. However, while the integrator works for
        an arbitrary direction, methods in :mod:`sarkas.tools.observables` work only for a magnetic field in the
        :math:`z` - direction. Hence, if you choose to use this integrator remember to change your physical observables.

        """
        # Half position update
        ptcls.pos += ptcls.vel * self.dt

        # Enforce boundary condition
        self.enforce_bc(ptcls)

        # Compute total potential energy and acceleration for second half step velocity update
        potential_energy = self.update_accelerations(ptcls)

        # Calculate the cross products
        b_cross_v = cross(self.magnetic_field_uvector, ptcls.vel)
        b_cross_b_cross_v = cross(self.magnetic_field_uvector, b_cross_v)
        b_cross_a = cross(self.magnetic_field_uvector, ptcls.acc)
        b_cross_b_cross_a = cross(self.magnetic_field_uvector, b_cross_a)

        # First half step of velocity update
        ptcls.vel += -self.sdt * b_cross_v + self.ccodt * b_cross_b_cross_v

        ptcls.vel += (
            ptcls.acc * self.dt - self.ccodt / self.omega_c * b_cross_a + self.dt * self.ssodt * b_cross_b_cross_a
        )

        # Second half position update
        ptcls.pos += ptcls.vel * self.dt

        # Enforce boundary condition
        self.enforce_bc(ptcls)

        return potential_energy

    def cyclotronic_zdir(self, ptcls):
        """
        Update particles' class using the cyclotronic algorithm in the case of a
        constant magnetic field along the :math:`z` axis.
        For more info see eqs. (16) - (17) of Ref. :cite:`Patacchini2009`

        Parameters
        ----------
        ptcls: :class:`sarkas.particles.Particles`
            Particles data.

        Returns
        -------
        potential_energy : float
             Total potential energy.

        """
        # Drift half step
        # Rotate Positions
        ptcls.pos[:, 0] += (
            ptcls.vel[:, 0] * self.sdt[:, 0] / self.omega_c[:, 0]
            + ptcls.vel[:, 1] * self.ccodt[:, 1] / self.omega_c[:, 1]
        )
        ptcls.pos[:, 1] += (
            ptcls.vel[:, 1] * self.sdt[:, 1] / self.omega_c[:, 1]
            - ptcls.vel[:, 0] * self.ccodt[:, 0] / self.omega_c[:, 0]
        )
        ptcls.pos[:, 2] += 0.5 * ptcls.vel[:, 2] * self.dt
        # Enforce boundary condition
        self.enforce_bc(ptcls)
        # Create rotated velocities
        self.v_B[:, 0] = self.cdt[:, 0] * ptcls.vel[:, 0] + self.sdt[:, 1] * ptcls.vel[:, 1]
        self.v_B[:, 1] = self.cdt[:, 1] * ptcls.vel[:, 1] - self.sdt[:, 0] * ptcls.vel[:, 0]
        ptcls.vel[:, :2] = self.v_B[:, :2].copy()
        # Compute total potential energy and accelerations
        potential_energy = self.update_accelerations(ptcls)

        # Kick full step
        ptcls.vel += ptcls.acc * self.dt

        # Drift half step
        # Rotate Positions
        ptcls.pos[:, 0] += (
            ptcls.vel[:, 0] * self.sdt[:, 0] / self.omega_c[:, 0]
            + ptcls.vel[:, 1] * self.ccodt[:, 1] / self.omega_c[:, 1]
        )
        ptcls.pos[:, 1] += (
            ptcls.vel[:, 1] * self.sdt[:, 1] / self.omega_c[:, 1]
            - ptcls.vel[:, 0] * self.ccodt[:, 0] / self.omega_c[:, 0]
        )
        ptcls.pos[:, 2] += 0.5 * ptcls.vel[:, 2] * self.dt
        # Enforce boundary condition
        self.enforce_bc(ptcls)
        # Create rotated velocities
        self.v_B[:, 0] = self.cdt[:, 0] * ptcls.vel[:, 0] + self.sdt[:, 1] * ptcls.vel[:, 1]
        self.v_B[:, 1] = self.cdt[:, 1] * ptcls.vel[:, 1] - self.sdt[:, 0] * ptcls.vel[:, 0]
        # Update final velocities
        ptcls.vel[:, :2] = self.v_B[:, :2].copy()

        return potential_energy

    def cyclotronic(self, ptcls):
        """
        Update particles' class using the cyclotronic algorithm in the case of a
        constant magnetic field along the :math:`z` axis.
        For more info see eqs. (16) - (17) of Ref. :cite:`Patacchini2009`

        Parameters
        ----------
        ptcls: :class:`sarkas.particles.Particles`
            Particles data.

        Returns
        -------
        potential_energy : float
             Total potential energy.

        """
        # Drift half step

        # Calculate the cross products
        b_cross_v = cross(self.magnetic_field_uvector, ptcls.vel)
        b_cross_b_cross_v = cross(self.magnetic_field_uvector, b_cross_v)
        # Rotate Positions
        ptcls.pos += (
            0.5 * ptcls.vel * self.dt
            - self.ccodt * b_cross_v / self.omega_c
            + 0.5 * self.dt * self.ssodt * b_cross_b_cross_v
        )
        # Enforce boundary condition
        self.enforce_bc(ptcls)
        # First half step of velocity update
        ptcls.vel += -self.sdt * b_cross_v + self.ccodt * b_cross_b_cross_v
        # Compute total potential energy and accelerations
        potential_energy = self.update_accelerations(ptcls)

        # Kick full step
        ptcls.vel += ptcls.acc * self.dt

        # Drift half step
        # Calculate the cross products
        b_cross_v = cross(self.magnetic_field_uvector, ptcls.vel)
        b_cross_b_cross_v = cross(self.magnetic_field_uvector, b_cross_v)
        # Rotate Positions
        ptcls.pos += (
            0.5 * ptcls.vel * self.dt
            - self.ccodt * b_cross_v / self.omega_c
            + 0.5 * self.dt * self.ssodt * b_cross_b_cross_v
        )
        # Enforce boundary condition
        self.enforce_bc(ptcls)
        # Second half step of velocity update
        ptcls.vel += -self.sdt * b_cross_v + self.ccodt * b_cross_b_cross_v

        return potential_energy

    def thermostate(self, ptcls):
        """
        Update particles' velocities according to the chosen thermostat

        Parameters
        ----------
        ptcls : :class:`sarkas.particles.Particles`
            Particles' data.

        """
        _, T = ptcls.kinetic_temperature()
        berendsen(ptcls.vel, self.species_temperatures, T, self.species_num, self.thermalization_rate)

    def periodic_bc(self, ptcls):
        """
        Applies periodic boundary conditions by calling enforce_pbc

        Parameters
        ----------
        ptcls: :class:`sarkas.particles.Particles`
            Particles data.

        """

        enforce_pbc(ptcls.pos, ptcls.pbc_cntr, self.box_lengths)

    def absorbing_bc(self, ptcls):
        """
        Applies absorbing boundary conditions by calling enforce_abc

        Parameters
        ----------
        ptcls: :class:`sarkas.particles.Particles`
            Particles data.

        """

        enforce_abc(ptcls.pos, ptcls.vel, ptcls.acc, ptcls.charges, self.box_lengths)

    def open_bc(self, ptcls):
        """
        Applies open boundary conditions. Basically it does nothing. pass

        Parameters
        ----------
        ptcls: :class:`sarkas.particles.Particles`
            Particles data.

        """

        pass

    def reflecting_bc(self, ptcls):
        """
        Applies reflective boundary conditions by calling enforce_rbc

        Parameters
        ----------
        ptcls: :class:`sarkas.particles.Particles`
            Particles data.

        """

        enforce_rbc(ptcls.pos, ptcls.vel, self.box_lengths, self.dt)

    def pretty_print(self):
        """Print integrator and thermostat information in a user-friendly way."""

        if self.thermalization:
            print("\nTHERMOSTAT: ")
            print(f"Type: {self.thermostat_type}")
            print(f"First thermostating timestep, i.e. thermalization_timestep = {self.thermalization_timestep}")
            print(f"Berendsen parameter tau: {self.berendsen_tau:.3f} [timesteps]")
            print(f"Berendsen relaxation rate: {self.thermalization_rate:.3f} [1/timesteps] ")
            print("Thermostating temperatures: ")
            for i, (t, t_ev) in enumerate(zip(self.thermostate_temperatures, self.thermostate_temperatures_eV)):
                print(f"Species ID {i}: T_eq = {t:.6e} [K] = {t_ev:.6e} [eV]")

        print("\nINTEGRATOR: ")
        print(f"Equilibration Integrator Type: {self.equilibration_type}")
        if self.magnetized:
            print(f"Magnetization Integrator Type: {self.magnetization_type}")
        print(f"Production Integrator Type: {self.production_type}")

        wp_tot = norm(self.species_plasma_frequencies)
        wp_dt = wp_tot * self.dt

        print(f"Time step = {self.dt:.6e} [s]")
        print(f"Total plasma frequency = {wp_tot:.6e} [rad/s]")
        print(f"w_p dt = {wp_dt:.4f} ~ 1/{int(1.0 / wp_dt)}")

        if self.potential_type == "qsp":
            wp_e = self.species_plasma_frequencies[0]
            wp_ions = norm(self.species_plasma_frequencies[1:])
            print(f"e plasma frequency = {wp_e:.6e} [rad/s]")
            print(f"total ion plasma frequency = {wp_ions:.6e} [rad/s]")
            print(f"w_pe dt = {self.dt * wp_e:.4f} ~ 1/{int(1.0 / (self.dt * wp_e))}")
            print(f"w_pi dt = {self.dt * wp_ions:.4f} ~ 1/{int(1.0 / (self.dt * wp_ions))}")

        elif self.potential_type == "lj":
            print(f"The plasma frequency is defined as w_p = sqrt( epsilon / (sigma^2 * mass) )")

        if self.magnetized:
            high_wc_dt = abs(self.species_cyclotron_frequencies).max() * self.dt
            low_wc_dt = abs(self.species_cyclotron_frequencies).min() * self.dt

            if high_wc_dt > low_wc_dt:
                print(f"Highest w_c dt = {high_wc_dt:2.4f} = {high_wc_dt / pi:.4f} pi")
                print(f"Smallest w_c dt = {low_wc_dt:2.4f} = {low_wc_dt / pi:.4f} pi")
            else:
                print(f"w_c dt = {high_wc_dt:2.4f} = {high_wc_dt / pi:.4f} pi")

        if self.equilibration_type == "langevin" or self.production_type == "langevin":
            print(f"langevin_gamma * dt = {self.langevin_gamma * self.dt:.4e}")
            print(f"langevin_gamma / wp = {self.langevin_gamma / wp_tot:.4e}")

@jit(void(float64[:, :], float64[:], float64[:], int64[:], float64), nopython=True)
def berendsen(vel, T_desired, T, species_np, tau):
    """
    Numba'd function to update particle velocity based on Berendsen thermostat :cite:`Berendsen1984`.

    Parameters
    ----------
    vel : numpy.ndarray
        Particles' velocities to rescale.

    T_desired : numpy.ndarray
        Target temperature of each species.

    T : numpy.ndarray
        Instantaneous temperature of each species.

    species_np : numpy.ndarray
        Number of each species.

    tau : float
        Scale factor.

    """

    # if it < therm_timestep:
    #     fact = sqrt(T_desired / T)
    # else:
    #     fact = sqrt(1.0 + (T_desired / T - 1.0) * tau)  # eq.(11)

    # branchless programming
    fact = sqrt(1.0 + (T_desired / T - 1.0) * tau)
    species_start = 0
    species_end = 0

    for i, num in enumerate(species_np):
        species_end += num
        vel[species_start:species_end, :] *= fact[i]
        species_start += num

@jit(void(float64[:, :], float64[:, :], float64[:]), nopython=True)
def enforce_pbc(pos, cntr, box_vector) -> None:
    """
    Numba'd function to enforce periodic boundary conditions.

    Parameters
    ----------
    pos : numpy.ndarray
        Particles' positions.

    cntr : numpy.ndarray
        Counter for the number of times each particle get folded back into the main simulation box

    box_vector : numpy.ndarray
        Box Dimensions.

    """

    # Loop over all particles
    for p in arange(pos.shape[0]):
        for d in arange(pos.shape[1]):

            # If particle is outside of box in positive direction, wrap to negative side
            if pos[p, d] > box_vector[d]:
                pos[p, d] -= box_vector[d]
                cntr[p, d] += 1
            # If particle is outside of box in negative direction, wrap to positive side
            if pos[p, d] < 0.0:
                pos[p, d] += box_vector[d]
                cntr[p, d] -= 1

@jit(void(float64[:, :], float64[:, :], float64[:, :], float64[:], float64[:]), nopython=True)
def enforce_abc(pos, vel, acc, charges, box_vector) -> None:
    """
    Numba'd function to enforce absorbing boundary conditions.

    Parameters
    ----------
    pos: numpy.ndarray
        Particles' positions.

    vel : numpy.ndarray
        Particles' velocities.

    acc : numpy.ndarray
        Particles' accelerations.

    charges : numpy.ndarray
        Charge of each particle. Shape = (``total_num_ptcls``).

    box_vector: numpy.ndarray
        Box Dimensions.

    """

    # Loop over all particles
    for p in arange(pos.shape[0]):
        for d in arange(pos.shape[1]):

            # If particle is outside of box in positive direction, remove charge, velocity and acceleration
            if pos[p, d] >= box_vector[d]:
                pos[p, d] = box_vector[d]
                vel[p, :] = zeros(3)
                acc[p, :] = zeros(3)
                charges[p] = 0.0
            # If particle is outside of box in negative direction, remove charge, velocity and acceleration
            if pos[p, d] <= 0.0:
                pos[p, d] = 0.0
                vel[p, :] = zeros(3)
                acc[p, :] = zeros(3)
                charges[p] = 0.0

@jit(void(float64[:, :], float64[:, :], float64[:], float64), nopython=True)
def enforce_rbc(pos, vel, box_vector, dt) -> None:
    """
    Numba'd function to enforce reflecting boundary conditions.

    Parameters
    ----------
    pos: numpy.ndarray
        Particles' positions.

    vel : numpy.ndarray
        Particles' velocities.

    box_vector: numpy.ndarray
        Box Dimensions.

    dt : float
        Timestep.

    """

    # Loop over all particles
    for p in arange(pos.shape[0]):
        for d in arange(pos.shape[1]):

            # If particle is outside of box in positive direction, wrap to negative side
            if pos[p, d] > box_vector[d] or pos[p, d] < 0.0:
                # Revert velocity
                vel[p, d] *= -1.0
                # Restore previous position assuming verlet algorithm
                pos[p, d] += vel[p, d] * dt

# @jit(void(float64[:, :], int64[:], float64[:]), nopython=True)
# def remove_drift(vel, nums, masses) -> None:
#     """
#     Numba'd function to enforce conservation of total linear momentum.
#     It updates :attr:`sarkas.particles.Particles.vel`.
#
#     Parameters
#     ----------
#     vel: numpy.ndarray
#         Particles' velocities.
#
#     nums: numpy.ndarray
#         Number of particles of each species.
#
#     masses: numpy.ndarray
#         Mass of each species.
#
#     """
#
#     P = zeros((len(nums), vel.shape[1]))
#     species_start = 0
#     for ic in range(len(nums)):
#         species_end = species_start + nums[ic]
#         P[ic, :] = vel[species_start:species_end, :].sum(axis=0) * masses[ic]
#         species_start = species_end
#
#     if P.sum(axis=0).any() > 1e-40:
#         # Remove tot momentum
#         species_start = 0
#         for ic in range(len(nums)):
#             species_end = species_start + nums[ic]
#             vel[species_start:species_end, :] -= P[ic, :] / (float(nums[ic]) * masses[ic])
#             species_start = species_end


import numpy as np
import random
from sub_func.get_iou import iou

__all__ = ['calc_rpn']

def calc_rpn(config, img_data, width, height, resized_width, resized_height, resize_func):
    """(Important part!) Calculate the rpn for all anchors
        If feature map has shape 38x50=1900, there are 1900x(3 x 3)=17100 potential anchors

    Args:
        config: Config instance
        img_data: augmented image data
        width: original image width (e.g. 600)
        height: original image height (e.g. 800)
        resized_width: resized image width according to config.im_size (e.g. 300)
        resized_height: resized image height according to config.im_size (e.g. 400)
        resize_func: function to calculate final layer's feature map (of base model) size according to input image size

    Returns:
        y_rpn_cls: list(num_bboxes, y_is_box_valid + y_rpn_overlap)
            y_is_box_valid: 0 or 1 (0 means the box is invalid, 1 means the box is valid)
            y_rpn_overlap: 0 or 1 (0 means the box is not an object, 1 means the box is an object)
        y_rpn_regr: list(num_bboxes, 4*y_rpn_overlap + y_rpn_regr)
            y_rpn_regr: x1,y1,x2,y2 bunding boxes coordinates
    """
    downscale = float(config.rpn_stride)
    anchor_sizes = config.anchor_box_scales  # 128, 256, 512
    anchor_ratios = config.anchor_box_ratios  # 1:1, 1:2*sqrt(2), 2*sqrt(2):1
    num_anchors = len(anchor_sizes) * len(anchor_ratios)  # 3x3=9

    # calculate the output map size based on the network architecture
    (output_width, output_height) = resize_func(resized_width, resized_height)

    n_anchratios = len(anchor_ratios)  # 3

    # initialise empty output objectives
    y_rpn_overlap = np.zeros((output_height, output_width, num_anchors))  # box is object?
    y_is_box_valid = np.zeros((output_height, output_width, num_anchors))  # box is valid?
    y_rpn_regr = np.zeros((output_height, output_width, num_anchors * 4))  # bbox coordinates

    num_bboxes = len(img_data['bboxes'])

    num_anchors_for_bbox = np.zeros(num_bboxes).astype(int)  # 각 bbox 마다 anchor 갯수
    best_anchor_for_bbox = -1 * np.ones((num_bboxes, 4)).astype(int)  # 각 bbox 마다 best anchor (x, y, w, h)
    best_iou_for_bbox = np.zeros(num_bboxes).astype(np.float32)  # 각 bbox 마다 best IOU
    best_x_for_bbox = np.zeros((num_bboxes, 4)).astype(int)  # 각 bbox 마다 (x1, x2, y1, y2)
    best_dx_for_bbox = np.zeros((num_bboxes, 4)).astype(np.float32)  # 각 bbox 마다 (tx, ty, tw, th)

    # get the GT box coordinates, and resize to account for image resizing
    gta = np.zeros((num_bboxes, 4))
    for bbox_num, bbox in enumerate(img_data['bboxes']):
        # get the GT box coordinates, and resize to account for image resizing
        gta[bbox_num, 0] = bbox['x1'] * (resized_width / float(width))
        gta[bbox_num, 1] = bbox['x2'] * (resized_width / float(width))
        gta[bbox_num, 2] = bbox['y1'] * (resized_height / float(height))
        gta[bbox_num, 3] = bbox['y2'] * (resized_height / float(height))

    # rpn ground truth

    for anchor_size_idx in range(len(anchor_sizes)):
        for anchor_ratio_idx in range(n_anchratios):
            anchor_x = anchor_sizes[anchor_size_idx] * anchor_ratios[anchor_ratio_idx][0]
            anchor_y = anchor_sizes[anchor_size_idx] * anchor_ratios[anchor_ratio_idx][1]
            # anchor_x, anchor_y : anchor size 에 가로 세율 비율 곱한 값

            for ix in range(output_width):
                # x-coordinates of the current anchor box
                x1_anc = downscale * (ix + 0.5) - anchor_x / 2
                x2_anc = downscale * (ix + 0.5) + anchor_x / 2

                # ignore boxes that go across image boundaries
                if x1_anc < 0 or x2_anc > resized_width:
                    continue

                for jy in range(output_height):
                    # y-coordinates of the current anchor box
                    y1_anc = downscale * (jy + 0.5) - anchor_y / 2
                    y2_anc = downscale * (jy + 0.5) + anchor_y / 2

                    # ignore boxes that go across image boundaries
                    if y1_anc < 0 or y2_anc > resized_height:
                        continue

                    # bbox_type indicates whether an anchor should be a target
                    # Initialize with 'negative'
                    bbox_type = 'neg'

                    # this is the best IOU for the (x,y) coord and the current anchor
                    # note that this is different from the best IOU for a GT bbox
                    best_iou_for_loc = 0.0

                    for bbox_num in range(num_bboxes):

                        # get IOU of the current GT box and the current anchor box
                        gt_coord = [
                            gta[bbox_num, 0],
                            gta[bbox_num, 2],
                            gta[bbox_num, 1],
                            gta[bbox_num, 3],
                        ]
                        anc_coord = [x1_anc, y1_anc, x2_anc, y2_anc]
                        curr_iou = iou(gt_coord, anc_coord)
                        # calculate the regression targets if they will be needed
                        if curr_iou > best_iou_for_bbox[bbox_num] or curr_iou > config.rpn_max_overlap:
                            # cx, cy : gt 의 중점좌표
                            cx = (gta[bbox_num, 0] + gta[bbox_num, 1]) / 2.0
                            cy = (gta[bbox_num, 2] + gta[bbox_num, 3]) / 2.0
                            # 현재 anchor 의 좌표
                            cxa = (x1_anc + x2_anc) / 2.0
                            cya = (y1_anc + y2_anc) / 2.0

                            # x,y are the center point of ground-truth bbox
                            # xa,ya are the center point of anchor bbox (xa=downscale * (ix + 0.5); ya=downscale * (iy+0.5))
                            # w,h are the width and height of ground-truth bbox
                            # wa,ha are the width and height of anchor bboxe
                            # tx = (x - xa) / wa
                            # ty = (y - ya) / ha
                            # tw = log(w / wa)
                            # th = log(h / ha)
                            wa = x2_anc - x1_anc
                            ha = y2_anc - y1_anc

                            tx = (cx - cxa) / wa
                            ty = (cy - cya) / ha
                            tw = np.log((gta[bbox_num, 1] - gta[bbox_num, 0]) / wa)
                            th = np.log((gta[bbox_num, 3] - gta[bbox_num, 2]) / ha)

                        if img_data['bboxes'][bbox_num]['class'] != 'bg':
                            # all GT boxes should be mapped to an anchor box, so we keep track of which anchor box was best
                            if curr_iou > best_iou_for_bbox[bbox_num]:
                                best_anchor_for_bbox[bbox_num] = [jy, ix, anchor_ratio_idx, anchor_size_idx]
                                best_iou_for_bbox[bbox_num] = curr_iou
                                best_x_for_bbox[bbox_num,:] = [x1_anc, x2_anc, y1_anc, y2_anc]
                                best_dx_for_bbox[bbox_num,:] = [tx, ty, tw, th]

                            # we set the anchor to positive if the IOU is >0.7 (it does not matter if there was another better box, it just indicates overlap)
                            if curr_iou > config.rpn_max_overlap:
                                bbox_type = 'pos'
                                num_anchors_for_bbox[bbox_num] += 1
                                # we update the regression layer target if this IOU is the best for the current (x,y) and anchor position
                                if curr_iou > best_iou_for_loc:
                                    best_iou_for_loc = curr_iou
                                    best_regr = (tx, ty, tw, th)

                            # if the IOU is >0.3 and <0.7, it is ambiguous and no included in the objective
                            if config.rpn_min_overlap < curr_iou < config.rpn_max_overlap:
                                # gray zone between neg and pos
                                if bbox_type != 'pos':
                                    bbox_type = 'neutral'

                    # turn on or off outputs depending on IOUs
                    if bbox_type == 'neg':
                        y_is_box_valid[jy, ix, anchor_ratio_idx + n_anchratios * anchor_size_idx] = 1
                        y_rpn_overlap[jy, ix, anchor_ratio_idx + n_anchratios * anchor_size_idx] = 0
                    elif bbox_type == 'neutral':
                        y_is_box_valid[jy, ix, anchor_ratio_idx + n_anchratios * anchor_size_idx] = 0
                        y_rpn_overlap[jy, ix, anchor_ratio_idx + n_anchratios * anchor_size_idx] = 0
                    elif bbox_type == 'pos':
                        y_is_box_valid[jy, ix, anchor_ratio_idx + n_anchratios * anchor_size_idx] = 1
                        y_rpn_overlap[jy, ix, anchor_ratio_idx + n_anchratios * anchor_size_idx] = 1
                        start = 4 * (anchor_ratio_idx + n_anchratios * anchor_size_idx)
                        y_rpn_regr[jy, ix, start:start+4] = best_regr

    # we ensure that every bbox has at least one positive RPN region
    for idx in range(num_anchors_for_bbox.shape[0]):
        if num_anchors_for_bbox[idx] == 0:
            # no box with an IOU greater than zero ...
            if best_anchor_for_bbox[idx, 0] == -1:
                continue
            y_is_box_valid[
                best_anchor_for_bbox[idx, 0], best_anchor_for_bbox[idx, 1], best_anchor_for_bbox[idx, 2] + n_anchratios *
                best_anchor_for_bbox[idx, 3]] = 1
            y_rpn_overlap[
                best_anchor_for_bbox[idx, 0], best_anchor_for_bbox[idx, 1], best_anchor_for_bbox[idx, 2] + n_anchratios *
                best_anchor_for_bbox[idx, 3]] = 1
            start = 4 * (best_anchor_for_bbox[idx,2] + n_anchratios * best_anchor_for_bbox[idx, 3])
            y_rpn_regr[
                best_anchor_for_bbox[idx, 0], best_anchor_for_bbox[idx, 1], start:start+4] = best_dx_for_bbox[idx, :]

    y_rpn_overlap = np.transpose(y_rpn_overlap, (2, 0, 1))
    y_rpn_overlap = np.expand_dims(y_rpn_overlap, axis=0)
    # y_rpn_overlap: (1, num_anchors, output_height, output_width)

    y_is_box_valid = np.transpose(y_is_box_valid, (2, 0, 1))
    y_is_box_valid = np.expand_dims(y_is_box_valid, axis=0)
    # y_is_box_valid: (1, num_anchors, output_height, output_width)

    y_rpn_regr = np.transpose(y_rpn_regr, (2, 0, 1))
    y_rpn_regr = np.expand_dims(y_rpn_regr, axis=0)
    # y_rpn_regr: (1, num_anchros * 4, output_height, output_width)

    pos_locs = np.where(np.logical_and(y_rpn_overlap[0, :, :, :] == 1, y_is_box_valid[0, :, :, :] == 1))
    neg_locs = np.where(np.logical_and(y_rpn_overlap[0, :, :, :] == 0, y_is_box_valid[0, :, :, :] == 1))

    num_pos = len(pos_locs[0])

    # one issue is that the RPN has many more negative than positive regions,
    # so we turn off some of the negative regions.
    # We also limit it to 256 regions.
    num_regions = 256

    if len(pos_locs[0]) > num_regions / 2:
        val_locs = random.sample(
            range(len(pos_locs[0])),
            len(pos_locs[0])-num_regions/2,
        )
        y_is_box_valid[0, pos_locs[0][val_locs], pos_locs[1][val_locs], pos_locs[2][val_locs]] = 0
        num_pos = num_regions / 2

    if len(neg_locs[0]) + num_pos > num_regions:
        val_locs = random.sample(
            range(len(neg_locs[0])),
            len(neg_locs[0])-num_pos,
        )
        y_is_box_valid[0, neg_locs[0][val_locs], neg_locs[1][val_locs], neg_locs[2][val_locs]] = 0

    y_rpn_cls = np.concatenate([y_is_box_valid, y_rpn_overlap], axis=1)
    y_rpn_regr = np.concatenate([np.repeat(y_rpn_overlap, 4, axis=1), y_rpn_regr], axis=1)
    return np.copy(y_rpn_cls), np.copy(y_rpn_regr), num_pos


"""
15N–1HN DQ/ZQ CPMG
==================

Analyzes 15N and 1H chemical exchange by applying CPMG pulses on 15N and 1H
simultaneously. The spin system is maintained as DQ or ZQ during Trelax, and
is calculated using the (15n)×(15n), two-spin matrix, where n is the number
of states::

    {        Ix(a),   Iy(a),   Iz(a),   Sx(a), IxSx(a), IySx(a), IzSx(a),
      Sy(a), IxSy(a), IySy(a), IzSy(a), Sz(a), IxSz(a), IySz(a), IzSz(a),
             Ix(b),   Iy(b),   Iz(b),   Sx(b), IxSx(b), IySx(b), IzSx(b),
      Sy(b), IxSy(b), IySy(b), IzSy(b), Sz(b), IxSz(b), IySz(b), IzSz(b), ... }

The phase cycle of CPMG pulses is chosen based on νCPMG as described in the
reference, which is a mixture of constant-phase and XY-family phase cycles.

References
----------

Orekhov, Korzhnev and Kay. J Am Chem Soc (2004) 126:1886-1891

Note
----

A sample configuration file for this module is available using the command::

    $ chemex config cpmg_hn_dq_zq

"""
import functools as ft

import numpy as np

import chemex.experiments.helper as ceh
import chemex.helper as ch
import chemex.nmr.liouvillian as cnl

_SCHEMA = {
    "type": "object",
    "properties": {
        "experiment": {
            "type": "object",
            "properties": {
                "time_t2": {"type": "number"},
                "carrier_h": {"type": "number"},
                "carrier_n": {"type": "number"},
                "pw90_h": {"type": "number"},
                "pw90_n": {"type": "number"},
                "dq_flg": {"type": "boolean"},
                "observed_state": {
                    "type": "string",
                    "pattern": "[a-z]",
                    "default": "a",
                },
            },
            "required": [
                "time_t2",
                "carrier_h",
                "carrier_n",
                "pw90_h",
                "pw90_n",
                "dq_flg",
            ],
        }
    },
}

def read(config):
    ch.validate(config, _SCHEMA)
    config["basis"] = cnl.Basis(type="ixyzsxyz", spin_system="nh")
    config["fit"] = _fit_this()
    return ceh.load_experiment(config=config, pulse_seq_cls=PulseSeq)

def _fit_this():
    return {
        "rates": ["r2mq_is_{observed_state}", "mu_is_{observed_state}"],
        "model_free": ["tauc_{observed_state}", "s2_{observed_state}"],
    }

class PulseSeq:
    def __init__(self, config, propagator):
        self.prop = propagator
        settings = config["experiment"]
        self.time_t2 = settings["time_t2"]
        self.prop.carrier_i = settings["carrier_n"]
        self.prop.carrier_s = settings["carrier_h"]
        self.pw90_i = settings["pw90_n"]
        self.pw90_s = settings["pw90_h"]
        self.prop.b1_i = 1 / (4.0 * self.pw90_i)
        self.prop.b1_s = 1 / (4.0 * self.pw90_s)
        self.dq_flg = settings["dq_flg"]
        self.prop.detection = self._get_detection(settings["observed_state"])

    @ft.lru_cache(maxsize=10000)
    def calculate(self, ncycs, params_local):
        self.prop.update(params_local)

        # Calculation of the propagators corresponding to all the delays
        tau_cps = self._get_tau_cps(ncycs)
        tau_cp_list = list(tau_cps.values())
        delays = dict(zip(tau_cp_list, self.prop.delays(tau_cp_list)))
        d_cp = {ncyc: delays[delay] for ncyc, delay in tau_cps.items()}

        # Calculation of the propagators corresponding to all the pulses
        p9024090_1 = self.prop.p9024090_is_1[[0, 1], [0, 1]]
        p9024090_2 = self.prop.p9024090_is_2[[0, 1], [0, 1]]

        # Getting the starting magnetization
        start = self.prop.get_start_magnetization(["2ixsx"])

        # Calculating the cpmg trains
        intst = {0: self.prop.detect(start)}
        for ncyc in set(ncycs) - {0}:
            phases1, phases2 = self._get_phases(ncyc)
            echo1 = d_cp[ncyc] @ p9024090_1 @ d_cp[ncyc]
            echo2 = d_cp[ncyc] @ p9024090_2 @ d_cp[ncyc]
            cpmg1 = ft.reduce(np.matmul, echo1[phases1])
            cpmg2 = ft.reduce(np.matmul, echo2[phases2])
            intst[ncyc] = self.prop.detect(0.5 * (cpmg1 + cpmg2) @ start)
        return np.array([intst[ncyc] for ncyc in ncycs])

    @ft.lru_cache()
    def _get_tau_cps(self, ncycs):
        ncycs_ = np.asarray(ncycs)
        ncycs_ = ncycs_[ncycs_ > 0]
        return dict(
            zip(ncycs_, self.time_t2 / (4.0 * ncycs_) - 7.0 / 3.0 * self.pw90_i)
        )

    @ft.lru_cache()
    def _get_phases(self, ncyc):
        nu_cpmg = self.ncycs_to_nu_cpmgs(ncyc)
        if nu_cpmg < 51.0:
            cp_phases1 = [0, 1, 0, 1]
            cp_phases2 = [1, 0, 1, 0]
        elif nu_cpmg < 255.0:
            cp_phases1 = [0]
            cp_phases2 = [1]
        else:
            cp_phases1 = [0, 1, 0, 1, 1, 0, 1, 0]
            cp_phases2 = [1, 0, 1, 0, 0, 1, 0, 1]
        phases1 = np.take(cp_phases1, np.flip(np.arange(2 * ncyc)), mode="wrap")
        phases2 = np.take(cp_phases2, np.flip(np.arange(2 * ncyc)), mode="wrap")
        return phases1, phases2

    def _get_detection(self, state):
        if self.dq_flg:
            return f"[2ixsx_{state}] - [2iysy_{state}]"
        return f"[2ixsx_{state}] + [2iysy_{state}]"

    def ncycs_to_nu_cpmgs(self, ncycs):
        ncycs_ = np.asarray(ncycs)
        ncycs_ = ncycs_[ncycs_ > 0]
        return ncycs_ / self.time_t2


# -*- coding: utf-8 -*-
'''This module provides a user-friendly pythonic wrapper for the low-level C interface functions.'''

from __future__ import division, print_function, absolute_import, unicode_literals

import datetime as dt
import calendar
import warnings

import numpy as np

from aacgm2._aacgmv2 import A2G, TRACE, BADIDEA, ALLOWTRACE, GEOCENTRIC, setDateTime, aacgmConvert
from aacgm2 import IGRF_12_COEFFS

aacgmConvert_vectorized = np.vectorize(aacgmConvert)

def convert(lat, lon, alt, date=None, a2g=False, trace=False, allowtrace=False, badidea=False, geocentric=False):
    '''Converts to/from geomagnetic coordinates.

    This is a user-friendly pythonic wrapper for the low-level C interface
    functions available in :mod:`aacgm2._aacgmv2`.

    Parameters
    ==========
    lat,lon,alt : array_like
        Input latitude(s), longitude(s) and altitude(s). They must be
        `broadcastable to the same shape <http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html>`_.
    date : :class:`datetime.date`/:class:`datetime.datetime`, optional
        The date/time to use for the magnetic field model, default ``None`` (uses
        current time). Must be between 1900 and 2020.
    a2g : bool, optional
        Convert from AACGM-v2 to geographic coordinates, default ``False``
        (converts geographic to AACGM-v2).
    trace : bool, optional
        Use field-line tracing, default ``False`` (uses coefficients). Tracing
        is more precise and needed at altitudes > 2000 km, but significantly
        slower.
    allowtrace : bool, optional
        Automatically use field-line tracing above 2000 km, default ``False``
        (raises an exception for these altitudes unless ``trace=True`` or
        ``badidea=True``).
    badidea : bool, optional
        Allow use of coefficients above 2000 km (bad idea!)
    geocentric : bool, optional
        Assume inputs are geocentric with Earth radius 6371.2 km.

    Returns
    =======

    lat_out : ``numpy.ndarray``
        Converted latitude
    lon_out : ``numpy.ndarray``
        Converted longitude
    alt_out : ``numpy.ndarray``
        Converted altitude

    Raises
    ======

    ValueError
        if max(alt) > 2000 and neither of `trace`, `allowtrace`, or `badidea` is ``True``
    ValueError
        if latitude is outside the range -90 to +90 degrees
    RuntimeError
        if there was a problem in the C extension

    Notes
    =====

    This function exclusively relies on the `AACGM-v2 C library
    <https://engineering.dartmouth.edu/superdarn/aacgm.html>`_. Specifically,
    it calls the functions :func:`_aacgmv2.setDateTime` and
    :func:`_aacgmv2.aacgmConvert`, which are simple interfaces to the
    C library functions :func:`AACGM_v2_SetDateTime` and
    :func:`AACGM_v2_Convert`. Details of the techniques used to derive the
    AACGM-v2 coefficients are described by Shepherd, 2014 [1]_.

    .. [1] Shepherd, S. G. (2014), Altitude-adjusted corrected geomagnetic
       coordinates: Definition and functional approximations,
       J. Geophys. Res. Space Physics, 119, 7501--7521,
       doi:`10.1002/2014JA020264 <http://dx.doi.org/10.1002/2014JA020264>`_.

   '''

    # check values
    if np.min(alt) < 0:
        warnings.warn('Coordinate transformations are not intended for altitudes < 0 km', UserWarning)

    if np.max(alt) > 2000 and not (trace or allowtrace or badidea):
        raise ValueError('Coefficients are not valid for altitudes above 2000 km. You must either use field-line '
                         'tracing (trace=True or allowtrace=True) or indicate you know this is a bad idea '
                         '(badidea=True)')

    # check if latitudes are > 90.1 (to allow some room for rounding errors, which will be clipped)
    if np.max(np.abs(lat)) > 90.1:
        raise ValueError('Latitude must be in the range -90 to +90 degrees')
    np.clip(lat, -90, 90)

    # constrain longitudes between -180 and 180
    lon = ((np.asarray(lon) + 180) % 360) - 180

    # set to current date if none is given
    if date is None:
        date = dt.datetime.now()

    # add time info if only date is given
    if isinstance(date, dt.date):
        date = dt.datetime.combine(date, dt.time(0))

    # set current date and time
    setDateTime(date.year, date.month, date.day, date.hour, date.minute, date.second)

    # make flag
    flag = A2G*a2g + TRACE*trace + ALLOWTRACE*allowtrace + BADIDEA*badidea + GEOCENTRIC*geocentric

    # convert
    lat_out, lon_out, alt_out = aacgmConvert_vectorized(lat, lon, alt, flag)

    return lat_out, lon_out, alt_out

def convert_mlt(arr, datetime, m2a=False):
    '''Converts between magnetic local time (MLT) and AACGM-v2 longitude.

    .. note:: This function is not related to the AACGM-v2 C library, but is provided as
              a convenience in the hopes that it might be useful for some purposes.

    Parameters
    ==========
    arr : array_like or float
        Magnetic longitudes or MLTs to convert.
    datetime : :class:`datetime.datetime`
        Date and time for MLT conversion in Universal Time (UT).
    m2a : bool
        Convert MLT to AACGM-v2 longitude (default is ``False``, which implies
        conversion from AACGM-v2 longitude to MLT).

    Returns
    =======
    out : numpy.ndarray
        Converted coordinates/MLT

    Notes
    =====

    The MLT conversion is not part of the AACGM-v2 C library and is instead based
    on Laundal et al., 2016 [1]_. A brief summary of the method is provided below.

    MLT is defined as

        MLT = (magnetic longitude - magnetic noon meridian longitude) / 15 + 12

    where the magnetic noon meridian longitude is the centered dipole longitude
    of the subsolar point.

    There are two important reasons for using centered dipole instead of AACGM for
    this calculation. One reason is that the AACGM longitude of the subsolar point
    is often undefined (being at low latitudes). More importantly, if the subsolar
    point close to ground was used, the MLT at polar latitudes would be affected
    by non-dipole features at low latitudes, such as the South Atlantic Anomaly.
    This is not desirable; since the Sun-Earth interaction takes place at polar
    field lines, it is these field lines the MLT should describe.

    In calculating the centered dipole longitude of the subsolar point, we use
    the first three IGRF Gauss coefficients, using linear interpolation between
    the model updates every five years.

    Both input and output MLON are taken modulo 360 to ensure they are between
    0 and 360 degrees. Similarly, input/output MLT are taken modulo 24.
    For implementation of the subsolar point calculation, see :func:`subsol`.

    .. [1] Laundal, K. M. and A. D. Richmond (2016), Magnetic Coordinate Systems,
       Space Sci. Rev., doi:`10.1007/s11214-016-0275-y <http://dx.doi.org/10.1007/s11214-016-0275-y>`_.

    '''
    d2r = np.pi/180

    # find subsolar point
    yr = datetime.year
    doy = datetime.timetuple().tm_yday
    ssm = datetime.hour*3600 + datetime.minute*60 + datetime.second
    subsol_lon, subsol_lat = subsol(yr, doy, ssm)

    # unit vector pointing at subsolar point:
    s = np.array([np.cos(subsol_lat * d2r) * np.cos(subsol_lon * d2r),
                  np.cos(subsol_lat * d2r) * np.sin(subsol_lon * d2r),
                  np.sin(subsol_lat * d2r)])

    # convert subsolar coordinates to centered dipole coordinates
    z = igrf_dipole_axis(datetime)  # Cartesian axis pointing at Northern dipole pole
    y = np.cross(np.array([0, 0, 1]), z)
    y = y/np.linalg.norm(y)
    x = np.cross(y, z)
    R = np.vstack((x, y, z))
    s_cd = R.dot(s)

    # centered dipole longitude of subsolar point:
    mlon_subsol = np.arctan2(s_cd[1], s_cd[0])/d2r

    # convert the input array
    if m2a:  # MLT to AACGM
        mlt = np.asarray(arr) % 24
        mlon = (15*(mlt - 12) + mlon_subsol) % 360
        return mlon
    else:  # AACGM to MLT
        mlon = np.asarray(arr) % 360
        mlt = ((mlon - mlon_subsol)/15 + 12) % 24
        return mlt

def subsol(year, doy, ut):
    '''Finds subsolar geocentric longitude and latitude.

    Helper function for :func:`convert_mlt`.

    Parameters
    ==========
    year : int [1601, 2100]
        Calendar year
    doy : int [1, 365/366]
        Day of year
    ut : float
        Seconds since midnight on the specified day

    Returns
    =======
    sbsllon : float
        Subsolar longitude for the given date/time
    sbsllat : float
        Subsolar latitude for the given date/time

    Notes
    =====

    Based on formulas in Astronomical Almanac for the year 1996, p. C24.
    (U.S. Government Printing Office, 1994). Usable for years 1601-2100,
    inclusive. According to the Almanac, results are good to at least 0.01
    degree latitude and 0.025 degrees longitude between years 1950 and 2050.
    Accuracy for other years has not been tested. Every day is assumed to have
    exactly 86400 seconds; thus leap seconds that sometimes occur on December
    31 are ignored (their effect is below the accuracy threshold of the
    algorithm).

    After Fortran code by A. D. Richmond, NCAR. Translated from IDL
    by K. Laundal.

    '''

    from numpy import sin, cos, pi, arctan2, arcsin

    yr = year - 2000

    if year >= 2101:
        print('subsol.py: subsol invalid after 2100. Input year is:', year)

    nleap = np.floor((year-1601)/4)
    nleap = nleap - 99
    if year <= 1900:
        if year <= 1600:
            print('subsol.py: subsol invalid before 1601. Input year is:', year)
        ncent = np.floor((year-1601)/100)
        ncent = 3 - ncent
        nleap = nleap + ncent

    l0 = -79.549 + (-0.238699*(yr-4*nleap) + 3.08514e-2*nleap)

    g0 = -2.472 + (-0.2558905*(yr-4*nleap) - 3.79617e-2*nleap)

    # Days (including fraction) since 12 UT on January 1 of IYR:
    df = (ut/86400 - 1.5) + doy

    # Addition to Mean longitude of Sun since January 1 of IYR:
    lf = 0.9856474*df

    # Addition to Mean anomaly since January 1 of IYR:
    gf = 0.9856003*df

    # Mean longitude of Sun:
    l = l0 + lf

    # Mean anomaly:
    g = g0 + gf
    grad = g*pi/180

    # Ecliptic longitude:
    lmbda = l + 1.915*sin(grad) + 0.020*sin(2*grad)
    lmrad = lmbda*pi/180
    sinlm = sin(lmrad)

    # Days (including fraction) since 12 UT on January 1 of 2000:
    n = df + 365*yr + nleap

    # Obliquity of ecliptic:
    epsilon = 23.439 - 4e-7*n
    epsrad = epsilon*pi/180

    # Right ascension:
    alpha = arctan2(cos(epsrad)*sinlm, cos(lmrad)) * 180/pi

    # Declination:
    delta = arcsin(sin(epsrad)*sinlm) * 180/pi

    # Subsolar latitude:
    sbsllat = delta

    # Equation of time (degrees):
    etdeg = l - alpha
    nrot = round(etdeg/360)
    etdeg = etdeg - 360*nrot

    # Apparent time (degrees):
    aptime = ut/240 + etdeg    # Earth rotates one degree every 240 s.

    # Subsolar longitude:
    sbsllon = 180 - aptime
    nrot = round(sbsllon/360)
    sbsllon = sbsllon - 360*nrot

    return sbsllon, sbsllat

def gc2gd_lat(gc_lat):
    '''Convert geocentric latitude to geodetic latitude using WGS84.

    Parameters
    ==========
    gc_lat : array_like or float
        Geocentric latitude

    Returns
    =======
    gd_lat : same as input
        Geodetic latitude

    '''
    WGS84_e2 = 0.006694379990141317
    return np.rad2deg(-np.arctan(np.tan(np.deg2rad(gc_lat))/(WGS84_e2 - 1)))

def igrf_dipole_axis(date):
    '''Get Cartesian unit vector pointing at dipole pole in the north, according to IGRF

    Parameters
    ==========
    date : :class:`datetime.datetime`
        Date and time

    Returns
    =======
    m: numpy.ndarray
        Cartesian 3 element vector pointing at dipole pole in the north (geocentric coords)

    Notes
    =====
    IGRF coefficients are read from the igrf12coeffs.txt file. It should also work after IGRF updates.
    The dipole coefficients are interpolated to the date, or extrapolated if date > latest IGRF model
    '''

    # get time in years, as float:
    year = date.year
    doy = date.timetuple().tm_yday
    year = year + doy/(365 + calendar.isleap(year))

    # read the IGRF coefficients
    with open(IGRF_12_COEFFS, 'r') as f:
        lines = f.readlines()

    years = lines[3].split()[3:][:-1]
    years = np.array(years, dtype=float)  # time array

    g10 = lines[4].split()[3:]
    g11 = lines[5].split()[3:]
    h11 = lines[6].split()[3:]

    # secular variation coefficients (for extrapolation)
    g10sv = np.float32(g10[-1])
    g11sv = np.float32(g11[-1])
    h11sv = np.float32(h11[-1])

    # model coefficients:
    g10 = np.array(g10[:-1], dtype=float)
    g11 = np.array(g11[:-1], dtype=float)
    h11 = np.array(h11[:-1], dtype=float)

    # get the gauss coefficient at given time:
    if year <= years[-1]:  # regular interpolation
        g10 = np.interp(year, years, g10)
        g11 = np.interp(year, years, g11)
        h11 = np.interp(year, years, h11)
    else:  # extrapolation
        dt = year - years[-1]
        g10 = g10[-1] + g10sv * dt
        g11 = g11[-1] + g11sv * dt
        h11 = h11[-1] + h11sv * dt

    # calculate pole position
    B0 = np.sqrt(g10**2 + g11**2 + h11**2)

    return -np.array([g11, h11, g10])/B0


"""
 Copyright (c) 2019 Intel Corporation

 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.
"""

import numpy as np
from scipy.optimize import linear_sum_assignment

class Detection:
    """Class that stores detected object"""

    def __init__(self, obj_id, roi, conf, waiting=0, duration=1):
        """Constructor"""

        self.id = obj_id
        self.roi = roi
        self.conf = conf
        self.waiting = waiting
        self.duration = duration

    @property
    def roi(self):
        """Returns ROI of detected object"""

        return self._roi

    @roi.setter
    def roi(self, roi):
        """Sets ROI of detected object"""

        self._roi = np.copy(roi.reshape(1, -1))

class Tracker:  # pylint: disable=too-few-public-methods
    """Class that carries out tracking of persons using Hungarian algorithm"""

    def __init__(self, detector, score_threshold, iou_threshold, smooth_weight=0.5, max_waiting=5):
        """Constructor"""

        self._detector = detector
        self._score_threshold = score_threshold
        self._iou_threshold = iou_threshold
        self._smooth_weight = smooth_weight
        self._max_waiting = max_waiting

        self._last_detections = []
        self._cur_req_id, self._next_req_id = 0, 1
        self._last_id = 0

    @staticmethod
    def _matrix_iou(set_a, set_b):
        """Computes IoU metric for the two sets of vectors"""

        intersect_ymin = np.maximum(set_a[:, 0].reshape([-1, 1]), set_b[:, 0].reshape([1, -1]))
        intersect_xmin = np.maximum(set_a[:, 1].reshape([-1, 1]), set_b[:, 1].reshape([1, -1]))
        intersect_ymax = np.minimum(set_a[:, 2].reshape([-1, 1]), set_b[:, 2].reshape([1, -1]))
        intersect_xmax = np.minimum(set_a[:, 3].reshape([-1, 1]), set_b[:, 3].reshape([1, -1]))

        intersect_heights = np.maximum(0.0, intersect_ymax - intersect_ymin)
        intersect_widths = np.maximum(0.0, intersect_xmax - intersect_xmin)
        intersect_areas = intersect_heights * intersect_widths

        areas_set_a = ((set_a[:, 2] - set_a[:, 0]) * (set_a[:, 3] - set_a[:, 1])).reshape([-1, 1])
        areas_set_b = ((set_b[:, 2] - set_b[:, 0]) * (set_b[:, 3] - set_b[:, 1])).reshape([1, -1])

        union_areas = areas_set_a + areas_set_b - intersect_areas

        return intersect_areas / union_areas

    @staticmethod
    def filter_rois(new_rois, score_threshold):
        """Filters input ROIs by valid height/width and score threshold values"""

        heights = new_rois[:, 2] - new_rois[:, 0]
        widths = new_rois[:, 3] - new_rois[:, 1]
        valid_sizes_mask = np.logical_and(heights > 0.0, widths > 0.0)
        valid_conf_mask = new_rois[:, 4] > score_threshold

        valid_roi_ids = np.where(np.logical_and(valid_sizes_mask, valid_conf_mask))[0]
        filtered_rois = new_rois[valid_roi_ids, :4]
        filtered_conf = new_rois[valid_roi_ids, 4]

        return filtered_rois, filtered_conf

    def _track(self, last_detections, new_rois):
        """Updates current tracks according new observations"""

        filtered_rois, filtered_conf = self.filter_rois(new_rois, self._score_threshold)

        if filtered_rois.shape[0] == 0:
            out_detections = []
            for det in last_detections:
                det.waiting = 1
                det.duration = 0
                out_detections.append(det)

            return out_detections

        if last_detections is None or len(last_detections) == 0:
            out_detections = []
            for roi, conf in zip(filtered_rois, filtered_conf):
                out_detections.append(Detection(self._last_id, roi.reshape(1, -1), conf))
                self._last_id += 1

            return out_detections

        last_rois = np.concatenate([det.roi for det in last_detections], axis=0)
        affinity_matrix = self._matrix_iou(last_rois, filtered_rois)
        cost_matrix = 1.0 - affinity_matrix

        row_ind, col_ind = linear_sum_assignment(cost_matrix)
        affinity_values = 1.0 - cost_matrix[row_ind, col_ind]

        valid_matches = affinity_values > self._iou_threshold
        row_ind = row_ind[valid_matches]
        col_ind = col_ind[valid_matches]

        out_detections = []
        for src_id, trg_id in zip(row_ind, col_ind):
            det = last_detections[src_id]
            det.waiting = 0
            det.duration += 1
            new_roi = filtered_rois[trg_id]
            det.roi = self._smooth_roi(det.roi, new_roi.reshape(1, -1), self._smooth_weight)
            det.conf = filtered_conf[trg_id]
            out_detections.append(det)

        unmatched_src_ind = set(range(len(last_detections))) - set(row_ind.tolist())
        for src_id in unmatched_src_ind:
            det = last_detections[src_id]
            det.waiting += 1
            det.duration = 0
            if det.waiting < self._max_waiting:
                out_detections.append(det)

        unmatched_trg_ind = set(range(len(filtered_rois))) - set(col_ind.tolist())
        for trg_id in unmatched_trg_ind:
            new_roi = filtered_rois[trg_id]
            new_roi_conf = filtered_conf[trg_id]
            out_detections.append(Detection(self._last_id, new_roi.reshape(1, -1), new_roi_conf))
            self._last_id += 1

        return out_detections

    @staticmethod
    def _smooth_roi(prev_roi, new_roi, weight):
        """Smooths tracking ROI"""

        if prev_roi is None:
            return new_roi

        return weight * prev_roi + (1.0 - weight) * new_roi

    @staticmethod
    def _clip_roi(roi, frame_size):
        """Clips ROI limits according frame sizes"""

        frame_height, frame_width = frame_size

        old_roi = roi.reshape(-1)
        new_roi = [np.maximum(0, int(old_roi[0])),
                   np.maximum(0, int(old_roi[1])),
                   np.minimum(frame_width, int(old_roi[2])),
                   np.minimum(frame_height, int(old_roi[3]))]

        return np.array(new_roi)

    def _get_last_detections(self, frame_size, max_num_detections, labels_map):
        """Returns active detections"""

        if self._last_detections is None or len(self._last_detections) == 0:
            return [], {}

        out_detections = []
        for det in self._last_detections:
            if det.waiting > 0 or det.duration <= 1:
                continue

            clipped_roi = self._clip_roi(det.roi, frame_size)
            out_det = Detection(det.id, clipped_roi, det.conf, det.waiting, det.duration)
            out_detections.append(out_det)

        if len(out_detections) > max_num_detections:
            out_detections.sort(key=lambda x: x.conf, reverse=True)
            out_detections = out_detections[:max_num_detections]

        matched_det_ids = {det.id for det in out_detections} & labels_map.keys()
        unused_det_ids = sorted(set(range(max_num_detections)) - matched_det_ids)

        out_labels_map = {}
        for det in out_detections:
            if det.id in matched_det_ids:
                out_labels_map[det.id] = labels_map[det.id]
            else:
                new_local_det_id = unused_det_ids[0]
                unused_det_ids = unused_det_ids[1:]

                out_labels_map[det.id] = new_local_det_id
                det.id = new_local_det_id

        return out_detections, labels_map

    def add_frame(self, frame, max_num_detections, labels_map):
        """Adds new detections and returns active tracks"""

        self._detector.async_infer(frame, self._next_req_id)
        new_rois = self._detector.wait_request(self._cur_req_id)
        self._cur_req_id, self._next_req_id = self._next_req_id, self._cur_req_id

        if new_rois is not None:
            self._last_detections = self._track(self._last_detections, new_rois)

        frame_size = frame.shape[:2]
        out_detections, out_labels_map = self._get_last_detections(
            frame_size, max_num_detections, labels_map)

        return out_detections, out_labels_map


import numpy as np
import numpy.linalg as la
import scipy.sparse as SP
import scipy
from scipy.sparse import linalg as SPla

def spdot(a, b, array_out=True):
    """
    Matrix multiplication function to deal with sparse and dense objects

    Parameters
    ----------

    a           : array
                  first multiplication factor. Can either be sparse or dense.
    b           : array
                  second multiplication factor. Can either be sparse or dense.
    array_out   : boolean
                  If True (default) the output object is always a np.array

    Returns
    -------

    ab : array
         product of a times b. Sparse if a and b are sparse. Dense otherwise.
    """
    if type(a).__name__ == 'ndarray' and type(b).__name__ == 'ndarray':
        ab = np.dot(a, b)
    elif type(a).__name__ == 'csr_matrix' or type(b).__name__ == 'csr_matrix' \
            or type(a).__name__ == 'csc_matrix' or type(b).__name__ == 'csc_matrix':
        ab = a * b
        if array_out:
            if type(ab).__name__ == 'csc_matrix' or type(ab).__name__ == 'csr_matrix':
                ab = ab.toarray()
    else:
        raise Exception, "Invalid format for 'spdot' argument: %s and %s" % (
            type(a).__name__, type(b).__name__)
    return ab

def spmultiply(a, b, array_out=True):
    """
    Element-wise multiplication function to deal with sparse and dense
    objects. Both objects must be of the same type.

    Parameters
    ----------

    a           : array
                  first multiplication factor. Can either be sparse or dense.
    b           : array
                  second multiplication factor. Can either be sparse or dense.
                  integer.
    array_out   : boolean
                  If True (default) the output object is always a np.array

    Returns
    -------

    ab : array
         elementwise multiplied object. Sparse if a is sparse. Dense otherwise.
    """
    if type(a).__name__ == 'ndarray' and type(b).__name__ == 'ndarray':
        ab = a * b
    elif (type(a).__name__ == 'csr_matrix' or type(a).__name__ == 'csc_matrix') \
            and (type(b).__name__ == 'csr_matrix' or type(b).__name__ == 'csc_matrix'):
        ab = a.multiply(b)
        if array_out:
            if type(ab).__name__ == 'csc_matrix' or type(ab).__name__ == 'csr_matrix':
                ab = ab.toarray()
    else:
        raise Exception, "Invalid format for 'spmultiply' argument: %s and %s" % (
            type(a).__name__, type(b).__name__)
    return ab

def sphstack(a, b, array_out=False):
    """
    Horizontal stacking of vectors (or matrices) to deal with sparse and dense objects

    Parameters
    ----------

    a           : array or sparse matrix
                  First object.
    b           : array or sparse matrix
                  Object to be stacked next to a
    array_out   : boolean
                  If True the output object is a np.array; if False (default)
                  the output object is an np.array if both inputs are
                  arrays or CSR matrix if at least one input is a CSR matrix

    Returns
    -------

    ab          : array or sparse matrix
                  Horizontally stacked objects
    """
    if type(a).__name__ == 'ndarray' and type(b).__name__ == 'ndarray':
        ab = np.hstack((a, b))
    elif type(a).__name__ == 'csr_matrix' or type(b).__name__ == 'csr_matrix':
        ab = SP.hstack((a, b), format='csr')
        if array_out:
            if type(ab).__name__ == 'csr_matrix':
                ab = ab.toarray()
    else:
        raise Exception, "Invalid format for 'sphstack' argument: %s and %s" % (
            type(a).__name__, type(b).__name__)
    return ab

def spbroadcast(a, b, array_out=False):
    """
    Element-wise multiplication of a matrix and vector to deal with sparse
    and dense objects

    Parameters
    ----------

    a           : array or sparse matrix
                  Object with one or more columns.
    b           : array
                  Object with only one column
    array_out   : boolean
                  If True the output object is a np.array; if False (default)
                  the output object is an np.array if both inputs are
                  arrays or CSR matrix if at least one input is a CSR matrix

    Returns
    -------

    ab          : array or sparse matrix
                  Element-wise multiplication of a and b
    """
    if type(a).__name__ == 'ndarray' and type(b).__name__ == 'ndarray':
        ab = a * b
    elif type(a).__name__ == 'csr_matrix':
        b_mod = SP.lil_matrix((b.shape[0], b.shape[0]))
        b_mod.setdiag(b)
        ab = (a.T * b_mod).T
        if array_out:
            if type(ab).__name__ == 'csr_matrix':
                ab = ab.toarray()
    else:
        raise Exception, "Invalid format for 'spbroadcast' argument: %s and %s" % (
            type(a).__name__, type(b).__name__)
    return ab

def spmin(a):
    """
    Minimum value in a matrix or vector to deal with sparse and dense objects

    Parameters
    ----------

    a           : array or sparse matrix
                  Object with one or more columns.

    Returns
    -------

    min a       : int or float
                  minimum value in a
    """
    return a.min()

def spmax(a):
    """
    Maximum value in a matrix or vector to deal with sparse and dense objects

    Parameters
    ----------

    a           : array or sparse matrix
                  Object with one or more columns.

    Returns
    -------

    max a       : int or float
                  maximum value in a
    """
    return a.max()

def splogdet(a):
    """
    Compute the log determinant of a large matrix.

    Parameters
    ----------

    a       :   array or sparse matrix
                Object with one or more columns

    Returns
    -------

    log determinant of a    :   int or float
                                logged determinant of a
    """
    if SP.issparse(a):
        LU = SPla.splu(a)
        det = np.sum(np.log(np.abs(LU.U.diagonal())))
    else:
        sgn, ldet = la.slogdet(a)
        det = sgn * ldet
    return det

def spfill_diagonal(a, val):
    """
    Fill the diagonal of a sparse or dense matrix

    Parameters
    ----------

    a       :   array or sparse matrix
                Object with one or more columns
    val     :   int or float
                value with which to fill the diagonal of a

    Returns
    -------

    a with val on each element of the diagonal
    """
    if SP.issparse(a):
        a.setdiag(val)
    else:
        np.fill_diagonal(a, val)
    return a

def spinv(a):
    """
    Compute the inverse of a sparse or dense matrix

    Parameters
    ----------

    a       :   array or sparse matrix
                Object with one or more columns

    Returns
    -------

    ai, the inverse of a
    """
    if SP.issparse(a):
        ai = SPla.inv(a)
    else:
        ai = la.inv(a)
    return ai

def spisfinite(a):
    """
    Determine whether an array has nan or inf values

    Parameters
    ----------
    a   :   array or sparse matrix
            Object with one or more columns

    Returns
    -------
    bool denoting whether or not the array contains any NaN or inf
    """
    return np.isfinite(a.sum())

def _test():
    import doctest
    doctest.testmod()

if __name__ == '__main__':
    _test()