import numpy as np
import gradio as gr
import pickle

import plotly.graph_objects as go
import plotly.express as px

import pandas as pd

# Remove the global variable and instead use gr.State
# define a Dataframe styler to highlight the Feasible? column to be green and red
def highlight_success(val):
    color = 'lightgreen' if val == 'Success' else 'lightcoral'
    return f'color:white;background-color: {color}'

def sim_espray_constrained(x, noise_se=None):
    # Define the equations
    conc = x[:, 0]
    flow_rate = x[:, 1]
    voltage = x[:, 2]
    solvent = x[:, 3]
    diameter = (np.sqrt(conc) * np.sqrt(flow_rate)) / np.log2(voltage) * 10 + 0.4 + solvent  # Diameter in micrometers
    if noise_se is not None:
        diameter = diameter + noise_se * np.random.randn(*diameter.shape)
    exp_con = (np.log(flow_rate) * (solvent - 0.5) + 1.40 >= 0).astype(float)
    return np.column_stack((diameter, exp_con))

X_init = np.array([[0.5, 15, 10, 0],
                   [0.5, 0.1, 10, 1],
                   [3, 20, 15, 0],
                   [1, 20, 10, 1],
                   [0.2, 0.02, 10, 1]])

Y_init = sim_espray_constrained(X_init)
exp_record_df = pd.DataFrame(X_init, columns=['Concentration (%w/v)', 'Flow Rate (mL/h)', 'Voltage (kV)', 'Solvent'])
exp_record_df['Size (um)'] = Y_init[:, 0]
# map 1 to CHCl3 and 0 to DMAc
exp_record_df['Solvent'] = ['DMAc' if x == 0 else 'CHCl3' for x in exp_record_df['Solvent']]
exp_record_df['Feasible?'] = ['Success' if x == 1 else 'Failed' for x in Y_init[:, 1]]

gr_exp_record_df = gr.DataFrame(value=exp_record_df.style.map(highlight_success, subset=['Feasible?']).format(precision=3), label="Prior Experiments")

def import_results():
    strategies = ['qEI', 'qEI_vi_mixed_con', 'qEICF_vi_mixed_con', 'rnd']
    
    # Load results from pickle file
    with open('best_distances.pkl', 'rb') as f:
        best_distances = pickle.load(f)

    # vstack all values in best_distances
    best_distances_vstack = {k: np.vstack(best_distances[k]) for k in strategies}
    best_distances_all_trials = -np.vstack([best_distances_vstack[k] for k in strategies])
    best_distances_all_trials_df = pd.DataFrame(best_distances_all_trials)
    best_distances_all_trials_df['strategy'] = np.repeat(['Vanilla BO', 'Constrained BO', 'CCBO', 'Random'], 20)
    best_distances_all_trials_df['trial'] = list(range(20)) * len(strategies)

    best_distances_df_long = pd.melt(best_distances_all_trials_df, id_vars=['strategy', 'trial'], var_name='iteration', value_name='regret')
    return best_distances_df_long

def calc_human_performance(df):
    # convert back solvent to 0 and 1
    df['Solvent'] = [0 if x == 'DMAc' else 1 for x in df['Solvent']]
    
    TARGET_SIZE = 3.0  # Example target size
    ROUNDS = len(df) // 2

    X_human = df[['Concentration (%w/v)', 'Flow Rate (mL/h)', 'Voltage (kV)', 'Solvent']].values

    X_human_init = X_init.copy()
    Y_human_init = Y_init.copy()
    
    best_human_distance = []

    for iter in range(ROUNDS + 1):
        Y_distance = -np.abs(Y_human_init[:, 0] - TARGET_SIZE)
        best_human_distance.append(np.ma.masked_array(Y_distance, mask=~Y_human_init[:, 1].astype(bool)).max())
        new_x = X_human[2 * iter:2 * (iter + 1)]
        X_human_init = np.vstack([X_human_init, new_x])
        Y_human_init = np.vstack([Y_human_init, sim_espray_constrained(new_x)])

    return -np.array(best_human_distance)

