import numpy as np
from scipy.integrate import solve_ivp
from pylcp.common import (progressBar, random_vector)
from pylcp.integration_tools import solve_ivp_random
from scipy.interpolate import interp1d

"""
Here, the function "rateeq" in the rateeq module of pylcp is adjusted. Typically, there will be no need to change anything here.

Goal:
    Be able to interrupt the calculation of a trajectory in a more useful way.

Implementation:
    Set a max boundary and a radius in phase space that are both able
    to interrupt the calculation. If the atom leaves a usefull region or is trapped
    in the MOT, the calculation will be stopped. Thus, a long timespan can be specified
    and the calculation will stop automatically, when reasonable.
"""

def evolve_motion(self, t_span, freeze_axis=[False, False, False],
                    random_recoil=False, random_force=False,
                    max_scatter_probability=0.1, progress_bar=False,
                    record_force=False, rng=np.random.default_rng(),
                    boundary=None, MOTradius=None, MOTvelocity=None, **kwargs):
    """
    Evolve the populations :math:`N` and the motion of the atom in time.

    This function evolves the rate equations, moving the atom through space,
    given the instantaneous force, for some period of time.

    Parameters
    ----------
    t_span : list or array_like
        A two element list or array that specify the initial and final time
        of integration.
    freeze_axis : list of boolean
        Freeze atomic motion along the specified axis.
        Default: [False, False, False]
    random_recoil : boolean
        Allow the atom to randomly recoil from scattering events.
        Default: False
    random_force : boolean
        Rather than calculating the force using the rateeq.force() method,
        use the calculated scattering rates from each of the laser beam
        (combined with the instantaneous populations) to randomly add photon
        absorption events that cause the atom to recoil randomly from the
        laser beam(s).
        Default: False
    max_scatter_probability : float
        When undergoing random recoils and/or force, this sets the maximum
        time step such that the maximum scattering probability is less than
        or equal to this number during the next time step.  Default: 0.1
    progress_bar : boolean
        If true, show a progress bar as the calculation proceeds.
        Default: False
    record_force : boolean
        If true, record the instantaneous force and store in the solution.
        Default: False
    rng : numpy.random.Generator()
        A properly-seeded random number generator.  Default: calls
        ``numpy.random.default.rng()``
    boundary : float
        Define a boundary of the simulation. This is the maximal position
        the atoms can reach before the calculation is interrupted by an event.
        If boundary = 0, no interruption will happen. The calculation is run 
        until t_spn[1] is reached. sry for the stupid coding. (default: 0)
    MOTradius : float
        Same as boundary, but specifies a small radius of "MOT region".
        When the atoms reach this radius and have a velocity smaller than
        MOTvelocity, the caculation will be interrupted by an event.
        If 0, no interruption. (default: 0)
    MOTvelocity : float
        See above. (default : 0)
    **kwargs :
        Additional keyword arguments get passed to solve_ivp_random, which
        is what actually does the integration.

    Returns
    -------
    sol : OdeSolution
        Bunch object that contains the following fields:

            * t: integration times found by solve_ivp
            * N: population vs. time
            * v: atomic velocity
            * r: atomic position

        It contains other important elements, which can be discerned from
        scipy's solve_ivp documentation.
    """
    free_axes = np.bitwise_not(freeze_axis)

    if progress_bar:
        progress = progressBar()

    if record_force:
        ts = []
        Fs = []

    def motion(t, y):
        N = y[:-6]
        v = y[-6:-3]
        r = y[-3:]

        Rev, Rijl = self.construct_evolution_matrix(r, v, t)
        if not random_force:
            if record_force:
                F = self.force(r, t, N, return_details=True)

                ts.append(t)
                Fs.append(F)

                F = F[0]
            else:
                F = self.force(r, t, N, return_details=False)

            dydt = np.concatenate((Rev @ N,
                                    F*free_axes/self.hamiltonian.mass+
                                    self.constant_accel,
                                    v))
        else:
            dydt = np.concatenate((Rev @ N,
                                    self.constant_accel,
                                    v))

        if np.any(np.isnan(dydt)):
            raise ValueError('Enountered a NaN!')

        if progress_bar:
            progress.update(t/t_span[-1])

        return dydt

    def random_force_func(t, y, dt):
        total_P = 0
        num_of_scatters = 0

        # Go over all available keys:
        for key in self.laserBeams:
            # Extract the pumping rate from each laser:
            Rl = np.sum(self.Rijl[key], axis=(1,2))

            # Calculate the probability to scatter a photon from the laser:
            P = Rl*dt

            # Roll the dice N times, where $N=\sum(lasers)
            dice = rng.random(len(P))

            # Give them kicks!
            for ii in np.arange(len(Rl))[dice<P]:
                num_of_scatters += 1
                y[-6:-3] += self.laserBeams[key].beam_vector[ii].kvec(y[-3:], t)/\
                            self.hamiltonian.mass
                # Can branch to a differe, lower state, but let's ignore that
                # for the moment.
                y[-6:-3] += self.recoil_velocity[key]*(random_vector(rng, free_axes))

            total_P += np.sum(P)

        # Calculate a new maximum dt to make sure we evolve while not
        # exceeding dt max:
        new_dt_max = (max_scatter_probability/total_P)*dt

        return (num_of_scatters, new_dt_max)


    def random_recoil_func(t, y, dt):
        num_of_scatters = 0
        total_P = 0.

        # Go over each block in the Hamiltonian and compute the decay:
        for key in self.decay_rates:
            P = dt*self.decay_rates[key]*y[self.decay_N_indices[key]]

            # Roll the dice N times, where $N=\sum(n_i)
            dice = rng.random(len(P))

            # For any random number that is lower than P_i, add a
            # recoil velocity.
            for ii in range(np.sum(dice<P)):
                num_of_scatters += 1
                y[-6:-3] += self.recoil_velocity[key]*(random_vector(rng, free_axes)+
                                                        random_vector(rng, free_axes))

            # Save the total probability of a scatter:
            total_P += np.sum(P)

        # Calculate a new maximum dt to make sure we evolve while not
        # exceeding dt max:
        new_dt_max = (max_scatter_probability/total_P)*dt

        return (num_of_scatters, new_dt_max)

    eventList=[]

    # Define the event function to stop the simulation when r > boundary
    def boundary_event(boundary):
        def event(t, y):
            return 0 if y[-3]**2+y[-2]**2+y[-1]**2 > boundary**2 else 1 # Trigger event when x > boundary
        return event
    
    # Define the event function to stop the simulation when position < 0.01 and velocity < 1
    def MOT_event(MOTradius, MOTvelocity):
        def event(t, y):
            radius = y[-3]**2+y[-2]**2+y[-1]**2
            velocity = y[-6]**2+y[-5]**2+y[-4]**2
            return radius - MOTradius**2 if velocity < MOTvelocity**2 else 1
        return event
    
    if boundary:
        if type(boundary) == float:
            boundary_event_func = boundary_event(boundary)
            boundary_event_func.terminal = True  # Stop the integration on condition
            boundary_event_func.direction = -1    # trigger when decreasing
            eventList.append(boundary_event_func)
        else:
            raise(TypeError(f"Boundary is type {type(boundary)}. Expected float."))
    
    if MOTradius and MOTvelocity:
        if type(MOTradius) == float and type(MOTvelocity) == float:
            MOT_event_func = MOT_event(MOTradius, MOTvelocity)
            MOT_event_func.terminal = True  # Stop the integration on condition
            MOT_event_func.direction = -1   # Trigger when decreasing
            eventList.append(MOT_event_func)
        else:
            raise(TypeError(f"Motvelocity is type {type(MOTvelocity)} and MOTradius is type {type(MOTradius)}. Expected float for both."))

    y0 = np.concatenate((self.N0, self.v0, self.r0))
    if random_force:
        self.sol = solve_ivp_random(motion, random_force_func, t_span, y0,
                                    initial_max_step=max_scatter_probability,
                                    **kwargs)
    elif random_recoil:
        self.sol = solve_ivp_random(motion, random_recoil_func, t_span, y0,
                                    initial_max_step=max_scatter_probability,
                                    **kwargs)
    else:
        self.sol = solve_ivp(motion, t_span, y0, events=eventList, **kwargs)

    if progress_bar:
        # Just in case the solve_ivp_random terminated due to an event.
        progress.update(1.)

    # Rearrange the solution:
    self.sol.N = self.sol.y[:-6]
    self.sol.v = self.sol.y[-6:-3]
    self.sol.r = self.sol.y[-3:]

    if record_force:
        f = interp1d(ts[:-1], np.array([f[0] for f in Fs[:-1]]).T)
        self.sol.F = f(self.sol.t)

        f = interp1d(ts[:-1], np.array([f[2] for f in Fs[:-1]]).T)
        self.sol.fmag = f(self.sol.t)

        self.sol.f = {}
        for key in Fs[0][1]:
            f = interp1d(ts[:-1], np.array([f[1][key] for f in Fs[:-1]]).T)
            self.sol.f[key] = f(self.sol.t)
            self.sol.f[key] = np.swapaxes(self.sol.f[key], 0, 1)

    del self.sol.y

    return self.sol