import gradio as gr
import numpy as np
from scipy.ndimage import convolve1d


def flip_lightness(X):
    print(f"In flip_lightness, min: {X.min()}, Max: {X.max()}")
    return -1*X + 255
    

def scale_lightness(X: np.ndarray, min:float = 0, max:float = 255):
    X_std = (X - X.min()) / (X.max() - X.min())
    print(f"Min: {X.min()}, Max: {X.max()}")
    return np.clip(X_std * (max - min) + min, min, max)


def convolve_vertically(a, b):
    conv = np.zeros(a.shape)
    for i, (trc_a, trc_b) in enumerate(np.nditer([a,b],['external_loop'], order='F')):
        conv[:,i] += convolve1d(trc_a, trc_b,
                                mode="constant", # pad with zeroes
                                origin=-(len(trc_b) // 2) # trc_b starts at t=0
                               )
    # normalize so it can be shown as image
    conv_scaled = scale_lightness(conv).astype(np.uint8)
    return flip_lightness(conv_scaled)


canvas_size = (800,500)


with gr.Blocks() as demo:
    with gr.Row():
        sketchpad_a = gr.Sketchpad(
            type="numpy",
            canvas_size=canvas_size,
        )

        sketchpad_b = gr.Sketchpad(
            type="numpy",
            canvas_size=canvas_size,
        )

        im_preview = gr.Image()

    sketchpad_b.change(convolve_vertically, outputs=im_preview, inputs=[sketchpad_a, sketchpad_b], show_progress="hidden")

if __name__ == "__main__":
    demo.launch()






# import gradio as gr
# import numpy as np
# from scipy.ndimage import convolve1d

# def flip_lightness(X):
#     print(f"In flip_lightness, min: {X.min()}, Max: {X.max()}")
#     return -1*X + 255

# def binary_lightness(X):
#     print(f"In binary_lightness, min: {X.min()}, Max: {X.max()}")
#     return (X > (254/2)).astype(int)

# def scale_lightness(X: np.ndarray, min:float = 0, max:float = 255):
#     X_std = (X - X.min()) / (X.max() - X.min())
#     print(f"Min: {X.min()}, Max: {X.max()}")
#     return np.clip(X_std * (max - min) + min, min, max)

# def convolve_vertically(a, b):
#     conv = np.zeros(a.shape)
#     for i, (trc_a, trc_b) in enumerate(np.nditer([a,b],['external_loop'], order='F')):
#         conv[:,i] += convolve1d(trc_a, trc_b,
#                                 mode="constant", # pad with zeroes
#                                 origin=-(len(trc_b) // 2) # trc_b starts at t=0
#                                )
#     # normalize so it can be shown as image
#     conv_scaled = scale_lightness(conv).astype(int)
#     return flip_lightness(conv_scaled)


# def convolve_if_ready(sketchpad_a, sketchpad_b):
#     '''sketchpad now returns dict of layers'''
#     a = sketchpad_a["composite"]
#     b = sketchpad_b["composite"]

#     print(a.keys())

#     if sketchpad_b is None:
#         # b = np.zeros(sketch_a.shape)
#         return flip_lightness(a)
#     else:
#         # print(f"In read, min: {np.asarray(a).min()}, Max: {np.asarray(a).max()}")
#         # nan_count = np.count_nonzero(np.isnan(np.asarray(a)))
#         # print(f"In read, nan count: {nan_count}")
#         # a = binary_lightness(a)
#         # b = binary_lightness(b)
        
#         return convolve_vertically(a, b)

# def passthrough(a, b):
#     return a["composite"]


# canvas_size = (800,500)
# white_bg = np.full((256, 256), 255, dtype=np.uint8)

# sketch_src = gr.Sketchpad(value=white_bg, canvas_size=canvas_size, image_mode="L", type="numpy")
# sketch_rcv = gr.Sketchpad(value=white_bg, canvas_size=canvas_size, image_mode="L", type="numpy")

# # sketch_src = gr.ImageEditor(sources=(), brush=gr.Brush(colors=["#000000"], color_mode="fixed"),
# #                             canvas_size=(30,40), image_mode="L", type="numpy")                    
# # sketch_rcv = gr.ImageEditor(sources=(), brush=gr.Brush(colors=["#000000"], color_mode="fixed"),
# #                             canvas_size=(30,40), image_mode="L", type="numpy")   

# iface = gr.Interface(
#     fn=convolve_if_ready,
#     # fn=passthrough,
#     inputs=[sketch_src, sketch_rcv],
#     # outputs=gr.Textbox(),
#     outputs="image",
#     live=True,
#     )
# iface.launch()