# %% MoleculeConfiguration

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

# Define coordinates
cartesian_coordinates = [[  0.            ,   0.            ,   0.            ],
                         [  1.513         ,   0.            ,   0.            ],
                         [ -0.363333307269,   0.            ,  -1.027661781602],
                         [ -0.363333307269,  -0.889981209366,   0.513830890801],
                         [ -0.363333307269,   0.889981209366,   0.513830890801],
                         [  1.924666666667,   1.008373277446,  -0.582184583177],
                         [  2.017333333333,   0.            ,   1.426470079914],
                         [  1.876333307269,  -0.889981209366,  -0.513830890801],
                         [  3.530333333333,   0.            ,   1.426470079914],
                         [  1.654000026065,   0.889981209366,   1.940300970715],
                         [  1.654000026065,  -0.889981209366,   1.940300970715],
                         [  3.893666640602,   0.889981209366,   0.912639189113],
                         [  4.034666666667,   0.            ,   2.852940159827],
                         [  3.942         ,  -1.008373277446,   0.844285496737],
                         [  5.547666666667,   0.            ,   2.852940159827],
                         [  3.671333359398,   0.889981209366,   3.366771050628],
                         [  3.671333359398,  -0.889981209366,   3.366771050628],
                         [  5.959333333333,   1.008373277446,   2.27075557665 ],
                         [  6.052         ,   0.            ,   4.279410239741],
                         [  5.910999973935,  -0.889981209366,   2.339109269026],
                         [  7.565         ,   0.            ,   4.279410239741],
                         [  5.688666692731,   0.889981209366,   4.793241130542],
                         [  5.688666692731,  -0.889981209366,   4.793241130542],
                         [  7.928333307269,   0.889981209366,   3.76557934894 ],
                         [  8.069333333333,   0.            ,   5.705880319655],
                         [  7.976666666667,  -1.008373277446,   3.697225656564],
                         [  9.582333333333,   0.            ,   5.705880319655],
                         [  7.706000026065,   0.889981209366,   6.219711210456],
                         [  7.706000026065,  -0.889981209366,   6.219711210456],
                         [  9.994         ,   1.008373277446,   5.123695736478],
                         [ 10.086666666667,   0.            ,   7.132350399568],
                         [  9.945666640602,  -0.889981209366,   5.192049428854],
                         [ 11.599666666667,   0.            ,   7.132350399568],
                         [  9.723333359398,   0.889981209366,   7.646181290369],
                         [  9.723333359398,  -0.889981209366,   7.646181290369],
                         [ 11.962999973935,   0.889981209366,   6.618519508767],
                         [ 12.104         ,   0.            ,   8.558820479482],
                         [ 12.011333333333,  -1.008373277446,   6.550165816391],
                         [ 13.617         ,   0.            ,   8.558820479482],
                         [ 11.740666692731,   0.889981209366,   9.072651370283],
                         [ 11.740666692731,  -0.889981209366,   9.072651370283],
                         [ 14.028666666667,   1.008373277446,   7.976635896305],
                         [ 14.121333333333,   0.            ,   9.985290559396],
                         [ 13.980333307269,  -0.889981209366,   8.044989588681],
                         [ 15.634333333333,   0.            ,   9.985290559396],
                         [ 13.758000026065,   0.889981209366,  10.499121450197],
                         [ 13.758000026065,  -0.889981209366,  10.499121450197],
                         [ 15.997666640602,   0.889981209366,   9.471459668595],
                         [ 16.138666666667,   0.            ,  11.411760639309],
                         [ 16.046         ,  -1.008373277446,   9.403105976219],
                         [ 17.651666666667,   0.            ,  11.411760639309],
                         [ 15.775333359398,   0.889981209366,  11.92559153011 ],
                         [ 15.775333359398,  -0.889981209366,  11.92559153011 ],
                         [ 18.063333333333,   1.008373277446,  10.829576056132],
                         [ 18.156         ,   0.            ,  12.838230719223],
                         [ 18.014999973935,  -0.889981209366,  10.897929748508],
                         [ 19.669         ,   0.            ,  12.838230719223],
                         [ 17.792666692732,   0.889981209366,  13.352061610024],
                         [ 17.792666692732,  -0.889981209366,  13.352061610024],
                         [ 20.032333307269,   0.889981209366,  12.324399828422],
                         [ 20.032333307269,   0.            ,  13.865892500825],
                         [ 20.080666666667,  -1.008373277446,  12.256046136046],
                         [  2.974666591342,   1.008373277446,  -0.582184583177],
                         [  4.991999924675,  -1.008373277446,   0.844285496737],
                         [  7.009333258008,   1.008373277446,   2.27075557665 ],
                         [  9.026666591342,  -1.008373277446,   3.697225656564],
                         [ 11.043999924675,   1.008373277446,   5.123695736478],
                         [ 13.061333258009,  -1.008373277446,   6.550165816391],
                         [ 15.078666591342,   1.008373277446,   7.976635896305],
                         [ 17.095999924675,  -1.008373277446,   9.403105976219],
                         [ 19.113333258009,   1.008373277446,  10.829576056132],
                         [ 21.130666591342,  -1.008373277446,  12.256046136046]]*Angstrom

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

