import numpy as np
import nglview
from bokeh.plotting import figure, output_notebook, show
from bokeh.layouts import row, column
import simtk.unit as unit
import mdtraj as md
import plotting_tools
output_notebook()

"""
Explanation: Nonequilibrium Switching Free Energy Analysis
This notebook is useful for examining the results of nonequilibrium switching calculations performed with the Perses tool. It enables examination of work traces as well as trajectories.
Import required libraries
End of explanation
"""


trajectory_directory = "/Users/grinawap/solvent_test_4"
trajectory_prefix = "cdk02" #these were set by the yaml file

data_reader = plotting_tools.NonequilibriumSwitchingAnalysis(trajectory_directory, trajectory_prefix)
cum_work_shape = np.shape(data_reader.cumulative_work)
n_steps_per_protocol = cum_work_shape[2]
protocol_steps = np.linspace(0, 1, n_steps_per_protocol)
n_protocols = cum_work_shape[1]

"""
Explanation: Set the appropriate variables to the directories where the output files are stored
End of explanation
"""


#create a convenience function to add lines to a bokeh plot
def plot_line(fig, x, y):
    fig.line(x, y)

def plot_work_values(protocol_steps, data_reader, n_protocols):
    """
    Convenience function to plot work values from nonequilibrium switching calculations.
    
    Arguments
    ---------
    protocol_steps : [n_steps] array
        Indexes the lambda value at each point in the corresponding work trajectory
    data_reader : NonequilibriumSwitchingAnalysis object
        object that contains the switching data
    n_protocols : int
        The total number of protocols in each direction (should be symmetric)
    """
    x_axis_label = "lambda value"
    y_axis_label = "Work (kT)"
    
    work_figure = figure(title="Nonequilibrium Switching Work", x_axis_label=x_axis_label,
                        y_axis_label=y_axis_label)
    
    for protocol_index in range(n_protocols):
        work_figure.line(protocol_steps, data_reader.cumulative_work[0, protocol_index, :], line_color="green")
        work_figure.line(protocol_steps, data_reader.cumulative_work[1, protocol_index, :], line_color="red")
    
    show(work_figure)

plot_work_values(protocol_steps, data_reader, n_protocols)

"""
Explanation: Now we can plot the forward and reverse work traces
These are the cumulative work going from lambda=0 to lambda=1, and the reverse
End of explanation
"""


nonequilibrium_trajectory = data_reader.get_nonequilibrium_trajectory("forward", 0)
view = nglview.NGLWidget()
view.add_trajectory(nonequilibrium_trajectory)
view.add_ball_and_stick()

view

"""
Explanation: What should I be looking for in the work traces?
Is there a place in the trajectory where the work suddenly becomes anomalously high? This might indicate either a bad protocol or a bug in the code
In general, these plots are giving an illustration of how fast the system is changing as lambda is changed. The faster that change is, the more unfavorable the variance of the ultimate free energy calculation
Let's take a closer look at the trajectories
See something weird? Sometimes it helps to just visualize the trajectory of the atoms. We can do that fairly straightforwardly:
End of explanation
"""