def plot_results(exp_data_df):
    # Extract human performance
    best_human_distance = calc_human_performance(exp_data_df)
    
    # Import results
    best_distances_df_long = import_results()
    
    fig = go.Figure()
    
    strategies = best_distances_df_long['strategy'].unique()
    
    for strategy in strategies:
        strategy_data = best_distances_df_long[best_distances_df_long['strategy'] == strategy]
        
        # Calculate mean and standard deviation
        mean_regret = strategy_data.groupby('iteration')['regret'].mean()
        std_regret = strategy_data.groupby('iteration')['regret'].std()
        
        iterations = mean_regret.index
        color = px.colors.qualitative.Set2[strategies.tolist().index(strategy)]
        
        # Add trace for mean line
        mean_trace = go.Scatter(
            x=iterations,
            y=mean_regret,
            mode='lines',
            name=strategy,
            line=dict(width=2, color=color)
        )
        fig.add_trace(mean_trace)
        
        # Add trace for shaded area (standard deviation)
        fig.add_trace(go.Scatter(
            x=list(iterations) + list(iterations[::-1]),
            y=list(mean_regret + std_regret) + list((mean_regret - std_regret)[::-1]),
            fill='toself',
            fillcolor=mean_trace.line.color.replace('rgb', 'rgba').replace(')', ',0.2)'),
            line=dict(color='rgba(255,255,255,0)'),
            showlegend=False,
            name=f'{strategy} (std dev)'
        ))
    # Add trace for human performance
    fig.add_trace(go.Scatter(
        x=list(range(len(best_human_distance))),
        y=best_human_distance,
        mode='lines+markers',
        name='Human',
        line=dict(width=2, color='brown')
    ))

    fig.update_layout(
        title='Performance Comparison',
        xaxis_title='Iteration',
        yaxis_title='Regret (μm)',
        legend_title='Strategy',
        template='plotly_white',
        legend=dict(
            x=0.01,
            y=0.01,
            bgcolor='rgba(255, 255, 255, 0.5)',
            bordercolor='rgba(0, 0, 0, 0.5)',
            borderwidth=1
        )
    )
    
    return fig

# Update predict function to use state
def predict(state, text1, conc1, flow_rate1, voltage1, solvent1, text2, conc2, flow_rate2, voltage2, solvent2):
    # Get current results storage from state or initialize if None
    if state is None:
        results_storage = pd.DataFrame(columns=['Concentration (%w/v)', 'Flow Rate (mL/h)', 'Voltage (kV)', 'Solvent', 'Size (um)', 'Feasible?'])
    else:
        results_storage = state.copy()
        
    solvent_value1 = 0 if solvent1 == 'DMAc' else 1
    solvent_value2 = 0 if solvent2 == 'DMAc' else 1
    
    # Convert inputs to numpy array
    inputs1 = np.array([[conc1, flow_rate1, voltage1, solvent_value1]])
    inputs2 = np.array([[conc2, flow_rate2, voltage2, solvent_value2]])

    # Get predictions
    results1 = sim_espray_constrained(inputs1)
    results2 = sim_espray_constrained(inputs2)
    
    # Format output
    diameter1 = results1[0, 0]
    exp_con1 = results1[0, 1]
    
    diameter2 = results2[0, 0]
    exp_con2 = results2[0, 1]

    # create a dataframe to store the results
    results_df = pd.DataFrame(np.array([
        [conc1, flow_rate1, voltage1, solvent_value1, diameter1, exp_con1],
        [conc2, flow_rate2, voltage2, solvent_value2, diameter2, exp_con2]
    ]), columns=['Concentration (%w/v)', 'Flow Rate (mL/h)', 'Voltage (kV)', 'Solvent', 'Size (um)', 'Feasible?'])

    results_df['Solvent'] = ['DMAc' if x == 0 else 'CHCl3' for x in results_df['Solvent']]
    results_df['Feasible?'] = ['Success' if x == 1 else 'Failed' for x in results_df['Feasible?']]
    
    results_storage = pd.concat([results_storage, results_df], ignore_index=True)
    results_display = results_storage.style.map(highlight_success, subset=['Feasible?']).format(precision=3)
   
    return (results_storage, gr_exp_record_df, gr.DataFrame(value=results_display, label="Your Results"), plot_results(results_storage))

# Update reset_results to use state
def reset_results(state):
    results_storage = pd.DataFrame(columns=['Concentration (%w/v)', 'Flow Rate (mL/h)', 'Voltage (kV)', 'Solvent', 'Size (um)', 'Feasible?'])
    return (results_storage, gr_exp_record_df, gr.DataFrame(value=results_storage.style.map(highlight_success, subset=['Feasible?']).format(precision=3), label="Your Results"), plot_results(results_storage))

