from QuantumATK import *
from AddOns.AlloyPlugins.SubstitutionalAlloy.SubstitutionalAlloyWidget import substitutionalAlloy

number_of_runs = 10

# -------------------------------------------------------------
# Calculator
# -------------------------------------------------------------
#----------------------------------------
# Basis Set
#----------------------------------------
basis_set = [
    LDABasis.Gallium_SingleZetaPolarized,
    LDABasis.Arsenic_SingleZetaPolarized,
    LDABasis.Indium_SingleZetaPolarized,
    ]

#----------------------------------------
# Exchange-Correlation
#----------------------------------------
exchange_correlation = MGGA.TB09LDA(c=1.01)

numerical_accuracy_parameters = NumericalAccuracyParameters(
    k_point_sampling=(3, 3, 3),
    )

calculator = LCAOCalculator(
    basis_set=basis_set,
    exchange_correlation=exchange_correlation,
    numerical_accuracy_parameters=numerical_accuracy_parameters,
    )


# Set up lattice with exp. lattice const. of In0.53Ga0.47As
lattice = FaceCenteredCubic(5.8687*Angstrom)

# Define elements
elements = [Indium, Arsenic]

# Define coordinates
fractional_coordinates = [[ 0.  ,  0.  ,  0.  ],
                          [ 0.25,  0.25,  0.25]]

# Set up configuration
primitive_configuration = BulkConfiguration(
    bravais_lattice=lattice,
    elements=elements,
    fractional_coordinates=fractional_coordinates
    )

# Attach calculator to primitive bulk_configuration
primitive_configuration.setCalculator(calculator)

# Repeat InAs
super_inas = primitive_configuration.repeat(3,3,3)

# List of indices corresponding to Indium atoms
indium_indices = []
symbols = super_inas.symbols()

for index,atom in enumerate(symbols):
    if atom=='In':
        indium_indices.append(index)

# Repeat creation of random alloy and effective band structure calculation.
for run in range(number_of_runs):

    # Call the substitutional alloy generator.
    super_ingaas = substitutionalAlloy(super_inas,indium_indices, 'fixed_fraction', 47, Gallium)

    # Set calculator and update.
    super_ingaas.setCalculator(calculator)
    super_ingaas.update()
    nlsave('InGaAs_3x3x3.nc', super_ingaas)

    # -------------------------------------------------------------
    # EffectiveBandstructure
    # -------------------------------------------------------------
    effective_bandstructure = EffectiveBandstructure(
        super_ingaas,
        primitive_configuration,
        energies = numpy.linspace(-2,2.75,500)*eV,
        route=['L','G','X'],
        points_per_segment=101,
        )

    nlsave('InGaAs_3x3x3.nc',effective_bandstructure)