from os import path
import pickle
import json
import autofit as af
from autogalaxy.pipeline.phase import abstract
from autogalaxy.pipeline.phase import extensions

from os import path
import os
import math
from scipy.stats import norm
import numpy as np

# noinspection PyAbstractClass
class StochasticPhase(extensions.HyperPhase):
    def __init__(
        self,
        phase: abstract.AbstractPhase,
        hyper_search,
        model_classes=tuple(),
        stochastic_method="gaussian",
        stochastic_sigma=0.0,
    ):

        self.is_stochastic = True
        self.stochastic_method = stochastic_method
        self.stochastic_sigma = stochastic_sigma

        super().__init__(
            phase=phase, hyper_search=hyper_search, model_classes=model_classes
        )

    @property
    def hyper_name(self):
        return "stochastic"

    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.
        """

        # TODO : This is a horror show. It will be cleaner once autofit doesnt use .zip files anymore.

        self.results = results

        self.search.paths.restore()

        stochastic_log_evidences_json_file = path.join(
            self.paths.output_path, "stochastic_log_evidences.json"
        )

        try:
            with open(stochastic_log_evidences_json_file, "r") as f:
                stochastic_log_evidences = np.asarray(json.load(f))
        except FileNotFoundError:
            stochastic_log_evidences = results.last.stochastic_log_evidences
            self.save_stochastic_log_evidences_to_json(
                stochastic_log_evidences=stochastic_log_evidences
            )

        self.search.paths.zip_remove()

        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 = 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(instance=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

        # TODO : Nasty hack to get log evidnees to copy, do something bettter in future.

        phase.modify_search_paths()

        phase.search.paths.restore()

        try:
            os.makedirs(self.paths.output_path)
        except FileExistsError:
            pass

        stochastic_log_evidences_json_file = path.join(
            phase.paths.output_path, "stochastic_log_evidences.json"
        )

        try:
            with open(stochastic_log_evidences_json_file, "w") as outfile:
                json.dump(
                    [float(evidence) for evidence in stochastic_log_evidences], outfile
                )
        except FileExistsError:
            pass

        stochastic_log_evidences_pickle_file = path.join(
            phase.paths.pickle_path, "stochastic_log_evidences.pickle"
        )

        with open(stochastic_log_evidences_pickle_file, "wb") as f:
            pickle.dump(stochastic_log_evidences, f)

        phase.search.paths.zip_remove()

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