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

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

# Set up lattice
vector_a = [1.3325, -2.30795770109, 0.0]*Angstrom
vector_b = [1.3325, 2.30795770109, 0.0]*Angstrom
vector_c = [0.0, 0.0, 4.947]*Angstrom
left_electrode_lattice = UnitCell(vector_a, vector_b, vector_c)

# Define elements
left_electrode_elements = [Zinc, Zinc]

# Define coordinates
left_electrode_coordinates = [[ 1.3325    ,  0.76931925,  1.23675   ],
                              [ 1.3325    , -0.76931925,  3.71025   ]]*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 = [1.3325, -2.30795770109, 0.0]*Angstrom
vector_b = [1.3325, 2.30795770109, 0.0]*Angstrom
vector_c = [0.0, 0.0, 4.947]*Angstrom
right_electrode_lattice = UnitCell(vector_a, vector_b, vector_c)

# Define elements
right_electrode_elements = [Zinc, Zinc]

# Define coordinates
right_electrode_coordinates = [[ 1.0280478 ,  0.59354366,  1.23675   ],
                               [ 1.0280478 , -0.94509484,  3.71025   ]]*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 = [1.3325, -2.30795770109, 0.0]*Angstrom
vector_b = [1.3325, 2.30795770109, 0.0]*Angstrom
vector_c = [0.0, 0.0, 40.5741587317]*Angstrom
central_region_lattice = UnitCell(vector_a, vector_b, vector_c)

# Define elements
central_region_elements = [Zinc, Zinc, Zinc, Zinc, Zinc, Zinc, Oxygen, Zinc, Oxygen, Zinc,
                           Oxygen, Zinc, Oxygen, Zinc, Oxygen, Zinc, Zinc, Zinc, Zinc, Zinc,
                           Zinc]

# Define coordinates
central_region_coordinates = [[  1.3325    ,   0.76931925,   1.23675   ],
                              [  1.3325    ,  -0.76931925,   3.71025   ],
                              [  1.3325    ,   0.76931925,   6.18375   ],
                              [  1.3325    ,  -0.76931925,   8.65725   ],
                              [  1.3325    ,   0.76931925,  11.13075   ],
                              [  1.3325    ,  -0.76931925,  13.60425   ],
                              [  1.3325    ,  -0.76931925,  15.60909881],
                              [  2.3605478 ,   0.17577554,  16.27738175],
                              [  2.3605478 ,   0.17577554,  18.28223056],
                              [  1.0280478 ,   0.59354366,  18.95051349],
                              [  1.0280478 ,   0.59354366,  20.9553623 ],
                              [  1.3325    ,  -0.76931925,  21.62364524],
                              [  1.3325    ,  -0.76931925,  23.62849405],
                              [  2.3605478 ,   0.17577554,  24.29677699],
                              [  2.3605478 ,   0.17577554,  26.3016258 ],
                              [  1.0280478 ,   0.59354366,  26.96990873],
                              [  1.0280478 ,  -0.94509484,  29.44340873],
                              [  1.0280478 ,   0.59354366,  31.91690873],
                              [  1.0280478 ,  -0.94509484,  34.39040873],
                              [  1.0280478 ,   0.59354366,  36.86390873],
                              [  1.0280478 ,  -0.94509484,  39.33740873]]*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],
    equivalent_electrode_lengths=[4.947, 4.947]*Angstrom,
    )

# -------------------------------------------------------------
# Calculator
# -------------------------------------------------------------
#----------------------------------------
# Hamiltonian Parametrization
#----------------------------------------
hamiltonian_parametrization = SlaterKosterHamiltonianParametrization(
    basis_set=DFTBDirectory("dftb/znorg-0-1/"))

#----------------------------------------
# Pair Potentials
#----------------------------------------
pair_potentials = DFTBDirectory("dftb/znorg-0-1/")

#----------------------------------------
# Numerical Accuracy Settings
#----------------------------------------
left_electrode_k_point_sampling = MonkhorstPackGrid(
    na=5,
    nb=5,
    nc=100,
    )
left_electrode_numerical_accuracy_parameters = NumericalAccuracyParameters(
    k_point_sampling=left_electrode_k_point_sampling,
    density_mesh_cutoff=10.0*Hartree,
    )

right_electrode_k_point_sampling = MonkhorstPackGrid(
    na=5,
    nb=5,
    nc=100,
    )
right_electrode_numerical_accuracy_parameters = NumericalAccuracyParameters(
    k_point_sampling=right_electrode_k_point_sampling,
    density_mesh_cutoff=10.0*Hartree,
    )

device_k_point_sampling = MonkhorstPackGrid(
    na=5,
    nb=5,
    nc=100,
    )
device_numerical_accuracy_parameters = NumericalAccuracyParameters(
    k_point_sampling=device_k_point_sampling,
    density_mesh_cutoff=10.0*Hartree,
    )

#----------------------------------------
# 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 = SemiEmpiricalCalculator(
    hamiltonian_parametrization=hamiltonian_parametrization,
    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 = SemiEmpiricalCalculator(
    hamiltonian_parametrization=hamiltonian_parametrization,
    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 = DeviceSemiEmpiricalCalculator(
    hamiltonian_parametrization=hamiltonian_parametrization,
    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)

# -------------------------------------------------------------
# Initial State
# -------------------------------------------------------------
old_calculation = nlread('zero-bias.hdf5', DeviceConfiguration)[0]

device_configuration.setCalculator(
    calculator,
    initial_state=old_calculation,
)
nlsave('iv-curve.hdf5', device_configuration)
nlprint(device_configuration)

# -------------------------------------------------------------
# IV Characteristics
# -------------------------------------------------------------
# Kpoint sampling
kpoint_grid = MonkhorstPackGrid(
    na=20,
    nb=20,
    )

# Gate-source voltages
gate_source_voltages = numpy.linspace(0, 0, 1)*Volt

# Drain-source voltages
drain_source_voltages = numpy.linspace(0, 1, 5)*Volt

# File name.
filename = u'iv-curve.hdf5'

iv_characteristics = IVCharacteristics(
    configuration=device_configuration,
    filename=filename,
    object_id='ivcharacteristics',
    gate_regions=[],
    gate_source_voltages=gate_source_voltages,
    drain_source_voltages=drain_source_voltages,
    energies=None,
    kpoints=kpoint_grid,
    self_energy_calculator=RecursionSelfEnergy(),
    energy_zero_parameter=AverageFermiLevel,
    infinitesimal=1e-06*eV,
    log_filename_prefix='ivcharacteristics_',
    number_of_processes_per_task=None,
)
iv_characteristics.update()
