# -*- coding: utf-8 -*-
# -------------------------------------------------------------
# Bulk Configuration
# -------------------------------------------------------------

# Set up lattice
vector_a = [16.646425, 0.0, 0.0]*Angstrom
vector_b = [-8.32321249999, 14.4162269322, 0.0]*Angstrom
vector_c = [0.0, 0.0, 21.326457581]*Angstrom
lattice = UnitCell(vector_a, vector_b, vector_c)

# Define elements
elements = [Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum]

# Define coordinates
fractional_coordinates = [[ 0.055555555556,  0.111111111111,  0.10621977455 ],
                          [ 0.222222222222,  0.111111111111,  0.10621977455 ],
                          [ 0.388888888889,  0.111111111111,  0.10621977455 ],
                          [ 0.555555555556,  0.111111111111,  0.10621977455 ],
                          [ 0.722222222222,  0.111111111111,  0.10621977455 ],
                          [ 0.888888888889,  0.111111111111,  0.10621977455 ],
                          [ 0.055555555556,  0.277777777778,  0.10621977455 ],
                          [ 0.222222222222,  0.277777777778,  0.10621977455 ],
                          [ 0.388888888889,  0.277777777778,  0.10621977455 ],
                          [ 0.555555555556,  0.277777777778,  0.10621977455 ],
                          [ 0.722222222222,  0.277777777778,  0.10621977455 ],
                          [ 0.888888888889,  0.277777777778,  0.10621977455 ],
                          [ 0.055555555556,  0.444444444444,  0.10621977455 ],
                          [ 0.222222222222,  0.444444444444,  0.10621977455 ],
                          [ 0.388888888889,  0.444444444444,  0.10621977455 ],
                          [ 0.555555555556,  0.444444444444,  0.10621977455 ],
                          [ 0.722222222222,  0.444444444444,  0.10621977455 ],
                          [ 0.888888888889,  0.444444444444,  0.10621977455 ],
                          [ 0.055555555556,  0.611111111111,  0.10621977455 ],
                          [ 0.222222222222,  0.611111111111,  0.10621977455 ],
                          [ 0.388888888889,  0.611111111111,  0.10621977455 ],
                          [ 0.555555555556,  0.611111111111,  0.10621977455 ],
                          [ 0.722222222222,  0.611111111111,  0.10621977455 ],
                          [ 0.888888888889,  0.611111111111,  0.10621977455 ],
                          [ 0.055555555556,  0.777777777778,  0.10621977455 ],
                          [ 0.222222222222,  0.777777777778,  0.10621977455 ],
                          [ 0.388888888889,  0.777777777778,  0.10621977455 ],
                          [ 0.555555555556,  0.777777777778,  0.10621977455 ],
                          [ 0.722222222222,  0.777777777778,  0.10621977455 ],
                          [ 0.888888888889,  0.777777777778,  0.10621977455 ],
                          [ 0.055555555556,  0.944444444444,  0.10621977455 ],
                          [ 0.222222222222,  0.944444444444,  0.10621977455 ],
                          [ 0.388888888889,  0.944444444444,  0.10621977455 ],
                          [ 0.555555555556,  0.944444444444,  0.10621977455 ],
                          [ 0.722222222222,  0.944444444444,  0.10621977455 ],
                          [ 0.888888888889,  0.944444444444,  0.10621977455 ],
                          [ 0.            ,  0.            ,  0.212439549099],
                          [ 0.166666666667,  0.            ,  0.212439549099],
                          [ 0.333333333333,  0.            ,  0.212439549099],
                          [ 0.5           ,  0.            ,  0.212439549099],
                          [ 0.666666666667,  0.            ,  0.212439549099],
                          [ 0.833333333333,  0.            ,  0.212439549099],
                          [-0.            ,  0.166666666667,  0.212439549099],
                          [ 0.166666666667,  0.166666666667,  0.212439549099],
                          [ 0.333333333333,  0.166666666667,  0.212439549099],
                          [ 0.5           ,  0.166666666667,  0.212439549099],
                          [ 0.666666666667,  0.166666666667,  0.212439549099],
                          [ 0.833333333333,  0.166666666667,  0.212439549099],
                          [-0.            ,  0.333333333333,  0.212439549099],
                          [ 0.166666666667,  0.333333333333,  0.212439549099],
                          [ 0.333333333333,  0.333333333333,  0.212439549099],
                          [ 0.5           ,  0.333333333333,  0.212439549099],
                          [ 0.666666666667,  0.333333333333,  0.212439549099],
                          [ 0.833333333333,  0.333333333333,  0.212439549099],
                          [-0.            ,  0.5           ,  0.212439549099],
                          [ 0.166666666667,  0.5           ,  0.212439549099],
                          [ 0.333333333333,  0.5           ,  0.212439549099],
                          [ 0.5           ,  0.5           ,  0.212439549099],
                          [ 0.666666666667,  0.5           ,  0.212439549099],
                          [ 0.833333333333,  0.5           ,  0.212439549099],
                          [ 0.            ,  0.666666666667,  0.212439549099],
                          [ 0.166666666667,  0.666666666667,  0.212439549099],
                          [ 0.333333333333,  0.666666666667,  0.212439549099],
                          [ 0.5           ,  0.666666666667,  0.212439549099],
                          [ 0.666666666667,  0.666666666667,  0.212439549099],
                          [ 0.833333333333,  0.666666666667,  0.212439549099],
                          [-0.            ,  0.833333333333,  0.212439549099],
                          [ 0.166666666667,  0.833333333333,  0.212439549099],
                          [ 0.333333333333,  0.833333333333,  0.212439549099],
                          [ 0.5           ,  0.833333333333,  0.212439549099],
                          [ 0.666666666667,  0.833333333333,  0.212439549099],
                          [ 0.833333333333,  0.833333333333,  0.212439549099],
                          [ 0.111111111111,  0.055555555556,  0.318659323649],
                          [ 0.277777777778,  0.055555555556,  0.318659323649],
                          [ 0.444444444444,  0.055555555556,  0.318659323649],
                          [ 0.611111111111,  0.055555555556,  0.318659323649],
                          [ 0.777777777778,  0.055555555556,  0.318659323649],
                          [ 0.944444444444,  0.055555555556,  0.318659323649],
                          [ 0.111111111111,  0.222222222222,  0.318659323649],
                          [ 0.277777777778,  0.222222222222,  0.318659323649],
                          [ 0.444444444444,  0.222222222222,  0.318659323649],
                          [ 0.611111111111,  0.222222222222,  0.318659323649],
                          [ 0.777777777778,  0.222222222222,  0.318659323649],
                          [ 0.944444444444,  0.222222222222,  0.318659323649],
                          [ 0.111111111111,  0.388888888889,  0.318659323649],
                          [ 0.277777777778,  0.388888888889,  0.318659323649],
                          [ 0.444444444444,  0.388888888889,  0.318659323649],
                          [ 0.611111111111,  0.388888888889,  0.318659323649],
                          [ 0.777777777778,  0.388888888889,  0.318659323649],
                          [ 0.944444444444,  0.388888888889,  0.318659323649],
                          [ 0.111111111111,  0.555555555556,  0.318659323649],
                          [ 0.277777777778,  0.555555555556,  0.318659323649],
                          [ 0.444444444444,  0.555555555556,  0.318659323649],
                          [ 0.611111111111,  0.555555555556,  0.318659323649],
                          [ 0.777777777778,  0.555555555556,  0.318659323649],
                          [ 0.944444444444,  0.555555555556,  0.318659323649],
                          [ 0.111111111111,  0.722222222222,  0.318659323649],
                          [ 0.277777777778,  0.722222222222,  0.318659323649],
                          [ 0.444444444444,  0.722222222222,  0.318659323649],
                          [ 0.611111111111,  0.722222222222,  0.318659323649],
                          [ 0.777777777778,  0.722222222222,  0.318659323649],
                          [ 0.944444444444,  0.722222222222,  0.318659323649],
                          [ 0.111111111111,  0.888888888889,  0.318659323649],
                          [ 0.277777777778,  0.888888888889,  0.318659323649],
                          [ 0.444444444444,  0.888888888889,  0.318659323649],
                          [ 0.611111111111,  0.888888888889,  0.318659323649],
                          [ 0.777777777778,  0.888888888889,  0.318659323649],
                          [ 0.944444444444,  0.888888888889,  0.318659323649],
                          [ 0.055555555556,  0.111111111111,  0.424879098199],
                          [ 0.222222222222,  0.111111111111,  0.424879098199],
                          [ 0.388888888889,  0.111111111111,  0.424879098199],
                          [ 0.555555555556,  0.111111111111,  0.424879098199],
                          [ 0.722222222222,  0.111111111111,  0.424879098199],
                          [ 0.888888888889,  0.111111111111,  0.424879098199],
                          [ 0.055555555556,  0.277777777778,  0.424879098199],
                          [ 0.222222222222,  0.277777777778,  0.424879098199],
                          [ 0.388888888889,  0.277777777778,  0.424879098199],
                          [ 0.555555555556,  0.277777777778,  0.424879098199],
                          [ 0.722222222222,  0.277777777778,  0.424879098199],
                          [ 0.888888888889,  0.277777777778,  0.424879098199],
                          [ 0.055555555556,  0.444444444444,  0.424879098199],
                          [ 0.222222222222,  0.444444444444,  0.424879098199],
                          [ 0.388888888889,  0.444444444444,  0.424879098199],
                          [ 0.555555555556,  0.444444444444,  0.424879098199],
                          [ 0.722222222222,  0.444444444444,  0.424879098199],
                          [ 0.888888888889,  0.444444444444,  0.424879098199],
                          [ 0.055555555556,  0.611111111111,  0.424879098199],
                          [ 0.222222222222,  0.611111111111,  0.424879098199],
                          [ 0.388888888889,  0.611111111111,  0.424879098199],
                          [ 0.555555555556,  0.611111111111,  0.424879098199],
                          [ 0.722222222222,  0.611111111111,  0.424879098199],
                          [ 0.888888888889,  0.611111111111,  0.424879098199],
                          [ 0.055555555556,  0.777777777778,  0.424879098199],
                          [ 0.222222222222,  0.777777777778,  0.424879098199],
                          [ 0.388888888889,  0.777777777778,  0.424879098199],
                          [ 0.555555555556,  0.777777777778,  0.424879098199],
                          [ 0.722222222222,  0.777777777778,  0.424879098199],
                          [ 0.888888888889,  0.777777777778,  0.424879098199],
                          [ 0.055555555556,  0.944444444444,  0.424879098199],
                          [ 0.222222222222,  0.944444444444,  0.424879098199],
                          [ 0.388888888889,  0.944444444444,  0.424879098199],
                          [ 0.555555555556,  0.944444444444,  0.424879098199],
                          [ 0.722222222222,  0.944444444444,  0.424879098199],
                          [ 0.888888888889,  0.944444444444,  0.424879098199],
                          [ 0.833333333333,  0.            ,  0.531098872748],
                          [ 0.666666666667,  0.166666666667,  0.531098872748],
                          [ 0.5           ,  0.333333333333,  0.531098872748],
                          [ 0.333333333333,  0.5           ,  0.531098872748],
                          [ 0.166666666667,  0.666666666667,  0.531098872748],
                          [ 0.666666666667,  0.666666666667,  0.531098872748],
                          [-0.            ,  0.833333333333,  0.531098872748]]

