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

# -------------------------------------------------------------
# Left Electrode
# -------------------------------------------------------------

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

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

# Define coordinates
left_electrode_coordinates = [[ 1.441879115187,  1.441879115187,  1.0197125     ],
                              [ 4.325637345561,  1.441879115187,  1.0197125     ],
                              [ 7.209395575935,  1.441879115187,  1.0197125     ],
                              [ 1.441879115187,  4.325637345561,  1.0197125     ],
                              [ 4.325637345561,  4.325637345561,  1.0197125     ],
                              [ 7.209395575935,  4.325637345561,  1.0197125     ],
                              [ 1.441879115187,  7.209395575935,  1.0197125     ],
                              [ 4.325637345561,  7.209395575935,  1.0197125     ],
                              [ 7.209395575935,  7.209395575935,  1.0197125     ],
                              [ 0.            ,  0.            ,  3.0588375     ],
                              [ 2.883758230374,  0.            ,  3.0588375     ],
                              [ 5.767516460748,  0.            ,  3.0588375     ],
                              [ 0.            ,  2.883758230374,  3.0588375     ],
                              [ 2.883758230374,  2.883758230374,  3.0588375     ],
                              [ 5.767516460748,  2.883758230374,  3.0588375     ],
                              [ 0.            ,  5.767516460748,  3.0588375     ],
                              [ 2.883758230374,  5.767516460748,  3.0588375     ],
                              [ 5.767516460748,  5.767516460748,  3.0588375     ]]*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 = [8.65127469112, 0.0, 0.0]*Angstrom
vector_b = [0.0, 8.65127469112, 0.0]*Angstrom
vector_c = [0.0, 0.0, 4.07825]*Angstrom
right_electrode_lattice = UnitCell(vector_a, vector_b, vector_c)

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

# Define coordinates
right_electrode_coordinates = [[ 1.441879115187,  1.441879115187,  1.0195625     ],
                               [ 4.325637345561,  1.441879115187,  1.0195625     ],
                               [ 7.209395575935,  1.441879115187,  1.0195625     ],
                               [ 1.441879115187,  4.325637345561,  1.0195625     ],
                               [ 4.325637345561,  4.325637345561,  1.0195625     ],
                               [ 7.209395575935,  4.325637345561,  1.0195625     ],
                               [ 1.441879115187,  7.209395575935,  1.0195625     ],
                               [ 4.325637345561,  7.209395575935,  1.0195625     ],
                               [ 7.209395575935,  7.209395575935,  1.0195625     ],
                               [ 0.            ,  0.            ,  3.0586875     ],
                               [ 2.883758230374,  0.            ,  3.0586875     ],
                               [ 5.767516460748,  0.            ,  3.0586875     ],
                               [ 0.            ,  2.883758230374,  3.0586875     ],
                               [ 2.883758230374,  2.883758230374,  3.0586875     ],
                               [ 5.767516460748,  2.883758230374,  3.0586875     ],
                               [ 0.            ,  5.767516460748,  3.0586875     ],
                               [ 2.883758230374,  5.767516460748,  3.0586875     ],
                               [ 5.767516460748,  5.767516460748,  3.0586875     ]]*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 = [8.65127469112, 0.0, 0.0]*Angstrom
vector_b = [0.0, 8.65127469112, 0.0]*Angstrom
vector_c = [0.0, 0.0, 8.15665]*Angstrom
central_region_lattice = UnitCell(vector_a, vector_b, vector_c)

# Define elements
central_region_elements = [Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold,
                           Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold,
                           Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold,
                           Gold, Gold, Gold]

