from QuantumATK import *

def random_alloy(random_seed, na, nb, nc, num_vacancies, si_ge_ratio=0.5):
    """
    Generate a random SiGe alloy.

    @param random_seed   : The random seed used to generate this crystal.
    @type                : int

    @param na            : The number of times to repeat the primitive cell
                           along the a direction.
    @type                : int

    @param nb            : The number of times to repeat the primitive cell
                           along the b direction.
    @type                : int

    @param nc            : The number of times to repeat the primitive cell
                           along the c direction.
    @type                : int

    @param num_vacancies : The number of vacancies in the crystal.
    @type                : int

    @param si_ge_ratio   : The ratio of Si to Ge atoms.
    @type                : float

    @return BulkConfiguration
    """

    rng = numpy.random.RandomState(random_seed)

    # Define a Silicon lattice.
    lattice = FaceCenteredCubic(5.4306*Angstrom)
    elements = [Silicon, Silicon]
    fractional_coordinates = [[ 0.  ,  0.  ,  0.  ],
                              [ 0.25,  0.25,  0.25]]

    bulk_configuration = BulkConfiguration(
        bravais_lattice=lattice,
        elements=elements,
        fractional_coordinates=fractional_coordinates
        )

    # Repeat the structure the requested number of times.
    bulk_configuration = bulk_configuration.repeat(na, nb, nc)


    n = len(bulk_configuration)
    keep_indicies = rng.choice(range(n), size=n-num_vacancies, replace=False)

    elements = numpy.array(bulk_configuration.elements())[keep_indicies]
    coordinates = bulk_configuration.cartesianCoordinates()[keep_indicies]
    lattice = bulk_configuration.bravaisLattice()

    # Make the alloy.
    n = len(elements)
    num_germanium = int(n * (1.0 - si_ge_ratio))
    germanium_indices = rng.choice(range(n), size=num_germanium, replace=False)
    for i in germanium_indices:
        elements[i] = Germanium

    # Set up configuration
    bulk_configuration = BulkConfiguration(
        bravais_lattice=lattice,
        elements=elements,
        cartesian_coordinates=coordinates
        )

    return bulk_configuration

# Setup the calculator.
potentialSet = StillingerWeber_SiGe_1995()
calculator = TremoloXCalculator(parameters=potentialSet)

# Generate a random alloy.
bulk_configuration = random_alloy(
        random_seed=1,
        na=4,
        nb=4,
        nc=4,
        num_vacancies=1,
        si_ge_ratio=0.5)

# Set the calculator.
bulk_configuration.setCalculator(calculator)

# Optimize the geometry.
bulk_configuration = OptimizeGeometry(
        bulk_configuration,
        max_forces=0.01*eV/Ang,
        max_stress=0.01*eV/Ang**3,
        max_steps=500,
        max_step_length=0.2*Ang,
        trajectory_filename='alloy.hdf5',
        optimizer_method=LBFGS(),
        )

# Save the result to alloy.hdf5.
nlsave('alloy.hdf5', bulk_configuration, 'minimized')