# -*- coding: utf-8 -*-
setVerbosity(MinimalLog)

# -------------------------------------------------------------
# Molecule Configuration
# -------------------------------------------------------------
# Define elements
elements = [Carbon, Oxygen, Oxygen, Oxygen, Carbon, Carbon, Hydrogen, Hydrogen,
            Hydrogen, Hydrogen]

# Define coordinates
cartesian_coordinates = [[ 1.609431413032,  0.050029424266, -0.200214465477],
                         [ 2.821872471776,  0.074345609581, -0.33257452873 ],
                         [ 0.786730069642,  0.285844748805, -1.270661995358],
                         [ 1.046027032925, -0.213547522321,  1.021790323208],
                         [-0.320536198865, -0.108625807588,  0.763421170161],
                         [-0.484549258604,  0.072564433479, -0.73758247059 ],
                         [-0.839575976013, -1.027933189299,  1.111755833993],
                         [-0.732343636487,  0.769867954492,  1.307351219363],
                         [-0.918484269576, -0.844342167145, -1.192644130163],
                         [-1.139565310547,  0.94179651573 , -0.964104219255]]*Angstrom

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

# Add tags
molecule_configuration.addTags('OPLS_CT_226', [4, 5])
molecule_configuration.addTags('OPLS_C_32',   [0])
molecule_configuration.addTags('OPLS_HC_63',  [6, 7, 8, 9])
molecule_configuration.addTags('OPLS_OS_22',  [2, 3])
molecule_configuration.addTags('OPLS_O_23',   [1])

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

# -------------------------------------------------------------
# Calculator
# -------------------------------------------------------------
numerical_accuracy_parameters = NumericalAccuracyParameters(
    density_mesh_cutoff=125.0*Hartree,
)

calculator = LCAOCalculator(
    numerical_accuracy_parameters=numerical_accuracy_parameters,
)

molecule_configuration.setCalculator(calculator)
nlprint(molecule_configuration)
molecule_configuration.update()
nlsave('ethylene_earbonate_charges.hdf5', molecule_configuration)

# -------------------------------------------------------------
# Optimize Geometry
# -------------------------------------------------------------
molecule_configuration = OptimizeGeometry(
    molecule_configuration,
    max_forces=0.05*eV/Ang,
    max_steps=200,
    max_step_length=0.2*Ang,
    trajectory_filename='ethylene_earbonate_charges_trajectory.hdf5',
    trajectory_interval=5.0*Minute,
    restart_strategy=RestartFromTrajectory(),
    optimize_cell=False,
    optimizer_method=LBFGS(),
    enable_optimization_stop_file=True,
)
nlsave('ethylene_earbonate_charges.hdf5', molecule_configuration)
nlprint(molecule_configuration)

# -------------------------------------------------------------
# Electrostatic Difference Potential
# -------------------------------------------------------------
electrostatic_difference_potential = ElectrostaticDifferencePotential(molecule_configuration)
nlsave('ethylene_earbonate_charges.hdf5', electrostatic_difference_potential)

# -------------------------------------------------------------
# Repeat ESP Partial Charges
# -------------------------------------------------------------
repeat_esp_partial_charges = RepeatESPPartialCharges(
    molecule_configuration,
    electrostatic_difference_potential,
    total_charge=0*elementary_charge,
    inner_vdw_shell=1.0,
    outer_vdw_shell=3.0,
    stride=1,
    tag_prefix='OPLS_',
    )
molecule_configuration.setPartialCharges(repeat_esp_partial_charges.evaluate())
nlsave('ethylene_earbonate_charges.hdf5', molecule_configuration)
nlsave('ethylene_earbonate_charges.hdf5', repeat_esp_partial_charges)