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

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

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

# Define elements
left_electrode_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
left_electrode_coordinates = [[  5.            ,   5.750915      ,   0.615250427611],
                              [  5.            ,   8.592635      ,   0.615250427611],
                              [  5.            ,  17.117795      ,   0.615250427611],
                              [  5.            ,  18.538655      ,   0.615250427611],
                              [  5.            ,   6.461345      ,   1.845751282832],
                              [  5.            ,   7.882205      ,   1.845751282832],
                              [  5.            ,  16.407365      ,   1.845751282832],
                              [  5.            ,  19.249085      ,   1.845751282832],
                              [  5.            ,   5.750915      ,   3.076252138053],
                              [  5.            ,   8.592635      ,   3.076252138053],
                              [  5.            ,  17.117795      ,   3.076252138053],
                              [  5.            ,  18.538655      ,   3.076252138053],
                              [  5.            ,   6.461345      ,   4.306752993274],
                              [  5.            ,   7.882205      ,   4.306752993274],
                              [  5.            ,  16.407365      ,   4.306752993274],
                              [  5.            ,  19.249085      ,   4.306752993274],
                              [  5.            ,   5.750915      ,   5.537253848495],
                              [  5.            ,   8.592635      ,   5.537253848495],
                              [  5.            ,  17.117795      ,   5.537253848495],
                              [  5.            ,  18.538655      ,   5.537253848495],
                              [  5.            ,   6.461345      ,   6.767754703716],
                              [  5.            ,   7.882205      ,   6.767754703716],
                              [  5.            ,  16.407365      ,   6.767754703716],
                              [  5.            ,  19.249085      ,   6.767754703716]]*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('lead1', [2, 3, 6, 7, 10, 11, 14, 15, 18, 19, 22, 23])
left_electrode.addTags('lead2', [0, 1, 4, 5, 8, 9, 12, 13, 16, 17, 20, 21])

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

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

# Define elements
right_electrode_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
right_electrode_coordinates = [[  5.            ,   5.750915      ,   0.615250427611],
                               [  5.            ,   8.592635      ,   0.615250427611],
                               [  5.            ,  17.117795      ,   0.615250427611],
                               [  5.            ,  18.538655      ,   0.615250427611],
                               [  5.            ,   6.461345      ,   1.845751282832],
                               [  5.            ,   7.882205      ,   1.845751282832],
                               [  5.            ,  16.407365      ,   1.845751282832],
                               [  5.            ,  19.249085      ,   1.845751282832],
                               [  5.            ,   5.750915      ,   3.076252138053],
                               [  5.            ,   8.592635      ,   3.076252138053],
                               [  5.            ,  17.117795      ,   3.076252138053],
                               [  5.            ,  18.538655      ,   3.076252138053],
                               [  5.            ,   6.461345      ,   4.306752993274],
                               [  5.            ,   7.882205      ,   4.306752993274],
                               [  5.            ,  16.407365      ,   4.306752993274],
                               [  5.            ,  19.249085      ,   4.306752993274],
                               [  5.            ,   5.750915      ,   5.537253848495],
                               [  5.            ,   8.592635      ,   5.537253848495],
                               [  5.            ,  17.117795      ,   5.537253848495],
                               [  5.            ,  18.538655      ,   5.537253848495],
                               [  5.            ,   6.461345      ,   6.767754703716],
                               [  5.            ,   7.882205      ,   6.767754703716],
                               [  5.            ,  16.407365      ,   6.767754703716],
                               [  5.            ,  19.249085      ,   6.767754703716]]*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('lead3', [2, 3, 6, 7, 10, 11, 14, 15, 18, 19, 22, 23])
right_electrode.addTags('lead4', [0, 1, 4, 5, 8, 9, 12, 13, 16, 17, 20, 21])

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

