# AUTOGENERATED! DO NOT EDIT! File to edit: source/91_bokeh.ipynb (unless otherwise specified).


from __future__ import annotations


__all__ = ['ALPHA', 'RED', 'GREEN', 'BLUE', 'C0', 'C1', 'C2', 'C3', 'C4', 'C5', 'C6', 'C7', 'C8', 'C9', 'get_lyp_path',
           'LayerProperty', 'LayerProperties', 'read_lyp', 'get_lyp', 'new_plot', 'adjust_plot', 'draw_polys',
           'draw_poly', 'draw_path', 'draw_point', 'draw_vector', 'draw_box', 'draw_cell', 'draw_inst', 'draw_layout']

# Internal Cell
#nbdev_comment from __future__ import annotations

import glob
import os
from functools import lru_cache
from typing import Dict, List, Optional, Tuple, Union

import bokeh.io as bio
import bokeh.models as bm
import bokeh.plotting as bp
import numpy as np
import pkg_resources
import pya
from lxml import etree

# Cell
ALPHA = 0.3
RED = "#FF0000"
GREEN = "#00FF00"
BLUE = "#0000FF"
C0 = '#1f77b4'
C1 = '#ff7f0e'
C2 = '#2ca02c'
C3 = '#d62728'
C4 = '#9467bd'
C5 = '#8c564b'
C6 = '#e377c2'
C7 = '#7f7f7f'
C8 = '#bcbd22'
C9 = '#17becf'

# Cell
LayerProperty = Dict[str, Union[str, bool]]
LayerProperties = Dict[Tuple[int, int], LayerProperty]

def get_lyp_path(path: Optional[str] = None):
    # first, let's try "~/.klayout"
    if path is None:
        path = ""
    path = os.path.abspath(os.path.expanduser(path))
    if os.path.isdir(path):
        possible_paths = glob.glob(f"{path}/*.lyp")
        if not possible_paths:
            path = get_lyp_path(pkg_resources.resource_filename("flayout", "layers.lyp"))
        else:
            path = possible_paths[0]
    return path

# Cell
@lru_cache
def read_lyp(path: Optional[str] = None) -> LayerProperties:
    """Load layer properties from a file

    Args:
        path: the path where to load the layer properties from

    Returns:
        a dictionary of layer property dictionaries
    """
    path = get_lyp_path(path)
    xml = etree.parse(path)
    root = xml.getroot()
    parsed: LayerProperties = {
        (0, 0): {
            "name": "",
            "frame-color": "#000000",
            "fill-color": "#000000",
            "visible": True,
        }
    }
    for properties in root.iter("properties"):
        name = properties.find("name")
        if name is not None:
            name, *_ = name.text.split("(")
            name = name.strip()
        else:
            name = ""

        layerstr = properties.find("source")
        if layerstr is not None:
            layerstr, *_ = layerstr.text.split("@")
            lr, dt = layerstr.strip().split("/")
            lr, dt = int(lr), int(dt)
            parsed[lr, dt] = {
                "name": name,
                "frame-color": properties.find("frame-color").text,
                "fill-color": properties.find("fill-color").text,
                "visible": bool(properties.find("visible").text),
            }

    return parsed

# Cell
@lru_cache
def get_lyp(
    layer_info: Union[pya.LayerInfo, Tuple[int, int]], path: Optional[str] = None
) -> LayerProperty:
    """Load layer properties for a specific layer from a file

    Args:
        layer_info: the layer info tuple to load the layer properties for
        path: the path where to load the layer properties from

    Returns:
        a layer property dictionary
    """
    if isinstance(layer_info, pya.LayerInfo):
        layer, datatype = layer_info.layer, layer_info.datatype # type: ignore
    else:
        layer, datatype, *_ = layer_info
    lyps = read_lyp(path=path)
    lyp = lyps.get((layer, datatype), lyps[0, 0])
    return lyp

