import numpy as np

from bokeh.plotting import figure
from bokeh.models import ColumnDataSource, Band, Whisker
from bokeh.palettes import Category10_10

from buki import Object


BLUE = Category10_10[0]
RED  = Category10_10[3]


class Plot1D(Object):

    def __init__(self, name=None):
        data = {
            "x": [],
            "y": [],
#            "y-std": [],
#            "y+std": [],
            "avg": [],
#            "avg-err": [],
#            "avg+err": [],
        }

        self.source = source = ColumnDataSource(data=data)

        fig = figure(name=name)

        line1, circle1, errband1, errbars1 = add_curve(fig, source, "x", "y", "y-std", "y+std", BLUE)
        line2, circle2, errband2, errbars2 = add_curve(fig, source, "x", "avg", "avg-err", "avg+err", RED)

#        line1.visible = line2.visible = False
#        circle1.visible = circle2.visible = False
#        errband1.visible = errband2.visible = False
#        errbars1.visible = errbars2.visible = False

        super().__init__(fig)


    def set(self, times, values):
        y = values[-1]

        if y is None: #TODO: is this needed?
            return

        n = len(values)
        values = np.asarray(values)

        x = np.arange(len(y))
        std = values.std(axis=0)

        #TODO: skip avg for n=1 ?
        avg = values.mean(axis=0)
        err = std / np.sqrt(n)

        data = {
            "x": x,
            "y": y,
            "y-std": y - std,
            "y+std": y + std,
            "avg": avg,
            "avg-err": avg - err,
            "avg+err": avg + err,
        }

        self.source.data.update(data)



def add_curve(fig, source, x, y, lower, upper, color):
    line = fig.line(x=x, y=y, source=source, color=color)
    circle = fig.circle(x=x, y=y, source=source, color=color)

    errband = Band(base=x, lower=lower, upper=upper, source=source, level="underlay", line_width=1, line_color=color, fill_color=color, fill_alpha=0.5)

    errbars = Whisker(base=x, lower=lower, upper=upper, source=source, line_color=color)
    errbars.lower_head.line_color = color
    errbars.upper_head.line_color = color

    fig.add_layout(errband)
    fig.add_layout(errbars)

    return line, circle, errband, errbars