# Copyright (c) 2018-2020, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
#
"""The documentation for Layouts"""

from omni import ui
from omni.ui import color as cl
from .doc_page import DocPage
import inspect
import asyncio
import omni.kit.app

SPACING = 5


def collapsable_simple():
    with ui.CollapsableFrame("Header"):
        with ui.VStack(height=0):
            ui.Button("Hello World")
            ui.Button("Hello World")


def collapsable_custom(style_system):
    def custom_header(collapsed, title):
        with ui.HStack(style=style_system):
            with ui.ZStack(width=30):
                ui.Circle(name="title")
                with ui.HStack():
                    ui.Spacer()
                    align = ui.Alignment.V_CENTER
                    ui.Line(name="title", width=6, alignment=align)
                    ui.Spacer()
                if collapsed:
                    with ui.VStack():
                        ui.Spacer()
                        align = ui.Alignment.H_CENTER
                        ui.Line(name="title", height=6, alignment=align)
                        ui.Spacer()

            ui.Label(title, name="title")

    style = {
        "CollapsableFrame": {
            "background_color": cl.transparent,
            "secondary_color": cl.transparent,
            "border_radius": 0,
            "border_color": cl("#B2B2B2"),
            "border_width": 0.5,
        },
        "CollapsableFrame:hovered": {"secondary_color": cl("#E5F1FB")},
        "CollapsableFrame:pressed": {"secondary_color": cl("#CCE4F7")},
        "Label::title": {"color": cl.black},
        "Circle::title": {
            "color": cl.black,
            "background_color": cl.transparent,
            "border_color": cl("#B2B2B2"),
            "border_width": 0.75,
        },
        "Line::title": {"color": cl("#666666"), "border_width": 1},
    }

    style.update(style_system)
    with ui.CollapsableFrame("Header", build_header_fn=custom_header, style=style):
        with ui.VStack(height=0):
            ui.Button("Hello World")
            ui.Button("Hello World")


def collapsable_layout(style_system):
    style = {
        "CollapsableFrame": {
            "border_color": cl("#005B96"),
            "border_radius": 4,
            "border_width": 2,
            "padding": 0,
            "margin": 0,
        }
    }
    frame = ui.CollapsableFrame("Header", style=style)
    with frame:
        with ui.VStack(height=0):
            ui.Button("Hello World")
            ui.Button("Hello World")

    def set_style(field, model, style=style, frame=frame):
        frame_style = style["CollapsableFrame"]
        frame_style[field] = model.get_value_as_float()
        frame.set_style(style)

    with ui.HStack():
        ui.Label("Padding:", width=ui.Percent(10), name="text")
        model = ui.FloatSlider(min=0, max=50).model
    model.add_value_changed_fn(lambda m: set_style("padding", m))

    with ui.HStack():
        ui.Label("Margin:", width=ui.Percent(10), name="text")
        model = ui.FloatSlider(min=0, max=50).model
    model.add_value_changed_fn(lambda m: set_style("margin", m))


def direction():
    def rotate(dirs, stack, label):
        dirs[0] = (dirs[0] + 1) % len(dirs[1])
        stack.direction = dirs[1][dirs[0]]
        label.text = str(stack.direction)

    dirs = [
        0,
        [
            ui.Direction.LEFT_TO_RIGHT,
            ui.Direction.RIGHT_TO_LEFT,
            ui.Direction.TOP_TO_BOTTOM,
            ui.Direction.BOTTOM_TO_TOP,
        ],
    ]
    stack = ui.Stack(ui.Direction.LEFT_TO_RIGHT, width=0, height=0)
    with stack:
        for name in ["One", "Two", "Three", "Four"]:
            ui.Button(name)
    with ui.HStack():
        with ui.HStack():
            ui.Label("Current direcion is ", name="text", width=0)
            label = ui.Label("", name="text")

        button = ui.Button("Change")
        button.set_clicked_fn(lambda d=dirs, s=stack, l=label: rotate(d, s, l))
        rotate(dirs, stack, label)


