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

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

# Set up lattice
vector_a = [8.65127469112, -6.75261456523e-33, 0.0]*Angstrom
vector_b = [-4.32563734556, 7.49222365763, 4.09015759516e-16]*Angstrom
vector_c = [0.0, 0.0, 7.06373620597]*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, Gold, Gold, Gold, Gold,
                           Gold, Gold, Gold, Gold, Gold]

# Define coordinates
left_electrode_coordinates = [[ 0.72107876,  1.24846289,  1.17728937],
                              [ 3.60483699,  1.24846289,  1.17728937],
                              [ 6.48859522,  1.24846289,  1.17728937],
                              [-0.72080035,  3.74587078,  1.17728937],
                              [ 5.04671611,  3.74587078,  1.17728937],
                              [ 2.16295788,  3.74587078,  1.17728937],
                              [ 3.60483699,  6.24327867,  1.17728937],
                              [-2.16267947,  6.24327867,  1.17728937],
                              [ 0.72107876,  6.24327867,  1.17728937],
                              [ 2.16295788,  2.08093219,  3.5318681 ],
                              [-0.72080035,  2.08093219,  3.5318681 ],
                              [ 5.04671611,  2.08093219,  3.5318681 ],
                              [-3.60455858,  7.07574796,  3.5318681 ],
                              [-2.16267947,  4.57834008,  3.5318681 ],
                              [ 2.16295788,  7.07574796,  3.5318681 ],
                              [ 3.60483699,  4.57834008,  3.5318681 ],
                              [ 0.72107876,  4.57834008,  3.5318681 ],
                              [-0.72080035,  7.07574796,  3.5318681 ],
                              [ 7.93047434,  0.4159936 ,  5.88644684],
                              [ 0.72107876,  2.91340148,  5.88644684],
                              [ 2.16295788,  5.41080937,  5.88644684],
                              [-0.72080035,  5.41080937,  5.88644684],
                              [ 5.04671611,  5.41080937,  5.88644684],
                              [ 2.16295788,  0.4159936 ,  5.88644684],
                              [ 3.60483699,  2.91340148,  5.88644684],
                              [ 5.04671611,  0.4159936 ,  5.88644684],
                              [ 6.48859522,  2.91340148,  5.88644684]]*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, -6.75261456523e-33, 0.0]*Angstrom
vector_b = [-4.32563734556, 7.49222365763, 4.09015759516e-16]*Angstrom
vector_c = [0.0, 0.0, 7.06373620597]*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, Gold, Gold, Gold, Gold,
                            Gold, Gold, Gold, Gold, Gold]

# Define coordinates
right_electrode_coordinates = [[ 0.72080035,  4.57882217,  1.17728937],
                               [-0.72107876,  7.07623006,  1.17728937],
                               [ 2.16267947,  7.07623006,  1.17728937],
                               [-2.16295788,  4.57882217,  1.17728937],
                               [-0.72107876,  2.08141429,  1.17728937],
                               [ 3.60455858,  4.57882217,  1.17728937],
                               [-3.60483699,  7.07623006,  1.17728937],
                               [ 5.0464377 ,  2.08141429,  1.17728937],
                               [ 2.16267947,  2.08141429,  1.17728937],
                               [ 7.93019593,  0.4164757 ,  3.5318681 ],
                               [ 3.60455858,  2.91388358,  3.5318681 ],
                               [ 6.48831681,  2.91388358,  3.5318681 ],
                               [ 2.16267947,  5.41129147,  3.5318681 ],
                               [ 5.0464377 ,  0.4164757 ,  3.5318681 ],
                               [ 5.0464377 ,  5.41129147,  3.5318681 ],
                               [-0.72107876,  5.41129147,  3.5318681 ],
                               [ 0.72080035,  2.91388358,  3.5318681 ],
                               [ 2.16267947,  0.4164757 ,  3.5318681 ],
                               [ 2.16267947,  3.74635288,  5.88644684],
                               [-2.16295788,  6.24376076,  5.88644684],
                               [ 0.72080035,  6.24376076,  5.88644684],
                               [-0.72107876,  3.74635288,  5.88644684],
                               [ 6.48831681,  1.24894499,  5.88644684],
                               [ 0.72080035,  1.24894499,  5.88644684],
                               [ 5.0464377 ,  3.74635288,  5.88644684],
                               [ 3.60455858,  1.24894499,  5.88644684],
                               [ 3.60455858,  6.24376076,  5.88644684]]*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, -6.75261456523e-33, 0.0]*Angstrom
