'''
Illustration of iterative self-consistent modelling machinery reimplemented in Python.
The two classes Component and SelfConsistentModel are pure Python equivalents of
the same classes provided by the Agama Python interface, and the iterative model results
should be exactly equivalent when using N-body units (unfortunately, when setting
non-trivial dimensional units, the results diverge after first iteration due to
minuscule roundoff effects accumulated in unit conversion).
The concept of a user-defined Python density function that computes moments of
a DF-based GalaxyModel and is used to initialize a cheaper density interpolator
(DensitySphericalHarmonic/DensityAzimuthalHarmonic) can be used in other contexts,
not necessarily within a self-consistent modelling loop.
For instance, one may construct a density expansion for the given DF, and subsequently
use its surfaceDensity() method for many input points -- the result is equivalent
to GalaxyModel.projectedMoments up to integration errors, but is much cheaper to compute.
'''
# if the module has been installed to the globally known directory, just import it
try: import agama
except ImportError:  # otherwise load the shared library from the parent folder
    import sys
    sys.path += ['../']
    try: import agama
    except ImportError as ex: sys.exit("\033[1;31mFAILED TO IMPORT AGAMA: %s\033[0m" % ex)

class Component(object):
    '''
    A pure Python analog of agama.Component class (slightly simplified)
    '''
    def __init__(self, density=None, potential=None, **kwargs):
        # these checks are a bit more stringent than performed by the C++ extension module:
        # the latter accepts either instances of built-in Density and Potential classes,
        # or arbitrary callable Python functions, whereas the code below ignores the second option
        if density is not None and not isinstance(density, agama.Density):
            raise TypeError("Argument 'density' must be a valid Density object")
        if potential is not None and not isinstance(potential, agama.Potential):
            raise TypeError("Argument 'potential' must be a valid Potential object")
        self._density = density
        self._potential = potential
        for k,v in kwargs.items(): self.__dict__[k.lower()] = v   # make parameters case-insensitive

    @property
    def density(self): return self._density

    @property
    def potential(self): return self._potential

    def update(self, potential, af):
        # the main procedure for recomputing the density generated by DF
        # and representing it by either a spherical- or azimuthal-harmonic expansion.
        # fixOrder argument instructs the density expansion to use exactly the minimum necessary
        # number of sampling points in angles for the requested expansion order,
        # thereby reducing the number of calls to the expensive density computation routine;
        # the same approach is used in the C++ code for the SelfConsistentModel class.
        if not hasattr(self, 'df'): return   # can only update DF-based components
        gm = agama.GalaxyModel(potential, self.df, af)
        densfnc = lambda x: gm.moments(x, dens=True, vel=False, vel2=False)
        if self.disklike:
            self._density = agama.Density(density=densfnc, type='DensityAzimuthalHarmonic',
                gridsizer=self.sizeradialcyl,   rmin=self.rmincyl, rmax=self.rmaxcyl,
                gridsizez=self.sizeverticalcyl, zmin=self.zmincyl, zmax=self.zmaxcyl,
                mmax=0, symmetry='a', fixOrder=True)
        else:
            self._density = agama.Density(density=densfnc, type='DensitySphericalHarmonic',
                gridsizer=self.sizeradialsph, rmin=self.rminsph, rmax=self.rmaxsph,
                lmax=self.lmaxangularsph, mmax=0, symmetry='a', fixOrder=True)

