# -*- coding: utf-8 -*-
setVerbosity(MinimalLog, ELECTRIC_FIELD_CONSTRAINT=True)

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

# Set up lattice
lattice = SimpleTetragonal(3.97816645571*Angstrom, 4.07906435409*Angstrom)

# Define elements
elements = [Barium, Titanium, Oxygen, Oxygen, Oxygen]

# Define coordinates
fractional_coordinates = [[ 0.00000000008 ,  0.000000000035, -0.010189450414],
                          [ 0.500000000154,  0.500000000008,  0.465308176297],
                          [ 0.500000000067,  0.50000000004 ,  1.011196206321],
                          [ 0.500000000058,  0.000000000049,  0.498397341501],
                          [ 0.000000000038,  0.500000000045,  0.498397341501]]

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

# -------------------------------------------------------------
# Calculator
# -------------------------------------------------------------
#----------------------------------------
# Basis Set
#----------------------------------------
basis_set = [
    LDABasis.Oxygen_SingleZetaPolarized,
    LDABasis.Titanium_SingleZetaPolarized,
    LDABasis.Barium_SingleZetaPolarized,
    ]

#----------------------------------------
# Exchange-Correlation
#----------------------------------------
exchange_correlation = LDA.PZ

k_point_sampling = MonkhorstPackGrid(
    na=2,
    nb=2,
    nc=2,
    )
numerical_accuracy_parameters = NumericalAccuracyParameters(
    density_mesh_cutoff=80.0*Hartree,
    k_point_sampling=k_point_sampling
    )

iteration_control_parameters = IterationControlParameters(
    tolerance=1e-02,
    )

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

bulk_configuration.setCalculator(calculator)
nlprint(bulk_configuration)
bulk_configuration.update()
nlsave('bulk_configuration.hdf5', bulk_configuration)

# -------------------------------------------------------------
# Optimize Geometry
# -------------------------------------------------------------

# Use default parameters for the Born effective charges and
# Polarization.
born_effective_charge_parameters = BornEffectiveChargeParameters()
polarization_parameters = PolarizationParameters()

# Setup an electric field constraint object.
efield = ElectricFieldConstraint(
    [0.0, 0.0, 0.1] * Volt / Ang,
    born_effective_charge_parameters=born_effective_charge_parameters,
    polarization_parameters=polarization_parameters,
    update_strategy=UpdateElectricFieldCorrection(max_displacement=0.*Angstrom))

constraints = [efield]

bulk_configuration = OptimizeGeometry(
    bulk_configuration,
    max_forces=0.05*eV/Ang,
    max_steps=200,
    max_step_length=0.2*Ang,
    constraints=constraints,
    trajectory_filename='optimize_bulk_configuration.hdf5',
    trajectory_interval=1*Second,
    restart_strategy=RestartFromTrajectory(),
    optimize_cell=False,
    enable_optimization_stop_file=True,
    )
nlsave('optimize_bulk_configuration.hdf5', bulk_configuration)
nlprint(bulk_configuration)