# Copyright 2018 Olivier Scholder <o.scholder@gmail.com>

"""
Helper functions to handle spectras.
"""

from .misc import alias, dec_debug, do_debug


def get_substance_peaks(substance, negative=True):
    import os
    import sqlite3

    DB_PATH = os.path.join(
        os.path.abspath(os.path.join(__file__, "../..")), "data", "elements.db"
    )
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()
    c.execute(
        "SELECT Peaks.Fragment from Peaks where Peaks.Substance==(SELECT ID from substance where Name LIKE '%{name}%') and Polarity{pol}=0".format(
            name=substance, pol="><"[negative]
        )
    )
    return [x[0] for x in c.fetchall()]


def get_dm(m, sf, k0, dsf, dk0):
    import numpy as np

    return 2 * np.sqrt(m) * np.sqrt((dk0**2 / (sf**2)) + m * (dsf**2 / (sf**2)))


@alias("showPeak")
def show_peak(
    m,
    D,
    m0,
    delta=None,
    errors=False,
    dm0=0,
    dofit=False,
    show_elts=True,
    debug=False,
    Aredux=1,
    label=None,
    include=None,
    include_only=None,
    exclude=None,
    polarity="+",
    colors="rgb",
    pretty=True,
    formula=True,
    auto_scale=True,
    fakefit=False,
    zero_axis=True,
    **kargs,
):
    """
    Plot a spectrum (given by its mass vector m and intensity vector D) around a mass m0 (Δm=delta).

    Parameters
    ----------
    m : numpy.ndarray
        mass vector
    D : numpy.ndarray
        Intensity vector
    m0 : float
        Central mass of the plot's x-axis
    delta: None or float
        the half-span of the plot's x-axis.
        If None will try to determine automatically the peaks region
    errors : boolean
        If true will include the fitting errors (if dofit is True)
    dofit : boolean
        If True will fit the peaks for the elements found
    show_elts : boolean
        If True will display a list of all elements for the plotted region. The basic elements comes from a database, but they can be adjusted with the include, include_only and exclude parameters.
    include : list
        List of elements to be included to the plotting elements
    include_only : list
        Only those elements will be displayed and none from the database
    exclude : list
        Exclude some elements from display
    fakefit : boolean
        If the fit fails, fakefit can be used in order to see and adjust the first fitting guess
    dm0 : float
        The mass shift used for the fitting first guess.
    Aredux : float
        The amplitude reduction used to the first fitting guess.
    label : string
        Plot label. Will be displayed if you use the legend() function from matplotlib.
    polarity : string ("+" or "-")
        will switch the elements chosen from the database
    colors : string
        A series of letters defining matplotlib color used to display the elements.
    pretty : boolean
        If True will display the elements in a pretty way such that the elements text are not overlapping (or at least will avoid it as much as possible)
    formula : boolean
        Will convert the formula in a pretty manner. Exponents and indices will be written correctly
    auto_scale : boolean
        If True will try to rescale the plot in order to have all labels inside.
    zero_axis : boolean or dict
        Will draw an horizontal line for the zero intensities.
        If a dict, the parameters of the lines can be tuned. See the parameters passed to axhline from matplotlib
    """
    import copy

    import matplotlib.pyplot as plt
    import numpy as np
    from scipy.optimize import curve_fit

    from . import LG, get_mass, get_peaklist
    from .elts import formulafy

    if exclude is None:
        exclude = []
    if "showElts" in kargs:
        from warnings import warn

        warn("Parameter showElts is deprecated. Please use show_elts", stacklevel=2)
        show_elts = kargs.pop("showElts")

    if include is None:
        include = []

    if isinstance(include, str):
        include = include.split(",")

    assert isinstance(include, list)

    if isinstance(m0, str):
        include.append(m0)
        m0 = get_mass(m0)

    if do_debug(debug):
        import time

        time.time()

    if isinstance(exclude, str):
        exclude = exclude.split(",")

    if fakefit:
        dofit = True

    if delta is None:
        delta0 = kargs.pop("delta0", 0.5)
        mask = (m >= m0 - delta0) * (m <= m0 + delta0)
        thresh = max(5, np.max(D[mask]) * kargs.get("min_perc", 0.01))
        li = np.argmax(mask) + np.argmax(D[mask] > thresh)
        ui = np.argmax(mask) + np.sum(mask) - np.argmax(D[mask][::-1] > thresh)
        upper_mass = m[ui] + 0.01
        lower_mass = m[li] - 0.01

        delta = (upper_mass - lower_mass) / 2.0 + 0.01
    else:
        lower_mass = m0 - delta
        upper_mass = m0 + delta
    if do_debug(debug):
        print(f"Lower mass: {lower_mass:.2f}\nUpper mass: {upper_mass:.2f}")
    mask = (m >= lower_mass) * (m <= upper_mass)

    negative = False
    if polarity in ["-", "Negative", "negative", "neg", "Neg", "NEG"]:
        negative = True
    if show_elts:
        if include_only is not None:
            E = (
                include_only.split(",")
                if isinstance(include_only, str)
                else include_only
            )
        else:
            E = [
                x
                for NM in range(int(round(lower_mass)), int(round(upper_mass)) + 1)
                for x in get_peaklist(NM, negative)
            ]
            E = [x for x in E if x not in exclude] + include
        E = list({x + [["+", "-"][negative], ""]["+" in x or "-" in x] for x in E})
    else:
        E = []
    if negative:
        E = [x + ["-", ""]["-" in x] for x in E]
    if do_debug(debug):
        print("Selected Elements: " + ", ".join(E))
    m0s = [get_mass(x) for x in E]
    E = [x for x, y in zip(E, m0s) if y >= lower_mass and y <= upper_mass]
    m0s = [get_mass(x) for x in E]
    E_labels = [formulafy(x) for x in E] if formula else E
    if do_debug(debug):
        print("Displayed elements:", ", ".join(E))

    Dt = np.copy(D[mask])
    mt = m[mask]
    if do_debug(debug):
        print(" ; ".join(E))
    ax = kargs.pop("ax", plt.gca())
    if dofit or fakefit:
        mp = m[mask][np.argmax(D[mask])]  # mass of max peak
        dm = dm0
        if dm is None:
            dm = 0
        if len(E) > 0:
            i = np.argmin(
                abs(
                    np.array(
                        [get_mass(x) for x in E if isinstance(x, str)]
                        + [x for x in E if isinstance(x, float)]
                    )
                    - mp
                )
            )  # which element is the closest to max_peak
            if dm0 is None:
                dm = mp - get_mass(E[i])
        p0 = [kargs.get("asym0", 1), dm] + [0, 0] * len(
            E
        )  # delta m is estimated from the deviation of the highest peak
        Et = copy.deepcopy(E)  # copy element list

        while len(Et) > 0:
            mp = mt[np.argmax(Dt)]  # mass of max peak
            ms = [get_mass(x) for x in Et]
            i = np.argmin(abs(ms - mp))
            idx = E.index(Et[i])
            j = np.argmin(abs(mt - ms[i] - dm))
            if do_debug(debug):
                print("Max element:", Et[i], mp, ms[i], Dt[j], idx)
            p0[2 + 2 * idx] = kargs.get("sig0", 0.002)
            p0[3 + 2 * idx] = Aredux * Dt[j]
            Dt -= LG(mt, ms[i], p0[2 + 2 * idx], amp=p0[3 + 2 * idx], asym=p0[0], lg=0)
            Dt[Dt < 0] = 0
            del Et[i]

        def fit(x, *p):
            y = x * 0
            for i in range((len(p) - 2) // 2):
                x0 = m0s[i] + p[1]
                y += LG(x, x0, p[2 + 2 * i], amp=p[3 + 2 * i], asym=p[0], lg=0)
            return y

        fit_type = None
        if dofit:
            try:
                assert not fakefit
                popt, pcov = curve_fit(
                    fit,
                    m[mask],
                    D[mask],
                    p0=p0,
                    bounds=(
                        [1 / kargs.get("asym_max", 10), -0.015]
                        + [0, 0] * ((len(p0) - 1) // 2),
                        [kargs.get("asym_max", 10), 0.015]
                        + [kargs.get("sig_max", 0.01), np.inf] * ((len(p0) - 1) // 2),
                    ),
                )
                fit_type = 0
            except:
                p0[1] = 0
                try:
                    assert not fakefit
                    popt, pcov = curve_fit(
                        fit,
                        m[mask],
                        D[mask],
                        p0=p0,
                        bounds=(
                            [1 / kargs.get("asym_max", 4), -0.015]
                            + [0, 0] * ((len(p0) - 1) // 2),
                            [kargs.get("asym_max", 4), 0.015]
                            + [kargs.get("sig_max", 0.01), np.inf]
                            * ((len(p0) - 1) // 2),
                        ),
                    )
                    fit_type = 1
                except Exception as e:
                    fit_type = 2
                    if do_debug(debug):
                        raise e
                    popt = p0
                    pcov = np.zeros((len(p0), len(p0)))
                    if ax is not None:
                        for x in ["right", "left", "top", "bottom"]:
                            ax.spines[x].set_color("red")
        else:
            popt = p0
            popt[1] = 0
            pcov = np.zeros((len(p0), len(p0)))
        if fakefit:
            popt[1] = dm0
        res = {}
        err = np.sqrt(np.diag(pcov))
        for i in range((len(popt) - 1) // 2):
            Y = LG(
                m[mask], m0s[i], popt[2 + 2 * i], popt[3 + 2 * i], lg=0, asym=popt[0]
            )
            Area = (
                popt[2 * i + 3]
                * popt[2 * i + 2]
                * np.sqrt(2 * np.pi)
                * (0.5 + 0.5 * popt[0])
            )
            Area_err = np.sqrt(
                (
                    err[2 * i + 3]
                    * popt[2 * i + 2]
                    * np.sqrt(2 * np.pi)
                    * (0.5 + 0.5 * popt[0])
                )
                ** 2
                + (
                    err[2 * i + 2]
                    * popt[2 * i + 3]
                    * np.sqrt(2 * np.pi)
                    * (0.5 + 0.5 * popt[0])
                )
                ** 2
                + (
                    err[0]
                    * popt[2 * i + 2]
                    * popt[2 * i + 3]
                    * np.sqrt(2 * np.pi)
                    * 0.5
                )
                ** 2
            )
            if not errors:
                res[E_labels[i]] = {
                    "m0": m0s[i],
                    "mass": m0s[i] + popt[1],
                    "Area": Area,
                    "Amp": popt[2 * i + 3],
                    "sig": popt[2 * i + 2],
                    "assym": popt[0],
                    "dm": popt[1] * 1e6 / m0s[i],
                    "fit": fit_type,
                }
            else:
                res[E_labels[i]] = {
                    "m0": m0s[i],
                    "mass": (m0s[i] + popt[1], err[1]),
                    "Area": (Area, Area_err),
                    "Amp": (popt[2 * i + 3], err[2 * i + 3]),
                    "sig": (popt[2 * i + 2], err[2 * i + 2]),
                    "assym": (popt[0], err[0]),
                    "dm": (popt[1] * 1e6 / m0s[i], err[1] * 1e6 / m0s[i]),
                    "fit": fit_type,
                }
    else:
        popt = [0, 0]
        res = None
    if ax is not None:
        if label is None:
            p = ax.plot(
                m[mask] - popt[1], D[mask], color=kargs.get("curve_color", None)
            )
        else:
            p = ax.plot(
                m[mask] - popt[1],
                D[mask],
                label=label,
                color=kargs.get("curve_color", None),
            )

    if show_elts and ax is not None:
        from . import put_Xlabels

        if dofit and ax is not None:
            ax.plot(m[mask], D[mask], color=p[0].get_color(), alpha=0.1)
            if pretty:
                put_Xlabels(
                    ax,
                    m0s,
                    [f"{E}: {res[E]['Area']:.2f}" for E in res],
                    color=colors,
                    debug=dec_debug(debug),
                    **kargs,
                )
            resO = [(res[x]["m0"], res[x]) for x in res]
            resO.sort(key=lambda x: x[0])
            for i, (_, r) in enumerate(resO):
                col = colors[i % len(colors)]
                Y = LG(m[mask], r["m0"], r["sig"], r["Amp"], lg=0, asym=r["assym"])
                ax.plot(m[mask], Y, "--", color=col)
        elif pretty:
            put_Xlabels(
                ax, m0s, E_labels, color=colors, debug=dec_debug(debug), **kargs
            )

        if "dsf" in kargs and "dk0" in kargs and "sf" in kargs and "k0" in kargs:
            sf, dsf, k0, dk0 = (kargs[x] for x in "sf,dsf,k0,dk0".split(","))
            P = list(zip(m0s, E))
            P.sort(key=lambda x: x[0])
            for i, (mi, Ei) in enumerate(P):
                col = colors[i % len(colors)]
                dmi = get_dm(mi, sf, k0, dsf, dk0)
                ax.axvline(mi - dmi, color=col, alpha=0.5, linestyle=":")
                ax.axvline(mi + dmi, color=col, alpha=0.5, linestyle=":")

        if not pretty:
            P = list(zip(m0s, E_labels))
            P.sort(key=lambda x: x[0])
            y = ax.get_ylim()[1]
            for i, (mi, Ei) in enumerate(P):
                col = colors[i % len(colors)]
                ax.axvline(mi, color=col, alpha=0.5)
                ax.annotate(
                    Ei,
                    (mi, y),
                    (5, 0),
                    rotation=90,
                    va="top",
                    ha="left",
                    textcoords="offset pixels",
                    color=col,
                )
        if auto_scale:
            fig = plt.gcf()
            renderer = fig.canvas.get_renderer()
            ymax = ax.get_ylim()[1]
            ymin = ax.get_ylim()[0]
            for child in ax.get_children():
                bbox = child.get_window_extent(renderer)
                bbox_data = bbox.transformed(ax.transData.inverted())
                if ymax < bbox_data.ymax:
                    ymax = bbox_data.ymax
            ax.set_ylim((ymin, ymax))
            ax.autoscale_view()

            # Add the Intensity=0 line in case there is negative intensities (happens when plotting PCA components)
    if zero_axis:
        if isinstance(zero_axis, dict):
            ax.achline(0, **zero_axis)
        else:
            ax.axhline(0, color="k", alpha=0.5, lw=0.5)

    if dofit and ax is not None:
        ax.plot(m[mask] - popt[1], fit(m[mask], *popt), "r:")
    return res


def plot_isotopes(
    elt,
    amp=None,
    main=None,
    ax=None,
    sig=None,
    asym=None,
    lg=0,
    limit=1,
    color="C1",
    show_elts=False,
    debug=False,
    dm=0,
    **kargs,
):
    """
    plot the isotopes of a given element on a spectral profile plot
    """
    # Old parameter name compatibility
    if "Amp" in kargs:
        from warnings import warn

        warn(
            "Parameter Amp is deprecated. Please use amp in order to set the amplitude!",
            stacklevel=2,
        )
        amp = kargs.pop("Amp")
    if "showElts" in kargs:
        from warnings import warn

        warn("Parameter showElts is deprecated. Please use show_elts", stacklevel=2)
        show_elts = kargs.pop("showElts")
    import re

    import matplotlib.pyplot as plt
    import numpy as np
    from scipy.optimize import curve_fit

    from . import LG, Molecule, get_abund, get_isotopes, get_mass

    if type(elt) is Molecule:
        elt = str(elt)
    if ax is None:
        ax = plt.gca()
    if main is None:
        main = ax
    main_isotope = re.sub("^([0-9]+)", "", elt)
    main_iso = (main_isotope, get_mass(main_isotope), get_abund(main_isotope))
    L = main.lines[0]
    m, y = L.get_xdata(), L.get_ydata()
    m -= dm
    if hasattr(ax, "log") and ax.log:
        y = 10**y
    mask = np.abs(m - main_iso[1]) < 0.1
    if amp is None:
        i = np.argmin(np.abs(m - main_iso[1]))
        amp = np.max(y[i - 10 : i + 10]) / main_iso[2]
    if sig is None and asym is None:

        def fit(x, s, a):
            return LG(x, main_iso[1], s, amp=amp * main_iso[2], lg=lg, asym=a)

        (sig, asym), _ = curve_fit(fit, m[mask], y[mask], (0.005, 1))
    elif sig is None:

        def fit(x, s):
            return LG(x, main_iso[1], s, amp=amp * main_iso[2], lg=lg, asym=asym)

        s, _ = curve_fit(fit, m[mask], y[mask], (0.005))
        sig = s[0]
    elif asym is None:

        def fit(x, a):
            return LG(x, main_iso[1], sig, amp=amp * main_iso[2], lg=lg, asym=a)

        a, _ = curve_fit(fit, m[mask], y[mask], (1))
        asym = a[0]
    if debug:
        print(sig, asym)
    isos = get_isotopes(elt, min_abund=limit / amp)
    L = ax.lines[0]
    m, y = L.get_xdata(), L.get_ydata()
    if hasattr(ax, "log") and ax.log:
        y = 10**y
    main.get_xlim()
    s = m * 0

    for iso in isos:
        s += LG(m, iso[1], sig, amp=amp * iso[2], lg=lg, asym=asym)
        if show_elts:
            ax.annotate(
                iso[0],
                (iso[1], amp * iso[2]),
                (0, 5),
                textcoords="offset pixels",
                ha="center",
            )
    if hasattr(ax, "log") and ax.log:
        s[s >= 1] = np.log10(s[s >= 1])
        s[s < 1] = 0
    ax.plot(m + dm, s, color, **kargs)
    return m, s
