import torch

from ...param import forward
from ...backend_obj import ArrayLike
from ...utils.decorators import ignore_numpy_warnings
from .._shared_methods import parametric_initialize, parametric_segment_initialize
from ...utils.parametric_profiles import gaussian_np
from .. import func


def _x0_func(model_params, R, F):
    return R[4], 10 ** F[0]


class GaussianMixin:
    """Gaussian radial light profile.

    The Gaussian profile is a simple and widely used model for extended objects.
    The functional form of the Gaussian profile is defined as:

    .. math::

       I(R) = \\frac{{\\rm flux}}{2\\pi\\sigma^2} \\exp(-R^2 / (2 \\sigma^2))

    where ``flux`` is the total flux of the profile and ``sigma`` is the
    standard deviation which controls the width of the profile.

    :param sigma: Standard deviation of the Gaussian profile in arcseconds.
    :param flux: Total flux of an unperturbed Gaussian profile (use model.total_flux() for general results).
    """

    _model_type = "gaussian"
    _parameter_specs = {
        "sigma": {
            "units": "arcsec",
            "valid": (0, None),
            "shape": (),
            "dynamic": True,
            "description": "Standard deviation of the Gaussian profile in arcseconds.",
        },
        "flux": {
            "units": "flux",
            "valid": (0, None),
            "shape": (),
            "dynamic": True,
            "description": "Total flux of an unperturbed Gaussian profile (use model.total_flux() for general results).",
        },
    }

    @torch.no_grad()
    @ignore_numpy_warnings
    def initialize(self):
        super().initialize()

        parametric_initialize(
            self,
            self.target[self.window],
            gaussian_np,
            ("sigma", "flux"),
            _x0_func,
        )

    @forward
    def radial_model(self, R: ArrayLike, sigma: ArrayLike, flux: ArrayLike) -> ArrayLike:
        return func.gaussian(R, sigma, flux)


class iGaussianMixin:
    """Gaussian radial light profile.

    The Gaussian profile is a simple and widely used model for extended objects.
    The functional form of the Gaussian profile is defined as:

    .. math::

       I(R) = \\frac{{\\rm flux}}{2\\pi\\sigma^2} \\exp(-R^2 / (2 \\sigma^2))

    where ``sigma`` is the standard deviation which controls the width of the
    profile and ``flux`` gives the total flux of the profile (assuming no
    perturbations).

    ``sigma`` and ``flux`` are batched by their first dimension, allowing for
    multiple Gaussian profiles to be defined at once.

    :param sigma: Standard deviation of the Gaussian profile in arcseconds.
    :param flux: Total flux of an unperturbed Gaussian profile (use model.total_flux() for general results).
    """

    _model_type = "gaussian"
    _parameter_specs = {
        "sigma": {
            "units": "arcsec",
            "valid": (0, None),
            "shape": (None,),
            "dynamic": True,
            "description": "Standard deviation of the Gaussian profile in arcseconds.",
        },
        "flux": {
            "units": "flux",
            "valid": (0, None),
            "shape": (None,),
            "dynamic": True,
            "description": "Total flux of an unperturbed Gaussian profile (use model.total_flux() for general results).",
        },
    }

    @torch.no_grad()
    @ignore_numpy_warnings
    def initialize(self):
        super().initialize()

        parametric_segment_initialize(
            model=self,
            target=self.target[self.window],
            prof_func=gaussian_np,
            params=("sigma", "flux"),
            x0_func=_x0_func,
            segments=self.segments,
        )

    @forward
    def iradial_model(self, i: int, R: ArrayLike, sigma: ArrayLike, flux: ArrayLike) -> ArrayLike:
        return func.gaussian(R, sigma[i], flux[i])


class GaussianPSFMixin:
    """Gaussian radial light profile.

    The Gaussian profile is a simple and widely used model for extended objects.
    The functional form of the Gaussian profile is defined as:

    .. math::

       I(R) = \\frac{{\\rm flux}}{2\\pi\\sigma^2} \\exp(-R^2 / (2 \\sigma^2))

    where ``flux`` is the total flux of the profile and ``sigma`` is the
    standard deviation which controls the width of the profile.

    :param sigma: Standard deviation of the Gaussian profile in pixels.
    :param flux: Total flux of an unperturbed Gaussian profile (use model.total_flux() for general results).
    """

    _model_type = "gaussian"
    _parameter_specs = {
        "sigma": {
            "units": "pix",
            "valid": (0, None),
            "shape": (),
            "dynamic": True,
            "description": "Standard deviation of the Gaussian profile in pixels.",
        },
        "flux": {
            "units": "flux",
            "valid": (0, None),
            "shape": (),
            "dynamic": False,
            "value": 1.0,
            "description": "Total flux of an unperturbed Gaussian profile (use model.total_flux() for general results).",
        },
    }

    @torch.no_grad()
    @ignore_numpy_warnings
    def initialize(self):
        super().initialize()

        parametric_initialize(
            self,
            self.target[self.window],
            gaussian_np,
            ("sigma", "flux"),
            _x0_func,
        )

    @forward
    def radial_model(self, R: ArrayLike, sigma: ArrayLike, flux: ArrayLike) -> ArrayLike:
        return func.gaussian(R, sigma, flux)
