# Licensed under a 3-clause BSD style license - see LICENSE.rst
import astropy.units as u
import numpy as np

from .core import _prefit, lnprobmodel
from .extern.validator import validate_array
from .plot import _plot_data_to_ax, color_cycle
from .utils import sed_conversion, validate_data_table

__all__ = ["InteractiveModelFitter"]


def _process_model(model):
    if (
        isinstance(model, tuple) or isinstance(model, list)
    ) and not isinstance(model, np.ndarray):
        return model[0]
    else:
        return model


class InteractiveModelFitter:
    """
    Interactive model fitter using matplotlib widgets


    Parameters
    ----------
    modelfn : function
        A function that takes a vector in the parameter space and the data
        table, and returns the expected fluxes at the energies in the
        spectrum.
    p0 : array
        Initial position vector.
    data : `~astropy.table.Table` or list of `~astropy.table.Table`
        Table or tables with observed spectra. Must follow the format described
        in `naima.run_sampler`.
    e_range : list of `~astropy.units.Quantity`, length 2, optional
        Limits in energy for the computation of the model.
        Note that setting this parameter will mean that the model output is
        computed twice when `data` is provided: once for display using
        `e_range` and once for computation of the log-probability using the
        energy values of the spectra.
    e_npoints : int, optional
        How many points to compute for the model if `e_range` is set. Default
        is 100.
    labels : iterable of strings, optional
        Labels for the parameters included in the position vector ``p0``. If
        not provided ``['par1','par2', ... ,'parN']`` will be used.
    sed : bool, optional
        Whether to plot SED or differential spectrum.
    auto_update : bool, optional
        Whether to update the model plot when parameter sliders are changed.
        Default is True and can also be changed through the GUI.
    """

    def __init__(
        self,
        modelfn,
        p0,
        data=None,
        e_range=None,
        e_npoints=100,
        labels=None,
        sed=True,
        auto_update=True,
    ):

        import matplotlib.pyplot as plt
        from matplotlib.widgets import Button, CheckButtons, Slider

        self.pars = p0
        self.P0_IS_ML = False
        npars = len(p0)
        if labels is None:
            labels = ["par{0}".format(i) for i in range(npars)]
        elif len(labels) < npars:
            labels += ["par{0}".format(i) for i in range(len(labels), npars)]

        self.hasdata = data is not None
        self.data = None
        if self.hasdata:
            self.data = validate_data_table(data, sed=sed)

        self.modelfn = modelfn
        self.fig = plt.figure()
        modelax = plt.subplot2grid((10, 4), (0, 0), rowspan=4, colspan=4)

        if e_range is not None:
            e_range = validate_array(
                "e_range", u.Quantity(e_range), physical_type="energy"
            )
            energy = (
                np.logspace(
                    np.log10(e_range[0].value),
                    np.log10(e_range[1].value),
                    e_npoints,
                )
                * e_range.unit
            )
            if self.hasdata:
                energy = energy.to(self.data["energy"].unit)
            else:
                e_unit = e_range.unit
        else:
            energy = np.logspace(-4, 2, e_npoints) * u.TeV
            e_unit = u.TeV

        # Bogus flux array to send to model if not using data
        if sed:
            flux = np.zeros(e_npoints) * u.Unit("erg/(cm2 s)")
        else:
            flux = np.zeros(e_npoints) * u.Unit("1/(TeV cm2 s)")

        if self.hasdata:
            e_unit = self.data["energy"].unit
            _plot_data_to_ax(self.data, modelax, sed=sed, e_unit=e_unit)
            if e_range is None:
                # use data for model
                energy = self.data["energy"]
                flux = self.data["flux"]

        self.data_for_model = {"energy": energy, "flux": flux}

        model = _process_model(self.modelfn(p0, self.data_for_model))

        if self.hasdata:
            if not np.array_equal(
                self.data_for_model["energy"].to(u.eV).value,
                self.data["energy"].to(u.eV).value,
            ):
                # this will be slow, maybe interpolate already computed model?
                model_for_lnprob = _process_model(
                    self.modelfn(self.pars, self.data)
                )
            else:
                model_for_lnprob = model
            lnprob = lnprobmodel(model_for_lnprob, self.data)
            if isinstance(lnprob, u.Quantity):
                lnprob = lnprob.decompose().value
            self.lnprobtxt = modelax.text(
                0.05,
                0.05,
                r"",
                ha="left",
                va="bottom",
                transform=modelax.transAxes,
                size=20,
            )
            self.lnprobtxt.set_text(
                r"$\ln\mathcal{{L}} = {0:.1f}$".format(lnprob)
            )

        self.f_unit, self.sedf = sed_conversion(energy, model.unit, sed)

        if self.hasdata:
            datamin = (
                self.data["energy"][0] - self.data["energy_error_lo"][0]
            ).to(e_unit).value / 3
            xmin = min(datamin, energy[0].to(e_unit).value)
            datamax = (
                self.data["energy"][-1] + self.data["energy_error_hi"][-1]
            ).to(e_unit).value * 3
            xmax = max(datamax, energy[-1].to(e_unit).value)
            modelax.set_xlim(xmin, xmax)
        else:
            # plot_data_to_ax has not set ylabel
            unit = self.f_unit.to_string("latex_inline")
            if sed:
                modelax.set_ylabel(r"$E^2 dN/dE$ [{0}]".format(unit))
            else:
                modelax.set_ylabel(r"$dN/dE$ [{0}]".format(unit))
            modelax.set_xlim(energy[0].value, energy[-1].value)

        (self.line,) = modelax.loglog(
            energy.to(e_unit),
            (model * self.sedf).to(self.f_unit),
            lw=2,
            c="k",
            zorder=10,
        )

        modelax.set_xlabel(
            "Energy [{0}]".format(energy.unit.to_string("latex_inline"))
        )

        paraxes = []
        for n in range(npars):
            paraxes.append(
                plt.subplot2grid((2 * npars, 10), (npars + n, 0), colspan=7)
            )
        self.parsliders = []
        slider_props = {"facecolor": color_cycle[-1], "alpha": 0.5}
        for label, parax, valinit in zip(labels, paraxes, p0):
            # Attempt to estimate reasonable parameter ranges from label
            pmin, pmax = valinit / 10, valinit * 3
            if "log" in label:
                span = 2
                if "norm" in label or "amplitude" in label:
                    # give more range for normalization parameters
                    span *= 2
                pmin, pmax = valinit - span, valinit + span
            elif ("index" in label) or ("alpha" in label):
                if valinit > 0.0:
                    pmin, pmax = 0, 5
                else:
                    pmin, pmax = -5, 0
            elif "norm" in label or "amplitude" in label:
                # norm without log, it will not be pretty because sliders are
                # only linear
                pmin, pmax = valinit / 100, valinit * 100

            slider = Slider(
                parax,
                label,
                pmin,
                pmax,
                valinit=valinit,
                valfmt="%.4g",
                **slider_props
            )
            slider.on_changed(self.update_if_auto)
            self.parsliders.append(slider)

        autoupdateax = plt.subplot2grid((8, 4), (4, 3), colspan=1, rowspan=1)
        auto_update_check = CheckButtons(
            autoupdateax, ("Auto update",), (auto_update,)
        )
        auto_update_check.on_clicked(self.update_autoupdate)
        self.autoupdate = auto_update

        updateax = plt.subplot2grid((8, 4), (5, 3), colspan=1, rowspan=1)
        update_button = Button(updateax, "Update model")
        update_button.on_clicked(self.update)

        if self.hasdata:
            fitax = plt.subplot2grid((8, 4), (6, 3), colspan=1, rowspan=1)
            fit_button = Button(fitax, "Do Nelder-Mead fit")
            fit_button.on_clicked(self.do_fit)

        closeax = plt.subplot2grid((8, 4), (7, 3), colspan=1, rowspan=1)
        close_button = Button(closeax, "Close window")
        close_button.on_clicked(self.close_fig)

        self.fig.subplots_adjust(top=0.98, right=0.98, bottom=0.02, hspace=0.2)

        plt.show()

    def update_autoupdate(self, label):
        self.autoupdate = not self.autoupdate

    def update_if_auto(self, val):
        if self.autoupdate:
            self.update(val)

    def update(self, event):
        # If we update, values have changed and P0 is not ML anymore
        self.P0_IS_ML = False
        self.pars = [slider.val for slider in self.parsliders]
        model = _process_model(self.modelfn(self.pars, self.data_for_model))
        self.line.set_ydata((model * self.sedf).to(self.f_unit))
        if self.hasdata:
            if not np.array_equal(
                self.data_for_model["energy"].to(u.eV).value,
                self.data["energy"].to(u.eV).value,
            ):
                # this will be slow, maybe interpolate already computed model?
                model = _process_model(self.modelfn(self.pars, self.data))
            lnprob = lnprobmodel(model, self.data)
            if isinstance(lnprob, u.Quantity):
                lnprob = lnprob.decompose().value
            self.lnprobtxt.set_text(
                r"$\ln\mathcal{{L}} = {0:.1f}$".format(lnprob)
            )
        self.fig.canvas.draw_idle()

    def do_fit(self, event):
        self.pars = [slider.val for slider in self.parsliders]
        self.pars, P0_IS_ML = _prefit(self.pars, self.data, self.modelfn, None)
        autoupdate = self.autoupdate
        self.autoupdate = False
        if P0_IS_ML:
            for slider, val in zip(self.parsliders, self.pars):
                slider.set_val(val)
            self.update("after_fit")
        self.autoupdate = autoupdate
        self.P0_IS_ML = P0_IS_ML

    def close_fig(self, event):
        import matplotlib.pyplot as plt

        plt.close(self.fig)
