# -------------------------------------------------------------
# Left electrode
# -------------------------------------------------------------
# Set up lattice
vector_a = [6.0, 0.0, 0.0]*Angstrom
vector_b = [0.0, 6.0, 0.0]*Angstrom
vector_c = [0.0, 0.0, 5.8]*Angstrom
left_electrode_lattice = UnitCell(vector_a, vector_b, vector_c)
# Define elements
left_electrode_elements = [Carbon, Carbon]
# Define coordinates
left_electrode_coordinates = [[ 3.  ,  3.  ,  1.45],
                              [ 3.  ,  3.  ,  4.35]]*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 = [6.0, 0.0, 0.0]*Angstrom
vector_b = [0.0, 6.0, 0.0]*Angstrom
vector_c = [0.0, 0.0, 5.8]*Angstrom
right_electrode_lattice = UnitCell(vector_a, vector_b, vector_c)
# Define elements
right_electrode_elements = [Carbon, Carbon]
# Define coordinates
right_electrode_coordinates = [[ 3.  ,  3.  ,  1.45],
                               [ 3.  ,  3.  ,  4.35]]*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 = [6.0, 0.0, 0.0]*Angstrom
vector_b = [0.0, 6.0, 0.0]*Angstrom
vector_c = [0.0, 0.0, 71.5]*Angstrom
central_region_lattice = UnitCell(vector_a, vector_b, vector_c)
# Define elements
central_region_elements = [Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
                           Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
                           Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon]
# Define coordinates
central_region_coordinates = [[  3.  ,   3.  ,   1.45],
                              [  3.  ,   3.  ,   4.35],
                              [  3.  ,   3.  ,   7.25],
                              [  3.  ,   3.  ,  10.15],
                              [  3.  ,   3.  ,  13.05],
                              [  3.  ,   3.  ,  15.95],
                              [  3.  ,   3.  ,  18.85],
                              [  3.  ,   3.  ,  21.75],
                              [  3.  ,   3.  ,  24.65],
                              [  3.  ,   3.  ,  27.55],
                              [  3.  ,   3.  ,  30.45],
                              [  3.  ,   3.  ,  33.35],
                              [  3.  ,   3.  ,  38.15],
                              [  3.  ,   3.  ,  41.05],
                              [  3.  ,   3.  ,  43.95],
                              [  3.  ,   3.  ,  46.85],
                              [  3.  ,   3.  ,  49.75],
                              [  3.  ,   3.  ,  52.65],
                              [  3.  ,   3.  ,  55.55],
                              [  3.  ,   3.  ,  58.45],
                              [  3.  ,   3.  ,  61.35],
                              [  3.  ,   3.  ,  64.25],
                              [  3.  ,   3.  ,  67.15],
                              [  3.  ,   3.  ,  70.05]]*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]
    )
# -------------------------------------------------------------
# Calculator
# -------------------------------------------------------------
#----------------------------------------
# Basis Set
#----------------------------------------
basis_set = [
    LDABasis.Carbon_SingleZeta,
    ]

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

#----------------------------------------
# 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]]
    )

# Use the special noncollinear mixing scheme
iteration_control_parameters = IterationControlParameters(
    algorithm=PulayMixer(noncollinear_mixing=True)
    )

#----------------------------------------
# Electrode Calculators
#----------------------------------------
left_electrode_calculator = LCAOCalculator(
    basis_set=basis_set,
    exchange_correlation=exchange_correlation,
    poisson_solver=left_electrode_poisson_solver,
    )
right_electrode_calculator = LCAOCalculator(
    basis_set=basis_set,
    exchange_correlation=exchange_correlation,
    poisson_solver=right_electrode_poisson_solver,
    )

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

# Define the spin rotation
theta = 90*Degrees
left_spins = [(i, 1, 0*Degrees, 0*Degrees) for i in range(12)]
right_spins = [(i, 1, theta, 0*Degrees ) for i in range(12,24)]
spin_list = left_spins+right_spins
initial_spin = InitialSpin(scaled_spins=spin_list)

# Setup the initial state
device_configuration.setCalculator(
calculator,
    initial_spin=initial_spin,
)

# Calculate and save
device_configuration.update()
nlsave("STT.nc", device_configuration)

# -------------------------------------------------------------
# Mulliken population
# -------------------------------------------------------------
mulliken_population = MullikenPopulation(device_configuration)
nlsave("STT.nc", mulliken_population)
nlprint(mulliken_population)

#----------------------------------------
# Spin Transfer Torque
#----------------------------------------
spin_transfer_torque = SpinTransferTorque(
    configuration=device_configuration,
    energy=0.0*eV,
    kpoints=MonkhorstPackGrid(1,1),
    contributions=Left,
    self_energy_calculator=RecursionSelfEnergy(),
    energy_zero_parameter=AverageFermiLevel,
    infinitesimal=1.0e-6*eV
)
nlsave("STT.nc", spin_transfer_torque)