# "Production code"
import numpy as np
from interactive_pipe import Curve, SingleCurve

# -----------------
# Processing blocks
# -----------------


def gen_color(
    frequency: int = 0,  # discrete slider (int)
    isotropy: float = 0.,  # continuous slider (float)
    context={}  # global parameters (dict)
) -> np.ndarray:
    # context allows sharing information between processing blocks.
    # Let's share the frequency value for further use.
    context["freq"] = frequency
    print(context.get("initialized_param", "Wrongly initialized context!"))
    lin_coord = np.linspace(0, 1., 256)
    X, Y = np.meshgrid(lin_coord, isotropy*lin_coord)
    radius = 0.5+0.5*np.cos(frequency*np.sqrt(X**2 + Y**2))
    return np.stack([np.abs(X), np.abs(Y), radius], axis=-1).clip(0, 1)


def modify_geometry(
    img: np.ndarray,
    effect: str = "flip"  # dropdown menus (str)
) -> np.ndarray:
    img = img[::-1] if "flip" in effect else img
    img = img[:, ::-1] if "mirror" in effect else img
    return img


def change_color(
    img: np.ndarray,
    bnw: bool = True  # checkboxes (bool)
):
    if bnw:
        return np.mean(img, axis=-1, keepdims=True).repeat(3, axis=-1)
    return img


def compare_by_splitting(
    img_1: np.ndarray,
    img_2: np.ndarray,
    ratio: float = 0.5,  # continuous slider (float)
) -> np.ndarray:
    out = np.zeros_like(img_1)
    split = int(ratio*img_1.shape[1])
    out[:, :split] = img_2[:, :split]
    out[:, split+5:] = img_1[:, split+5:]
    return out


def extract_profile(img: np.ndarray, context={}) -> Curve:
    # context allows sharing information between processing blocks
    frequency = context.get('freq', 0)
    init_shared_param = context.get(
        'initialized_param', "Wrongly initialized context!")
    luma = img.mean(axis=-1)
    h_profile = SingleCurve(
        y=luma[0, :], label=f"H profile f={frequency}")
    v_profile = SingleCurve(y=luma[:, 0], label="V profile")
    return Curve(
        [h_profile, v_profile],
        xlabel="Position", ylabel="Luminance",
        grid=True,
        title=f"Luminance profiles\n(global context parameter {init_shared_param})"
    )


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


def tutorial_pipeline():
    inp = gen_color()
    out_geometry = modify_geometry(inp)
    out_bnw = change_color(inp)
    out_image = compare_by_splitting(out_geometry, out_bnw)
    out_profile = extract_profile(out_image)
    return [[inp, out_geometry], [out_profile, out_image]]

# -----------------------------------------------------------------------------------------
# Other layout you can try! None will create an empty space.
# return [[inp, out_geometry], [out_profile, out_image], [out_bnw, inp]]
# return [out_profile, out_image, out_bnw, inp]
# return out_profile, out_image, out_bnw, inp, None
# return [[inp, out_geometry], [out_profile, out_image], [out_bnw, None], [None, inp]]
# return inp
# return [[out_image], [inp]]