import sys
import numpy
import networkx

from NL.IO.NLSaveUtilities import nlread

from NL.CommonConcepts.PhysicalQuantity import atomic_mass_unit

# Open the trajectory
traj = nlread(sys.argv[1])[-1]

# Allocate necessary data arrays
reaction_percent = numpy.zeros(len(traj))
avg_mass = numpy.zeros(len(traj)) * atomic_mass_unit
red_mass = numpy.zeros(len(traj)) * atomic_mass_unit

# For each image in the trajectory
for i in range(len(traj)):
    conf = traj.image(i)
    graph = conf._bondGraph()

    # Get the total number of molecules
    total_index = []
    for sub_graph in networkx.connected_component_subgraphs(graph):
        total_index.append(sub_graph.nodes())

    # Get the molecular mass of each molecule
    mass_array = numpy.zeros(len(total_index)) * atomic_mass_unit
    for j, idx in enumerate(total_index):
        mass_array[j] = (conf.atomicMasses()[idx]).sum()
    mass_array.sort()

    # Save the data to ploy
    avg_mass[i] = mass_array.mean()
    red_mass[i] = mass_array[:-1].mean()
    reaction_percent[i] = traj._getQuantity('Reaction_Complete', i)

# Plot the graph
line_a = Plot.Line(reaction_percent, avg_mass)
line_a.setLabel('Avg. molecular mass')
line_a.setColor('red')
line_a.setLineWidth(1)

line_b = Plot.Line(reaction_percent, red_mass)
line_b.setLabel('Avg. reduced mass')
line_b.setColor('blue')
line_b.setLineWidth(1)

model_a = Plot.PlotModel(y_unit=atomic_mass_unit)
model_a.framing().setTitle('Thermoset molecular mass`')
model_a.xAxis().setLabel('Reaction completion')
model_a.yAxis().setLabel('Molecular mass')
model_a.legend().setVisible(True)
model_a.legend().setLocation('upper left')
model_a.addItem(line_a)
model_a.setLimits()

model_b = Plot.PlotModel(y_unit=atomic_mass_unit)
model_b.xAxis().setLabel('')
model_b.yAxis().setLabel('Reduced mass')
model_b.yAxis().setMirrored(True)
model_b.legend().setVisible(True)
model_b.legend().setLocation('center left')
model_b.addItem(line_b)
model_b.setLimits()

layout = Plot.OverlayLayout()
layout.setMode(Plot.LAYOUT_MODES.SHARE_X)

# Add main model.
frame_a = Plot.PlotFrame(use_frame_model=False)
frame_a.addModel(model_a)
layout.addItem(frame_a)

frame_b = Plot.PlotFrame(use_frame_model=False)
frame_b.addModel(model_b)
layout.addItem(frame_b)

# Show the plot.
Plot.show(layout)

# Save plot.
Plot.save(layout, 'molecular_mass.png')