src/heatmap.py

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import scipy as sp
import sklearn as sk
import plotly.express as px
from plotly.subplots import make_subplots
import plotly.graph_objects as go
import streamlit as st


def clean_data(df):

    df["date"] = pd.to_datetime(df["date"], format="%Y-%m-%d")
    df["day"] = df["date"].dt.day_name()
    df["hour"] = df["time"].str[:2] + ":00"
    df.drop(columns=["motorcycle"], axis=1, inplace=True)
    df["vehicle"] = df["car"] + df["large_vehicle"]

    return df


class HeatMap:
    def __init__(self, counts_df):
        self.df = clean_data(counts_df)
        new = (
            self.df.groupby(["hour", "day"])
            .sum()
            .drop(columns=["car", "large_vehicle"])
            .reset_index()
        )
        table = pd.pivot_table(
            new, values="vehicle", index=["day"], columns=["hour"]
        ).reset_index()
        self.table = table.reindex([1, 5, 6, 4, 0, 2, 3])


    def vehicle_count_bar(self):
        new_df = self.df.groupby(["day"]).sum().reset_index()
        new_df = new_df.reindex([1, 5, 6, 4, 0, 2, 3])

        veh_count_fig = px.bar(
            new_df,
            x="day",
            y="vehicle",
            color="day",
            text_auto=True,
            labels={"day": "Day of the Week", "vehicle": "Vehicle Count"},
        )

        return veh_count_fig

    def heatmap(self):

        new_table = self.table.iloc[:, 1:].to_numpy()

        hm_fig = px.imshow(
            new_table,
            labels=dict(
                x="Hour of the Day", y="Day of the Week", color="Causeway Traffic"
            ),
            x=[
                "12am",
                "1am",
                "2am",
                "3am",
                "4am",
                "5am",
                "6am",
                "7am",
                "8am",
                "9am",
                "10am",
                "11am",
                "12pm",
                "1pm",
                "2pm",
                "3pm",
                "4pm",
                "5pm",
                "6pm",
                "7pm",
                "8pm",
                "9pm",
                "10pm",
                "11pm",
            ],
            y=[
                "Monday",
                "Tuesday",
                "Wednesday",
                "Thursday",
                "Friday",
                "Saturday",
                "Sunday",
            ],
            text_auto=True,
        )
        hm_fig.update_xaxes(side="top")

        return hm_fig

    def update_hour_bar_chart(self, hour="08:00"):

        fig_hours = px.bar(
            self.table,
            x="day",
            y=str(hour),
            color="day",
            text_auto=True,
            labels={"day": "Day of the Week"},
        )
        return fig_hours

    def update_day_bar_chart(self, day="Saturday"):

        t = self.table.T
        t.drop("day", inplace=True)
        t.columns = [
            "Monday",
            "Tuesday",
            "Wednesday",
            "Thursday",
            "Friday",
            "Saturday",
            "Sunday",
        ]
        t = t.reset_index()

        fig_days = px.bar(
            t,
            x="hour",
            y=str(day),
            color=str(day),
            text_auto=True,
            labels={"hour": "Count of Each Hour"},
        )

        return fig_days