from SMW import *

# Define each layer in the multilayer structure.
layers = []
annealing_parameters = []
relaxation_calculator = []
amorphous_method = MolecularDynamicsMeltQuench(
    melt_temperature=3500*Kelvin,
    melt_time=40.0*ps,
    quench_temperature=300*Kelvin,
    quench_rate=2*37*Kelvin/ps,
    time_step=1.0*fs,
    max_retries=30,
    random_seed=1234,
)

material = MaterialSpecificationsDatabase.HKMG_MATERIALS['Silicon']()
layer_0 = CrystalLayer(
    material,
    length=10.0*Angstrom,
    miller_indices=(1, 0, 0),
    )
layers.append(layer_0)

material = MaterialSpecificationsDatabase.HKMG_MATERIALS['Silica'](
    amorphous_density=2.200*gram/cm**3,
    elements_with_stoichiometry={Silicon: 1, Oxygen: 2},
    )
layer_1 = AmorphousLayer(
    material,
    length=10.0*Angstrom,
    method=amorphous_method,
    )
layers.append(layer_1)

material = MaterialSpecificationsDatabase.HKMG_MATERIALS['Hafnia (monoclinic)'](
    amorphous_density=9.390*gram/cm**3,
    elements_with_stoichiometry={Hafnium: 1, Oxygen: 2},
    )
layer_2 = AmorphousLayer(
    material,
    length=10.0*Angstrom,
    method=amorphous_method,
    )
layers.append(layer_2)

material = MaterialSpecificationsDatabase.HKMG_MATERIALS['Titanium nitride'](
    amorphous_density=5.210*gram/cm**3,
    elements_with_stoichiometry={Titanium: 1, Nitrogen: 1},
    )
layer_3 = AmorphousLayer(
    material,
    length=10.0*Angstrom,
    method=amorphous_method,
    )
layers.append(layer_3)


# -------------------------------------------------------------
# Setup annealing parameter used for optimizing interfaces.
# -------------------------------------------------------------
initial_velocity = MaxwellBoltzmannDistribution(
    temperature=600.0*Kelvin,
    remove_center_of_mass_momentum=True,
    random_seed=None,
    enforce_temperature=True,
)

method = Langevin(
    time_step=1*femtoSecond,
    reservoir_temperature=600*Kelvin,
    friction=0.05*femtoSecond**-1,
    initial_velocity=initial_velocity,
)
md_steps = 40000
annealing_parameters.append(
    AnnealingParameters(
        md_method=method,
        md_steps=md_steps
    )
)

initial_velocity = MaxwellBoltzmannDistribution(
    temperature=600.0*Kelvin,
    remove_center_of_mass_momentum=True,
    random_seed=None,
    enforce_temperature=True,
)

method = Langevin(
    time_step=1*femtoSecond,
    reservoir_temperature=600*Kelvin,
    friction=0.05*femtoSecond**-1,
    initial_velocity=initial_velocity,
    heating_rate=-50*Kelvin/picoSecond,
)
md_steps = 10000
annealing_parameters.append(
    AnnealingParameters(
        md_method=method,
        md_steps=md_steps
    )
)


Si_SiO2_calculator = TremoloXCalculator(QuantumATK_MTP_Si_SiO2_2022())
relaxation_calculator.append(Si_SiO2_calculator)
HfO2_SiO2_calculator = TremoloXCalculator(QuantumATK_MTP_HfO2_SiO2_2022())
relaxation_calculator.append(HfO2_SiO2_calculator)
HfO2_TiN_calculator = TremoloXCalculator(QuantumATK_MTP_HfO2_TiN_2022())
relaxation_calculator.append(HfO2_TiN_calculator)

surface_vectors = [[ 11.520042257735,   0.            ],
                   [  7.680028171823,  -7.680028171823]]*Angstrom


multilayer_builder = MultilayerBuilder(
    filename='HKMG_builder.hdf5',
    object_id='multilayer_builder',
    layers=layers,
    surface_vectors=surface_vectors,
    relaxation_calculator=relaxation_calculator,
    annealing_parameters=annealing_parameters,
    max_iteration_annealing_parameters=1,
    select_from_different_cut_plane=True,
)
multilayer_builder.update()
nlsave('HKMG_builder.hdf5', multilayer_builder.configuration())