import numpy as np
import holoviews as hv

hv.extension("bokeh")

def BCE(y, y_hat):
    return y_hat, -y * np.log(y_hat) - (1 - y) * np.log(1 - y_hat)


def wBCE(y, y_hat):
    return y_hat, -y * np.log(y_hat) + y_hat

y_hat = np.linspace(1e-9, 1-1e-09, 1000)
plots = 6

hv.Layout(
    [
        (
            hv.Curve(BCE(y, y_hat), label="BCE")
            * hv.Curve(wBCE(y, y_hat), label="wBCE")
            * hv.VLine(y).opts(color="k", line_dash="dashed")
        ).opts(
            xlabel="y_hat",
            ylabel="loss",
            width=200,
            show_grid=True,
            title=f"y={np.round(y, 3)}",
            xlim=(-0.1, 1.1),
            ylim=(0, 2),
        )
        for y in np.linspace(0, 1, plots)
    ]
).cols(plots)

def d_dyhat(f):
    def d_dyhat_f(y, y_hat):
        f_x, f_y = f(y, y_hat)
        return f_x, np.diff(f_y)/np.diff(f_x)
    return d_dyhat_f

hv.Layout(
    [
        (
            hv.Curve(d_dyhat(BCE)(y, y_hat), label="BCE")
            * hv.Curve(d_dyhat(wBCE)(y, y_hat), label="wBCE")
            * hv.VLine(y).opts(color="k", line_dash="dashed")
            * hv.HLine(0).opts(color="k", line_dash="dashed")
        ).opts(
            xlabel="y_hat",
            ylabel="gradient",
            width=200,
            show_grid=True,
            title=f"y={np.round(y, 3)}",
            xlim=(-0.1, 1.1),
            ylim=(-3, 3),
        )
        for y in np.linspace(0, 1, plots)
    ]
).cols(plots)