import matplotlib

matplotlib.use("Agg")
# Libraries
import streamlit as st

# Utility
import time
import io
import argparse
import sys
import os.path
import subprocess
import tempfile
import logging

# Viz
import altair as alt
import av

# ML
import numpy as np
import pandas as pd
import cv2 as cv
from PIL import Image, ImageOps
from tqdm import tqdm

# Custom
import inference
from app_utils import *


def run_app():
    # Options
    st.set_option("deprecation.showfileUploaderEncoding", False)
    # App documentation
    st.title("MIT Count Fish Counter")
    st.text("Upload a video file to detect and count fish")
    video_url = "yolo2_out_py.mp4"  # Replace with the actual video URL or file path
    video_bytes = open(video_url, "rb").read()


    col1, col2 = st.columns(2)

    ## Col1 #########################################
    with col1:
        ## Initial visualizations
        # Historical values
        st.altair_chart(
            plot_historical_data(pd.read_csv("herring_count_all.csv")),
            use_container_width=True,
        )

        # Locations
        st.subheader("Map of Fishery Locations")
        st.map(
            pd.DataFrame(
                np.random.randn(5, 2) / [50, 50] + [42.41, -71.38],
                columns=["lat", "lon"],
            )
        )

    ## Col2 #########################################
    with col2:
        # Replace the loading message with the video
        st.subheader("Example of processed video")
        st.video(video_bytes)
        st.subheader("Upload your own video...")
        # Initialize
        img_types = ["jpg", "png", "jpeg"]
        video_types = ["mp4", "avi"]
        uploaded_file = st.file_uploader(
            "Select an image or video file...", type=img_types + video_types
        )

        # Display uploaded file
        if uploaded_file is not None:
            if str(uploaded_file.type).split("/")[-1] in img_types:
                image = Image.open(uploaded_file)
                st.image(image, caption="Uploaded image", use_column_width=True)

                # TBD: Inference code to run and display for single image

            elif str(uploaded_file.type).split("/")[-1] in video_types:
                # Display uploaded file
                st.video(uploaded_file)

                # Convert streamlit video object to OpenCV format to run inferences
                tfile = tempfile.NamedTemporaryFile(delete=False)
                tfile.write(uploaded_file.read())
                vf = cv.VideoCapture(tfile.name)

                # Run inference
                with st.spinner("Running inference..."):
                    frames, counts, timestamps = inference.main(vf)
                logging.info("INFO: Completed running inference on frames")
                st.balloons()

                # Convert OpenCV Numpy frames in-memory to IO Bytes for streamlit
                streamlit_video_file = frames_to_video(frames=frames, fps=11)

                st.video(streamlit_video_file)  # Show processed video\
                st.download_button(
                    label="Download processed video",
                    data=streamlit_video_file,
                    mime="mp4",
                    file_name="processed_video.mp4",
                )
                df_counts_time = pd.DataFrame(
                    data={"fish_count": counts, "timestamps": timestamps[1:]}
                )
                st.altair_chart(
                    plot_count_date(dataframe=df_counts_time),
                    use_container_width=True,
                )

        else:
            st.write("No file uploaded")


if __name__ == "__main__":
    run_app()