from SMW import *

random_seed = 667
sample_size = 10

# Define each layer in the multilayer structure.
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,
)

# -------------------------------------------------------------
# SiO2/Si parameter
# -------------------------------------------------------------
# Setup annealing parameter used for optimizing interfaces.
annealing_parameters_Si_SiO2 = []
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_Si_SiO2.append(
    AnnealingParameters(
        md_method=method,
        md_steps=md_steps
    )
)

initial_velocity = ConfigurationVelocities(
    remove_center_of_mass_momentum=False,
)
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_Si_SiO2.append(
    AnnealingParameters(md_method=method, md_steps=md_steps)
)
annealing_parameters.append(annealing_parameters_Si_SiO2)
# -------------------------------------------------------------
# HfO2/SiO2 parameter
# -------------------------------------------------------------
# Setup annealing parameter used for optimizing interfaces.
annealing_parameters_HfO2_SiO2 = []
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_HfO2_SiO2.append(
    AnnealingParameters(
        md_method=method,
        md_steps=md_steps
    )
)

initial_velocity = ConfigurationVelocities(
    remove_center_of_mass_momentum=False,
)

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_HfO2_SiO2.append(
    AnnealingParameters(md_method=method, md_steps=md_steps)
)
annealing_parameters.append(annealing_parameters_HfO2_SiO2)
# -------------------------------------------------------------
# HfO2/TiN parameter
# -------------------------------------------------------------
# Setup annealing parameter used for optimizing interfaces.
annealing_parameters_HfO2_TiN = []
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_HfO2_TiN.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_HfO2_TiN.append(
    AnnealingParameters(
        md_method=method,
        md_steps=md_steps
    )
)
annealing_parameters.append(annealing_parameters_HfO2_TiN)

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 = [[ 7.680028171823,  7.680028171823],
                   [-7.680028171823,  7.680028171823]]*Angstrom

for i in range(sample_size):
    layers = []
    material = MaterialSpecificationsDatabase.HKMG_MATERIALS['Silicon']()
    layer_0 = CrystalLayer(
        material,
        length=10.0*Angstrom,
        miller_indices=(1, 0, 0),
        top_atom_index=0,
        bottom_atom_index=0,
        displacement_vector=(0.0, 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,
        random_seed=random_seed+i
        )
    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,
        random_seed=random_seed+i
        )
    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,
        random_seed=random_seed+i
        )
    layers.append(layer_3)

    multilayer_builder = MultilayerBuilder(
        filename=f'HKMG_builder{i}.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,
        log_filename_prefix=f'hkmg_builder_sample{i}_'
    )
    multilayer_builder.update()
    nlsave(f'HKMG_builder{i}.hdf5', multilayer_builder.configuration())