from interactive_pipe import interactive_pipeline, interactive
import argparse
from typing import Tuple
import numpy as np

# --- Define processing blocks ---


def gen_color(
    green: float = 0.8,
    decay: int = 1
) -> np.ndarray:  # 1 image output
    coord = np.linspace(0., 1., 256)
    X, Y = np.meshgrid(np.exp(-coord*decay), coord)
    return np.stack([X, green*Y, Y[::-1]], -1)


def flip_img(
    img: np.ndarray,  # 1 image input
    flip: bool = True
) -> Tuple[np.ndarray, np.ndarray]:  # 2 image outputs
    out = img[::-1] if flip else img.copy()
    return out, img[:, ::-1]


def amplify(
    img: np.ndarray,  # 1 image input
    amp: str = "L"
) -> np.ndarray:  # 1 image output
    return 1/{"S": 4, "M": 3, "L": 2, "XL": 1}[amp] * img

# --- Define pipeline = a sequence of processing blocks ---


def pipe() -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    inp = gen_color()
    out_flip, out_mirror = flip_img(inp)
    out = amplify(out_mirror)
    return [inp, out_flip, out]  # Display 3 images side by side


# --- Add interactivity ---

def add_controls() -> None:
    # Add controls to the processing blocks
    interactive(
        green=(1., [0., 1.], "Green amount"),  # float <-> continuous slider
        decay=(2, [0, 10], "Color decay")  # int <-> discrete slider
    )(gen_color)
    interactive(
        amp=("L", ["S", "M", "L", "XL"], "Amplify")  # str <-> dropdown menu
    )(amplify)
    interactive(
        flip=(True, "Flip")  # bool <-> checkbox
    )(flip_img)


def launch(backend="gradio", markdown_description="") -> None:
    # Decorate pipeline which becomes interactive
    playable_pipe = interactive_pipeline(
        gui=backend,
        markdown_description=markdown_description
    )(pipe)
    playable_pipe()  # Launch the interactive pipeline


# --------------------------------------------------
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-b", "--backend", default="gradio",
                        choices=["gradio", "qt", "mpl"], type=str)
    args = parser.parse_args()
    markdown_description = "# Code to build this app on gradio \n\n"
    markdown_description += "```python\n"+open(__file__, 'r').read()+"```"
    add_controls()
    launch(args.backend, markdown_description)