inputs = [
    gr.Markdown("### Experiment 1"),
    gr.Number(value=1.2, label="Concentration (%w/v, range: 0.05-5.00)", minimum=0.05, maximum=5.0, precision=3),
    gr.Number(value=20.0, label="Flow Rate (mL/h, range: 0.01-60.00)", minimum=0.01, maximum=60.0, precision=3),
    gr.Number(value=15.0, label="Voltage (kV, range: 10.00-18.00)", minimum=10.0, maximum=18.0, precision=3),
    gr.Dropdown(['DMAc', 'CHCl3'], value='DMAc', label='Solvent'),
    gr.Markdown("### Experiment 2"),
    gr.Number(value=2.8, label="Concentration (%w/v, range: 0.05-5.00)", minimum=0.05, maximum=5.0, precision=3),
    gr.Number(value=20.0, label="Flow Rate (mL/h, range: 0.01-60.00)", minimum=0.01, maximum=60.0, precision=3),
    gr.Number(value=15.0, label="Voltage (kV, 10.00-18.00)", minimum=10.0, maximum=18.0, precision=3),
    gr.Dropdown(['DMAc', 'CHCl3'], value='CHCl3', label='Solvent')
]

outputs = [gr_exp_record_df, gr.DataFrame(label="Your Results"), gr.Plot(label="Performance Comparison")]

description = "<h3>Welcome, challenger!</h3><p> If you think you may perform better than <strong>CCBO</strong>, try this interactive game to optimize electrospray! Rules are simple: <ul><li>Examine carefully the initial experiments you have on the right (or below if you're using your phone), remember, your target size is <u><i><strong>3.000 um</strong></i></u> ----></li><li>Select your parameters, you have <strong>2</strong> experiments (chances) in each round, use them wisely! </li><li>Click <strong>Submit</strong> to see the results, reflect and improve your selection!</li><li>Repeat the process for <strong>5</strong> rounds to see if you can beat CCBO!</li></ul></p>"

# Create a Blocks interface instead of using Interface
with gr.Blocks() as demo:
    
    gr.Markdown("## Human vs CCBO Campaign - Simulated Electrospray")
    gr.Markdown(description)
    
    # Add state component to store user-specific results
    results_state = gr.State()
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("### Experiment 1")
            conc1 = gr.Number(value=1.2, label="Concentration (%w/v, range: 0.05-5.00)", minimum=0.05, maximum=5.0, precision=3)
            flow_rate1 = gr.Number(value=20.0, label="Flow Rate (mL/h, range: 0.01-60.00)", minimum=0.01, maximum=60.0, precision=3)
            voltage1 = gr.Number(value=15.0, label="Voltage (kV, range: 10.00-18.00)", minimum=10.0, maximum=18.0, precision=3)
            solvent1 = gr.Dropdown(['DMAc', 'CHCl3'], value='DMAc', label='Solvent')
            
            gr.Markdown("### Experiment 2")
            conc2 = gr.Number(value=2.8, label="Concentration (%w/v, range: 0.05-5.00)", minimum=0.05, maximum=5.0, precision=3)
            flow_rate2 = gr.Number(value=20.0, label="Flow Rate (mL/h, range: 0.01-60.00)", minimum=0.01, maximum=60.0, precision=3)
            voltage2 = gr.Number(value=15.0, label="Voltage (kV, 10.00-18.00)", minimum=10.0, maximum=18.0, precision=3)
            solvent2 = gr.Dropdown(['DMAc', 'CHCl3'], value='CHCl3', label='Solvent')
            
        with gr.Column():
            prior_experiments = gr.DataFrame(value=exp_record_df.style.map(highlight_success, subset=['Feasible?']).format(precision=3), label="Prior Experiments")
            results_df = gr.DataFrame(label="Your Results")
            perf_plot = gr.Plot(label="Performance Comparison")
    
    with gr.Row():
        submit_btn = gr.Button("Submit")
        reset_btn = gr.Button("Reset Results")
    
    # Add invisible text input components to match the predict function signature
    text1 = gr.Textbox(visible=False)
    text2 = gr.Textbox(visible=False)
    
    # Connect the submit button to the predict function
    submit_btn.click(
        fn=predict,
        inputs=[
            results_state,
            text1, conc1, flow_rate1, voltage1, solvent1,
            text2, conc2, flow_rate2, voltage2, solvent2
        ],
        outputs=[results_state, prior_experiments, results_df, perf_plot]
    )
    
    # Connect the reset button to the reset_results function
    reset_btn.click(
        fn=reset_results,
        inputs=[results_state],
        outputs=[results_state, prior_experiments, results_df, perf_plot]
    )

demo.launch()