from ATK.TwoProbe import *
from ATK.MPI import processIsMaster

# Generate time stamp
if processIsMaster():
    import platform, time
    print '#',time.ctime()
    print '#',platform.node(),platform.platform()+'\n'

# Opening vnlfile
if processIsMaster(): file = VNLFile('femgofe_antipara.vnl')

# Scattering elements and coordinates
scattering_elements = [Iron,      Iron,      Iron,      Iron,      
                       Oxygen,    Oxygen,    Oxygen,    Oxygen,    
                       Magnesium, Magnesium, Magnesium, Magnesium, 
                       Iron,      Iron,      Iron,      Iron]
scattering_coordinates = [[  0.7165 ,   0.7165 ,   1.433  ],
                          [  2.1495 ,   2.1495 ,   2.866  ],
                          [  0.7165 ,   0.7165 ,   4.299  ],
                          [  2.1495 ,   2.1495 ,   5.732  ],
                          [  2.1495 ,   2.1495 ,   7.932  ],
                          [  0.7165 ,   0.7165 ,  10.12735],
                          [  2.1495 ,   2.1495 ,  12.3227 ],
                          [  0.7165 ,   0.7165 ,  14.51805],
                          [  0.7165 ,   0.7165 ,   7.932  ],
                          [  2.1495 ,   2.1495 ,  10.12735],
                          [  0.7165 ,   0.7165 ,  12.3227 ],
                          [  2.1495 ,   2.1495 ,  14.51805],
                          [  0.7165 ,   0.7165 ,  16.71805],
                          [  2.1495 ,   2.1495 ,  18.15105],
                          [  0.7165 ,   0.7165 ,  19.58405],
                          [  2.1495 ,   2.1495 ,  21.01705]]*Angstrom
        

# Set up left electrode
left_electrode_elements = [Iron, Iron, Iron, Iron]
left_electrode_coordinates = [[ 0.7165,  0.7165,  0.7165],
                              [ 2.1495,  2.1495,  2.1495],
                              [ 0.7165,  0.7165,  3.5825],
                              [ 2.1495,  2.1495,  5.0155]]*Angstrom

left_electrode_cell = [[ 2.866,  0.   ,  0.   ],
                       [ 0.   ,  2.866,  0.   ],
                       [ 0.   ,  0.   ,  5.732]]*Angstrom

# Set up right electrode
right_electrode_elements = [Iron, Iron, Iron, Iron]
right_electrode_coordinates = [[ 0.7165,  0.7165,  0.7165],
                               [ 2.1495,  2.1495,  2.1495],
                               [ 0.7165,  0.7165,  3.5825],
                               [ 2.1495,  2.1495,  5.0155]]*Angstrom

right_electrode_cell = [[ 2.866,  0.   ,  0.   ],
                        [ 0.   ,  2.866,  0.   ],
                        [ 0.   ,  0.   ,  5.732]]*Angstrom

# Set up electrodes
left_electrode_configuration = PeriodicAtomConfiguration(
    left_electrode_cell,
    left_electrode_elements,
    left_electrode_coordinates
    )

right_electrode_configuration = PeriodicAtomConfiguration(
    right_electrode_cell,
    right_electrode_elements,
    right_electrode_coordinates
    )

# Set up two-probe configuration
twoprobe_configuration = TwoProbeConfiguration(
    (left_electrode_configuration,right_electrode_configuration),
    scattering_elements,
    scattering_coordinates,
    electrode_repetitions=[[1,1],[1,1]],
    equivalent_atoms=([0,0],[3,15])
    )
if processIsMaster(): nlPrint(twoprobe_configuration)
if processIsMaster(): file.addToSample(twoprobe_configuration, 'twoprobe_configuration')

######################################################################
# Central region parameters
######################################################################
exchange_correlation_type = LDA.PZ

iteration_mixing_parameters = iterationMixingParameters(
    algorithm = IterationMixing.Pulay,
    diagonal_mixing_parameter = 0.05,
    quantity = IterationMixing.Hamiltonian,
    history_steps = 6
)

electron_density_parameters = electronDensityParameters(
    mesh_cutoff = 150.0*Rydberg,
    initial_scaled_spin = [ 0.56,  0.56,  0.56,  0.56,  
        0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  
        -0.56,  -0.56,  -0.56,  -0.56 ]
)

basis_set_parameters_Common = basisSetParameters(
    type = DoubleZetaPolarized,
    radial_sampling_dr = 0.001*Bohr,
    energy_shift = 0.01*Rydberg,
    delta_rinn = 0.8,
    v0 = 40.0*Rydberg,
    charge = 0.0,
    split_norm = 0.15
)

basis_set_parameters_Iron = basisSetParameters(
    type = SingleZetaPolarized,
    energy_shift = 0.01*Rydberg,
    delta_rinn = 0.8,
    v0 = 40.0*Rydberg,
    charge = 0.0,
    split_norm = 0.15,
    element = Iron
)

basis_set_parameters = [
    basis_set_parameters_Common,
    basis_set_parameters_Iron
    ]