# Internal Cell
def _get_range(box: Union[pya.Box, pya.DBox]) -> Tuple[float, float, float, float, float, float]:
    """Get the plotting bbox for a klayout box

    Args:
        poly: the polygon to create the bbox for

    returns:
        x1, x2, y1, y2, w, h
    """
    if isinstance(box, pya.Box):
        box = box.to_dtype()
    x1, x2 = min(box.p1.x, box.p2.x), max(box.p1.x, box.p2.x)
    y1, y2 = min(box.p1.y, box.p2.y), max(box.p1.y, box.p2.y)
    w, h = x2 - x1, y2 - y1
    if w > 5 * h:
        y, h = 0.5 * (y1 + y2), w / 5
        y1, y2 = y - h / 2, y + h / 2
    if h > 5 * w:
        x, w = 0.5 * (x1 + x2), h / 5
        x1, x2 = x - w / 2, x + w / 2
    if w < 2.0:
        x1, x2 = (x2 + x1) / 2 - 1.0, (x2 + x1) / 2 + 1.0
        w = 2.0
    if h < 2.0:
        y1, y2 = (y2 + y1) / 2 - 1.0, (y2 + y1) / 2 + 1.0
        h = 2.0
    return x1, x2, y1, y2, w, h

# Cell
def new_plot(box: Union[pya.Box, pya.DBox], max_dim: Optional[float] = None) -> bp.Figure:
    """Create a new plot with limits determined by a bbox

    Args:
        box: the bbox of the polygon or cell to create the figure for

    Returns:
        a bokeh Figure.
    """
    if max_dim is None:
        max_dim = 500
    x0, x1, y0, y1, w, h = _get_range(box)
    if w > h:
        plot_width = max_dim
        plot_height = max_dim * h / w
    else:
        plot_height = max_dim
        plot_width = max_dim * w / h
    plot = bp.figure(
        plot_width=round(plot_width),
        plot_height=round(plot_height),
        x_range=(x0, x1),
        y_range=(y0, y1),
        match_aspect=True,
        toolbar_location=None,
        tools=[bm.PanTool(), zoom := bm.WheelZoomTool()],
    )
    plot.toolbar.active_scroll = zoom  # type: ignore
    return plot

# Cell
def adjust_plot(plot: bp.Figure, box: Union[pya.Box, pya.DBox], max_dim: Optional[float] = None) -> bp.Figure:
    """Adjust a plot with limits determined by a bbox

    Args:
        plot: the plot to adjust the limits for
        box: the bbox of the polygon or cell to create the figure for

    Returns:
        a bokeh Figure.
    """
    if plot is None:
        return new_plot(box, max_dim=max_dim)
    if max_dim is None:
        max_dim = max(plot.height, plot.width)
    assert max_dim is not None
    x0, x1, y0, y1, w, h = _get_range(box)
    if w > h:
        plot_width = max_dim
        plot_height = max_dim * h / w
    else:
        plot_height = max_dim
        plot_width = max_dim * w / h
    plot.plot_width = round(plot_width)
    plot.plot_height = round(plot_height)
    plot.x_range = bm.Range1d(start=x0, end=x1)
    plot.y_range = bm.Range1d(start=y0, end=y1)
    return plot

# Cell
def draw_polys(
    plot: bp.Figure,
    polys: List[Union[pya.Polygon, pya.SimplePolygon, pya.DPolygon, pya.DSimplePolygon]],
    layer=(0, 0),
    fill_color=None,
    line_color=None,
    fill_alpha=ALPHA,
):
    """draw polygons with bokeh

    Args:
        plot: the plot to draw the polygon in
        polys: the polygons to draw

    Returns:
        the (inplace) modified plot containing the polygons
    """
    # some typing definitions for our own sanity...

    # Array of single coordinates (x OR y): ndim==1, shape=(N,).
    SimplePolygon = np.ndarray

    # List of coordinate arrays (first element: hull, other elements: holes)
    Polygon = List[SimplePolygon]

    # List of individual polygons (bokeh will XOR those polygons -> usually 1 Polygon per MultiPolygon)
    MultiPolygon = List[Polygon]

    # List of multi polygons
    MultiPolygons = List[MultiPolygon]

    xs: MultiPolygons = []
    ys: MultiPolygons = []

    for poly in polys:
        if isinstance(poly, (pya.Polygon, pya.SimplePolygon)):
            poly = poly.to_dtype()

        if isinstance(poly, pya.DPolygon):
            hull = np.asarray([(p.x, p.y) for p in poly.each_point_hull()])
            holes = [
                np.asarray([(p.x, p.y) for p in poly.each_point_hole(i)])
                for i in range(poly.holes())
            ]
        elif isinstance(poly, pya.DSimplePolygon):
            hull = np.asarray([(p.x, p.y) for p in poly.each_point()])
            holes = []
        else:
            raise ValueError(
                f"Invalid polygon type. Got: {type(poly)}. "
                f"Expected 'DPolygon' or 'DSimplePolygon'"
            )

        if hull.shape[0] < 3:
            continue

        plot = adjust_plot(
            plot, pya.Polygon([pya.DPoint(x, y) for x, y in hull]).bbox()
        )

        xs_: MultiPolygon = [[hull[:, 0], *(h[:, 0] for h in holes)]]
        ys_: MultiPolygon = [[hull[:, 1], *(h[:, 1] for h in holes)]]
        xs.append(xs_)
        ys.append(ys_)

    source = bm.ColumnDataSource({"xs": xs, "ys": ys})
    lyp = get_lyp(layer)
    patch = bm.MultiPolygons(
        xs="xs",
        ys="ys",
        fill_color=fill_color or lyp["fill-color"],
        fill_alpha=fill_alpha,
        line_color=line_color or lyp["frame-color"],
    )
    plot.add_glyph(source, patch)
    return plot