# Add tags
polyvinyl_alcohol.addTags('H_C', [ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13,
                                  14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26,
                                  27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
                                  40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52,
                                  53, 54, 55, 56, 57, 58, 59, 60, 61])
polyvinyl_alcohol.addTags('H_O', [62, 63, 64, 65, 66, 67, 68, 69, 70, 71])

# Add bonds
bonds = [[ 0,  1],
         [ 0,  2],
         [ 0,  3],
         [ 0,  4],
         [ 1,  5],
         [ 1,  6],
         [ 1,  7],
         [ 5, 62],
         [ 6,  8],
         [ 6,  9],
         [ 6, 10],
         [ 8, 11],
         [ 8, 12],
         [ 8, 13],
         [12, 14],
         [12, 15],
         [12, 16],
         [13, 63],
         [14, 17],
         [14, 18],
         [14, 19],
         [17, 64],
         [18, 20],
         [18, 21],
         [18, 22],
         [20, 23],
         [20, 24],
         [20, 25],
         [24, 26],
         [24, 27],
         [24, 28],
         [25, 65],
         [26, 29],
         [26, 30],
         [26, 31],
         [29, 66],
         [30, 32],
         [30, 33],
         [30, 34],
         [32, 35],
         [32, 36],
         [32, 37],
         [36, 38],
         [36, 39],
         [36, 40],
         [37, 67],
         [38, 41],
         [38, 42],
         [38, 43],
         [41, 68],
         [42, 44],
         [42, 45],
         [42, 46],
         [44, 47],
         [44, 48],
         [44, 49],
         [48, 50],
         [48, 51],
         [48, 52],
         [49, 69],
         [50, 53],
         [50, 54],
         [50, 55],
         [53, 70],
         [54, 56],
         [54, 57],
         [54, 58],
         [56, 59],
         [56, 60],
         [56, 61],
         [61, 71]]
polyvinyl_alcohol.setBonds(bonds)

polyvinyl_alcohol_name = "polyvinyl_alcohol"


# %% Create united atom configuration


def create_united_atom_configuration(all_atom_configuration):
    # This custom script supports atkpython syntax
    # and can perform almost any procedure.
    united_atom_configuration = unitedAtomsConfiguration(
        all_atom_configuration, atom_list=[Carbon], min_hydrogen=1, max_hydrogen=4
    )

    return united_atom_configuration


united_atom_configuration = create_united_atom_configuration(polyvinyl_alcohol)

nlsave('UnitedAtom_Example_results.hdf5', united_atom_configuration)


# %% Set ForceFieldCalculator

# %% ForceFieldCalculator

DREIDING = DreidingPotentialBuilder()
DREIDING.assignAtomTypes(united_atom_configuration, overwrite=True)
atomic_partial_charges = QEqAtomicCharges().charges(united_atom_configuration)
calculator = DREIDING.createCalculator(
    united_atom_configuration,
    atomic_charges=atomic_partial_charges,
)


# %% Set Calculator

united_atom_configuration.setCalculator(calculator)

nlsave('UnitedAtom_Example_results.hdf5', united_atom_configuration)


# %% OptimizeGeometry

restart_strategy = RestartFromTrajectory(
    trajectory_filename='UnitedAtom_Example_results.hdf5', object_id='optimize_trajectory'
)

optimized_configuration = OptimizeGeometry(
    configuration=united_atom_configuration,
    trajectory_filename='UnitedAtom_Example_results.hdf5',
    trajectory_object_id='optimize_trajectory',
    restart_strategy=restart_strategy,
)

nlsave('UnitedAtom_Example_results.hdf5', optimized_configuration, object_id='optgeom')


# %% Create all-atom configuration


def create_allatom_configuration(united_atom_configuration):
    # This custom script supports atkpython syntax
    # and can perform almost any procedure.
    final_configuration = explicitHydrogenConfiguration(united_atom_configuration)

    return final_configuration


final_configuration = create_allatom_configuration(optimized_configuration)

nlsave('UnitedAtom_Example_results.hdf5', final_configuration)
