from autoconf import conf
import autofit as af
from autogalaxy.pipeline.phase import abstract
from autogalaxy.pipeline.phase import extensions

import math
from scipy.stats import norm
import pickle
import json
import numpy as np

# noinspection PyAbstractClass
class StochasticPhase(extensions.ModelFixingHyperPhase):
    def __init__(
        self,
        phase: abstract.AbstractPhase,
        search,
        model_classes=tuple(),
        histogram_samples=100,
        histogram_bins=10,
        stochastic_method="gaussian",
        stochastic_sigma=0.0,
    ):

        self.is_stochastic = True
        self.histogram_samples = histogram_samples
        self.histogram_bins = histogram_bins
        self.stochastic_method = stochastic_method
        self.stochastic_sigma = stochastic_sigma

        super().__init__(
            phase=phase,
            search=search,
            model_classes=model_classes,
            hyper_name="stochastic",
        )

    def make_model(self, instance):
        return instance.as_model(self.model_classes)

    def run_hyper(
        self,
        dataset,
        results: af.ResultsCollection,
        info=None,
        pickle_files=None,
        **kwargs,
    ):
        """
        Run the phase, overriding the search's model instance with one created to
        only fit pixelization hyperparameters.
        """

        self.results = results

        stochastic_log_evidences_file = f"{conf.instance.output_path}/{self.paths.path_prefix}/{self.paths.name}/stochastic_log_evidences.json"

        try:
            stochastic_log_evidences = self.stochastic_log_evidences_from_json(
                filename=stochastic_log_evidences_file
            )
        except FileNotFoundError:
            stochastic_log_evidences = results.last.stochastic_log_evidences(
                histogram_samples=self.histogram_samples
            )
            self.stochastic_log_evidences_to_json(
                filename=stochastic_log_evidences_file,
                stochastic_log_evidences=stochastic_log_evidences,
            )

        if self.stochastic_method in "gaussian":

            mean, sigma = norm.fit(stochastic_log_evidences)

            limit = math.erf(0.5 * np.abs(self.stochastic_sigma) * math.sqrt(2))

            if self.stochastic_sigma >= 0.0:
                log_likelihood_cap = mean + (sigma * limit)
            else:
                log_likelihood_cap = mean - (sigma * limit)

            stochastic_tag = f"{self.stochastic_method}_{str(self.stochastic_sigma)}"

        else:

            log_likelihood_cap = np.median(stochastic_log_evidences)

            stochastic_tag = f"{self.stochastic_method}"

        phase = self.make_hyper_phase()
        phase.hyper_name = f"{phase.hyper_name}_{stochastic_tag}"

        phase.settings.log_likelihood_cap = log_likelihood_cap
        #       phase.paths.tag = phase.settings.phase_tag_no_inversion

        phase.use_as_hyper_dataset = False

        phase.model = self.make_model(results.last.instance)

        # TODO : HACK

        from autogalaxy.profiles import mass_profiles as mp

        mass = af.PriorModel(mp.EllipticalPowerLaw)

        mass.centre = af.last[-1].model.galaxies.lens.mass.centre
        mass.elliptical_comps = af.last[-1].model.galaxies.lens.mass.elliptical_comps
        mass.einstein_radius = af.last[-1].model.galaxies.lens.mass.einstein_radius

        phase.model.galaxies.lens.mass = mass

        result = phase.run(
            dataset,
            mask=results.last.mask,
            results=results,
            info=info,
            pickle_files=pickle_files,
            log_likelihood_cap=log_likelihood_cap,
        )

        self.save_stochastic_log_evidences(
            stochastic_log_evidences=stochastic_log_evidences
        )

        return result

    def save_stochastic_log_evidences(self, stochastic_log_evidences):
        """
        Save the dataset associated with the phase
        """
        with open(
            "{}/stochastic_log_evidences.pickle".format(self.paths.pickle_path), "wb"
        ) as f:
            pickle.dump(stochastic_log_evidences, f)

    def stochastic_log_evidences_from_json(cls, filename):
        with open(filename, "r") as f:
            return np.asarray(json.load(f))

    def stochastic_log_evidences_to_json(self, filename, stochastic_log_evidences):
        """
        Save the dataset associated with the phase
        """
        with open(filename, "w") as outfile:
            json.dump(
                [float(evidence) for evidence in stochastic_log_evidences], outfile
            )