iteration_control_parameters = iterationControlParameters(
    tolerance = 1e-005,
    criterion = IterationControl.TotalEnergy,
    max_steps = 100
)

electrode_voltages = (0.0,0.0)*Volt

two_probe_algorithm_parameters = twoProbeAlgorithmParameters(
    electrode_constraint = ElectrodeConstraints.Off,
    initial_density_type = InitialDensityType.NeutralAtom
)

energy_contour_integral_parameters = energyContourIntegralParameters(
    circle_points = 50,
    integral_lower_bound = 7.0*Rydberg,
    fermi_line_points = 10,
    fermi_function_poles = 4,
    real_axis_infinitesimal = 0.01*electronVolt,
    real_axis_point_density = 0.005*electronVolt
)

two_center_integral_parameters = twoCenterIntegralParameters(
    cutoff = 2500.0*Rydberg,
    points = 1024
)

######################################################################
# Left electrode parameters
######################################################################
left_electrode_electron_density_parameters = electronDensityParameters(
    mesh_cutoff = 150.0*Rydberg,
    initial_scaled_spin = [ 0.56,  0.56,  0.56,  0.56 ]
)

left_electrode_iteration_control_parameters = iterationControlParameters(
    tolerance = 1e-005,
    criterion = IterationControl.TotalEnergy,
    max_steps = 100
)

left_electrode_brillouin_zone_integration_parameters = brillouinZoneIntegrationParameters(
    monkhorst_pack_parameters = (10, 10, 100)
)

left_electrode_iteration_mixing_parameters = iterationMixingParameters(
    algorithm = IterationMixing.Pulay,
    diagonal_mixing_parameter = 0.05,
    quantity = IterationMixing.Hamiltonian,
    history_steps = 6
)

left_electrode_eigenstate_occupation_parameters = eigenstateOccupationParameters(
    temperature = 1200.0*Kelvin
)

######################################################################
# Collect left electrode parameters
######################################################################
left_electrode_parameters = ElectrodeParameters(
    brillouin_zone_integration_parameters = left_electrode_brillouin_zone_integration_parameters,
    electron_density_parameters = left_electrode_electron_density_parameters,
    eigenstate_occupation_parameters = left_electrode_eigenstate_occupation_parameters,
    iteration_mixing_parameters = left_electrode_iteration_mixing_parameters,
    iteration_control_parameters = left_electrode_iteration_control_parameters
)

######################################################################
# Right electrode parameters
######################################################################
right_electrode_electron_density_parameters = electronDensityParameters(
    mesh_cutoff = 150.0*Rydberg,
    initial_scaled_spin = [-0.56, -0.56, -0.56, -0.56 ]
)

right_electrode_iteration_control_parameters = iterationControlParameters(
    tolerance = 1e-005,
    criterion = IterationControl.TotalEnergy,
    max_steps = 100
)

right_electrode_brillouin_zone_integration_parameters = brillouinZoneIntegrationParameters(
    monkhorst_pack_parameters = (10, 10, 100)
)

right_electrode_iteration_mixing_parameters = iterationMixingParameters(
    algorithm = IterationMixing.Pulay,
    diagonal_mixing_parameter = 0.05,
    quantity = IterationMixing.Hamiltonian,
    history_steps = 6
)

right_electrode_eigenstate_occupation_parameters = eigenstateOccupationParameters(
    temperature = 1200.0*Kelvin
)

######################################################################
# Collect right electrode parameters
######################################################################
right_electrode_parameters = ElectrodeParameters(
    brillouin_zone_integration_parameters = right_electrode_brillouin_zone_integration_parameters,
    electron_density_parameters = right_electrode_electron_density_parameters,
    eigenstate_occupation_parameters = right_electrode_eigenstate_occupation_parameters,
    iteration_mixing_parameters = right_electrode_iteration_mixing_parameters,
    iteration_control_parameters = right_electrode_iteration_control_parameters
)

######################################################################
# Initialize self-consistent field calculation
######################################################################
two_probe_method = TwoProbeMethod(
    electrode_parameters = (left_electrode_parameters,right_electrode_parameters),
    exchange_correlation_type = exchange_correlation_type,
    iteration_mixing_parameters = iteration_mixing_parameters,
    electron_density_parameters = electron_density_parameters,
    basis_set_parameters = basis_set_parameters,
    iteration_control_parameters = iteration_control_parameters,
    energy_contour_integral_parameters = energy_contour_integral_parameters,
    two_center_integral_parameters = two_center_integral_parameters,
    electrode_voltages = electrode_voltages,
    algorithm_parameters = two_probe_algorithm_parameters
)
if processIsMaster(): nlPrint(two_probe_method)

runtime_parameters = runtimeParameters(
    verbosity_level = 1,
    checkpoint_filename = 'femgofe_antipara.nc'
)

# Perform self-consistent field calculation
scf = executeSelfConsistentCalculation(
    twoprobe_configuration,
    two_probe_method,
    runtime_parameters = runtime_parameters
)