import autofit as af
from autofit_workspace.howtofit.chapter_1_introduction.tutorial_5_visualization_masking.src.dataset.dataset import (
    Dataset,
    MaskedDataset,
)
from autofit_workspace.howtofit.chapter_1_introduction.tutorial_5_visualization_masking.src.phase.result import (
    Result,
)
from autofit_workspace.howtofit.chapter_1_introduction.tutorial_5_visualization_masking.src.phase.analysis import (
    Analysis,
)


"""The Phase class ``.nit__ method is unchanged from the previous tutorial, however other methods are changed."""


class Phase(af.AbstractPhase):

    gaussian = af.PhaseProperty("gaussian")

    Result = Result

    def __init__(self, *, gaussian, search):
        """
        A phase which fits a `Gaussian` model using a `NonLinearSearch`.

        Parameters
        ----------
        gaussian : gaussians.Gaussian
            The model component `Gaussian` class fitted by this phase.
        search: class
            The class of a non_linear search
        """

        super().__init__(search=search)

        self.gaussian = gaussian

    """
    The run method is slightly different, as it now passed a mask in addition to the `Dataset`. These are used to set up
    the masked-dataset in the `analysis.py` module.
    """

    def run(self, dataset: Dataset, mask):
        """
        Pass a `Dataset` to the phase, running the phase and `NonLinearSearch`.

        Parameters
        ----------
        dataset: aa.Dataset
            The `Dataset` fitted by the phase, as defined in the `dataset.py` module.
        mask: Mask2D
            The mask used for the analysis.

        Returns
        -------
        result: AbstractPhase.Result
            A result object comprising information on the `NonLinearSearch` and the maximum likelihood model.
        """

        analysis = self.make_analysis(dataset=dataset, mask=mask)

        result = self.run_analysis(analysis=analysis)

        return self.make_result(result=result, analysis=analysis)

    def make_analysis(self, dataset, mask):
        """
        Returns an Analysis object, which creates the `Dataset` and contains the functions which perform the fit.

        Parameters
        ----------
        dataset: aa.Dataset
            The `Dataset` fitted by the phase, as defined in the `dataset.py` module.

        Returns
        -------
        analysis : Analysis
            An analysis object that the `NonLinearSearch` calls to determine the fit log_likelihood for a given model
            instance.
        """

        """To mask the `Dataset` we simply pass both to the MaskedDataset class."""

        masked_dataset = MaskedDataset(dataset=dataset, mask=mask)

        """
        The `image_path` is where visualizatiion of the model fit is output. Below, we direct it to the same path as
        the `NonLinearSearch` output, but with an additional folder `image` at the end. This path should be used
        for pretty much any project.
        """

        return Analysis(
            masked_dataset=masked_dataset, image_path=self.search.paths.image_path
        )

    def make_result(self, result, analysis):
        return self.Result(
            samples=result.samples,
            previous_model=self.model,
            search=self.search,
            analysis=analysis,
        )
