# Load in an equilibrated initial configuration with calculator
bulk_configuration = nlread('Battery_Electrolyte.hdf5')[-1]

# Calculate the trajectory, including charge current measurements
method = NVTNoseHoover(
    time_step=1*femtoSecond,
    reservoir_temperature=340*Kelvin,
    thermostat_timescale=100*femtoSecond,
    heating_rate=0*Kelvin/picoSecond,
    chain_length=3,
    initial_velocity=ConfigurationVelocities(remove_center_of_mass_momentum=True),
)

constraints = [FixCenterOfMass()]
measurement_names = ['charge_current', 'charge_current_components']

md_trajectory = MolecularDynamics(
    bulk_configuration,
    constraints=constraints,
    trajectory_filename='Battery_Electrolyte_Trajectory.hdf5',
    steps=10000000,
    trajectory_interval=10000,
    measurement_hook=MDMeasurement(measurement_names, call_interval=10),
    method=method
)

# Calculate the ionic conductivity
self_diffusion = IonicConductivity(
    md_trajectory,
    atom_selection=Lithium
)

# Get the times in ps and the conductivity values in Siemens/Meter.
times        = self_diffusion.times()
conductivity = self_diffusion.data()

# Plot the diffusion coefficient calculated with different time limits
model = Plot.PlotModel(ps, Siemens/Meter)
model.title().setText('Ionic Conductivity of Lithium Electrolyte')
model.xAxis().setLabel('Time')
model.yAxis().setLabel('Conductivity')
model.legend().setVisible(False)

# Add diffusion integral
line = Plot.Line(times, conductivity)
line.setColor('blue')

# Add average
average = Plot.Average()
line.addItem(average)
max_time = times[-1]
average.setBounds(max_time*0.1, max_time*0.9)
average.setColor('red')

# Add line to the plot
model.addItem(line)

# Set limits and show the plot
model.setLimits()
Plot.show(model)