# Set up lattice
vector_a = [10.0, 0.0, 0.0]*Angstrom
vector_b = [0.0, 25.0, 0.0]*Angstrom
vector_c = [0.0, 0.0, 29.5320205253]*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,
                           Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
                           Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
                           Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
                           Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
                           Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
                           Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
                           Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
                           Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
                           Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
                           Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
                           Carbon, 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 = [[  5.            ,   5.750915      ,   0.615250427611],
                              [  5.            ,   8.592635      ,   0.615250427611],
                              [  5.            ,  17.117795      ,   0.615250427611],
                              [  5.            ,  18.538655      ,   0.615250427611],
                              [  5.            ,   6.461345      ,   1.845751282832],
                              [  5.            ,   7.882205      ,   1.845751282832],
                              [  5.            ,  16.407365      ,   1.845751282832],
                              [  5.            ,  19.249085      ,   1.845751282832],
                              [  5.            ,   5.750915      ,   3.076252138053],
                              [  5.            ,   8.592635      ,   3.076252138053],
                              [  5.            ,  17.117795      ,   3.076252138053],
                              [  5.            ,  18.538655      ,   3.076252138053],
                              [  5.            ,   6.461345      ,   4.306752993274],
                              [  5.            ,   7.882205      ,   4.306752993274],
                              [  5.            ,  16.407365      ,   4.306752993274],
                              [  5.            ,  19.249085      ,   4.306752993274],
                              [  5.            ,   5.750915      ,   5.537253848495],
                              [  5.            ,   8.592635      ,   5.537253848495],
                              [  5.            ,  17.117795      ,   5.537253848495],
                              [  5.            ,  18.538655      ,   5.537253848495],
                              [  5.            ,   6.461345      ,   6.767754703716],
                              [  5.            ,   7.882205      ,   6.767754703716],
                              [  5.            ,  16.407365      ,   6.767754703716],
                              [  5.            ,  19.249085      ,   6.767754703716],
                              [  5.            ,   5.750915      ,   7.998255558938],
                              [  5.            ,   8.592635      ,   7.998255558938],
                              [  5.            ,  17.117795      ,   7.998255558938],
                              [  5.            ,  18.538655      ,   7.998255558938],
                              [  5.            ,   6.461345      ,   9.228756414159],
                              [  5.            ,   7.882205      ,   9.228756414159],
                              [  5.            ,  10.723925      ,   9.228756414159],
                              [  5.            ,  12.144785      ,   9.228756414159],
                              [  5.            ,  14.986505      ,   9.228756414159],
                              [  5.            ,  16.407365      ,   9.228756414159],
                              [  5.            ,  19.249085      ,   9.228756414159],
                              [  5.            ,   5.750915      ,  10.45925726938 ],
                              [  5.            ,   8.592635      ,  10.45925726938 ],
                              [  5.            ,  10.013495      ,  10.45925726938 ],
                              [  5.            ,  12.855215      ,  10.45925726938 ],
                              [  5.            ,  14.276075      ,  10.45925726938 ],
                              [  5.            ,  17.117795      ,  10.45925726938 ],
                              [  5.            ,  18.538655      ,  10.45925726938 ],
                              [  5.            ,   6.461345      ,  11.689758124601],
                              [  5.            ,   7.882205      ,  11.689758124601],
                              [  5.            ,  10.723925      ,  11.689758124601],
                              [  5.            ,  12.144785      ,  11.689758124601],
                              [  5.            ,  14.986505      ,  11.689758124601],
                              [  5.            ,  16.407365      ,  11.689758124601],
                              [  5.            ,  19.249085      ,  11.689758124601],
                              [  5.            ,  14.257883546055,  12.773287892691],
                              [  5.            ,  12.873406454041,  12.773287892982],
                              [  5.            ,   5.750915      ,  12.920258979822],
                              [  5.            ,   8.592635      ,  12.920258979822],
                              [  5.            ,  10.013495      ,  12.920258979822],
                              [  5.            ,  17.117795      ,  12.920258979822],
                              [  5.            ,  18.538655      ,  12.920258979822],
                              [  5.            ,  12.19456505321 ,  14.030708696767],
                              [  5.            ,  14.936724946796,  14.030708696776],
                              [  5.            ,   6.461345      ,  14.150759835043],
                              [  5.            ,   7.882205      ,  14.150759835043],
                              [  5.            ,  10.723925      ,  14.150759835043],
                              [  5.            ,  16.407365      ,  14.150759835043],
                              [  5.            ,  19.249085      ,  14.150759835043],
                              [  5.            ,  13.565644999904,  14.750691998909],
                              [  5.            ,   5.750915      ,  15.381260690264],
                              [  5.            ,   8.592635      ,  15.381260690264],
                              [  5.            ,  10.013495      ,  15.381260690264],
                              [  5.            ,  17.117795      ,  15.381260690264],
                              [  5.            ,  18.538655      ,  15.381260690264],
                              [  5.            ,  13.565645000211,  16.011829097332],
                              [  5.            ,   6.461345      ,  16.611761545486],
                              [  5.            ,   7.882205      ,  16.611761545486],
                              [  5.            ,  10.723925      ,  16.611761545486],
                              [  5.            ,  16.407365      ,  16.611761545486],
                              [  5.            ,  19.249085      ,  16.611761545486],
                              [  5.            ,  12.194565051441,  16.7318126837  ],
                              [  5.            ,  14.936724948547,  16.731812683718],
                              [  5.            ,   5.750915      ,  17.842262400707],
                              [  5.            ,   8.592635      ,  17.842262400707],
                              [  5.            ,  10.013495      ,  17.842262400707],
                              [  5.            ,  17.117795      ,  17.842262400707],
                              [  5.            ,  18.538655      ,  17.842262400707],
                              [  5.            ,  14.257883594282,  17.989233630584],
                              [  5.            ,  12.873406405504,  17.989233631238],
                              [  5.            ,   6.461345      ,  19.072763255928],
                              [  5.            ,   7.882205      ,  19.072763255928],
                              [  5.            ,  10.723925      ,  19.072763255928],
                              [  5.            ,  12.144785      ,  19.072763255928],
                              [  5.            ,  14.986505      ,  19.072763255928],
                              [  5.            ,  16.407365      ,  19.072763255928],
                              [  5.            ,  19.249085      ,  19.072763255928],
                              [  5.            ,   5.750915      ,  20.303264111149],
                              [  5.            ,   8.592635      ,  20.303264111149],
                              [  5.            ,  10.013495      ,  20.303264111149],
                              [  5.            ,  12.855215      ,  20.303264111149],
                              [  5.            ,  14.276075      ,  20.303264111149],
                              [  5.            ,  17.117795      ,  20.303264111149],
                              [  5.            ,  18.538655      ,  20.303264111149],
                              [  5.            ,   6.461345      ,  21.53376496637 ],
                              [  5.            ,   7.882205      ,  21.53376496637 ],
                              [  5.            ,  10.723925      ,  21.53376496637 ],
                              [  5.            ,  12.144785      ,  21.53376496637 ],
                              [  5.            ,  14.986505      ,  21.53376496637 ],
                              [  5.            ,  16.407365      ,  21.53376496637 ],
                              [  5.            ,  19.249085      ,  21.53376496637 ],
                              [  5.            ,   5.750915      ,  22.764265821591],
                              [  5.            ,   8.592635      ,  22.764265821591],
                              [  5.            ,  17.117795      ,  22.764265821591],
                              [  5.            ,  18.538655      ,  22.764265821591],
                              [  5.            ,   6.461345      ,  23.994766676813],
                              [  5.            ,   7.882205      ,  23.994766676813],
                              [  5.            ,  16.407365      ,  23.994766676813],
                              [  5.            ,  19.249085      ,  23.994766676813],
                              [  5.            ,   5.750915      ,  25.225267532034],
                              [  5.            ,   8.592635      ,  25.225267532034],
                              [  5.            ,  17.117795      ,  25.225267532034],
                              [  5.            ,  18.538655      ,  25.225267532034],
                              [  5.            ,   6.461345      ,  26.455768387255],
                              [  5.            ,   7.882205      ,  26.455768387255],
                              [  5.            ,  16.407365      ,  26.455768387255],
                              [  5.            ,  19.249085      ,  26.455768387255],
                              [  5.            ,   5.750915      ,  27.686269242476],
                              [  5.            ,   8.592635      ,  27.686269242476],
                              [  5.            ,  17.117795      ,  27.686269242476],
                              [  5.            ,  18.538655      ,  27.686269242476],
                              [  5.            ,   6.461345      ,  28.916770097697],
                              [  5.            ,   7.882205      ,  28.916770097697],
                              [  5.            ,  16.407365      ,  28.916770097697],
                              [  5.            ,  19.249085      ,  28.916770097697]]*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('lead1', [2, 3, 6, 7, 10, 11, 14, 15, 18, 19, 22, 23])
