# -------------------------------------------------------------
# Bulk Configuration
# -------------------------------------------------------------

# Set up lattice
vector_a = [2.84772899324, 0.0, 0.0]*Angstrom
vector_b = [0.0, 2.84772899324, 0.0]*Angstrom
vector_c = [0.0, 0.0, 23.0545939284]*Angstrom
lattice = UnitCell(vector_a, vector_b, vector_c)

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

# Define coordinates
fractional_coordinates = [[ 0.5           ,  0.5           ,  0.304219128903],
                          [ 0.            ,  0.            ,  0.391561742195],
                          [ 0.5           ,  0.5           ,  0.478552505581],
                          [ 0.            ,  0.            ,  0.565044244745],
                          [ 0.434513847759,  0.499904931099,  0.666569187641],
                          [ 0.437099981028,  0.49978846352 ,  0.610664319366]]

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

# Add tags
bulk_configuration.addTags('Selection 0', [0, 1])

# -------------------------------------------------------------
# Calculator
# -------------------------------------------------------------
#----------------------------------------
# Basis Set
#----------------------------------------
basis_set = [
    GGABasis.Carbon_DoubleZetaPolarized,
    GGABasis.Oxygen_DoubleZetaPolarized,
    GGABasis.Palladium_DoubleZetaPolarized,
    ]

#----------------------------------------
# Exchange-Correlation
#----------------------------------------
exchange_correlation = GGA.PBE

numerical_accuracy_parameters = NumericalAccuracyParameters(
    k_point_sampling=(9, 9, 1),
    )

poisson_solver = FastFourier2DSolver(
    boundary_conditions=[[PeriodicBoundaryCondition,PeriodicBoundaryCondition],
                         [PeriodicBoundaryCondition,PeriodicBoundaryCondition],
                         [NeumannBoundaryCondition,DirichletBoundaryCondition]]
    )

calculator = LCAOCalculator(
    basis_set=basis_set,
    exchange_correlation=exchange_correlation,
    numerical_accuracy_parameters=numerical_accuracy_parameters,
    poisson_solver=poisson_solver,
    )

bulk_configuration.setCalculator(calculator)
nlprint(bulk_configuration)
bulk_configuration.update()
nlsave('Pd100_CO.nc', bulk_configuration)

# -------------------------------------------------------------
# Optimize Geometry
# -------------------------------------------------------------
constraints = [0, 1]

bulk_configuration = OptimizeGeometry(
        bulk_configuration,
        max_forces=0.05*eV/Ang,
        max_steps=200,
        max_step_length=0.2*Ang,
        constraints=constraints,
        trajectory_filename='Pd100_CO.nc',
        disable_stress=True,
        optimizer_method=LBFGS(),
        )
nlsave('Pd100_CO.nc', bulk_configuration)
nlprint(bulk_configuration)

# -------------------------------------------------------------
# Total Energy
# -------------------------------------------------------------
total_energy = TotalEnergy(bulk_configuration)
nlsave('Pd100_CO.nc', total_energy)
nlprint(total_energy)