def recreate(self):
    self._recreate_ui = ui.Frame(height=40)

    def changed(model, recreate_ui=self._recreate_ui):
        with recreate_ui:
            with ui.HStack():
                for i in range(model.get_value_as_int()):
                    ui.Button(f"Button #{i}")

    model = ui.IntDrag(min=0, max=10).model
    self._sub_recreate = model.subscribe_value_changed_fn(changed)


def clipping(self):
    self._clipping_frame = ui.Frame()
    with self._clipping_frame:
        ui.Button("Hello World")

    def set_width(model):
        self._clipping_frame.width = ui.Pixel(model.as_float)

    def set_height(model):
        self._clipping_frame.height = ui.Pixel(model.as_float)

    def set_hclipping(model):
        self._clipping_frame.horizontal_clipping = model.as_bool

    def set_vclipping(model):
        self._clipping_frame.horizontal_clipping = model.as_bool

    width = ui.FloatDrag(min=50, max=1000).model
    height = ui.FloatDrag(min=50, max=1000).model
    hclipping = ui.ToolButton(text="HClipping").model
    vclipping = ui.ToolButton(text="VClipping").model
    self._width = width.subscribe_value_changed_fn(set_width)
    self._height = height.subscribe_value_changed_fn(set_height)
    self._hclipping = hclipping.subscribe_value_changed_fn(set_hclipping)
    self._vclipping = vclipping.subscribe_value_changed_fn(set_vclipping)
    width.set_value(100)
    height.set_value(100)


def visibility():
    def invisible(button):
        button.visible = False

    def visible(buttons):
        for button in buttons:
            button.visible = True

    buttons = []
    with ui.HStack():
        for n in ["One", "Two", "Three", "Four", "Five"]:
            button = ui.Button(n, width=0)
            button.set_clicked_fn(lambda b=button: invisible(b))
            buttons.append(button)

        ui.Spacer()
        button = ui.Button("Visible all", width=0)
        button.set_clicked_fn(lambda b=buttons: visible(b))


def vgrid():
    with ui.ScrollingFrame(
        height=425,
        horizontal_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_ALWAYS_OFF,
        vertical_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_ALWAYS_ON,
    ):
        with ui.VGrid(column_width=100, row_height=100):
            for i in range(100):
                with ui.ZStack():
                    ui.Rectangle(
                        style={
                            "border_color": cl.black,
                            "background_color": cl.white,
                            "border_width": 1,
                            "margin": 0,
                        }
                    )
                    ui.Label(f"{i}", style={"margin": 5})


def hgrid():
    with ui.ScrollingFrame(
        height=425,
        horizontal_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_ALWAYS_ON,
        vertical_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_ALWAYS_OFF,
    ):
        with ui.HGrid(column_width=100, row_height=100):
            for i in range(100):
                with ui.ZStack():
                    ui.Rectangle(
                        style={
                            "border_color": cl.black,
                            "background_color": cl.white,
                            "border_width": 1,
                            "margin": 0,
                        }
                    )
                    ui.Label(f"{i}", style={"margin": 5})