# Set up configuration
bulk_configuration = BulkConfiguration(
    bravais_lattice=lattice,
    elements=elements,
    fractional_coordinates=fractional_coordinates
    )

# Add tags
bulk_configuration.addTags('Selection 0', [ 0,  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])

# Morse Potential
potentialSet = TremoloXPotentialSet(name='Pt-Morse')
potentialSet.addParticleType(ParticleType.fromElement(Platinum))
potentialSet.addPotential(MorsePotential(
    'Pt',
    'Pt',
    r_0=2.897*Angstrom,
    k=1.6047*Angstrom**-1,
    E_0=0.7102*eV,
    r_i=6.5*Angstrom,
    r_cut=9.5*Angstrom)
)
calculator = TremoloXCalculator(parameters=potentialSet)
bulk_configuration.setCalculator(calculator)

# -------------------------------------------------------------
# Optimize Geometry
# -------------------------------------------------------------
fix_atom_indices_0 = [ 0,  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]
constraints = [FixAtomConstraints(fix_atom_indices_0)]

bulk_configuration = OptimizeGeometry(
    bulk_configuration,
    max_forces=0.01*eV/Ang,
    max_steps=200,
    max_step_length=0.2*Ang,
    constraints=constraints,
    trajectory_filename=None,
    disable_stress=True,
    optimizer_method=LBFGS(),
)
nlsave('initial.hdf5', bulk_configuration)
nlprint(bulk_configuration)