# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/interactive_show.ipynb.

# %% auto 0
__all__ = ['plot_image', 'plot_graph', 'tonp', 'tolist', 'convert_to_nx', 'viz2', 'df2graph_nodes', 'ishow']

# %% ../nbs/interactive_show.ipynb 1
from . import *
from .bokeh_loader import bshow
from bokeh.io import output_notebook, show as bokeh_show, output_file
from bokeh.plotting import figure, from_networkx
from bokeh.models import (
    Circle,
    Rect,
    WheelZoomTool,
    PanTool,
    BoxZoomTool,
    ResetTool,
    MultiLine,
    NodesAndLinkedEdges,
    EdgesAndLinkedNodes,
    HoverTool,
    TapTool,
    BoxSelectTool,
)
from bokeh.palettes import Spectral7
import networkx as nx
from torch_geometric.utils import to_networkx
import torch
import numpy as np
from fastcore.basics import ifnone

output_notebook()


# %% ../nbs/interactive_show.ipynb 2
def plot_image(p, image, sz):
    if isinstance(image, str):
        image = read(image, 1)
    imga = resize(image, sz)
    h, w = imga.shape[:2]
    img = np.empty((h, w), dtype=np.uint32)
    view = img.view(dtype=np.uint8).reshape((h, w, 4))
    for i in range(h):
        for j in range(w):
            view[i, j, :3] = imga[h - i - 1, j]
            view[i, j, 3] = 255

    p.x_range.range_padding = p.y_range.range_padding = 0
    p.image_rgba(image=[img], x=0, y=0, dw=1, dh=h / w)


def plot_graph(g, output, im=None, **kwargs):
    plot = figure(
        title=kwargs.get("title", "Networkx Integration Demonstration"),
        tools=[WheelZoomTool(), PanTool(), BoxZoomTool(), ResetTool()],
        plot_width=900,
        plot_height=900,
        match_aspect=True,
    )
    h, w = im.shape[:2]
    plot_image(plot, im, sz=kwargs.get("sz", 0.5))
    fill_color = {
        ix: "green" if correct else "red"
        for ix, correct in nx.get_node_attributes(g, "correct").items()
    }
    nx.set_node_attributes(g, fill_color, "node_color")

    plot.add_tools(
        HoverTool(
            tooltips=[("index", "$index")]
            + [(i, f"@{i}") for i in kwargs.get("tooltips", [])]
        ),
        TapTool(),
        BoxSelectTool(),
    )
    plot.toolbar.active_scroll = plot.select_one(WheelZoomTool)
    plot.toolbar.active_drag = plot.select_one(PanTool)

    if kwargs.get("pos"):
        pos = nx.get_node_attributes(g, kwargs.get("pos"))
        pos = {i: (x, y * h / w) for i, (x, y) in pos.items()}
        graph = from_networkx(g, pos)
    else:
        print("Using random")
        # graph = from_networkx(g, nx.layout.fruchterman_reingold_layout, scale=20, center=(0,0))
        graph = from_networkx(g, nx.random_layout)

    # graph.node_renderer.glyph = Circle(
    #     size=kwargs.get("size", 10),
    #     # fill_color=kwargs.get("color", "green"),
    #     fill_color="node_color",
    #     fill_alpha=kwargs.get("opacity", 0.9),
    # )
    graph.node_renderer.glyph = Rect(
        width="w",
        height="h",
        # fill_color=kwargs.get("color", "green"),
        fill_color="node_color",
        fill_alpha=kwargs.get("opacity", 0.3),
    )
    # graph.node_renderer.selection_glyph = Circle(
    #     size=15, fill_color=kwargs.get("color", "gray")
    # )
    # graph.node_renderer.hover_glyph = Circle(
    #     size=15, fill_color=kwargs.get("color", "green")
    # )

    # graph.edge_renderer.glyph = MultiLine(
    #     line_color="#CCCCCC", line_alpha=0.8, line_width=1
    # )
    # graph.edge_renderer.selection_glyph = MultiLine(line_color="blue", line_width=2)
    # graph.edge_renderer.hover_glyph = MultiLine(line_color="blue", line_width=2)
    # graph.selection_policy = NodesAndLinkedEdges()
    # graph.inspection_policy = EdgesAndLinkedNodes()
    plot.renderers.append(graph)
    # output_notebook(f"{output}.html")
    bshow(plot)


def tonp(i):
    if isinstance(i, torch.Tensor):
        i = i.cpu().detach().numpy()
        return i
    if isinstance(i, list):
        return np.array(i)


def tolist(i):
    if isinstance(i, torch.Tensor):
        i = tonp(i)
    if isinstance(i, np.ndarray):
        i = i.tolist()
    return i


def convert_to_nx(g, node_attrs=None, undirected=True):
    graph = to_networkx(g, to_undirected=undirected)
    if node_attrs:
        for attr in node_attrs:
            _attr = tolist(g[attr])
            nx.set_node_attributes(
                graph, {i: _attr[i] for i in range(graph.number_of_nodes())}, attr
            )
    return graph


def viz2(graph, node_attrs=None, undirected=True, **kwargs):
    if not isinstance(graph, nx.classes.graph.Graph):
        graph = convert_to_nx(graph, node_attrs=node_attrs, undirected=undirected)
    plot_graph(graph, "networkx_graph", **kwargs)


def df2graph_nodes(
    df, text_attr="text", additional_attrs=None,
):
    df = df.copy()
    bbs = df2bbs(df)
    pos = [(bb.xc, 1 - bb.yc) for bb in bbs]
    w = [bb.w for bb in bbs]
    h = [bb.h for bb in bbs]
    if additional_attrs:
        df = df[additional_attrs]
    df["pos"] = pos
    df["w"] = w
    df["h"] = h
    nodes = [(ix, node) for ix, node in enumerate(df.to_dict(orient="records"))]
    return nodes


def ishow(im, df, additional_attrs=None, **kwargs):
    df = df.copy()
    if isinstance(im, (str, P)):
        im = read(im, 1)
    elif isinstance(im, PIL.Image.Image):
        im = np.array(im)
    h, w = im.shape[:2]
    df = df.reset_index(drop=True)

    G = nx.Graph()

    if "conf" not in df.columns:
        df["conf"] = 0
    if "conf_of_correct_class" not in df.columns:
        df["conf_of_correct_class"] = 0

    truth_col = kwargs.pop("truth_col", "truth")
    df["truth"] = None if truth_col not in df.columns else df[truth_col]

    pred_col = kwargs.pop("pred_col", "pred")
    df["pred"] = None if pred_col not in df.columns else df[pred_col]

    df["text"] = None if "text" not in df.columns else df["text"]

    df["correct"] = df["pred"] == df["truth"]
    if "x" not in df.columns:
        x, y, X, Y = list(zip(*df2bbs(df)))
        df["x"] = x
        df["X"] = X
        df["y"] = y
        df["Y"] = Y
    if df["X"].max() > 1:
        df["x"] = df["x"] / w
        df["y"] = df["y"] / h
        df["X"] = df["X"] / w
        df["Y"] = df["Y"] / h
        if "bb" in df.columns:
            df.drop("bb", axis=1, inplace=True)
    additional_attrs = [
        "truth",
        "pred",
        "conf",
        "conf_of_correct_class",
        "correct",
        "text",
    ] + ifnone(additional_attrs, [])
    G.add_nodes_from(df2graph_nodes(df, additional_attrs=additional_attrs))
    viz2(G, im=im, tooltips=additional_attrs, pos="pos")

