# Third-party
import numpy as np

# This package
from .. import combine, PhaseSpacePosition
from ..nbody import DirectNBody
from ...potential import Hamiltonian, PotentialBase
from ...integrate.timespec import parse_time_specification
from ._mockstream import mockstream_dop853, mockstream_dop853_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):
            raise TypeError('The input distribution function (DF) instance '
                            'must be an instance of a subclass of '
                            'BaseStreamDF, not {}.'.format(type(df)))
        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.