class SelfConsistentModel(object):
    '''
    A pure Python analog of agama.SelfConsistentModel class (slightly simplified)
    '''
    def __init__(self, components=None, potential=None, **kwargs):
        self.components= components  # performs type checking
        self.potential = potential   # on assignment
        self._af = None
        for k,v in kwargs.items(): self.__dict__[k.lower()] = v
        if 'verbose' not in self.__dict__: self.verbose = True

    @property
    def components(self): return self._components

    @components.setter
    def components(self, components):
        if not isinstance(components, list):
            raise TypeError("Argument 'components' must be a list")
        self._components = components

    @property
    def potential(self): return self._potential

    @potential.setter
    def potential(self, potential):
        if not isinstance(potential, agama.Potential):
            raise TypeError("potential should be a valid Potential object")
        self._potential = potential

    @property
    def af(self): return self._af

    def iterate(self):
        if len(self.components)==0:
            raise TypeError("'components' is an empty list")
        # prepare ground: make sure the potential and the corresponding action finder are defined
        if self._potential is None:
            self._updatePotential()
        elif self._af is None:
            self._af = agama.ActionFinder(self.potential)
        # compute the density of all components
        for ic, component in enumerate(self.components):
            if self.verbose:
                print("Computing density for component %i..." % ic)
            component.update(self.potential, self._af)
        # update the total potential and the corresponding action finder
        self._updatePotential()

    def _updatePotential(self):
        if self.verbose:
            print("Updating potential...")
        # sort out density and potential components into several groups
        densitySph = []
        densityCyl = []
        potentials = []
        for component in self.components:
            dens = component.density
            if dens is not None:
                if component.disklike: densityCyl.append(dens)
                else: densitySph.append(dens)
            else:
                potentials.append(component.potential)
        # create a single Multipole potential for all non-disk-like density components
        if len(densitySph) > 0:
            potentials.append(agama.Potential(type='Multipole', density=agama.Density(*densitySph),
                gridsizer=self.sizeradialsph, rmin=self.rminsph, rmax=self.rmaxsph,
                lmax=self.lmaxangularsph, mmax=0, symmetry='a'))
        # create a single CylSpline potential representing all disk-like density components
        if len(densityCyl) > 0:
            potentials.append(agama.Potential(type='CylSpline', density=agama.Density(*densityCyl),
                gridsizer=self.sizeradialcyl,   rmin=self.rmincyl, rmax=self.rmaxcyl,
                gridsizez=self.sizeverticalcyl, zmin=self.zmincyl, zmax=self.zmaxcyl,
                mmax=0, symmetry='a'))
        # combine all potential components and reinitialize the action finder
        self._potential = agama.Potential(*potentials)
        if self.verbose:
            print("Updating action finder...")
        self._af = agama.ActionFinder(self._potential)


if __name__ == "__main__":
    # example of usage of both the original SCM classes
    # from the core Agama library and their Python analogues

    def test(Component, SelfConsistentModel):
        # test a two-component spherical model (to speed up things, do not use disky components);
        # the first component is an isotropic DF of a NFW halo with a cutoff,
        # and the second one represents a more concentrated baryonic component,
        # which will cause the adiabatic contraction of the halo
        # when both components are iterated to achieve equilibrium
        params_comp1 = dict(
            disklike = False,
            rminSph = 0.01,
            rmaxSph = 100.0,
            sizeRadialSph = 21,
            lmaxAngularSph = 0)
        params_comp2 = dict(
            disklike = False,
            rminSph = 0.01,
            rmaxSph = 10.0,
            sizeRadialSph = 16,
            lmaxAngularSph = 0)
        params_scm = dict(
            rminSph = 0.005,
            rmaxSph = 200,
            sizeRadialSph = 25,
            lmaxAngularSph = 0,
            verbose = False)

        potential_init = agama.Potential(type='spheroid',
            gamma=1, beta=3, mass=20.0, scaleRadius=5.0, outerCutoffRadius=40.0)
        df_comp1 = agama.DistributionFunction(type='quasispherical',
            density=potential_init, potential=potential_init)
        df_comp2 = agama.DistributionFunction(type='doublepowerlaw',
            norm=12.0, J0=1.0, coefJrIn=1., coefJzIn=1., coefJrOut=1., coefJzOut=1., slopeIn=1, slopeOut=6)
        model = SelfConsistentModel(
            components=[Component(df=df_comp1, **params_comp1), Component(df=df_comp2, **params_comp2)],
            potential=potential_init,
            **params_scm)
        print(dir(model)); print(dir(model.components[0]))
        for it in range(5):
            model.iterate()
        return agama.GalaxyModel(model.potential, df_comp1, model.af).moments([1,0.5,0.3])

    # run the model twice - first with built-in classes from agama, then with their Python analogues
    result1 = test(agama.Component, agama.SelfConsistentModel)
    result2 = test(Component, SelfConsistentModel)

    if result1[0]==result2[0] and all(result1[1]==result2[1]):   # should coincide exactly
        print("\033[1;32mALL TESTS PASSED\033[0m")
    else:
        print("%s != %s" % (result1, result2))
        print("\033[1;31mSOME TESTS FAILED\033[0m")
