import pandas as pd
import dash
from dash import dcc, html
import webbrowser
from src.viz import Visualizer as viz
import plotly.graph_objects as go
import os
import datetime
import plotly.express as px
import json

start_time = 0#240
end_time = 1600#600

start_date = datetime.datetime(2024, 6, 7, 12, 0, 0)
end_date = start_date + datetime.timedelta(seconds=end_time)

# Initialize the Dash app
app = dash.Dash(__name__, suppress_callback_exceptions=True)

# Global dictionary to store data
data_store = {}

def parse_action_trace(path, exclude = []):
    """
    Parses action data from CSV files into a dictionary.

    This function reads CSV files specified in a dictionary where each key-value pair
    corresponds to an condition ID and its associated file. It extracts columns related to
    time, action name, agent name, action end time, and attributes from each file and
    stores them in a nested dictionary structure.

    Args:
        files_dict (dict): A dictionary where keys are action IDs and values are file paths.

    Returns:
        dict: A dictionary where each key is an condition ID and the value is another dictionary
              containing parsed data from the corresponding CSV file.
    """

    df = pd.read_csv(path, engine = "python")
    
    # Convert to events list for timeline plotting
    events = []

    for i in range(len(df['time'])):
        if not any(excl in df['actionName'][i] for excl in exclude):
            event = {
                'agent': df['agentName'][i],
                'name': df['actionName'][i],
                'start': start_date + datetime.timedelta(seconds=df['time'][i]),
                'end': start_date + datetime.timedelta(seconds=df['time'][i] + max(df['actionEndTime'][i], 2)),
                'duration': max(df['actionEndTime'][i], 2),
                'attributes': df['attributes'][i]
            }
            events.append(event)

    return events

def parse_all_data(path):

    # Get the list of subfolders in the 'data' folder
    subfolders = [f.name for f in os.scandir(path) if f.is_dir()]

    # Load all data at the start of the program
    # Load the CSV files into pandas DataFrames from the selected subfolder
    data = {}
    exclude = ['Detect_crash','Detect_conflict','Flight_Dynamics','Reroute_flight','Show_radar','Direct_to_waypoint','Change_heading','fly']
    for subfolder in subfolders:
        data[subfolder] = {
            'wnd_action_trace': pd.read_csv(f'{path}/{subfolder}/actionTrace_Agent_PIC Blaze.csv'),
            'all_action_trace': parse_action_trace(f'{path}/{subfolder}/actionTrace.csv', exclude),
        }
        data[subfolder]['aircraft_data'] = {}
        for file in os.listdir(f'{path}/{subfolder}'):
            if file.endswith('_acstate.csv'):
                aircraft_id = file.split('_')[0]
                data[subfolder]['aircraft_data'][aircraft_id] = pd.read_csv(f'{path}/{subfolder}/{file}')

    subfolder_labels = [{'label': i, 'value': i} for i in subfolders]

    return subfolder_labels, data

# Create the layout for the app
app.layout = html.Div([
    html.Div([
        html.Div("Input the relative path to your results folder, then click 'Submit':", id='text1'),
        dcc.Input(
            id='path-input',
            value='data',
            type='text',
            placeholder='Type your path... For example, ../wmc5.1/Scenario/AAMv2/Results',
            style={'width': '90%'}  # Make the input wider
        ),
        html.Button('Submit', id='submit-button', n_clicks=0),
        # Create a dropdown menu with the subfolders as options
        html.Div("Once all subfolders are loaded (representing the different conditions you ran in WMC), you can select which one to plot:", id='text2'),
        dcc.Dropdown(
            id='subfolder-dropdown',
            value=None #subfolders[0]  # Default value
            )
    ], style={'width': '50%', 'border': '1px solid gray'}),

    html.Div([
        html.Div([
            dcc.Graph(id='map'),  # Second graph
            dcc.Dropdown(
                id='variable-dropdown',
                # options=[{'label': i, 'value': i} for i in subfolders],
                value=None #subfolders[0]  # Default value
            ),
            dcc.Graph(id='altitude-series') # Third graph
        ], style={'width': '39%', 'display': 'inline-block', 'border': '1px solid gray'}),      
        html.Div([
            dcc.Tabs(id='tabs-example-1', value='tab-1', children=[
                dcc.Tab(label='Traditional Action Trace', value='tab-1'),
                dcc.Tab(label='What`s Next Diagram', value='tab-2'),
            ]),
            html.Div(id='tabs-example-content-1'),
            html.Div([dcc.Graph(id='timeline')], id='tab-1-content', style={'display': 'none'}),  # Fourth graph
            html.Div([dcc.Graph(id='time-series')], id='tab-2-content', style={'display': 'none'}),  # First graph
        ], style={'width': '59%', 'display': 'inline-block', 'border': '1px solid gray'})
    ], style={'display': 'flex'}),

    html.Div(
        dcc.Slider(
            id='time-slider',
            min=start_time,
            max=end_time,
            value=start_time+100,
            marks={str(time): str(time) for time in range(start_time, end_time+1, 50)},
            updatemode='drag'
        ),
        style={'width': '100%'}  # Set the width of the div containing the slider to 80% of the container
    )

])

@app.callback(
    dash.dependencies.Output('tab-1-content', 'style'),
    dash.dependencies.Output('tab-2-content', 'style'),
    dash.dependencies.Input('tabs-example-1', 'value')
)
def render_content(tab):
    if tab == 'tab-1':
        return {'display': 'block'}, {'display': 'none'}
    elif tab == 'tab-2':
        return {'display': 'none'}, {'display': 'block'}
    
