#!/usr/bin/env python3
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
import subprocess

start_time = 0#240
end_time = 1700#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 Vista.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.B("This is an interactive demonstration of case study results to showcase the ability of modeling and simulation to help envision distributed work systems"),
        html.Br(),
        dcc.Markdown("This interactive demo was performed under NASA contract number 80NSSC23CA121.",  id="explanation")
    ], style={'width': '100%', 'border': '1px solid gray'}),
    html.Div([
        dcc.Markdown("It is envisioned that in future operatoins, each UAV is going to be controlled using a command and control (C2) link that allows communication with the remote pilot. This architecture can have potentially dangerous implications in the case of a failure of the command and control link. The UAV's are envisioned to be programmed with one or more contingency procedures that are to be executed in the case of a lost link. For example, to divert to and land at the nearest vertiport or continue along its preapproved route. This brings the question of how will the rest of the system need to react to the change in the aircraft trajectory. Is a system like this robust to an event like loss of link? ",  id="explanation2")
    ], style={'width': '100%', 'border': '1px solid gray'}),
    html.Div([
        dcc.Markdown("If the dropdown menu is empty please wait a few seconds to allow the data to load. If the data doesnt load then click the submit button to force reload the dropdown. The files that contain data will take a few seconds to load so please be patient.",  id="explanation"),
            html.Button('Submit', id='submit-button', n_clicks=0)
    ], style={'width': '100%', 'border': '1px solid gray'}),
    html.Div([
#        html.Div([
#            html.Div("Select the desired input conditions, then click 'Simulate!':", id='text1'),
#            html.Label("Work Allocation: "),
#            dcc.Input( id='work-allocation', type='text', placeholder='The path to the CSV file that has the work allocation.', style={'width': '80%'}),
#            html.Br(),
#            html.Label("Contingency Time: "),
#            dcc.Input( id='cont-time', type='text', placeholder='Time at which the lost link event occurs (set to "ALONGTIMEAWAY" (no quotes) if no lost link event needs to occur)', style={'width': '80%'}),
#            html.Br(),
#            html.Label("Contingency Aircraft Response: "),
#            dcc.RadioItems(
#                id='ConAircraftResponse',
#                options=[
#                    {'label': 'Divert', 'value': '1'},
#                    {'label': 'Continue', 'value': '2'}
#                ],
#                value=True,
#                labelStyle={'display': 'inline-block'}
#            ),
#            html.Br(),
#            html.Label("Other Aircraft Response: "),
#            dcc.RadioItems(
#                id='OtherAircraftResponse',
#                options=[
#                    {'label': 'Increase Separation', 'value': '1'},
#                    {'label': 'Continue', 'value': '2'}
#                ],
#                value=True,
#                labelStyle={'display': 'inline-block'}
#            ),
#            
#            html.Br(),
#            html.Button('Simulate!', id='button', n_clicks=0),
#        ], style={'width': '50%', 'border': '1px solid gray'}),
        html.Div([
#            html.Div("Next, input the relative path to your results folder, then click 'Submit'. This will load the data:", id='text2'),
#            dcc.Input(
#                id='path-input',
#                # value='data',
#                type='text',
#                placeholder='Type your path... For example, ../wmc5.1/Scenario/AAMv2/Results',
#                style={'width': '70%'}  # Make the input wider
#            ),
#            html.Button('Submit', id='submit-button', n_clicks=0),
            # Create a dropdown menu with the subfolders as options
            html.Div("Select which condition to load in the interactive graphs:", id='text2'),
            dcc.Dropdown(
                id='subfolder-dropdown',
                value=None #subfolders[0]  # Default value
                ),
            dcc.Interval(
                id='interval-component',  # Unique ID for the interval component
                interval=1*1000,  # In milliseconds (1 second)
                n_intervals=0,  # Start at 0, increment once
                max_intervals=1  # Fire only once
            )
        ], style={'width': '100%', 'border': '1px solid gray'}),
    ], style={'display': 'flex'}),

    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': '40%', 'display': 'inline-block', 'border': '1px solid gray'}),
        html.Div([
            dcc.Graph(id='timeline'),
        ], style={'width': '60%', '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('button', 'n_clicks'),  # Add this output to reset button clicks
    [dash.dependencies.Input('work-allocation', 'value')],
    [dash.dependencies.Input('cont-time', 'value')],
    [dash.dependencies.Input('ConAircraftResponse', 'value')],
    [dash.dependencies.Input('OtherAircraftResponse', 'value')],
    [dash.dependencies.Input('button', 'n_clicks')],    
)
def run_wmc(workAllocation, selected_time, ConAircraftResponse, OtherAircraftResponse, nc_clicks):
    if nc_clicks > 0:
        # try:
        print("Running WMC with input arguments",workAllocation,selected_time,ConAircraftResponse+OtherAircraftResponse)
        # Create a string representing the command to be run
        command = f"wmc/AAMv2 {workAllocation} {selected_time} {ConAircraftResponse + OtherAircraftResponse} 4000"

        # Write the command to a wrapper.sh file
        with open("wmc/wrapper.sh", "w") as file:
            file.write("#!/bin/bash\n")
            file.write(command + "\n")

        # Make the wrapper.sh script executable
        subprocess.run(["chmod", "+x", "wmc/wrapper.sh"])

        # Run the wrapper.sh file
        process = subprocess.Popen(["wmc/wrapper.sh"])
        #process = subprocess.Popen(['wmc/AAMv2', workAllocation, selected_time, ConAircraftResponse+OtherAircraftResponse, "4000"])        # return f"WMC was run!"
        # process = subprocess.Popen(['./UAM', 'Example', '240', 'True'])        # return f"WMC was run!"
        try:
            process.wait(timeout=50)  # Wait for the process to complete
        except subprocess.TimeoutExpired:
            process.terminate()  # Terminate the process after timeout
        # Continue with the rest of the code
        print("Completed running WMC...")
    
    return 0
        # except:
            # return f"error!"
    
# 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.Output('submit-button', 'n_clicks'),  # Add this output to reset submit button clicks
    [dash.dependencies.Input('interval-component', 'n_intervals'),
     dash.dependencies.Input('submit-button', 'n_clicks')],
#    [dash.dependencies.Input('path-input', 'value')],
)
def load_data(n_intervals, n_clicks):
    selected_path = "Results"
    print("selected path is",selected_path)
    global data_store

    if n_intervals > 0 or 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, 0
    else:
        return [], [], 0

# 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()#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=1200, height=1000))
    # 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=2100))
    fig4 = viz.timeline(fig4, filtered_all_action_trace, selected_time, start_date, end_date)
    fig4['layout']['uirevision'] = 'Hello world!'
    fig4.update_layout(height=900)  # Explicitly set the height again

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

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