"""
Plummer gas model generator

This module contains a function used to create Plummer (1911) models consisting 
of (SPH) gas particles.
Plummer models follow a spherically symmetric density profile of the form:
rho = c * (1 + r**2)**(-5/2)
"""

import numpy

from amuse.ext.evrard_test import uniform_unit_sphere
from amuse.units import nbody_system
from amuse.units import units
from amuse import datamodel

__all__ = ["new_plummer_gas_model"]


class MakePlummerGasModel:
    def __init__(
        self,
        targetN,
        convert_nbody=None,
        base_grid=None,
        rscale=1 / 1.695,
        mass_cutoff=0.999,
        do_scale=False,
        **kwargs
    ):
        self.kwargs = kwargs
        self.targetN = targetN
        self.convert_nbody = convert_nbody
        self.rscale = rscale
        self.mass_frac = mass_cutoff
        self.do_scale = do_scale
        self.internal_energy = 0.25 / self.rscale
        self.base_sphere = uniform_unit_sphere(targetN, base_grid)

    def new_model(self):
        x, y, z = self.base_sphere.make_xyz()
        self.actualN = len(x)
        r = numpy.sqrt(x**2 + y**2 + z**2) * self.mass_frac ** (1 / 3.0)
        rtarget = self.rscale * (r**2 / (1 - r**2)) ** 0.5
        mr = self.mass_frac ** (1 / 3.0)
        maxr = self.rscale * (mr**2 / (1 - mr**2)) ** 0.5
        mass = numpy.ones_like(x) / self.actualN
        internal_energy = self.internal_energy / (1 + (rtarget / self.rscale) ** 2) ** (
            1.0 / 2
        )
        r = r.clip(1.0e-8, maxr)
        x = rtarget * x / r
        y = rtarget * y / r
        z = rtarget * z / r
        vx = numpy.zeros_like(x)
        vy = numpy.zeros_like(x)
        vz = numpy.zeros_like(x)
        return (mass, x, y, z, vx, vy, vz, internal_energy)

    @property
    def result(self):
        mass, x, y, z, vx, vy, vz, u = self.new_model()
        result = datamodel.Particles(self.actualN, **self.kwargs)
        result.mass = nbody_system.mass.new_quantity(mass)
        result.x = nbody_system.length.new_quantity(x)
        result.y = nbody_system.length.new_quantity(y)
        result.z = nbody_system.length.new_quantity(z)
        result.vx = nbody_system.speed.new_quantity(vx)
        result.vy = nbody_system.speed.new_quantity(vy)
        result.vz = nbody_system.speed.new_quantity(vz)
        result.u = (nbody_system.speed**2).new_quantity(u)

        result.move_to_center()
        if self.do_scale:
            potential_energy = result.potential_energy(G=nbody_system.G)
            result.position *= potential_energy / (-0.5 | nbody_system.energy)

            internal_energy = result.thermal_energy()
            result.u *= (0.25 | nbody_system.energy) / internal_energy

        if self.convert_nbody is not None:
            result = datamodel.ParticlesWithUnitsConverted(
                result, self.convert_nbody.as_converter_from_si_to_generic()
            )
            result = result.copy()

        return result


def new_plummer_gas_model(number_of_particles, *list_arguments, **keyword_arguments):
    """
    Create a plummer gas model with the given number of particles. Returns
    a set of SPH particles with equal masses and positions distributed
    to fit a plummer distribution model. The model is centered around the
    origin. Velocities are set to zero, and internal energies are set to
    balance the gravitational forces between the gas particles.

    :argument number_of_particles: Number of particles to include in the plummer sphere
    :argument convert_nbody:  When given will convert the resulting set to SI units
    :argument mass_cutoff: Mass percentage inside radius of 1
    :argument do_scale: scale the result, similar to true nbody units (M=1, Q=0.25, U=-0.5)
    """
    uc = MakePlummerGasModel(number_of_particles, *list_arguments, **keyword_arguments)
    return uc.result
