import numpy as np

from ...integrate.timespec import parse_time_specification
from ...potential import Hamiltonian, PotentialBase
from ..core import PhaseSpacePosition
from ..nbody import DirectNBody
from ..util import combine
from ._mockstream import (
    mockstream_dop853,
    mockstream_dop853_animate,
    mockstream_leapfrog,
    # mockstream_leapfrog_animate,
)
from .core import MockStream

__all__ = ["MockStreamGenerator"]


class MockStreamGenerator:
    def __init__(self, df, hamiltonian, progenitor_potential=None):
        """Generate a mock stellar stream in the specified external potential.

        By default, you must pass in a specification of the stream distribution
        function (``df``), and the external gravitational potential and
        reference frame (via a `~gala.potential.Hamiltonian` object passed in
        through the ``hamiltonian`` argument).

        Also by default, the stream generation does not include the self-gravity
        of the progenitor system: star particles are generated using the ``df``
        object, and released into the external potential specified by the
        ``hamiltonian``. If you would like the star particles to feel the
        gravitational field of the progenitor system, you may pass in a
        potential object to represent the progenitor via the
        ``progenitor_potential`` argument. This can be any valid gala potential
        instance.

        Parameters
        ----------
        df : `~gala.dynamics.BaseStreamDF` subclass instance
            The stream distribution function (DF) object that specifies how to
            generate stream star particle initial conditions.
        hamiltonian : `~gala.potential.Hamiltonian`
            The external potential and reference frame to numerically integrate
            orbits in.
        progenitor_potential : `~gala.potential.PotentialBase` (optional)
            If specified, the self-gravity of the progenitor system is included
            in the force calculation and orbit integration. If not specified,
            self-gravity is not accounted for. Default: ``None``
        """
        from .df import BaseStreamDF

        if not isinstance(df, BaseStreamDF):
            msg = (
                "The input distribution function (DF) instance "
                "must be an instance of a subclass of "
                f"BaseStreamDF, not {type(df)}."
            )
            raise TypeError(msg)
        self.df = df

        # Validate the inpute hamiltonian
        self.hamiltonian = Hamiltonian(hamiltonian)

        if progenitor_potential is not None:
            # validate the potential class
            if not isinstance(progenitor_potential, PotentialBase):
                raise TypeError(
                    "If specified, the progenitor_potential must be a gala.potential "
                    "class instance."
                )

            self.self_gravity = True

        else:
            self.self_gravity = False

        self.progenitor_potential = progenitor_potential

    def _get_nbody(self, prog_w0, nbody):
        """
        Internal function that adds the progenitor to the list of nbody objects to
        integrate along with the test particles in the stream.
        """

        kwargs = {}
        if nbody is not None:
            if nbody.external_potential != self.hamiltonian.potential:
                raise ValueError(
                    "The external potential of the input nbody instance must match the "
                    "potential of the mock stream input hamiltonian! "
                    f"{nbody.external_potential} vs. {self.hamiltonian.potential}"
                )

            if nbody.frame != self.hamiltonian.frame:
                raise ValueError(
                    "The reference frame of the input nbody instance must match the "
                    "frame of the mock stream input hamiltonian! "
                    f"{nbody.frame} vs. {self.hamiltonian.frame}"
                )

            kwargs["w0"] = combine((prog_w0, nbody.w0))
            kwargs["particle_potentials"] = [
                self.progenitor_potential,
                *nbody.particle_potentials,
            ]
            kwargs["external_potential"] = self.hamiltonian.potential
            kwargs["frame"] = self.hamiltonian.frame
            kwargs["units"] = self.hamiltonian.units

        else:
            kwargs["w0"] = prog_w0
            kwargs["particle_potentials"] = [self.progenitor_potential]
            kwargs["external_potential"] = self.hamiltonian.potential
            kwargs["frame"] = self.hamiltonian.frame
            kwargs["units"] = self.hamiltonian.units

        return DirectNBody(**kwargs)

    def run(
        self,
        prog_w0,
        prog_mass,
        nbody=None,
        release_every=1,
        n_particles=1,
        output_every=None,
        output_filename=None,
        check_filesize=True,
        overwrite=False,
        progress=False,
        Integrator=None,
        Integrator_kwargs=None,
        **time_spec,
    ):
        """
        Run the mock stream generator with the specified progenitor initial conditions.

        This method generates the mock stellar stream for the specified progenitor
        system properties. The progenitor orbit is specified by passing in the initial
        or final conditions ``prog_w0`` and by specifying time-stepping information via
        the ``**time_spec`` keyword arguments. If the time-stepping specification
        proceeds forward in time, ``prog_w0`` is interpreted as initial conditions and
        the mock stream is generated forwards from this position. If the time-stepping
        proceeds backwards in time, the progenitor orbit is first numerically integrated
        backwards given the time-stepping information, then the stream is generated
        forward from the past such that ``prog_w0`` becomes the final position of the
        progenitor.

        Note that the stream generation also supports including other massive perturbers
        that can gravitationally influence the stream stars. These other massive bodies
        must be passed in as a `~gala.dynamics.DirectNBody` instance through the
        ``nbody`` argument. The phase-space coordinates of the bodies, ``nbody.w0``, are
        interpreted as initial or final conditions with the same logic as above.

        Parameters
        ----------
        prog_w0 : `~gala.dynamics.PhaseSpacePosition`
            The initial or final phase-space position of the progenitor system (see note
            above).
        prog_mass : `~astropy.units.Quantity` [mass]
            The mass of the progenitor system, passed in to the stream distribution
            function (df) ``.sample()`` method. This quantity sets the scale mass of the
            particle release df, but not the mass of the progenitor potential used to
            compute the self-gravity on the stream particles.
        nbody : `~gala.dynamics.DirectNBody` (optional)
            This allows specifying other massive perturbers (N-bodies) that can
            gravitationally influence the stream star orbits.
        release_every : int (optional)
            Controls how often to release stream particles from each tail. Default: 1,
            meaning release particles at each timestep.
        n_particles : int, array_like (optional)
            If an integer, this controls the number of particles to release in each tail
            at each release timestep. Alternatively, you can pass in an array with the
            same shape as the number of timesteps to release bursts of particles at
            certain times (e.g., pericenter).
        output_every : int (optional)
            Controls whether to output snapshots of the stream particle orbits. This is
            relative to the global time array.
        output_filename : str (optional)
            The path to the HDF5 file to be generated by the snapshotting.
        check_filesize : bool (optional)
            If True (the default value), this controls whether to check the estimated
            size of the  output file, and emits a warning if the file is >8GB in size.
        overwrite : bool (optional)
            Overwrite the output file if it exists.
        progress : bool (optional)
            Print a very basic progress bar while computing the stream.
        Integrator : `~gala.integrate.Integrator`, str (optional)
            Integrator class to use, or a string name like 'leapfrog', 'dopri853'.
            Currently, only the `~gala.integrate.DOPRI853Integrator` and
            `~gala.integrate.LeapfrogIntegrator` are supported.
        Integrator_kwargs : dict (optional)
            Any extra keyword arguments to pass to the integrator class
            when initializing. For example, you can pass in the
            ``atol`` and ``rtol`` keyword arguments to set the absolute and
            relative tolerances for the DOPRI853 integrator.
        **time_spec
            Specification of how long to integrate. Most commonly, this is a timestep
            ``dt`` and number of steps ``n_steps``, or a timestep ``dt``, initial time
            ``t1``, and final time ``t2``. You may also pass in a time array with ``t``.
            See documentation for `~gala.integrate.parse_time_specification` for more
            information.

        Returns
        -------
        stream_w : `~gala.dynamics.PhaseSpacePosition`
        nbody_w : `~gala.dynamics.PhaseSpacePosition`

        """
        from gala.integrate import (
            DOPRI853Integrator,
            LeapfrogIntegrator,
            get_integrator,
            # Ruth4Integrator,
        )

        if Integrator_kwargs is None:
            Integrator_kwargs = {}

        if Integrator is None:
            Integrator = DOPRI853Integrator

        # Validates and retrieves the integrator class from string name if needed
        Integrator = get_integrator(Integrator)

        units = self.hamiltonian.units
        t = parse_time_specification(units, **time_spec)

        prog_nbody = self._get_nbody(prog_w0, nbody)
        nbody_orbits = prog_nbody.integrate_orbit(
            t=t, Integrator=Integrator, Integrator_kwargs=Integrator_kwargs
        )

        # If the time stepping passed in is negative, assume this means that all
        # of the initial conditions are at *end time*, and we first need to
        # integrate them backwards before treating them as initial conditions
        if t[1] < t[0]:
            nbody_orbits = nbody_orbits[::-1]

            # TODO: this could be cleaned up...
            nbody0 = DirectNBody(
                nbody_orbits[0],
                prog_nbody.particle_potentials,
                external_potential=self.hamiltonian.potential,
                frame=self.hamiltonian.frame,
                units=units,
            )

        else:
            nbody0 = prog_nbody

        # Note: assumes that this is an orbit not a psp, i.e. that save_all is True
        prog_orbit = nbody_orbits[:, 0]  # Note: Progenitor must be idx 0!
        orbit_t = prog_orbit.t.decompose(units).value

        # Generate initial conditions from the DF
        stream_w0 = self.df.sample(
            prog_orbit,
            prog_mass,
            hamiltonian=self.hamiltonian,
            release_every=release_every,
            n_particles=n_particles,
        )
        w0 = np.vstack(
            (
                stream_w0.xyz.decompose(units).value,
                stream_w0.v_xyz.decompose(units).value,
            )
        ).T
        w0 = np.ascontiguousarray(w0)

        unq_t1s, nstream = np.unique(
            stream_w0.release_time.decompose(units).value, return_counts=True
        )

        all_nstream = np.zeros(prog_orbit.ntimes, dtype=int)
        for t1, n in zip(unq_t1s, nstream):
            all_nstream[np.isclose(orbit_t, t1)] = n

        nstream_idx = np.where(all_nstream != 0)[0]
        if 0 not in nstream_idx:
            nstream_idx = np.insert(nstream_idx, 0, 0)
            unq_t1s = np.insert(unq_t1s, 0, orbit_t[0])

        if Integrator == DOPRI853Integrator:
            if output_every is None:
                raw_nbody, raw_stream = mockstream_dop853(
                    nbody0,
                    orbit_t[nstream_idx],
                    w0,
                    unq_t1s,
                    orbit_t[-1],
                    all_nstream[nstream_idx].astype("i4"),
                    progress=int(progress),
                    **Integrator_kwargs,
                )
            else:  # store snapshots
                if output_filename is None:
                    raise ValueError(
                        "If output_every is specified, you must also pass in a "
                        "filename to store the snapshots in"
                    )

                raw_nbody, raw_stream = mockstream_dop853_animate(
                    nbody0,
                    orbit_t,
                    w0,
                    all_nstream.astype("i4"),
                    output_every=output_every,
                    output_filename=output_filename,
                    check_filesize=check_filesize,
                    overwrite=overwrite,
                    progress=int(progress),
                    **Integrator_kwargs,
                )

        elif Integrator == LeapfrogIntegrator:
            if output_every is None:
                raw_nbody, raw_stream = mockstream_leapfrog(
                    nbody0,
                    orbit_t,  # Pass full time array for N-body integration
                    orbit_t[nstream_idx],  # Spawn times only
                    w0,
                    unq_t1s,
                    orbit_t[-1],
                    all_nstream[nstream_idx].astype("i4"),
                    progress=int(progress),
                )
            else:  # store snapshots
                if output_filename is None:
                    raise ValueError(
                        "If output_every is specified, you must also pass in a filename to "
                        "store the snapshots in"
                    )
                raise NotImplementedError(
                    "Animation output for LeapfrogIntegrator is not implemented"
                )
                # raw_nbody, raw_stream = mockstream_leapfrog_animate(
                #     nbody0,
                #     orbit_t,
                #     w0,
                #     all_nstream.astype("i4"),
                #     output_every=output_every,
                #     output_filename=output_filename,
                #     check_filesize=check_filesize,
                #     overwrite=overwrite,
                #     progress=int(progress),
                # )
        else:
            raise ValueError(
                "Currently, only the DOPRI853Integrator and LeapfrogIntegrator "
                "are supported for mock stream generation."
            )

        x_unit = units["length"]
        v_unit = units["length"] / units["time"]
        stream_w = MockStream(
            pos=raw_stream[:, :3].T * x_unit,
            vel=raw_stream[:, 3:].T * v_unit,
            release_time=stream_w0.release_time,
            lead_trail=stream_w0.lead_trail,
            frame=self.hamiltonian.frame,
            copy=False,
        )
        nbody_w = PhaseSpacePosition(
            pos=raw_nbody[:, :3].T * x_unit,
            vel=raw_nbody[:, 3:].T * v_unit,
            frame=self.hamiltonian.frame,
            copy=False,
        )

        return stream_w, nbody_w
