from streamlit_js_eval import streamlit_js_eval
import torchvision.transforms as transforms
import streamlit as st
from PIL import Image
import torch
import time

st.set_page_config(page_title="Invertable Steganography", layout="centered", page_icon="📷")


# model
class HiddenNetwork:
    def __init__(self, public_key):
        super().__init__()
        self.public_key = public_key

    def noise(self, x):
        torch.manual_seed(self.public_key)
        noise = torch.randn_like(x)
        return noise

    def forward(self, x):
        return x - self.noise(x)

    def backward(self, x):
        return x + self.noise(x)


def load(encode_key=0, decode_key=0):
    with st.spinner('Getting Neruons in Order ...'):
        encode_model = HiddenNetwork(int(encode_key))
        decode_model = HiddenNetwork(int(decode_key))
        process = transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor()])
        time.sleep(1)
    return encode_model, decode_model, process


def main():
    # Set Streamlit theme to dark mode
    st.markdown(
        """
        <style>
        body {
            color: #f8f8f8;
            background-color: #121212;
        }
        </style>
        """,
        unsafe_allow_html=True,
    )
    st.markdown("<h1 style='text-align: center; color: white;'>📷 Invertable Steganography 📷</h1>", unsafe_allow_html=True)
    st.image("secret.png", use_column_width=True)
    st.write(
        """
        Invertible neural networks like FreiA and Glow are neural architectures designed for reversible data transformations. 
        They are often used in image steganography, allowing data to be encoded and perfectly reconstructed without any loss of 
        information. These networks utilize public keys as seeds, enhancing security and ensuring that only authorized parties can 
        access the hidden data within encoded images.
        """
    )

    # Create inputs
    st.markdown("""---""")
    input_image = st.file_uploader("Encode Image", type=["jpg", "jpeg", "png"], key="1")
    col1, col2 = st.columns(2)
    encode_key = col1.number_input("Encoding Key", value=42, key="2")
    decode_key = col2.number_input("Decoding Key", value=42, key="3")

    # core
    col1, col2, col3, col4 = st.columns(4)
    col1.write("Input")
    col2.write("Encode")
    col3.write("Decode Encode")
    col4.write("Decode Input")

    # columns
    if input_image is not None:
        encode_model, decode_model, process = load(encode_key, decode_key)
        image = process(Image.open(input_image).convert("RGB"))
        forward = encode_model.forward(image)
        backward = decode_model.backward(forward)
        backward_input = decode_model.backward(image)

        with col1:
            st.image(transforms.ToPILImage()(image), use_column_width=True)
            transforms.ToPILImage()(image).save("tmp/image.png")
            st.download_button(label='Download Image', data=open('tmp/image.png', 'rb').read(), file_name='image.png', mime='image/png', key="4")

        with col2:
            st.image(transforms.ToPILImage()(forward), use_column_width=True)
            transforms.ToPILImage()(forward).save("tmp/forward.png")
            st.download_button(label='Download Image', data=open('tmp/forward.png', 'rb').read(), file_name='forward.png', mime='image/png', key="5")

        with col3:
            st.image(transforms.ToPILImage()(backward), use_column_width=True)
            transforms.ToPILImage()(backward).save("tmp/back.png")
            st.download_button(label='Download Image', data=open('tmp/back.png', 'rb').read(), file_name='back.png', mime='image/png', key="6")

        with col4:
            st.image(transforms.ToPILImage()(backward_input), use_column_width=True)
            transforms.ToPILImage()(backward_input).save("tmp/input.png")
            st.download_button(label='Download Image', data=open('tmp/input.png', 'rb').read(), file_name='input.png', mime='image/png', key="7")

    # Create a button to reset the interface page
    st.markdown("""---""")
    if st.button("Reset", use_container_width=True):
        streamlit_js_eval(js_expressions="parent.window.location.reload()")


if __name__ == "__main__":
    main()