# Cell
def draw_poly(
    plot: bp.Figure,
    poly: Union[pya.Polygon, pya.SimplePolygon, pya.DPolygon, pya.DSimplePolygon],
    layer=(0, 0),
    fill_color=None,
    line_color=None,
    fill_alpha=ALPHA,
):
    """draw a polygon with bokeh

    Args:
        plot: the plot to draw the polygon in
        poly: the polygon to draw

    Returns:
        the (inplace) modified plot containing the polygon
    """
    plot = adjust_plot(plot, poly.bbox())
    return draw_polys(
        plot,
        [poly],
        layer=layer,
        fill_color=fill_color,
        line_color=line_color,
        fill_alpha=fill_alpha,
    )

# Cell
def draw_path(
    plot: bp.Figure,
    path: Union[pya.Path, pya.DPath],
    layer=(0, 0),
    fill_color=None,
    line_color=None,
    fill_alpha=ALPHA,
):
    """draw a path with bokeh

    Args:
        plot: the plot to draw the path in
        poly: the path to draw

    Returns:
        the (inplace) modified plot containing the path
    """
    return draw_polys(
        plot,
        [path.polygon()],
        layer=layer,
        fill_color=fill_color,
        line_color=line_color,
        fill_alpha=fill_alpha,
    )

# Cell
def draw_point(
    plot: bp.Figure,
    p: Union[pya.Point, pya.DPoint],
    layer=(0, 0),
    fill_color=None,
    line_color=None,
    fill_alpha=ALPHA,
):
    """draw a point with bokeh

    Args:
        plot: the plot to draw the point in
        p: the point to draw

    Returns:
        the (inplace) modified plot containing the point
    """
    if isinstance(p, pya.Point):
        p = p.to_dtype()
    v = pya.DVector(1.0, 1.0)
    box = pya.DBox(p - v, p + v)
    plot = adjust_plot(plot, box)
    *_, w, h = _get_range(box)
    radius = max(w, h) / 30
    lyp = get_lyp(layer)
    plot.circle(
        p.x,
        p.y,
        fill_alpha=fill_alpha,
        fill_color=fill_color or lyp["fill-color"],
        line_color=line_color or ["frame-color"],
        radius=radius,
    )
    return plot

# Cell
def draw_vector(
    plot: bp.Figure,
    v: Union[pya.Vector, pya.DVector],
    layer=(0, 0),
    fill_color=None,
    line_color=None,
    fill_alpha=ALPHA,
):
    """draw a vector as an arrow with bokeh

    Args:
        plot: the plot to draw the vector in
        v: the vector to draw

    Returns:
        the (inplace) modified plot containing the vector
    """
    if isinstance(v, pya.Vector):
        v = pya.DVector(v.x / 1000.0, v.y / 1000.0)
    box = pya.DBox(0, 0, v.x, v.y)
    plot = adjust_plot(plot, box, max_dim=250)
    lyp = get_lyp(layer)
    arrow_head = bm.VeeHead(
        fill_alpha=fill_alpha,
        fill_color=fill_color or lyp["fill-color"],
        line_color=line_color or lyp["frame-color"],
    )
    arrow = bm.Arrow(
        end=arrow_head,
        x_start=0,
        y_start=0,
        x_end=v.x,
        y_end=v.y,
        line_color=line_color or lyp["frame-color"],
    )
    plot.add_layout(arrow)
    return plot

