# -------------------------------------------------------------
# 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, 5.732]*Angstrom
left_electrode_lattice = UnitCell(vector_a, vector_b, vector_c)

# Define elements
left_electrode_elements = [Iron, Iron, Iron, Iron]

# Define coordinates
left_electrode_coordinates = [[ 0.716500000003,  0.716500000003,  0.716499999992],
                              [ 2.149500000003,  2.149500000003,  2.149500000008],
                              [ 0.716500000003,  0.716500000003,  3.582499999992],
                              [ 2.149500000003,  2.149500000003,  5.015500000008]]*Angstrom

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

# Add tags
left_electrode.addTags('Selection 0')

# -------------------------------------------------------------
# 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, 5.732]*Angstrom
right_electrode_lattice = UnitCell(vector_a, vector_b, vector_c)

# Define elements
right_electrode_elements = [Iron, Iron, Iron, Iron]

# Define coordinates
right_electrode_coordinates = [[ 0.716500000003,  0.716500000003,  0.716499999992],
                               [ 2.149500000003,  2.149500000003,  2.149500000008],
                               [ 0.716500000003,  0.716500000003,  3.582499999992],
                               [ 2.149500000003,  2.149500000003,  5.015500000008]]*Angstrom

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

# Add tags
right_electrode.addTags('Selection 0')

# -------------------------------------------------------------
# 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, 31.13975]*Angstrom
central_region_lattice = UnitCell(vector_a, vector_b, vector_c)

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

# Define coordinates
central_region_coordinates = [[  0.716500000003,   0.716500000003,   0.716499999992],
                              [  2.149500000003,   2.149500000003,   2.149500000008],
                              [  0.716500000003,   0.716500000003,   3.582499999992],
                              [  2.149500000003,   2.149500000003,   5.015500000008],
                              [  0.7165        ,   0.7165        ,   6.461176916299],
                              [  2.1495        ,   2.1495        ,   7.883348084818],
                              [  0.7165        ,   0.7165        ,  10.046913166035],
                              [  2.1495        ,   2.1495        ,  10.125854994682],
                              [  0.7165        ,   0.7165        ,  12.2742460054  ],
                              [  2.1495        ,   2.1495        ,  12.287565384199],
                              [  0.7165        ,   0.7165        ,  14.47251361208 ],
                              [  2.1495        ,   2.1495        ,  14.47501183558 ],
                              [  0.7165        ,   0.7165        ,  16.667197114   ],
                              [  2.1495        ,   2.1495        ,  16.669729913304],
                              [  0.7165        ,   0.7165        ,  18.854675912015],
                              [  2.1495        ,   2.1495        ,  18.867967064779],
                              [  0.7165        ,   0.7165        ,  21.016369375396],
                              [  2.1495        ,   2.1495        ,  21.095331583047],
                              [  0.7165        ,   0.7165        ,  23.258901803911],
                              [  2.1495        ,   2.1495        ,  24.681132523465],
                              [  0.716500000003,   0.716500000003,  26.124249999992],
                              [  2.149500000003,   2.149500000003,  27.557250000008],
                              [  0.716500000003,   0.716500000003,  28.990249999992],
                              [  2.149500000003,   2.149500000003,  30.423250000008]]*Angstrom

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

# Add tags
central_region.addTags('Selection 0', [0, 1, 2, 3, 20, 21, 22, 23])

device_configuration = DeviceConfiguration(
    central_region,
    [left_electrode, right_electrode]
    )

# -------------------------------------------------------------
# Calculator
# -------------------------------------------------------------
#----------------------------------------
# Basis Set
#----------------------------------------
basis_set = [
    GGABasis.Oxygen_DoubleZetaPolarized,
    GGABasis.Magnesium_DoubleZetaPolarized,
    GGABasis.Iron_SingleZetaPolarized,
    ]

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

#----------------------------------------
# Numerical Accuracy Settings
#----------------------------------------
left_electrode_k_point_sampling = MonkhorstPackGrid(
    na=7,
    nb=7,
    nc=100,
    )
