# -*- coding: utf-8 -*-
# -------------------------------------------------------------
# Polymer Sequences
# -------------------------------------------------------------

# ---------------------------------------------
# Monomer Molecule 1
# ---------------------------------------------

# Define elements
elements = [Carbon, Hydrogen, Hydrogen, Hydrogen, Carbon, Hydrogen, Hydrogen,
            Chlorine]

# Define coordinates
cartesian_coordinates = [[ 0.082573986725, -0.715837372657, -1.049741055856],
                         [ 1.200632880086, -0.741832973127, -1.084132698825],
                         [-0.307857413992, -0.784975505075, -2.095797344985],
                         [-0.281106687902, -1.604219748865, -0.47524813875 ],
                         [-0.418374973727,  0.632715291686, -0.357887724741],
                         [-0.129477342652,  1.524334868219, -0.972967818507],
                         [-1.538044977329,  0.661945827687, -0.309514512742],
                         [ 0.280758705109,  0.751777559344,  1.28562566904 ]]*Angstrom

# Set up configuration
monomer_1 = MoleculeConfiguration(
    elements=elements,
    cartesian_coordinates=cartesian_coordinates
    )

# Add tags
monomer_1.addTags('ENDH_OPLS_CT_4',   [0])
monomer_1.addTags('ENDT_OPLS_CT_168', [4])
monomer_1.addTags('HEAD_CONNECT',     [0])
monomer_1.addTags('HEAD_GROUP_A',     [1])
monomer_1.addTags('OPLS_CT_169',      [4])
monomer_1.addTags('OPLS_CT_5',        [0])
monomer_1.addTags('OPLS_Cl_9',        [7])
monomer_1.addTags('OPLS_HC_1',        [1, 2, 3])
monomer_1.addTags('OPLS_HC_54',       [5, 6])
monomer_1.addTags('TAIL_CONNECT',     [4])
monomer_1.addTags('TAIL_GROUP_A',     [6])
monomer_1.addTags('TAIL_GROUP_B',     [5])

# Add bonds
bonds = [[0, 1],
         [0, 2],
         [0, 3],
         [0, 4],
         [4, 5],
         [4, 6],
         [4, 7]]
monomer_1.setBonds(bonds)


# Create the list of configurations
monomer_list = [
    monomer_1,
]

# Create the polymer sequence object
polymer_sequence_a = PolymerSequence(
    monomer_configurations=monomer_list,
    number_of_monomers=20,
    number_of_chains=20,
    tactic_ratio=0.5,
)

polymer_sequences = [polymer_sequence_a]

# -------------------------------------------------------------
# Polymer Monte Carlo Builder
# -------------------------------------------------------------
polymer_builder = PolymerMonteCarloBuilder(
    polymer_sequence=polymer_sequences,
    density=1.3000*gram/cm**3,
    monte_carlo_temperature=300.0*Kelvin,
    angle_sampling_points=20,
    repack_molecules=False,
    )
bulk_configuration = polymer_builder.buildPolymerConfiguration()

# Repeat the system along the c-axis to increase the distance between the heat source and heat
# sink.
bulk_configuration = bulk_configuration.repeat(1, 1, 3)

# -------------------------------------------------------------
# Calculator
# -------------------------------------------------------------
potential_builder = DreidingPotentialBuilder(
    include_electrostatic=False,
    )
calculator = potential_builder.createCalculator(bulk_configuration)
bulk_configuration.setCalculator(calculator)
nlsave('nemd_polymer.hdf5', bulk_configuration)

# -------------------------------------------------------------
# Force Capped Equilibration
# -------------------------------------------------------------
equilibration_method = ForceCappedEquilibration(
    temperature=300.00*Kelvin,
    md_steps_per_force_capped_simulation=40000,
    md_time_step=0.50*fs,
    force_capped_simulations=4,
    starting_factor=1.040,
    ending_factor=0.800,
    fixed_indices=None,
    )
bulk_configuration = equilibration_method.runEquilibration(bulk_configuration)
nlsave('nemd_polymer.hdf5', bulk_configuration)

# -------------------------------------------------------------
# Molecular Dynamics
# -------------------------------------------------------------

initial_velocity = MaxwellBoltzmannDistribution(
    temperature=300.0*Kelvin,
    remove_center_of_mass_momentum=True
)

method = NVEVelocityVerlet(
    time_step=1*femtoSecond,
    initial_velocity=initial_velocity,
)

momentum_exchange_hook = NonEquilibriumMomentumExchange(
    configuration=bulk_configuration,
    exchange_interval=100,
    heat_source=(0.0, 10.0*Angstrom),
    heat_sink=(0.5, 10.0*Angstrom),
    update_profile_interval=0,
    profile_resolution=2.0*Ang
)

constraints = [FixCenterOfMass()]

md_trajectory = MolecularDynamics(
    bulk_configuration,
    constraints=constraints,
    trajectory_filename='nemd_polymer.hdf5',
    steps=50000,
    log_interval=100,
    measurement_hook=momentum_exchange_hook,
    method=method
)

bulk_configuration = md_trajectory.lastImage()
nlprint(momentum_exchange_hook)