vector_b = [-4.32563734556, 7.49222365763, 4.09015759516e-16]*Angstrom
vector_c = [0.0, 0.0, 21.4923367139]*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, Sulfur, Carbon, Hydrogen, Hydrogen,
                           Carbon, Carbon, Carbon, Carbon, Hydrogen, Hydrogen, Carbon, Sulfur,
                           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 = [[  0.72107876,   1.24846289,   1.17728937],
                              [  3.60483699,   1.24846289,   1.17728937],
                              [  6.48859522,   1.24846289,   1.17728937],
                              [ -0.72080035,   3.74587078,   1.17728937],
                              [  5.04671611,   3.74587078,   1.17728937],
                              [  2.16295788,   3.74587078,   1.17728937],
                              [  3.60483699,   6.24327867,   1.17728937],
                              [ -2.16267947,   6.24327867,   1.17728937],
                              [  0.72107876,   6.24327867,   1.17728937],
                              [  2.16295788,   2.08093219,   3.5318681 ],
                              [ -0.72080035,   2.08093219,   3.5318681 ],
                              [  5.04671611,   2.08093219,   3.5318681 ],
                              [ -3.60455858,   7.07574796,   3.5318681 ],
                              [ -2.16267947,   4.57834008,   3.5318681 ],
                              [  2.16295788,   7.07574796,   3.5318681 ],
                              [  3.60483699,   4.57834008,   3.5318681 ],
                              [  0.72107876,   4.57834008,   3.5318681 ],
                              [ -0.72080035,   7.07574796,   3.5318681 ],
                              [  7.93047434,   0.4159936 ,   5.88644684],
                              [  0.72107876,   2.91340148,   5.88644684],
                              [  2.16295788,   5.41080937,   5.88644684],
                              [ -0.72080035,   5.41080937,   5.88644684],
                              [  5.04671611,   5.41080937,   5.88644684],
                              [  2.16295788,   0.4159936 ,   5.88644684],
                              [  3.60483699,   2.91340148,   5.88644684],
                              [  5.04671611,   0.4159936 ,   5.88644684],
                              [  6.48859522,   2.91340148,   5.88644684],
                              [  2.16295788,   3.74587078,   7.59644684],
                              [  2.16295788,   3.74587078,   9.34644649],
                              [  1.07984884,   5.62136998,   9.49557904],
                              [  3.24594854,   1.87057656,   9.49607409],
                              [  1.55660877,   4.79581794,  10.04613074],
                              [  2.76904869,   2.69637089,  10.04634658],
                              [  1.55658866,   4.79585277,  11.44599014],
                              [  2.76902857,   2.69640571,  11.44620598],
                              [  1.0796888 ,   5.6216471 ,  11.99626262],
                              [  3.24578851,   1.87085367,  11.99675768],
                              [  2.16267947,   3.74635288,  12.14589023],
                              [  2.16267947,   3.74635288,  13.89588988],
                              [  0.72080035,   4.57882217,  15.60588988],
                              [ -0.72107876,   7.07623006,  15.60588988],
                              [  2.16267947,   7.07623006,  15.60588988],
                              [ -2.16295788,   4.57882217,  15.60588988],
                              [ -0.72107876,   2.08141429,  15.60588988],
                              [  3.60455858,   4.57882217,  15.60588988],
                              [ -3.60483699,   7.07623006,  15.60588988],
                              [  5.0464377 ,   2.08141429,  15.60588988],
                              [  2.16267947,   2.08141429,  15.60588988],
                              [  7.93019593,   0.4164757 ,  17.96046861],
                              [  3.60455858,   2.91388358,  17.96046861],
                              [  6.48831681,   2.91388358,  17.96046861],
                              [  2.16267947,   5.41129147,  17.96046861],
                              [  5.0464377 ,   0.4164757 ,  17.96046861],
                              [  5.0464377 ,   5.41129147,  17.96046861],
                              [ -0.72107876,   5.41129147,  17.96046861],
                              [  0.72080035,   2.91388358,  17.96046861],
                              [  2.16267947,   0.4164757 ,  17.96046861],
                              [  2.16267947,   3.74635288,  20.31504735],
                              [ -2.16295788,   6.24376076,  20.31504735],
                              [  0.72080035,   6.24376076,  20.31504735],
                              [ -0.72107876,   3.74635288,  20.31504735],
                              [  6.48831681,   1.24894499,  20.31504735],
                              [  0.72080035,   1.24894499,  20.31504735],
                              [  5.0464377 ,   3.74635288,  20.31504735],
                              [  3.60455858,   1.24894499,  20.31504735],
                              [  3.60455858,   6.24376076,  20.31504735]]*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 = [
    CerdaHuckelParameters.Gold_fcc_Basis,
    HoffmannHuckelParameters.Carbon_Basis,
    HoffmannHuckelParameters.Hydrogen_Basis,
    HoffmannHuckelParameters.Sulfur_Basis,
    ]

#----------------------------------------
# Numerical Accuracy Settings
#----------------------------------------
left_electrode_k_point_sampling = MonkhorstPackGrid(
    shift_to_gamma=(False, False, False),
    na=3,
    nb=3,
    nc=100,
    )
left_electrode_numerical_accuracy_parameters = NumericalAccuracyParameters(
    k_point_sampling=left_electrode_k_point_sampling,
    )

right_electrode_k_point_sampling = MonkhorstPackGrid(
    shift_to_gamma=(False, False, False),
    na=3,
    nb=3,
    nc=100,
    )
right_electrode_numerical_accuracy_parameters = NumericalAccuracyParameters(
    k_point_sampling=right_electrode_k_point_sampling,
    )

device_k_point_sampling = MonkhorstPackGrid(
    shift_to_gamma=(False, False, False),
    na=3,
    nb=3,
    nc=100,
    )
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()]]
    )

#----------------------------------------
# Device Algorithm Settings
#----------------------------------------
initial_density_type = EquivalentBulk(
    electrode_constraint_length=10.0*Angstrom,
    )

device_algorithm_parameters = DeviceAlgorithmParameters(
    initial_density_type=initial_density_type,
    )

#----------------------------------------
# Electrode Calculators
#----------------------------------------
left_electrode_calculator = HuckelCalculator(
    basis_set=basis_set,
    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 = HuckelCalculator(
    basis_set=basis_set,
    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 = DeviceHuckelCalculator(
    basis_set=basis_set,
    numerical_accuracy_parameters=device_numerical_accuracy_parameters,
    iteration_control_parameters=device_iteration_control_parameters,
    device_algorithm_parameters=device_algorithm_parameters,
    electrode_calculators=
        [left_electrode_calculator, right_electrode_calculator],
    )

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

# -------------------------------------------------------------
# Transmission Spectrum
# -------------------------------------------------------------
kpoint_grid = MonkhorstPackGrid(
    shift_to_gamma=(False, False, False),
    na=1,
    nb=1,
    nc=1,
    )

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