#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
#  app.py
#
#  Copyright 2023 KP
#
#  This program is free software; you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation; either version 2 of the License, or
#  (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, write to the Free Software
#  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
#  MA 02110-1301, USA.
#
#

import argparse

import streamlit as st

from lib.generate import generate
from lib.util import convert_to_video, get_device


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--device",
        choices=["cuda", "mps", "cpu"],
        help="Override device",
    )
    args = parser.parse_args()
    device = args.device if args.device is not None else get_device()

    st.set_page_config(
        page_title="ez-text2video",
        page_icon="🎥",
        layout="wide",
        menu_items={
            "Get Help": "https://github.com/kpthedev/ez-text2video",
            "Report a bug": "https://github.com/kpthedev/ez-text2video/issues",
            "About": "# ez-text2video 🎥 \n A Streamlit app to easily run the text-to-video diffusion model.",
        },
    )
    st.write("# ez-text2video 🎥")
    col_left, col_right = st.columns(2)

    with col_left:
        st.info(
            "The very first time you run this app, it will take some time to download all of the models (~5 mins).",
            icon="ℹ️",
        )
        prompt = st.text_area("Prompt")

        # Number inputs
        num_sub_col_1, num_sub_col_2, num_sub_col_3, num_sub_col_4 = st.columns(4)
        frames = num_sub_col_1.number_input(
            label="Number of total frames", min_value=1, max_value=999999, value=16
        )
        n_fps = num_sub_col_2.number_input(
            label="Frames per second (fps)", min_value=1, max_value=999999, value=8
        )
        steps = num_sub_col_3.number_input(
            label="Number of inference steps", min_value=1, max_value=999999, value=50
        )
        seed = num_sub_col_4.number_input(
            label="Seed", min_value=1, max_value=999999, value=42
        )

        # Dim inputs
        dim_sub_col_1, dim_sub_col_2 = st.columns(2)
        height = dim_sub_col_1.slider(
            label="Height", min_value=16, max_value=1024, value=256, step=1
        )
        width = dim_sub_col_2.slider(
            label="Width", min_value=16, max_value=1024, value=256, step=1
        )

        with st.expander("Optimizations", expanded=True):
            st.markdown(f"**Device:** `{device}`")
            cpu_offload = st.checkbox(
                "Enable CPU offloading",
                value=True if device == "cuda" else False,
                disabled=True if device == "cpu" else False,
            )
            attention_slice = st.checkbox(
                "Enable attention slicing (slow)",
                value=True if device == "mps" else False,
                disabled=True if device == "cpu" else False,
            )

        if st.button("Generate", use_container_width=True):
            with st.spinner("Generating..."):
                raw_video = generate(
                    prompt=prompt,
                    num_frames=int(frames),
                    num_steps=int(steps),
                    seed=int(seed),
                    height=height,
                    width=width,
                    device=device,
                    cpu_offload=cpu_offload,
                    attention_slice=attention_slice,
                )
                video = convert_to_video(
                    video_frames=raw_video,
                    fps=int(n_fps),
                    filename=f"{prompt.replace(' ', '_').lower()}-{seed}",
                )

            with col_right:
                st.video(video)


if __name__ == "__main__":
    main()