import streamlit as st
import pandas as pd
import plotly.express as px
from datasets import load_dataset
import os
from src.basic_plot import basic_chart
from src.map_viz import calling_map_viz
from src.data_ingestion import daily_average
from src.heatmap import HeatMap

def fetch_data():
    # comment out for local testing, but be sure to include after testing
    dataset = load_dataset("tappyness1/causion", use_auth_token=os.environ['TOKEN'])
    # print (dataset)
    # print (pd.DataFrame(dataset['train']))
    counts_df = pd.DataFrame(dataset['train'])

    # only use this part before for local testing
    # once local testing is completed, comment out and use the dataset above
    # counts_df = pd.read_csv("data/counts_dataset.csv")
    return counts_df

def main():

    counts_df = fetch_data()
    # st.set_page_config(layout="wide")
    height = 650

    st.markdown(""" <style>
                #MainMenu {visibility: hidden;}
                footer {visibility: hidden;}
                </style> """, 
                unsafe_allow_html=True
                )

    # Select Plot Option
    st.sidebar.markdown("Select Plots to show")
    checkbox_one = st.sidebar.checkbox('Overall Traffic', value = True) # rename as necessary
    checkbox_two = st.sidebar.checkbox('Traffic Map', value = True)
    checkbox_three = st.sidebar.checkbox('Heat Map', value = True)

    if checkbox_one:
        st.plotly_chart(basic_chart(counts_df),use_container_width=True)

    if checkbox_two:
        st.pyplot(calling_map_viz(counts_df))
    
    if checkbox_three:

        heatmap = HeatMap(counts_df)

        st.plotly_chart(heatmap.vehicle_count_bar())
        st.plotly_chart(heatmap.heatmap())

        hour_choice = st.selectbox(
            "Choose Hour",
            options=[
                "00:00", "01:00", "02:00", "03:00", "04:00", "05:00",
                "06:00", "07:00", "08:00", "09:00", "10:00", "11:00",
                "12:00", "13:00", "14:00", "15:00", "16:00", "17:00",
                "18:00", "19:00", "20:00", "21:00", "22:00", "23:00",
            ],
            key = "hour"
        )
        st.plotly_chart(heatmap.update_hour_bar_chart(hour_choice))

        day_choice = st.selectbox("Choose Day of the Week", ["Monday", "Tuesday", "Wednesday", 
                                                "Thursday", "Friday","Saturday", "Sunday"], key = "day")
        st.plotly_chart(heatmap.update_day_bar_chart(day_choice))
        

if __name__ == "__main__":
    main()