import bokeh.plotting as bklt
from bokeh.embed import file_html
from bokeh.resources import CDN
import numpy as np
import strax

export, __all__ = strax.exporter()


@export
def bokeh_to_wiki(fig, outputfile=None):
    """Function which converts bokeh HTML code to a wiki readable code.

    :param fig: Figure to be conerted
    :param outputfile: String of absolute file path. If specified output is writen to the file. Else
        output is print to the notebook and can be simply copied into the wiki.

    """
    # convert plot to wiki format:
    html = file_html(fig, CDN)
    html = "\n".join((["<html>"] + html.split("\n")[2:]))

    if outputfile:
        with open(outputfile, mode="w") as file:
            file.write(html)
    else:
        print(html)


def get_peaks_source(peaks, relative_start=0, time_scaler=1, keep_amplitude_per_sample=True):
    """Computes bokeh.plotting.ColumnDataSource for given peaks.

    :param peaks: Peaks to be plotted.
    :param relative_start: t0 from which on the peaks should be plotted.
    :param time_scaler: Factor to rescale the time from ns to other scale. E.g. =1000 scales to µs.
    :param keep_amplitude_per_sample: Boolean if true amplitude of the plotted peaks is in
        pe/sample. False pe/ns.
    :return: bokeh.plotting.ColumnDataSource instance which can be used to plot peaks.

    """
    if not (np.all(peaks["type"] == peaks[0]["type"])):
        raise ValueError("All peaks must be of the same type (S1, S2 or Unknown)!")

    x_p = []
    y_p = []
    for ind, p in enumerate(peaks):
        x, y = _patches_x_y(p, keep_amplitude_per_sample=keep_amplitude_per_sample)
        x -= relative_start  # relative to first peak
        x = x / time_scaler
        x_p.append(x)
        y_p.append(y)

    if peaks[0]["type"] == 2:
        scaler = 10**-3
    else:
        scaler = 1

    source = bklt.ColumnDataSource(
        data={
            "type": peaks["type"],
            "xs": x_p,  # Coordinates for peak patches
            "ys": y_p,
            "x": peaks["x"],  # XY-Pos in PMZ Hitpattern
            "y": peaks["y"],
            "dt": peaks["dt"],
            "time": peaks["time"],
            "center_time": peaks["center_time"],
            "endtime": strax.endtime(peaks),
            "area": peaks["area"],
            "nhits": peaks["n_hits"],
            "aft": peaks["area_fraction_top"],
            "width_50": peaks["range_50p_area"] * scaler,
            "width_90": peaks["range_90p_area"] * scaler,
            "rise": peaks["rise_time"] * scaler,
            "rel_center_time": (peaks["center_time"] - peaks["time"]) * scaler,
        }
    )
    return source


def _patches_x_y(peak, keep_amplitude_per_sample=False):
    """Creates x,y coordinates needed to draw peaks via bokeh.models.patches.

    :param peak: Peak for which we need the x/y samples
    :param keep_amplitude_per_sample: If y-data should be in units of "per sample" or "per ns".
    :return: Tuple of x, y

    """
    if keep_amplitude_per_sample:
        dt_a = 1
    else:
        dt_a = peak["dt"]

    x = np.arange(peak["length"])
    xx = np.zeros(2 * len(x), dtype=x.dtype)
    mx = 0.5 * (x[1::] + x[:-1])
    xx[1:-1:2] = mx
    xx[2::2] = mx
    xx[0] = 1.5 * x[0] - 0.5 * x[1]
    xx[-1] = 1.5 * x[-1] - 0.5 * x[-2]
    xx = np.array(xx) * peak["dt"] + peak["time"]
    y = peak["data"][: peak["length"]]
    yy = np.zeros(2 * len(x))
    yy[0::2] = y
    yy[1::2] = y
    yy = np.array(yy) / dt_a

    # baseline since we'll fill underneath
    xx = np.concatenate([[xx[0]], xx, [xx[-1]]])
    yy = np.concatenate([[0], yy, [0]])
    return xx, yy


def peak_tool_tip(peak_type):
    """Default mouseover tooltip for peaks.

    :param peak_type: If 2, all time variables are in µs else in ns.
    :return: dictionary of tooltips. Can be converted to a list for bokeh.models.HoverTool.

    """

    tool_tip = dict()
    tool_tip["type"] = ("type", "@type")

    # Add non-time parameters (results in an ordered tooltip)
    tool_tip["amplitude"] = ("amplitude [pe/ns]", "$y")
    tool_tip["area"] = ("area [pe]", "@area")
    tool_tip["x"] = ("x [cm]", "@x")
    tool_tip["y"] = ("y [cm]", "@y")

    # Now ns/µs parameters for S1 and S2
    tool_tip["time_dynamic"] = ("time [ns]", "$x")
    tool_tip["dt"] = ("dt [ns/sample]", "@dt")
    tool_tip["range_50p_width"] = ("50% width [ns]", "@width_50")
    tool_tip["range_90p_width"] = ("90% width [ns]", "@width_90")
    tool_tip["rise_time"] = ("rise time [ns]", "@rise")
    tool_tip["rel_center_time"] = ("center time [ns]", "@rel_center_time")

    # Misc parameters
    tool_tip["nhits"] = ("nhits", "@nhits")
    tool_tip["aft"] = ("AFT", "@aft")

    # Add static time parameters:
    tool_tip["time_static"] = ("time [ns]", "@time")
    tool_tip["center_time"] = ("center_time [ns]", "@center_time")
    tool_tip["endtime"] = ("endtime [ns]", "@endtime")

    if peak_type == 2:
        for k, i in tool_tip.items():
            if k not in ["time_static", "center_time", "endtime"]:
                tool_tip[k] = (i[0].replace("[ns]", "[µs]"), i[1])

    return tool_tip


def default_fig(width=400, height=400, title="", **kwargs):
    """Helper function which returns a bokeh.plotting.figure instance with sizing_mode set to
    'scale_both' and an aspect ratio set according to the specified width and height.

    :param width: Plot width in pixels
    :param height: PLot height in pixels.
    :param title: Title of the plot. Also allows for additional kwargs accepted by bokeh.plotting.
    :return: bokeh.plotting.figure instance.

    """
    fig = bklt.figure(
        width=width,
        height=height,
        sizing_mode="scale_width",
        title=title,
        tools="pan,box_zoom,reset,save",
        **kwargs,
    )
    return fig