# Load data when the app starts and store it in the hidden div
@app.callback(
    dash.dependencies.Output('subfolder-dropdown', 'options'),
    dash.dependencies.Output('variable-dropdown', 'options'),
    [dash.dependencies.Input('submit-button', 'n_clicks')],
    [dash.dependencies.Input('path-input', 'value')],
)
def load_data(n_clicks, selected_path):
    print("selected path is",selected_path)
    global data_store
    if n_clicks > 0:
        if selected_path is None or not os.path.isdir(selected_path):
            return [], []
        subfolder_labels, data = parse_all_data(selected_path)
        print('subfolders',subfolder_labels)
        if data is None:
            data = {}
        # Store data in the global dictionary
        data_store = data
        labels = list(data[list(data.keys())[0]]['aircraft_data'][list(data[list(data.keys())[0]]['aircraft_data'].keys())[0]].columns.tolist())  # Extracting column names for the first aircraft-specific dataframe
        return subfolder_labels, labels
    else:
        return [], []

# Define the callback to update the graphs based on the slider value
@app.callback(
    dash.dependencies.Output('time-series', 'figure'),
    dash.dependencies.Output('map', 'figure'),
    dash.dependencies.Output('altitude-series', 'figure'),
    dash.dependencies.Output('timeline', 'figure'),
    #  dash.dependencies.Input('path-input', 'value'),
    [dash.dependencies.Input('subfolder-dropdown', 'value'),
     dash.dependencies.Input('time-slider', 'value'),
     dash.dependencies.Input('variable-dropdown','value')]
)
# def update_graph(selected_subfolder, selected_time):
#     # Load the CSV files into pandas DataFrames from the selected subfolder
#     df1 = pd.read_csv(f'data/{selected_subfolder}/actionTrace_Agent_PIC Blaze.csv')
#     df2 = pd.read_csv(f'data/{selected_subfolder}/SNLBlaze_acstate.csv')
#     df3 = pd.read_csv(f'data/{selected_subfolder}/GCCRaven_acstate.csv')

#     exclude = ['Detect_crash','Detect_conflict','Flight_Dynamics','Reroute_flight','Show_radar','Direct_to_waypoint','Change_heading']
#     events = parse_action_trace(f'data/{selected_subfolder}/actionTrace.csv', exclude)

def update_graph(selected_subfolder, selected_time, selected_variable):
    
    if selected_subfolder not in data_store:
        return go.Figure(), go.Figure(), go.Figure(layout=dict(autosize=True, width=None, height=300)), go.Figure()

    wnd_action_trace = data_store[selected_subfolder].get('wnd_action_trace', pd.DataFrame())
    all_action_trace = data_store[selected_subfolder].get('all_action_trace', [])
    aircraft_data = data_store[selected_subfolder].get('aircraft_data', {})

    filtered_wnd_action_trace = wnd_action_trace[(wnd_action_trace['time'] >= start_time) & (wnd_action_trace['time'] <= selected_time)]
    
    filtered_all_action_trace = []
    for event in all_action_trace:
        if (event['start'] - start_date).total_seconds() <= selected_time:
            event_copy = event.copy()  # Create a local copy of the event
            if (event_copy['end'] - start_date).total_seconds() > selected_time:
                event_copy['end'] = start_date + datetime.timedelta(seconds=selected_time)
            filtered_all_action_trace.append(event_copy)

    ac_data = {}
    for aircraft_id, df in aircraft_data.items():
        filtered_df = df[(df['time'] >= start_time) & (df['time'] <= selected_time)]
        ac_data[aircraft_id] = filtered_df

    fig1 = go.Figure(layout=dict(autosize=True, width=None, height=500))
    fig1['layout']['uirevision'] = 'Hello world!'
    fig1 = viz.plot_trajectory(fig1, ac_data)

    fig2 = go.Figure(layout=dict(autosize=True, width=700, height=700))
    fig2 = viz.wnd_visualization(fig2, filtered_wnd_action_trace, selected_time, start_time, end_time)
    fig2['layout']['uirevision'] = 'Hello world!'

    if selected_variable is not None:
        fig3 = go.Figure(layout=dict(autosize=True, width=None, height=300))
        for aircraft_id, df in ac_data.items():
            fig3.add_trace(go.Scatter(x=df['time'], y=df[selected_variable], mode='lines', name=aircraft_id))
        fig3.update_layout(xaxis_title='Real Time', yaxis_title=selected_variable, showlegend=True, xaxis=dict(range=[start_time, end_time]))#, yaxis=dict(range=[0, 1200])
        fig3['layout']['uirevision'] = 'Hello world!'

    else:
        fig3 = go.Figure(layout=dict(autosize=True, width=None, height=300))
    fig4 = go.Figure(layout=dict(autosize=True, width=None, height=700))
    fig4 = viz.timeline(fig4, filtered_all_action_trace, selected_time, start_date, end_date)
    fig4['layout']['uirevision'] = 'Hello world!'

    return fig2, fig1, fig3, fig4
    
# Run the app
if __name__ == '__main__':

    # webbrowser.open_new("http://127.0.0.1:8050/")
    app.run_server(debug=False, host='0.0.0.0', port=7860)