# Import necessary libraries
import matplotlib

# Use Agg backend for Matplotlib
matplotlib.use("Agg")

# Libraries for the app
import streamlit as st
import time
import io
import argparse
import sys
import os.path
import subprocess
import tempfile
import logging

# Visualization libraries
import altair as alt
import av

# Machine Learning and Image Processing libraries
import numpy as np
import pandas as pd
import cv2 as cv
from PIL import Image, ImageOps
from tqdm import tqdm

# Custom modules
import inference
from app_utils import *

@st.cache_data
def load_video(video_url):
    video_bytes = open(video_url, "rb").read()
    return video_bytes

@st.cache_data
def load_historical(fpath):
    return pd.read_csv(fpath)

st.set_page_config(layout="wide")

# Define the main function to run the Streamlit app
def run_app():
    # Set Streamlit options
    st.set_option("deprecation.showfileUploaderEncoding", False)
    
    # App title and description
    st.title("MIT Count Fish Counter")
    st.text("Upload a video file to detect and count fish")
    
    # Example video URL or file path (replace with actual video URL or file path)
    video_url = "yolo2_out_py.mp4"
    video_bytes = load_video(video_url)

    # Load historical herring
    df_historical_herring = load_historical(fpath="herring_count_all.csv")

    tab1, map_tab = st.tabs(["📈 Chart", "Map of Fishery Locations"])

    # Create two columns for layout
    col1, col2 = st.columns(2)

    ## Col1 #########################################
    with col1:
        ## Initial visualizations
        # Plot historical data
        st.altair_chart(
            plot_historical_data(df_historical_herring),
            use_container_width=True,
        )

        # Display map of fishery 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:
        # Display example processed video
        st.subheader("Example of processed video")
        st.video(video_bytes)
        st.subheader("Upload your own video...")

        # Initialize accepted file types for upload
        img_types = ["jpg", "png", "jpeg"]
        video_types = ["mp4", "avi"]

        # Allow user to upload an image or video file
        uploaded_file = st.file_uploader("Select an image or video file...", type=img_types + video_types)

        # Display the uploaded file
        if uploaded_file is not None:
            if str(uploaded_file.type).split("/")[-1] in img_types:
                # Display uploaded image
                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 video
                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 on the uploaded video
                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)

                # Show processed video and provide download button
                st.video(streamlit_video_file)
                st.download_button(
                    label="Download processed video",
                    data=streamlit_video_file,
                    mime="mp4",
                    file_name="processed_video.mp4",
                )

                # Create dataframe for fish counts and timestamps
                df_counts_time = pd.DataFrame(
                    data={"fish_count": counts, "timestamps": timestamps[1:]}
                )

                # Display fish count vs. timestamp chart
                st.altair_chart(
                    plot_count_date(dataframe=df_counts_time),
                    use_container_width=True,
                )

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

# Run the app if the script is executed directly
if __name__ == "__main__":
    run_app()