from QuantumATK import *
from AddOns.FET2D.Study.FET2DAtomisticParameterExtraction import FET2DAtomisticParameterExtraction

# Channel configuration
# Set up lattice
lattice = Hexagonal(3.1881795204035988*Angstrom, 11.057859406824896*Angstrom)

# Define elements
elements = [Sulfur, Molybdenum, Sulfur, Sulfur, Molybdenum, Sulfur]

# Define coordinates
fractional_coordinates = [[ 0.333333333333,  0.666666666667,  0.609057255744],
                          [ 0.666666666667,  0.333333333333,  0.75          ],
                          [ 0.333333333333,  0.666666666667,  0.890942744256],
                          [ 0.666666666667,  0.333333333333,  0.109057255744],
                          [ 0.333333333333,  0.666666666667,  0.25          ],
                          [ 0.666666666667,  0.333333333333,  0.390942744256]]

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

# Contact configuration
# Set up lattice
lattice = FaceCenteredCubic(4.07825*Angstrom)

# Define elements
elements = [Gold]

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

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

# Define the Parameter Extraction Study Object.
parameter_extraction = FET2DAtomisticParameterExtraction(
    filename='fet2d_parameter_extraction_1642514327.22.hdf5',
    object_id='FET2DAtomisticParameterExtraction_0',
    channel_material_name='MoS2',
    contact_material_name='Gold',
    channel_thickness=1,
    maxE=1.1 * eV,
    contact_orientation=(0, 0, 1),
    channel_configuration=channel_configuration,
    contact_configuration=contact_configuration,
    log_filename_prefix_default='fet2d_parameter_extraction_',
)

# Start the calculations.
parameter_extraction.update()

# Extract the results.
parameter_extraction_results = parameter_extraction.extractionResults()