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.map_viz_pred import calling_pred_map_viz
from src.data_ingestion import daily_average
from src.heatmap import HeatMap
from src.data_ingestion import remove_previous_view, merge_volumes
from src.pred_plot import prep_data_pred_plot, data_split, train_model, predicted_figure, get_today, gen_fig, pred_bars
from datetime import date

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()
    pred_df = counts_df.copy()
    counts_df1 = counts_df.copy()
    counts_df = remove_previous_view(counts_df)
    counts_df = merge_volumes(counts_df)
    # st.set_page_config(layout="wide")
    height = 650

    st.markdown(""" <style>
                #MainMenu {visibility: hidden;}
                footer {visibility: hidden;}
                </style> """, 
                unsafe_allow_html=True
                )
    hours = [
        "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",
    ]
    st.write("<- Click on the sidebar to select your plot preferences")
    # Select Plot Option
    plot_type = st.sidebar.selectbox("Choose Plot Type", options = ['Historical', 'Predictive'])
    if plot_type == 'Historical':
        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)
        view_options = list(counts_df["view"].unique())
        view_options.append('All')
        view = st.sidebar.selectbox("Choose View", options=view_options, index = view_options.index("Woodlands - to Johor"))
        if view != 'All':
            st.header(f"Showing Traffic for {view}")
            counts_df = counts_df[counts_df['view'] == view]

        if checkbox_one:
            st.subheader("Overall Traffic")
            plot = st.selectbox("Choose Plot", options=["Day", "Hour", "Raw"], index = 0)
            st.plotly_chart(basic_chart(counts_df, plot = plot),use_container_width=True)

        if checkbox_two:
            st.subheader("Traffic Map")
            st.pyplot(calling_map_viz(counts_df1))
        
        if checkbox_three:

            heatmap = HeatMap(counts_df)
            
            # st.header("Mean Vehicle Count by Day of Week")
            # st.plotly_chart(heatmap.vehicle_count_bar())
            st.subheader("Heatmap")
            st.plotly_chart(heatmap.heatmap())

            hour_choice = st.selectbox(
                "Choose Hour",
                options= hours,
                key = "hour", index = hours.index("08:00")
            )
            st.subheader(f"Traffic Volume of Each Day at {hour_choice}")
            st.plotly_chart(heatmap.update_hour_bar_chart(hour_choice))
            days = ["Monday", "Tuesday", "Wednesday", 
                    "Thursday", "Friday","Saturday", "Sunday"]
            day_choice = st.selectbox("Choose Day of the Week", options = days, key = "day", index = days.index("Saturday"))
            st.subheader(f"Traffic Volume of Each Hour on {day_choice}")
            st.plotly_chart(heatmap.update_day_bar_chart(day_choice))

    else: 

        st.sidebar.markdown("Select Plots to show")
        checkbox_two_pred = st.sidebar.checkbox('Predictive Traffic Map', value = True)
        
        figs = gen_fig()
        today = get_today()
        final_table = prep_data_pred_plot(pred_df)
        x_train, _, y_train, _ = data_split(final_table)
        clf = train_model(x_train, y_train)
        col1, col2, col3 = st.columns(3)
        with col1:
            d = st.date_input(
                "Choose Your Planned Date",
                date(today[0],today[1], today[2]))

        with col2: 
            pred_view_choice = st.selectbox(
                "Choose View",
                options= ['Johor-Tuas','Johor-Woodlands', 'Tuas-Johor', 'Woodlands-Johor'],
                key = "pred_view"
            )
        with col3: 
            pred_hour_choice = st.selectbox(
                "Choose Your Planned Hour",
                options= hours,
                key = "pred_hour", index = hours.index("08:00")
            )

        starter_variables = [x_train, str(d), pred_hour_choice, pred_view_choice]
        st.plotly_chart(predicted_figure(clf, starter_variables, figs))
        st.plotly_chart(pred_bars(d, final_table))
        

        if checkbox_two_pred:
            st.subheader("Predictive Traffic Map")


            hour_choice = st.selectbox(
                "Choose Hour",
                options= hours,
                key = "hour", index = hours.index("08:00")
            )

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

            st.pyplot(calling_pred_map_viz(counts_df, day_choice, hour_choice))

if __name__ == "__main__":
    main()