def placer_track(self, id):
    # Initial size
    BEGIN = 50 + 100 * id
    END = 120 + 100 * id

    HANDLE_WIDTH = 10

    class EditScope:
        """The class to avoid circular event calling"""

        def __init__(self):
            self.active = False

        def __enter__(self):
            self.active = True

        def __exit__(self, type, value, traceback):
            self.active = False

        def __bool__(self):
            return not self.active

    class DoLater:
        """A helper to collect data and process it one frame later"""

        def __init__(self):
            self.__task = None
            self.__data = []

        def do(self, data):
            # Collect data
            self.__data.append(data)

            # Update in the next frame. We need it because we want to accumulate the affected prims
            if self.__task is None or self.__task.done():
                self.__task = asyncio.ensure_future(self.__delayed_do())

        async def __delayed_do(self):
            # Wait one frame
            await omni.kit.app.get_app().next_update_async()

            print(f"In the previous frame the user clicked the rectangles: {self.__data}")
            self.__data.clear()

    self.edit = EditScope()
    self.dolater = DoLater()

    def start_moved(start, body, end):
        if not self.edit:
            # Something already edits it
            return

        with self.edit:
            body.offset_x = start.offset_x
            rect.width = ui.Pixel(end.offset_x - start.offset_x + HANDLE_WIDTH)

    def body_moved(start, body, end, rect):
        if not self.edit:
            # Something already edits it
            return

        with self.edit:
            start.offset_x = body.offset_x
            end.offset_x = body.offset_x + rect.width.value - HANDLE_WIDTH

    def end_moved(start, body, end, rect):
        if not self.edit:
            # Something already edits it
            return

        with self.edit:
            body.offset_x = start.offset_x
            rect.width = ui.Pixel(end.offset_x - start.offset_x + HANDLE_WIDTH)

    with ui.ZStack(height=30):
        # Body
        body = ui.Placer(draggable=True, drag_axis=ui.Axis.X, offset_x=BEGIN)
        with body:
            rect = ui.Rectangle(width=END - BEGIN + HANDLE_WIDTH)
            rect.set_mouse_pressed_fn(lambda x, y, b, m, id=id: self.dolater.do(id))
        # Left handle
        start = ui.Placer(draggable=True, drag_axis=ui.Axis.X, offset_x=BEGIN)
        with start:
            ui.Rectangle(width=HANDLE_WIDTH, style={"background_color": cl("#FF660099")})
        # Right handle
        end = ui.Placer(draggable=True, drag_axis=ui.Axis.X, offset_x=END)
        with end:
            ui.Rectangle(width=HANDLE_WIDTH, style={"background_color": cl("#FF660099")})

    # Connect them together
    start.set_offset_x_changed_fn(lambda _, s=start, b=body, e=end: start_moved(s, b, e))
    body.set_offset_x_changed_fn(lambda _, s=start, b=body, e=end, r=rect: body_moved(s, b, e, r))
    end.set_offset_x_changed_fn(lambda _, s=start, b=body, e=end, r=rect: end_moved(s, b, e, r))


def placer_percents(self):
    # The size of the rectangle
    SIZE = 20.0

    with ui.ZStack(height=300):
        # Background
        ui.Rectangle(style={"background_color": cl("#999999")})

        # Small rectangle
        p = ui.Percent(50)
        placer = ui.Placer(draggable=True, offset_x=p, offset_y=p)
        with placer:
            ui.Rectangle(width=SIZE, height=SIZE)

    def clamp_x(offset):
        if offset.value < 0:
            placer.offset_x = ui.Percent(0)
        max_per = 100.0 - SIZE / placer.computed_width * 100.0
        if offset.value > max_per:
            placer.offset_x = ui.Percent(max_per)

    def clamp_y(offset):
        if offset.value < 0:
            placer.offset_y = ui.Percent(0)
        max_per = 100.0 - SIZE / placer.computed_height * 100.0
        if offset.value > max_per:
            placer.offset_y = ui.Percent(max_per)

    # Calbacks
    placer.set_offset_x_changed_fn(clamp_x)
    placer.set_offset_y_changed_fn(clamp_y)


