import streamlit as st
import pandas as pd
import numpy as np
import sklearn
import time
import plotly.express as px
from Univariate import Univariate_Models
from Multivariate import Multivariate_Models

st.set_page_config(layout="wide")


def GetFile():
    uploaded_file = st.file_uploader("", type=["csv", 'json'])
    file_type = st.selectbox("Please select file type", ['csv', 'json'])
    if uploaded_file is not None:
        uploaded_file.seek(0)
        if file_type == 'json':
            data = pd.read_json(uploaded_file)
            return data
        elif file_type == 'csv':
            seperator = st.text_input("Input csv seperator here ", ",")
            data = pd.read_csv(uploaded_file, sep=seperator)
            # data = data.sample(n=1000)
            return data

def header():
    col1, col2, col3, col4, col5 = st.columns(5)
    with col1: st.image('./assets/yash_logo.png')
    with col2: st.write(' ')
    with col3: st.write(' ')
    with col4: st.write(' ')
    with col5: st.write(' ')
    st.markdown("<h1 style='text-align: center;'>Time Series Accelerator - AutoML</h1>",
                unsafe_allow_html=True)


def univaiate_modelling(data, date_column, target_column):
        
        # Data Pre-Processing
        data[date_column] = pd.to_datetime(data[date_column])
        data = data.set_index(date_column)
        data = data[[target_column]]
        
        forcasting_horizon = round(data.shape[0] * 0.067)
        st.write("Forecasting horizon selected as: ", forcasting_horizon)

        # fh = 12
        with st.spinner('Initializing modeling setup...'):
            uni_mod = Univariate_Models(data, forcasting_horizon)
        
        with st.spinner('Getting Stats...'): 
            st.write("Statistics of the Data")
            st.dataframe(uni_mod.get_stats())
            

        ### Plots ##
        with st.spinner('Creating Plots...'):
            stats_plots = uni_mod.get_plots()

        for name, plt in stats_plots.items():
            st.write(name)
            st.plotly_chart(plt)

        #### Train Models ####
        with st.spinner('Training Multiple models'):
            model_metric_df, best_mod_param_df, best_model = uni_mod.train_models()
        # st.success('Done!')

        st.write("Models trained and their metrics")
        st.write(model_metric_df)

        st.write("Parameters of Best Model")
        st.dataframe(best_mod_param_df)

        #### Predictions/Forecasts ###
        st.write("Predicted values")
        st.dataframe(uni_mod.get_predictions())

        insam, outsam = uni_mod.plot_insample_and_forecast()
        st.plotly_chart(insam)
        st.plotly_chart(outsam)


def multivaiate_modelling(data, target_column, date_column, time_interval):
        forcasting_horizon = round(data.shape[0] * 0.067)
        st.write("Forecasting horizon selected as: ", forcasting_horizon)

        # fh = 12
        with st.spinner('Initializing modeling setup...'):
            multi_mod = Multivariate_Models(data, forcasting_horizon, target_column, date_column,
                                            time_interval)
        
        with st.spinner('Getting Stats...'): 
            st.write("Statistics of the Data")
            cat_features, num_features, desc_df = multi_mod.get_stats()
            st.dataframe(desc_df)
            st.write("Inferred numerical columns :", ", ".join(num_features))
            st.write("Inferred categorical columns :", ", ".join(cat_features))

        #### Train Models ####
        with st.spinner('Training Multiple models...'):
            model_metric_df, best_mod_param_df, _ = multi_mod.train_models()
        
        st.write("Models trained and their metrics")
        st.write(model_metric_df)

        st.write("Parameters of Best Model")
        st.dataframe(best_mod_param_df)

        ### Plots ##
        with st.spinner('Creating Plots...'):
            stats_plots = multi_mod.get_plots()

        for name, plt in stats_plots.items():
            try: 
                st.write(name)
                # st.plotly_chart(plt)
                st.image(plt)
            except Exception as e:
                print(f"[PLOT EXCEPTION] ==> {type(e).__name__} : {e}")

        #### Predictions/Forecasts ###
        st.write("Predicted values")
        st.dataframe(multi_mod.get_predictions())

        insam, outsam = multi_mod.plot_insample_and_forecast()
        st.plotly_chart(insam)
        st.plotly_chart(outsam)


def main():
    data = GetFile()
    # date_column = "date"

    if (data is not None):
        data_shape = data.shape
        st.write("Data head", data.head())
        st.write("Data shape :", data_shape)

        ts_type_radio = st.radio("Please select the type of analysis required: ",
                                 ('Univariate', 'Multivariate'))

        all_cols = set(data.columns)
        target_column = st.selectbox('Select the dependent variable', all_cols)
        date_column = st.selectbox('Select the date column',
                                    all_cols - set([target_column]))
        time_interval_radio = st.radio("Please select the type of analysis required: ",
                                           ('daily', 'weekly', 'monthly', 'yearly'), index=0)  
        
        ts_type = ts_type_radio.lower().strip()
        start_btn = st.button('Start Modelling')

        if start_btn and ts_type == "univariate":
            univaiate_modelling(data, date_column, target_column)

        elif start_btn and ts_type == "multivariate":
            multivaiate_modelling(data, target_column, date_column, time_interval_radio)



if __name__=='__main__':
    header()
    main()