"""
Description: Functions that directly support the Streamlit app
"""
import pandas as pd
import altair as alt
import io
import av
from tqdm import tqdm
import numpy as np
import logging


def frames_to_video(frames=None, fps=12):
    """
    Convert frames to video for Streamlit

    Args:
        frames: frame from cv2.VideoCapture as numpy. E.g. frame.astype(np.uint8)
        fps: Frames per second. Usefull if inference video is compressed to slow down for analysis
    """
    height, width, layers = frames[0].shape  # grab info from first frame
    output_memory_file = io.BytesIO()  # Create BytesIO "in memory file".

    output = av.open(
        output_memory_file, "w", format="mp4"
    )  # Open "in memory file" as MP4 video output
    stream = output.add_stream(
        "h264", str(fps)
    )  # Add H.264 video stream to the MP4 container, with framerate = fps.
    stream.width = width  # Set frame width
    stream.height = height  # Set frame height
    stream.pix_fmt = "yuv420p"  # NOTE: yuv444p doesn't work on mac. Select yuv444p pixel format (better quality than default yuv420p).
    stream.options = {
        "crf": "17"
    }  # Select low crf for high quality (the price is larger file size).

    # Iterate the created images, encode and write to MP4 memory file.
    logging.info("INFO: Encoding frames and writing to MP4 format.")
    for frame in tqdm(frames):
        frame = av.VideoFrame.from_ndarray(frame.astype(np.uint8), format="bgr24")
        packet = stream.encode(frame)  # Encode video frame
        output.mux(
            packet
        )  # "Mux" the encoded frame (add the encoded frame to MP4 file).

    packet = stream.encode(None)  # Flush the encoder
    output.mux(packet)
    output.close()

    output_memory_file.seek(0)
    return output_memory_file


def plot_historical_data(dataframe):
    """Returns altair plot of historical counts to be rendered on main dashboard."""
    dataframe["Date"] = pd.to_datetime(dataframe["Date"])
    s = (
        dataframe.resample(rule="D", on="Date")["Count"].sum().reset_index()
    )  # Resample on day
    return (
        alt.Chart(s, title="Historical Video Counts of Herring")
        .mark_bar()
        .transform_window(
            # The field to average
            rolling_mean="mean(Count)",
            # The number of values before and after the current value to include.
            frame=[-9, 0],
        )
        .encode(x="Date", y="Count", tooltip=["Count", "Date"])
        .interactive()
    )


def plot_count_date(dataframe):
    """Plots counts vs relative time for uploaded video."""
    dataframe["seconds"] = dataframe["timestamps"] / 1000
    dataframe["class"] = "Herring"  # TBD: Hard-coded for now
    return (
        alt.Chart(dataframe, title="Processed video detected fish")
        .mark_line()
        .encode(x="seconds", y="fish_count", color="class")
        .interactive()
    )