# -------------------------------------------------------------
# Initial Configuration
# -------------------------------------------------------------
bulk_configuration = nlread('Polystyrene.hdf5')[-1]

# -------------------------------------------------------------
# OPLS-AA Calculator
# -------------------------------------------------------------
potential_builder = OPLSPotentialBuilder()
calculator = potential_builder.createCalculator(bulk_configuration)
bulk_configuration.setCalculator(calculator)

# -------------------------------------------------------------
# Molecular Dynamics
# -------------------------------------------------------------
initial_velocity = ConfigurationVelocities(
    remove_center_of_mass_momentum=True
)

method = NPTMartynaTobiasKlein(
    time_step=1*femtoSecond,
    reservoir_temperature=300*Kelvin,
    reservoir_pressure=1*bar,
    thermostat_timescale=100*femtoSecond,
    barostat_timescale=500*femtoSecond,
    initial_velocity=initial_velocity,
    heating_rate=0*Kelvin/picoSecond,
)

constraints = [FixCenterOfMass()]

md_trajectory = MolecularDynamics(
    bulk_configuration,
    constraints=constraints,
    trajectory_filename='Polystyrene_Trajectory.hdf5',
    steps=100000,
    log_interval=10,
    trajectory_interval=1000,
    method=method
)

# -------------------------------------------------------------
# Calculate End to End Distance
# -------------------------------------------------------------
analyzer = EndToEndDistances(
      md_trajectory,
      start_time=0*ps,
      end_time=100*ps,
      end_tags=['HEAD_CONNECT', 'TAIL_CONNECT']
)
end_end_distance = analyzer.distances()
end_end_distance_average = end_end_distance.mean(axis=1)
end_end_correlation = analyzer.autocorrelation().mean(axis=0)
time = analyzer.times().convertTo(ps)

# -------------------------------------------------------------
# Plot The Results
# -------------------------------------------------------------
import pylab
pylab.figure()
pylab.plot(time, end_end_distance[:, 0], label='Polymer 1')
pylab.plot(time, end_end_distance[:, 5], label='Polymer 6')
pylab.plot(time, end_end_distance_average, label='Average')
pylab.xlabel('Time / ps')
pylab.ylabel('End to end distance / Angstrom')
pylab.legend()
pylab.savefig('End_To_End_Distances_Plot.png')

pylab.figure()
pylab.plot(time[:51], end_end_correlation[:51])
pylab.xlabel('Time / ps')
pylab.ylabel('End to end autocorrelation')
pylab.savefig('End_To_End_Autocorrelation_Plot.png')