central_region.addTags('lead2', [0, 1, 4, 5, 8, 9, 12, 13, 16, 17, 20, 21])
central_region.addTags('lead3', [107, 108, 111, 112, 115, 116, 119, 120, 123, 124,
                                 127, 128])
central_region.addTags('lead4', [105, 106, 109, 110, 113, 114, 117, 118, 121, 122,
                                 125, 126])

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

# -------------------------------------------------------------
# Calculator
# -------------------------------------------------------------
#----------------------------------------
# Basis Set
#----------------------------------------
basis_set = Hancock.C_Basis

#----------------------------------------
# Numerical Accuracy Settings
#----------------------------------------
left_electrode_k_point_sampling = MonkhorstPackGrid(
    shift_to_gamma=(False, False, False),
    na=1,
    nb=1,
    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=1,
    nb=1,
    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=1,
    nb=1,
    nc=100,
    )
device_numerical_accuracy_parameters = NumericalAccuracyParameters(
    k_point_sampling=device_k_point_sampling,
    )

#----------------------------------------
# Iteration Control Settings
#----------------------------------------
left_electrode_iteration_control_parameters = IterationControlParameters(
    tolerance=0.0001,
    )

right_electrode_iteration_control_parameters = IterationControlParameters(
    tolerance=0.0001,
    )

#----------------------------------------
# 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 = SlaterKosterCalculator(
    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 = SlaterKosterCalculator(
    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 = DeviceSlaterKosterCalculator(
    basis_set=basis_set,
    numerical_accuracy_parameters=device_numerical_accuracy_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('4probe_pi_model.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(0,0,1)*eV,
    kpoints=kpoint_grid,
    energy_zero_parameter=AverageFermiLevel,
    infinitesimal=1e-06*eV,
    self_energy_calculator=RecursionSelfEnergy(),
    )
nlsave('4probe_pi_model.nc', transmission_spectrum)
nlprint(transmission_spectrum)
