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
from src.data_ingestion import remove_previous_view, merge_volumes


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):

        # counts_df = counts_df[counts_df['view'] == view]
        
        self.df = clean_data(counts_df)
        
        new = (
            self.df.groupby(["hour", "day"])
            .mean(numeric_only = True)
            .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]).round(1)


    def vehicle_count_bar(self):
        new_df = self.df.groupby(["day"]).mean(numeric_only = True).round(1).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"},
        )
        veh_count_fig.update_layout(showlegend=False)

        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="Traffic"
            ),
            x=[
                "00", "01", "02", "03",
                "04", "05", "06", "07",
                "08", "09", "10", "11",
                "12", "13", "14", "15",
                "16",
                "17",
                "18",
                "19",
                "20",
                "21",
                "22",
                "23",
            ],
            y=[
                "Mon",
                "Tue",
                "Wed",
                "Thu",
                "Fri",
                "Sat",
                "Sun",
            ],
            # text_auto=True,
        )
        hm_fig.update_xaxes(side="top", tickmode = 'linear', type = 'category')
        hm_fig.update_layout(yaxis_visible=True, 
                             yaxis_showticklabels=True,
                             yaxis_title = ''
        )
        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"},
        )
        fig_hours.update_layout(showlegend=False)
        fig_hours.update_layout(yaxis_visible=False, 
                             yaxis_showticklabels=False,
                             yaxis_title = False
        )
        fig_hours.update_layout(xaxis_visible=True, 
                             xaxis_showticklabels=True,
                             xaxis_title = ''
        )
        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()

        t['hour'] = t['hour'].str[:2]

        fig_days = px.bar(
            t,
            x="hour",
            y=str(day),
            color=str(day),
            # text_auto=True,
            labels={"hour": "Count of Each Hour"},
        )
        fig_days.update_layout(showlegend=False)
        fig_days.update_layout(yaxis_visible=False, 
                             yaxis_showticklabels=False,
                             yaxis_title = False
        )
        fig_days.update_layout(xaxis_visible=True, 
                             xaxis_showticklabels=True,
                             xaxis_title = '',
        )
        fig_days.update_xaxes(tickangle=0, tickmode = 'linear', type = 'category', categoryorder='category ascending')
        return fig_days