import streamlit as st
from matplotlib import pyplot as plt
import albumentations as A
import cv2

# https://github.com/phrasenmaeher/image_augmentations_visualization/blob/master/Start.py
def create_pipeline(transformations: list):
    pipeline = []
    for index, transformation in enumerate(transformations):
        if transformation:
            pipeline.append(index_to_transformation(index))

    return pipeline


def load_image(filename):
    img = cv2.imread(filename)
    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

def index_to_transformation(index: int):
    if index == 0:
        return A.GaussNoise(p=1.0, var_limit=(0.25, 0.5))
    elif index == 1:
        return A.HorizontalFlip(p=1.0)
    elif index == 2:
        return A.VerticalFlip(p=1.0)
    elif index == 3:
        return A.RandomBrightness(p=1.0, limit=(0.5, 1.5))
    # elif index == 4:
    #     return A.AdvancedBlur(p=1.0, blur_limit=3)
    elif index == 4:
        return A.ChannelShuffle(p=1.0)
    elif index == 5:
        return A.ChannelDropout(p=1.0)
    elif index == 6:
        return A.RandomContrast(p=1.0, limit=(0.5, 1.5))

placeholder = st.empty()
placeholder2 = st.empty()
placeholder.markdown(
        "# Visualize an image augmentation pipeline\n"
        "### Select the components of the pipeline in the sidebar.\n"
        "Once you have chosen the augmentation techniques, select or upload an image.\n"
        "Then click 'Apply' to start!\n"
    )
placeholder2.markdown(
        "After clicking start, the individual steps of the pipeline are visualized. The ouput of the previous step is the input to the next step."
    )
    # placeholder.write("Create your audio pipeline by selecting augmentations in the sidebar.")
st.markdown("Choose the transformations here:")
gaussian_noise = st.sidebar.checkbox("GaussianNoise")

horizontal_flip = st.sidebar.checkbox("HorizontalFlip")
vertical_flip = st.sidebar.checkbox("VerticalFlip")
random_brightness = st.sidebar.checkbox("RandomBrightness")
# advanced_blur = st.sidebar.checkbox("AdvancedBlur")
channel_shuffle = st.sidebar.checkbox("ChannelShuffle")
channel_dropout = st.sidebar.checkbox("ChannelDropout")
random_contrast = st.sidebar.checkbox("RandomContrast")

st.markdown("---")
st.markdown("(Optional) Upload an image file here:")
file_uploader = st.sidebar.file_uploader(label="", type=[".png", ".jpg", ".jpeg"])
st.markdown("Or select a sample file here:")


st.markdown("---")
transformations = [
            gaussian_noise,
            horizontal_flip,
            vertical_flip,
            random_brightness,
            # advanced_blur,
            channel_shuffle,
            channel_dropout,
            random_contrast,

]

pipeline = A.Compose(create_pipeline(transformations))
print(pipeline)
tmp_img = load_image("/home/hodor/dev/Learning/XAI/streamlit_demo/multipage-app/figures/cnv.png")
# apply the transformation to the image
data = pipeline(image=tmp_img)["image"]
# modified_image = individual_transformation(image=tmp_img)["image"]
st.image(data)