# -*- coding: utf-8 -*-
# -------------------------------------------------------------
# Two-probe Configuration
# -------------------------------------------------------------
# -------------------------------------------------------------
# Left Electrode
# -------------------------------------------------------------

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

# Define elements
left_electrode_elements = [Iron, Iron]

# Define coordinates
left_electrode_coordinates = [[ 0.7165,  0.7165,  0.7165],
                              [ 2.1495,  2.1495,  2.1495]]*Angstrom

# Set up configuration
left_electrode = BulkConfiguration(
    bravais_lattice=left_electrode_lattice,
    elements=left_electrode_elements,
    cartesian_coordinates=left_electrode_coordinates
    )

# -------------------------------------------------------------
# Right Electrode
# -------------------------------------------------------------

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

# Define elements
right_electrode_elements = [Iron, Iron]

# Define coordinates
right_electrode_coordinates = [[ 2.1495,  2.1495,  0.7165],
                               [ 0.7165,  0.7165,  2.1495]]*Angstrom

# Set up configuration
right_electrode = BulkConfiguration(
    bravais_lattice=right_electrode_lattice,
    elements=right_electrode_elements,
    cartesian_coordinates=right_electrode_coordinates
    )

# -------------------------------------------------------------
# Central Region
# -------------------------------------------------------------

# Set up lattice
vector_a = [2.866, 0.0, 0.0]*Angstrom
vector_b = [0.0, 2.866, 0.0]*Angstrom
vector_c = [0.0, 0.0, 39.06710000000001]*Angstrom
central_region_lattice = UnitCell(vector_a, vector_b, vector_c)

# Define elements
central_region_elements = [Iron, Iron, Iron, Iron, Iron, Iron, Iron, Iron, Magnesium, Oxygen,
                           Oxygen, Magnesium, Magnesium, Oxygen, Oxygen, Magnesium, Magnesium,
                           Oxygen, Oxygen, Magnesium, Magnesium, Oxygen, Iron, Iron, Iron,
                           Iron, Iron, Iron, Iron, Iron]

# Define coordinates
central_region_coordinates = [[  0.7165 ,   0.7165 ,   0.7165 ],
                              [  2.1495 ,   2.1495 ,   2.1495 ],
                              [  0.7165 ,   0.7165 ,   3.5825 ],
                              [  2.1495 ,   2.1495 ,   5.0155 ],
                              [  0.7165 ,   0.7165 ,   6.4485 ],
                              [  2.1495 ,   2.1495 ,   7.8815 ],
                              [  0.7165 ,   0.7165 ,   9.3145 ],
                              [  2.1495 ,   2.1495 ,  10.7475 ],
                              [  0.7165 ,   0.7165 ,  12.9475 ],
                              [  2.1495 ,   2.1495 ,  12.9475 ],
                              [  0.7165 ,   0.7165 ,  15.14285],
                              [  2.1495 ,   2.1495 ,  15.14285],
                              [  0.7165 ,   0.7165 ,  17.3382 ],
                              [  2.1495 ,   2.1495 ,  17.3382 ],
                              [  0.7165 ,   0.7165 ,  19.53355],
                              [  2.1495 ,   2.1495 ,  19.53355],
                              [  0.7165 ,   0.7165 ,  21.7289 ],
                              [  2.1495 ,   2.1495 ,  21.7289 ],
                              [  0.7165 ,   0.7165 ,  23.92425],
                              [  2.1495 ,   2.1495 ,  23.92425],
                              [  0.7165 ,   0.7165 ,  26.1196 ],
                              [  2.1495 ,   2.1495 ,  26.1196 ],
                              [  2.1495 ,   2.1495 ,  28.3196 ],
                              [  0.7165 ,   0.7165 ,  29.7526 ],
                              [  2.1495 ,   2.1495 ,  31.1856 ],
                              [  0.7165 ,   0.7165 ,  32.6186 ],
                              [  2.1495 ,   2.1495 ,  34.0516 ],
                              [  0.7165 ,   0.7165 ,  35.4846 ],
                              [  2.1495 ,   2.1495 ,  36.9176 ],
                              [  0.7165 ,   0.7165 ,  38.3506 ]]*Angstrom

# Set up configuration
central_region = BulkConfiguration(
    bravais_lattice=central_region_lattice,
    elements=central_region_elements,
    cartesian_coordinates=central_region_coordinates
    )

device_configuration = DeviceConfiguration(
    central_region,
    [left_electrode, right_electrode],
    equivalent_electrode_lengths=[5.732, 5.732]*Angstrom,
    transverse_electrode_repetitions=[[1, 1], [1, 1]],
    )

# -------------------------------------------------------------
# Calculator
# -------------------------------------------------------------
#----------------------------------------
# Basis Set
#----------------------------------------
basis_set = [
    LDABasis.Oxygen_SingleZeta,
    LDABasis.Iron_SingleZetaPolarized,
    LDABasis.Magnesium_SingleZetaPolarized,
    ]

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

#----------------------------------------
# Numerical Accuracy Settings
#----------------------------------------
device_k_point_sampling = MonkhorstPackGrid(
    na=21,
    nb=21,
    nc=300,)
device_numerical_accuracy_parameters = NumericalAccuracyParameters(
    k_point_sampling=device_k_point_sampling,
    )

#----------------------------------------
# Contour Integral Settings
#----------------------------------------
equilibrium_contour = SemiCircleContour(
    integral_lower_bound=2.55775293896*Hartree,
    )

contour_parameters = ContourParameters(
    equilibrium_contour=equilibrium_contour,
    )

device_algorithm_parameters = DeviceAlgorithmParameters(
    initial_density_type=EquivalentBulk(electrode_constraint_length=3.*Ang)
    )

#----------------------------------------
# Device Calculator
#----------------------------------------
calculator = DeviceLCAOCalculator(
    basis_set=basis_set,
    exchange_correlation=exchange_correlation,
    numerical_accuracy_parameters=device_numerical_accuracy_parameters,
    contour_parameters=contour_parameters,
    device_algorithm_parameters=device_algorithm_parameters,
    )

device_configuration.setCalculator(calculator)

# -------------------------------------------------------------
# Initial State
# -------------------------------------------------------------
scaled_spins = [1] * 8 + [0] * 14 + [1] * 8
initial_spin = InitialSpin(scaled_spins=scaled_spins)

device_configuration.setCalculator(
    calculator,
    initial_spin=initial_spin,
)
device_configuration.update()
filename = 'polarized-0.hdf5'
nlsave(filename, device_configuration)