from interactive_pipe import interactive_pipeline, interactive, Image, TextPrompt
import cv2
from typing import Tuple
import numpy as np


# Processing blocks
# -----------------
def processing_block(
    inp: np.ndarray,
    amplify: float = 1.,
    context: dict = {}
) -> np.ndarray:
    context["amplify"] = amplify
    return amplify*inp


def draw_text(img: np.ndarray, free_text: str = "Hello", preset_text: str = "choice") -> np.ndarray:
    for txt, pos in [(free_text, (10, 50)), (preset_text, (10, 100))]:
        out = cv2.putText(
            img,
            txt,
            pos,
            cv2.FONT_HERSHEY_SIMPLEX,
            1,
            (255, 255, 255), 2
        )
    return out


def add_extra_caption(img: np.ndarray, caption: str = "Extra caption") -> np.ndarray:
    out = cv2.putText(
        img,
        caption,
        (10, 200),
        cv2.FONT_HERSHEY_SIMPLEX,
        1,
        (255, 0, 255), 2
    )
    return out


def move_circle(img: np.ndarray, x: int = 100, y: int = 100) -> np.ndarray:
    out = img.copy()
    cv2.circle(out, (x, y), 10, (255, 0, 0), -1)
    return out


def remove_caps(txt_prefix: str, txt_in: str = "TUTU") -> str:
    return txt_prefix + "\n" + txt_in.lower()


# Pipeline definition
# -------------------


def pipe(inp: np.ndarray, txt_prefix) -> Tuple[np.ndarray, np.ndarray]:
    out = add_extra_caption(inp)
    out = processing_block(out)
    out = draw_text(out)
    out = move_circle(out)
    txt_out = remove_caps(txt_prefix)
    return inp, out, txt_out

# Add interactivity to the processing blocks
# ------------------------------------------


def add_controls() -> None:

    # 2 ways to add prompts:
    # 1 = free_text=TextPrompt("DEFAULT TEXT", name="Caption custom")
    # 2 = free_text=("DEFAULT TEXT", None, "Write Free text")
    interactive(
        caption=TextPrompt("Created with interactive pipe",
                           name="Caption custom")
    )(add_extra_caption)
    interactive(
        free_text=("Hello world!", None, "Write Free text"),
        preset_text=("Hello", ["Hello", "World", "Goodbye"])
    )(draw_text)
    # Keyboards shortcuts: up, down, left, right -> only supported when using the Qt/MPL backend
    # Replaced by sliders when using the Gradio backend
    interactive(
        x=(100, [0, 200], "X", ["left", "right"]),
        y=(100, [0, 200], "Y", ["up", "down"])
    )(move_circle)
    interactive(
        amplify=(0.5, [0., 1.], "Amplify")
    )(processing_block)

    interactive(
        txt_in=TextPrompt("OUT TEXT",
                          name="Out text no caps")
    )(remove_caps)
# Add interactivity to the processing blocks
# ------------------------------------------


def launch(img: np.ndarray, backend: str = "gradio", markdown_description: str = "") -> None:
    add_controls()
    pipe_interactive = interactive_pipeline(
        gui=backend,
        cache=True,
        markdown_description=markdown_description
    )(pipe)
    pipe_interactive(img, "REMOVE CAPS! ")


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("-b", "--backend", default="gradio",
                        choices=["gradio", "qt", "mpl"], type=str)
    parser.add_argument(
        "-d", "--debug", action="store_true",
        help="Debug mode (to tune difficulty and tolerance)"
    )
    args = parser.parse_args()
    markdown_description = "# Code to build this app on gradio \n\n"
    markdown_description += "```python\n"+open(__file__, 'r').read()+"```"
    img = Image.load_image("sample.jpg")
    launch(img, backend=args.backend, markdown_description=markdown_description)