# Cell
def _box_to_poly(box: Union[pya.Box, pya.DBox]):
    """convert a box into a polygon

    Args:
        box: the box to convert into a polygon

    Returns:
        the polygon
    """
    if isinstance(box, pya.Box):
        box = box.to_dtype()
    x0, y0 = box.p1.x, box.p1.y
    x1, y1 = box.p2.x, box.p2.y
    return pya.DPolygon(
        [pya.DPoint(x0, y0), pya.DPoint(x1, y0), pya.DPoint(x1, y1), pya.DPoint(x0, y1)]
    )


def draw_box(
    plot: bp.Figure,
    box: Union[pya.Box, pya.DBox],
    fill_alpha=0.0,
    fill_color="#000000",
    line_color=GREEN,
):
    """draw a box with bokeh

    Args:
        plot: the plot to draw the box in
        box: the box to draw

    Returns:
        the (inplace) modified plot containing the box
    """
    plot = adjust_plot(plot, box)
    poly = _box_to_poly(box)
    return draw_poly(
        plot, poly, fill_alpha=fill_alpha, fill_color=fill_color, line_color=line_color
    )

# Cell
def _draw_shapes(plot, shapes, layer=(0, 0)):
    """draw shapes with bokeh

    Args:
        plot: the plot to draw the shape in
        shapes: the shapes to draw

    Returns:
        the (inplace) modified plot containing the shape
    """
    polys = []
    for shape in shapes:
        if shape.is_box():
            polys.append(_box_to_poly(shape.dbbox()))
        elif shape.is_path():
            polys.append(shape.dpath.polygon())
        elif shape.is_polygon():
            polys.append(shape.dpolygon)
        elif shape.is_simple_polygon():
            polys.append(shape.dsimple_polygon)
    #for poly in polys:
    #    plot = draw_poly(plot, poly, layer=layer)
    return draw_polys(plot, polys, layer=layer)

def draw_cell(plot, cell, draw_bbox=True):
    """draw a cell with bokeh

    Args:
        plot: the plot to draw the cell in
        cell: the cell to draw

    Returns:
        the (inplace) modified plot containing the cell
    """
    layout = pya.Layout()
    new_cell = layout.create_cell(cell.name)
    new_cell.copy_tree(cell)

    cell = new_cell
    cell = cell.flatten(-1, True)
    box = cell.dbbox()
    plot = adjust_plot(plot, box, max_dim=500)
    for lr in layout.layer_infos():
        shapes = [*cell.shapes(layout.layer(lr)).each()]
        plot = _draw_shapes(plot, shapes, layer=(lr.layer, lr.datatype))
    if draw_bbox:
        draw_box(plot, box)
    return plot

# Cell
def draw_inst(plot, inst, draw_bbox=True, draw_arrow=True):
    """draw a instance with bokeh

    Args:
        plot: the plot to draw the instance in
        inst: the instance to draw

    Returns:
        the (inplace) modified plot with containing the instance
    """
    _layout = pya.Layout()
    _cell = _layout.create_cell(inst.cell.name)
    _cell.copy_tree(inst.cell)
    _refcell = _layout.create_cell(f"ref_{inst.cell.name}")
    _refcell.insert(pya.CellInstArray(_layout.cell_by_name(inst.cell.name), inst.trans))

    plot = new_plot(_refcell.dbbox() + pya.Point(0, 0))
    plot = draw_vector(plot, inst.trans.disp)
    plot = draw_cell(plot, _refcell)
    plot = adjust_plot(plot, _refcell.dbbox() + pya.Point(0, 0))
    return plot

# Cell
def draw_layout(plot, layout):
    """draw a layout with bokeh

    Args:
        plot: the plot to draw the layout in
        layout: the layout to draw

    Returns:
        the (inplace) modified plot with containing the layout
    """
    plots = bp.Column(*[draw_cell(plot, cell) for cell in layout.top_cells()])
    return plots