"""Streamlit demo to explore ChangeIt dataset."""
from os.path import join, exists, dirname, abspath
from glob import glob

import numpy as np
import pandas as pd
import streamlit as st

import warnings
warnings.simplefilter(action='ignore')

curr_filepath = abspath(__file__)
repo_path = dirname(curr_filepath)


categories = [
    "ball",
    "beer",
    "boil_milk",
    "butter",
    "cake",
    "champagne",
    "coffee",
    "cream",
    "drill",
    "eggs",
    "juice",
    "milk",
    "potatoes",
    "tea",
    "tree",
    "whisk_eggs",
]


def load_category_df(path, desc):
    df = pd.read_csv(path, header=None)
    df.columns = ["video_id", "param"]
    df["category"] = path.split("/")[-1].split(".")[0]
    df["action"] = desc
    return df


def make_grid(cols,rows):
    grid = [0]*cols
    for i in range(cols):
        with st.container():
            grid[i] = st.columns(rows)
    return grid


if __name__ == "__main__":
    # Streamlit app code
    st.set_page_config(layout="wide")
    st.title("Clips from ChangeIt dataset (possibly of temporal nature) 🎬")
    

    # Load all the data
    if "video_ids" not in st.session_state:
        cdf = pd.read_csv(join(repo_path, "data", "categories.csv"))
        cdf = cdf[cdf.DIR_NAME.isin(categories)]
        action_names = cdf.ACTION_NAME.values
        paths = [join(repo_path, "data/videos", c + ".csv") for c in cdf.DIR_NAME]
        dfs = [load_category_df(p, a) for p, a in zip(paths, action_names)]
        df = pd.concat(dfs, ignore_index=True)
        
        # df["annot_path"] = df[["video_id", "category"]].apply(lambda x: join(repo_path, f"data/annotations/{x[1]}/{x[0]}.fps1.csv"), axis=1)
        # df = df[df["annot_path"].apply(exists)]

        st.session_state.df = df
    else:
        df = st.session_state.df
    st.markdown(f"**Total number of relevant clips**: {len(df)}", unsafe_allow_html=True)

    # Select a random subset of clips
    reload_button = st.button("Reload")
    NUM = 9
    indices = np.random.randint(0, len(st.session_state.df), NUM)
    if reload_button:
        indices = np.random.randint(0, len(st.session_state.df), NUM)

    # Show videos
    grid = make_grid(3, 3)
    per_video_width = 360
    per_video_height = 240
    for i, idx in enumerate(indices):
        row = i // 3
        col = i % 3
        video_id = df.iloc[idx].video_id
        action = df.iloc[idx].action

        url = f"https://www.youtube.com/embed/{video_id}"
        html_code = f"""
        <iframe height="{per_video_height}" width="{per_video_width}" src="{url}" frameborder="0" allowfullscreen></iframe>
        """
        grid[row][col].markdown(html_code, unsafe_allow_html=True)
        grid[row][col].markdown(f"**Action**: {action}", unsafe_allow_html=True)