# Define coordinates
central_region_coordinates = [[ 1.441879115187,  1.441879115187,  1.0197125     ],
                              [ 4.325637345561,  1.441879115187,  1.0197125     ],
                              [ 7.209395575935,  1.441879115187,  1.0197125     ],
                              [ 1.441879115187,  4.325637345561,  1.0197125     ],
                              [ 4.325637345561,  4.325637345561,  1.0197125     ],
                              [ 7.209395575935,  4.325637345561,  1.0197125     ],
                              [ 1.441879115187,  7.209395575935,  1.0197125     ],
                              [ 4.325637345561,  7.209395575935,  1.0197125     ],
                              [ 7.209395575935,  7.209395575935,  1.0197125     ],
                              [ 0.            ,  0.            ,  3.0588375     ],
                              [ 2.883758230374,  0.            ,  3.0588375     ],
                              [ 5.767516460748,  0.            ,  3.0588375     ],
                              [ 0.            ,  2.883758230374,  3.0588375     ],
                              [ 2.883758230374,  2.883758230374,  3.0588375     ],
                              [ 5.767516460748,  2.883758230374,  3.0588375     ],
                              [ 0.            ,  5.767516460748,  3.0588375     ],
                              [ 2.883758230374,  5.767516460748,  3.0588375     ],
                              [ 5.767516460748,  5.767516460748,  3.0588375     ],
                              [ 1.441879115187,  1.441879115187,  5.0979625     ],
                              [ 4.325637345561,  1.441879115187,  5.0979625     ],
                              [ 7.209395575935,  1.441879115187,  5.0979625     ],
                              [ 1.441879115187,  4.325637345561,  5.0979625     ],
                              [ 4.325637345561,  4.325637345561,  5.0979625     ],
                              [ 7.209395575935,  4.325637345561,  5.0979625     ],
                              [ 1.441879115187,  7.209395575935,  5.0979625     ],
                              [ 4.325637345561,  7.209395575935,  5.0979625     ],
                              [ 7.209395575935,  7.209395575935,  5.0979625     ],
                              [ 0.            ,  0.            ,  7.1370875     ],
                              [ 2.883758230374,  0.            ,  7.1370875     ],
                              [ 5.767516460748,  0.            ,  7.1370875     ],
                              [ 0.            ,  2.883758230374,  7.1370875     ],
                              [ 2.883758230374,  2.883758230374,  7.1370875     ],
                              [ 5.767516460748,  2.883758230374,  7.1370875     ],
                              [ 0.            ,  5.767516460748,  7.1370875     ],
                              [ 2.883758230374,  5.767516460748,  7.1370875     ],
                              [ 5.767516460748,  5.767516460748,  7.1370875     ]]*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 = HotbitDirectory("hotbit/standard/")

#----------------------------------------
# Pair Potentials
#----------------------------------------
pair_potentials = HotbitDirectory("hotbit/standard/")

#----------------------------------------
# Numerical Accuracy Settings
#----------------------------------------
left_electrode_k_point_sampling = MonkhorstPackGrid(
    na=2,
    nb=2,
    nc=42,
    )
left_electrode_numerical_accuracy_parameters = NumericalAccuracyParameters(
    k_point_sampling=left_electrode_k_point_sampling,
    )

right_electrode_k_point_sampling = MonkhorstPackGrid(
    na=2,
    nb=2,
    nc=42,
    )
right_electrode_numerical_accuracy_parameters = NumericalAccuracyParameters(
    k_point_sampling=right_electrode_k_point_sampling,
    )

device_k_point_sampling = MonkhorstPackGrid(
    na=2,
    nb=2,
    nc=42,
    )
device_numerical_accuracy_parameters = NumericalAccuracyParameters(
    k_point_sampling=device_k_point_sampling,
    )

#----------------------------------------
# Iteration Control Settings
#----------------------------------------
left_electrode_iteration_control_parameters = IterationControlParameters()

right_electrode_iteration_control_parameters = IterationControlParameters()

device_iteration_control_parameters = IterationControlParameters()

#----------------------------------------
# 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 = SlaterKosterCalculator(
    basis_set=basis_set,
    pair_potentials=pair_potentials,
    numerical_accuracy_parameters=left_electrode_numerical_accuracy_parameters,
    iteration_control_parameters=left_electrode_iteration_control_parameters,
    poisson_solver=left_electrode_poisson_solver,
    )

right_electrode_calculator = SlaterKosterCalculator(
    basis_set=basis_set,
    pair_potentials=pair_potentials,
    numerical_accuracy_parameters=right_electrode_numerical_accuracy_parameters,
    iteration_control_parameters=right_electrode_iteration_control_parameters,
    poisson_solver=right_electrode_poisson_solver,
    )

#----------------------------------------
# Device Calculator
#----------------------------------------
calculator = DeviceSlaterKosterCalculator(
    basis_set=basis_set,
    pair_potentials=pair_potentials,
    numerical_accuracy_parameters=device_numerical_accuracy_parameters,
    iteration_control_parameters=device_iteration_control_parameters,
    electrode_calculators=
        [left_electrode_calculator, right_electrode_calculator],
    )

device_configuration.setCalculator(calculator)
nlprint(device_configuration)
device_configuration.update()
nlsave('Au100_L0_T0.nc', device_configuration)

# -------------------------------------------------------------
# Transmission Spectrum
# -------------------------------------------------------------
kpoint_grid = MonkhorstPackGrid(
    na=10,
    nb=10,
    )

transmission_spectrum = TransmissionSpectrum(
    configuration=device_configuration,
    energies=numpy.linspace(-0.2,0.2,51)*eV,
    kpoints=kpoint_grid,
    energy_zero_parameter=AverageFermiLevel,
    infinitesimal=1e-06*eV,
    self_energy_calculator=KrylovSelfEnergy(),
    )
nlsave('Au100_L0_T0.nc', transmission_spectrum)
nlprint(transmission_spectrum)