import streamlit as st
import numpy as np
from tqdm.auto import trange
import torch

from source import Model, Block


def cvtImg(img):
    img = img - img.min()
    img = (img / img.max())
    return img.numpy().astype(np.float32)


def load_DM():
    PATH = "./Diffusion-cuda"
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    st.write("Devise used:", device)
    model = Model().to(device)
    st.write("Diffusion Model's architecture created.")

    try:
        # Load the entire model object
        model = torch.load(PATH, map_location=device)
        st.write("Diffusion Model loaded.")
    except FileNotFoundError:
        st.error("Model file not found at specified path.")
    except Exception as e:
        st.error(f"Error occurred while loading model: {e}")

    return model, device


def predict_step_new(model, device):
    xs = []
    IMG_SIZE = 32
    timesteps = 16

    progress_text = "Denoising the image."
    my_bar = st.progress(0, text=progress_text)
    x = torch.randn(size=(1, 3, IMG_SIZE, IMG_SIZE), device=device)
    placeholder = st.empty()
    li = []
    with torch.no_grad():
        for i in trange(timesteps):
            my_bar.progress(i*100//15, text=progress_text)
            t = i
            x = model(x, torch.full([8, 1], t, dtype=torch.float, device=device))
            xs = x[0].permute(1, 2, 0)
            xs = torch.clip(xs, -1, 1)
            xs = cvtImg(xs.cpu())
            images = np.transpose(xs, (2, 0, 1))
            images = np.transpose(images, (1, 2, 0))
            placeholder.image(images, caption=f"Time step {t}", width=500)
            li.append(images)
    st.markdown("Brand new **car** generated!")
    return li


def main():
    st.set_page_config(page_title="Diffusion Model App Demo")
    st.header("Diffusion Model App 🖼️ 🚕", divider="rainbow")
    st.markdown("""Made by:
    Marijan$^\dag$, Thomas, Mauritz, Amanda and Baris for **COGITO** Diffusion Model""")

    if "model" not in st.session_state:
        st.session_state.model = []

    if "generation" not in st.session_state:
        st.session_state.generation = []

    if "device" not in st.session_state:
        st.session_state.device = None

    placeholder_1 = st.empty()
    isclick = placeholder_1.button('Import Diffusion Model')
    if isclick:
        st.session_state.model, st.session_state.device = load_DM()
        placeholder_1.empty()

    if st.button('Generate!'):
        li = predict_step_new(st.session_state.model, st.session_state.device)
        st.session_state.generation.append(li)
        for num, gen in enumerate(st.session_state.generation):
            st.write("Generation number", num+1)
            st.image(gen, width=40)


if __name__ == "__main__":
    main()