class LayoutDoc(DocPage):
    """ document Layout classes"""

    def create_doc(self, navigate_to=None):
        self._section_title("Layout")
        self._text("We have three main components: VStack, HStack, and ZStack.")

        self._caption("HStack", navigate_to)
        self._text("This class is used to construct horizontal layout objects.")

        with ui.HStack():
            ui.Button("One")
            ui.Button("Two")
            ui.Button("Three")
            ui.Button("Four")
            ui.Button("Five")

        self._text("The simplest use of the class is like this:")
        self._code(
            """
        with ui.HStack():
            ui.Button("One")
            ui.Button("Two")
            ui.Button("Three")
            ui.Button("Four")
            ui.Button("Five")
            """
        )

        self._caption("VStack", navigate_to)
        self._text("The VStack class lines up widgets vertically.")

        with ui.VStack(width=100.0):
            with ui.VStack():
                ui.Button("One")
                ui.Button("Two")
                ui.Button("Three")
                ui.Button("Four")
                ui.Button("Five")

        self._text("The simplest use of the class is like this:")
        self._code(
            """
        with ui.VStack():
            ui.Button("One")
            ui.Button("Two")
            ui.Button("Three")
            ui.Button("Four")
            ui.Button("Five")
            """
        )

        self._caption("ZStack", navigate_to)
        self._text("A view that overlays its children, aligning them on top of each other.")

        with ui.VStack(width=100.0):
            with ui.ZStack():
                ui.Button("Very Long Text to See How Big it Can Be", height=0)
                ui.Button("Another\nMultiline\nButton", width=0)

        self._code(
            """
        with ui.ZStack():
            ui.Button("Very Long Text to See How Big it Can Be", height=0)
            ui.Button("Another\\nMultiline\\nButton", width=0)
            """
        )

        self._caption("Spacing", navigate_to)
        self._text("Spacing is a non-stretchable space in pixels between child items of the layout.")

        def set_spacing(stack, spacing):
            stack.spacing = spacing

        spacing_stack = ui.HStack(style={"margin": 0})
        with spacing_stack:
            for name in ["One", "Two", "Three", "Four"]:
                ui.Button(name)

        with ui.HStack(spacing=SPACING):
            with ui.HStack(width=100):
                ui.Spacer()
                ui.Label("spacing", width=0, name="text")
            with ui.HStack(width=ui.Percent(20)):
                field = ui.FloatField(width=50)
                slider = ui.FloatSlider(min=0, max=50, style={"color": cl.transparent})
                # Link them together
                slider.model = field.model
                slider.model.add_value_changed_fn(lambda m, s=spacing_stack: set_spacing(s, m.get_value_as_float()))

        self._code(
            """
        def set_spacing(stack, spacing):
            stack.spacing = spacing

        ui.Spacer(height=SPACING)
        spacing_stack = ui.HStack(style={"margin": 0})
        with spacing_stack:
            for name in ["One", "Two", "Three", "Four"]:
                ui.Button(name)

        ui.Spacer(height=SPACING)
        with ui.HStack(spacing=SPACING):
            with ui.HStack(width=100):
                ui.Spacer()
                ui.Label("spacing", width=0, name="text")
            with ui.HStack(width=ui.Percent(20)):
                field = ui.FloatField(width=50)
                slider = ui.FloatSlider(min=0, max=50, style={"color": cl.transparent})
                # Link them together
                slider.model = field.model
                slider.model.add_value_changed_fn(
                    lambda m, s=spacing_stack: set_spacing(s, m.get_value_as_float()))
            """
        )

        self._caption("VGrid", navigate_to)

        self._text(
            "Grid is a container that arranges its child views in a grid. The grid vertical/horizontal direction is "
            "the direction the grid size growing with creating more children."
        )

        self._text("It has two modes for cell width.")
        self._text(" - If the user set column_count, the column width is computed from the grid width.")
        self._text(" - If the user sets column_width, the column count is computed.")

        self._text("It also has two modes for height.")
        self._text(
            " - If the user sets row_height, VGrid uses it to set the height for all the cells. It's the fast mode "
            "because it's considered that the cell height never changes. VGrid easily predict which cells are visible."
        )
        self._text(
            " - If the user sets nothing, VGrid computes the size of the children. This mode is slower than the "
            "previous one, but the advantage is that all the rows can be different custom sizes. VGrid still draws "
            "only visible items, but to predict it, it uses cache, which can be big if VGrid has hundreds of "
            "thousands of items."
        )

        vgrid()
        self._code(inspect.getsource(vgrid).split("\n", 1)[-1])

        self._caption("HGrid", navigate_to)

        self._text("HGrid works exactly like VGrid, but with swapped width and height..")

        hgrid()
        self._code(inspect.getsource(hgrid).split("\n", 1)[-1])

        self._caption("Frame", navigate_to)

        self._text(
            "Frame is a container that can keep only one child. Each child added to Frame overrides the previous one. "
            "This feature is used for creating dynamic layouts. The whole layout can be easily recreated with a "
            "simple callback. In the following example, the buttons are recreated each time the slider changes."
        )

        recreate(self)
        self._code(inspect.getsource(recreate).split("\n", 1)[-1])

        self._text(
            "Another feature of Frame is the ability to clip its child. When the content of Frame is bigger than "
            "Frame itself, the exceeding part is not drawn if the clipping is on."
        )

        clipping(self)
        self._code(inspect.getsource(clipping).split("\n", 1)[-1])

        self._caption("CollapsableFrame", navigate_to)
        self._text(
            "CollapsableFrame is a frame widget that can hide or show its content. It has two states expanded and "
            "collapsed. When it's collapsed, it looks like a button. If it's expanded, it looks like a button and a "
            "frame with the content. It's handy to group properties, and temporarily hide them to get more space for "
            "something else."
        )
        collapsable_simple()
        self._code(inspect.getsource(collapsable_simple).split("\n", 1)[-1])

        self._text("It's possible to use a custom header.")

        collapsable_custom(self._style_system)
        self._code(inspect.getsource(collapsable_custom).split("\n", 1)[-1])

        self._text("The next example demonstrates how padding and margin work in the collapsable frame.")

        collapsable_layout(self._style_system)
        self._code(inspect.getsource(collapsable_layout).split("\n", 1)[-1])

        self._caption("Examples", navigate_to)
        with ui.VStack(style=self._style_system):
            for i in range(2):
                with ui.HStack():
                    ui.Spacer(width=50)

                    with ui.VStack(height=0):
                        ui.Button("Left {}".format(i), height=0)
                        ui.Button("Vertical {}".format(i), height=50)

                    with ui.HStack(width=ui.Fraction(2)):
                        ui.Button("Right {}".format(i))
                        ui.Button("Horizontal {}".format(i), width=ui.Fraction(2))

                    ui.Spacer(width=50)

        self._code(
            """
        with ui.VStack():
            for i in range(2):
                with ui.HStack():
                    ui.Spacer()

                    with ui.VStack(height=0):
                        ui.Button("Left {}".format(i), height=0)
                        ui.Button("Vertical {}".format(i), height=50)

                    with ui.HStack(width=ui.Fraction(2)):
                        ui.Button("Right {}".format(i))
                        ui.Button("Horizontal {}".format(i), width=ui.Fraction(2))

                    ui.Spacer()
            """
        )

        self._caption("Placer", navigate_to)
        self._text(
            "Enable you to place a widget presisely with offset. Placer's property `draggable` allows changing "
            "the position of the child widget with dragging it with the mouse."
        )

        with ui.ScrollingFrame(
            height=170,
            horizontal_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_ALWAYS_OFF,
            vertical_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_ALWAYS_OFF,
        ):
            with ui.ZStack():
                with ui.HStack():
                    for index in range(60):
                        ui.Line(width=10, style={"color": cl.black, "border_width": 0.5}, alignment=ui.Alignment.LEFT)
                with ui.VStack():
                    ui.Line(
                        height=10,
                        width=600,
                        style={"color": cl.black, "border_width": 0.5},
                        alignment=ui.Alignment.TOP,
                    )
                    for index in range(15):
                        ui.Line(
                            height=10,
                            width=600,
                            style={"color": cl.black, "border_width": 0.5},
                            alignment=ui.Alignment.TOP,
                        )
                    ui.Line(
                        height=10,
                        width=600,
                        style={"color": cl.black, "border_width": 0.5},
                        alignment=ui.Alignment.TOP,
                    )

                with ui.Placer(offset_x=100, offset_y=10):
                    ui.Button("moved 100px in X, and 10px in Y", width=0, height=20, name="placed")

                with ui.Placer(offset_x=200, offset_y=50):
                    ui.Button("moved 200px X , and 50 Y", width=0, height=0)

                def set_text(widget, text):
                    widget.text = text

                with ui.Placer(draggable=True, offset_x=300, offset_y=100):
                    ui.Circle(radius=50, width=50, height=50, size_policy=ui.CircleSizePolicy.STRETCH, name="placed")

                placer = ui.Placer(draggable=True, drag_axis=ui.Axis.Y, offset_x=400, offset_y=120)

                with placer:
                    with ui.ZStack(width=150, height=40):
                        ui.Rectangle(name="placed")
                        with ui.HStack(spacing=5):
                            ui.Circle(
                                radius=3,
                                width=15,
                                size_policy=ui.CircleSizePolicy.FIXED,
                                style={"background_color": cl.white},
                            )
                            ui.Label("UP / Down", style={"color": cl.white, "font_size": 16.0})
                            offset_label = ui.Label("120", style={"color": cl.white})

                placer.set_offset_y_changed_fn(lambda o: set_text(offset_label, str(o)))

        self._code(
            """
with ui.Placer(
    draggable=True, offset_x=300, offset_y=100):
    ui.Circle(
        radius=50,
        width=50,
        height=50,
        size_policy=ui.CircleSizePolicy.STRETCH,
        name="placed"
    )

placer = ui.Placer(
    draggable=True,
    drag_axis=ui.Axis.Y,
    offset_x=400,
    offset_y=120)

with placer:
    with ui.ZStack(
        width=150,
        height=40,
    ):
        ui.Rectangle(name="placed")
        with ui.HStack(spacing=5):
            ui.Circle(
                radius=3,
                width=15,
                size_policy=ui.CircleSizePolicy.FIXED,
                style={"background_color": cl.white},
            )
            ui.Label(
                "UP / Down",
                style={
                    "color": cl.white,
                    "font_size": 16.0
                }
            )
            offset_label = ui.Label(
                "120", style={"color": cl.white})

placer.set_offset_y_changed_fn(
    lambda o: set_text(offset_label, str(o)))
            """
        )

        self._text(
            "The following example shows the way to interact between three Placers to create a resizable rectangle. "
            "The rectangle can be moved on X-axis and can be resized with small side handles."
        )
        self._text(
            "When multiple widgets fire the callbacks simultaneously, it's possible to collect the event data and "
            "process them one frame later using asyncio."
        )
        with ui.ZStack():
            placer_track(self, 0)
            placer_track(self, 1)

        self._code(inspect.getsource(placer_track).split("\n", 1)[-1])

        self._text(
            "It's possible to set `offset_x` and `offset_y` in percents. It allows stacking the children to the "
            "proportions of the parent widget. If the parent size is changed, then the offset is updated accordingly."
        )

        placer_percents(self)
        self._code(inspect.getsource(placer_percents).split("\n", 1)[-1])

        self._caption("Visibility", navigate_to)
        self._text(
            "This property holds whether the widget is visible. Invisible widget is not rendered, and it doesn't take "
            "part in the layout. The layout skips it."
        )
        self._text("In the following example click the button to hide it.")
        visibility()
        self._code(inspect.getsource(visibility).split("\n", 1)[-1])

        self._caption("Direction", navigate_to)
        self._text(
            "It's possible to determine the direction of a stack with the property direction. The stack is able to "
            "change its direction dynamically."
        )
        direction()
        self._code(inspect.getsource(direction).split("\n", 1)[-1])
        ui.Spacer(height=50)