left_electrode_numerical_accuracy_parameters = NumericalAccuracyParameters(
    electron_temperature=1200.0*Kelvin,
    k_point_sampling=left_electrode_k_point_sampling,
    )

right_electrode_k_point_sampling = MonkhorstPackGrid(
    na=7,
    nb=7,
    nc=100,
    )
right_electrode_numerical_accuracy_parameters = NumericalAccuracyParameters(
    electron_temperature=1200.0*Kelvin,
    k_point_sampling=right_electrode_k_point_sampling,
    )

device_k_point_sampling = MonkhorstPackGrid(
    na=7,
    nb=7,
    nc=100,
    )
device_numerical_accuracy_parameters = NumericalAccuracyParameters(
    electron_temperature=1200.0*Kelvin,
    k_point_sampling=device_k_point_sampling,
    )

#----------------------------------------
# Poisson Solver Settings
#----------------------------------------
left_electrode_poisson_solver = FastFourier2DSolver(
    boundary_conditions=[[PeriodicBoundaryCondition(),PeriodicBoundaryCondition()],
                         [PeriodicBoundaryCondition(),PeriodicBoundaryCondition()],
                         [PeriodicBoundaryCondition(),PeriodicBoundaryCondition()]]
    )

right_electrode_poisson_solver = FastFourier2DSolver(
    boundary_conditions=[[PeriodicBoundaryCondition(),PeriodicBoundaryCondition()],
                         [PeriodicBoundaryCondition(),PeriodicBoundaryCondition()],
                         [PeriodicBoundaryCondition(),PeriodicBoundaryCondition()]]
    )

#----------------------------------------
# Electrode Calculators
#----------------------------------------
left_electrode_calculator = LCAOCalculator(
    basis_set=basis_set,
    exchange_correlation=exchange_correlation,
    numerical_accuracy_parameters=left_electrode_numerical_accuracy_parameters,
    poisson_solver=left_electrode_poisson_solver,
    )

right_electrode_calculator = LCAOCalculator(
    basis_set=basis_set,
    exchange_correlation=exchange_correlation,
    numerical_accuracy_parameters=right_electrode_numerical_accuracy_parameters,
    poisson_solver=right_electrode_poisson_solver,
    )

#----------------------------------------
# Device Calculator
#----------------------------------------
calculator = DeviceLCAOCalculator(
    basis_set=basis_set,
    exchange_correlation=exchange_correlation,
    numerical_accuracy_parameters=device_numerical_accuracy_parameters,
    electrode_calculators=
        [left_electrode_calculator, right_electrode_calculator],
    )

device_configuration.setCalculator(calculator)

# -------------------------------------------------------------
# Initial State
# -------------------------------------------------------------
initial_spin = InitialSpin(scaled_spins=[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0])
device_configuration.setCalculator(
    calculator,
    initial_spin=initial_spin,
)
device_configuration.update()
nlsave('mgo_para.nc', device_configuration)
nlprint(device_configuration)

# -------------------------------------------------------------
# Mulliken Population
# -------------------------------------------------------------
mulliken_population = MullikenPopulation(device_configuration)
nlsave('mgo_para.nc', mulliken_population)
nlprint(mulliken_population)

# -------------------------------------------------------------
# Forces
# -------------------------------------------------------------
forces = Forces(device_configuration)
nlsave('mgo_para.nc', forces)
nlprint(forces)

# -------------------------------------------------------------
# Transmission Spectrum
# -------------------------------------------------------------
transmission_spectrum = TransmissionSpectrum(
    configuration=device_configuration,
    energies=numpy.linspace(0,0,1)*eV,
    kpoints=MonkhorstPackGrid(151,151),
    energy_zero_parameter=AverageFermiLevel,
    infinitesimal=1e-06*eV,
    self_energy_calculator=RecursionSelfEnergy(),
    )
nlsave('mgo_para.nc', transmission_spectrum)
nlprint(transmission_spectrum)
