# %% Common parameters


def common_parameters():
    # Defined variables.
    md_steps = 20000
    temperature = 300.0 * Kelvin
    pressure = 1.0 * bar

    # Script.
    # This custom script supports atkpython syntax
    # and can perform almost any procedure.
    auto_log_interval = md_steps // 100
    auto_trajectory_interval = md_steps // 10
    test_string = 'pineapple'

    return md_steps, temperature, auto_log_interval, auto_trajectory_interval, pressure


md_steps, temperature, auto_log_interval, auto_trajectory_interval, pressure = (
    common_parameters()
)

nlsave('Au_Ag_MonteCarlo_results.hdf5', md_steps)

nlsave('Au_Ag_MonteCarlo_results.hdf5', temperature)

nlsave('Au_Ag_MonteCarlo_results.hdf5', auto_log_interval)

nlsave('Au_Ag_MonteCarlo_results.hdf5', auto_trajectory_interval)

nlsave('Au_Ag_MonteCarlo_results.hdf5', pressure)


# %% Gold

# Set up lattice
vector_a = [12.234749999999998, 0.0, 0.0]*Angstrom
vector_b = [0.0, 12.234749999999998, 0.0]*Angstrom
vector_c = [0.0, 0.0, 12.234749999999998]*Angstrom
lattice = UnitCell(vector_a, vector_b, vector_c)

# Define elements
elements = [Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold,
            Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold,
            Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold,
            Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold,
            Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold,
            Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold,
            Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold,
            Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold,
            Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold,
            Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold, Gold]

# Define coordinates
fractional_coordinates = [[ 0.            ,  0.            ,  0.            ],
                          [ 0.            ,  0.            ,  0.333333333333],
                          [ 0.            ,  0.            ,  0.666666666667],
                          [ 0.            ,  0.333333333333,  0.            ],
                          [ 0.            ,  0.333333333333,  0.333333333333],
                          [ 0.            ,  0.333333333333,  0.666666666667],
                          [ 0.            ,  0.666666666667,  0.            ],
                          [ 0.            ,  0.666666666667,  0.333333333333],
                          [ 0.            ,  0.666666666667,  0.666666666667],
                          [ 0.333333333333,  0.            ,  0.            ],
                          [ 0.333333333333,  0.            ,  0.333333333333],
                          [ 0.333333333333,  0.            ,  0.666666666667],
                          [ 0.333333333333,  0.333333333333,  0.            ],
                          [ 0.333333333333,  0.333333333333,  0.333333333333],
                          [ 0.333333333333,  0.333333333333,  0.666666666667],
                          [ 0.333333333333,  0.666666666667,  0.            ],
                          [ 0.333333333333,  0.666666666667,  0.333333333333],
                          [ 0.333333333333,  0.666666666667,  0.666666666667],
                          [ 0.666666666667,  0.            ,  0.            ],
                          [ 0.666666666667,  0.            ,  0.333333333333],
                          [ 0.666666666667,  0.            ,  0.666666666667],
                          [ 0.666666666667,  0.333333333333,  0.            ],
                          [ 0.666666666667,  0.333333333333,  0.333333333333],
                          [ 0.666666666667,  0.333333333333,  0.666666666667],
                          [ 0.666666666667,  0.666666666667,  0.            ],
                          [ 0.666666666667,  0.666666666667,  0.333333333333],
                          [ 0.666666666667,  0.666666666667,  0.666666666667],
                          [ 0.166666666667,  0.166666666667,  0.            ],
                          [ 0.166666666667,  0.166666666667,  0.333333333333],
                          [ 0.166666666667,  0.166666666667,  0.666666666667],
                          [ 0.166666666667,  0.5           ,  0.            ],
                          [ 0.166666666667,  0.5           ,  0.333333333333],
                          [ 0.166666666667,  0.5           ,  0.666666666667],
                          [ 0.166666666667,  0.833333333333,  0.            ],
                          [ 0.166666666667,  0.833333333333,  0.333333333333],
                          [ 0.166666666667,  0.833333333333,  0.666666666667],
                          [ 0.5           ,  0.166666666667,  0.            ],
                          [ 0.5           ,  0.166666666667,  0.333333333333],
                          [ 0.5           ,  0.166666666667,  0.666666666667],
                          [ 0.5           ,  0.5           ,  0.            ],
                          [ 0.5           ,  0.5           ,  0.333333333333],
                          [ 0.5           ,  0.5           ,  0.666666666667],
                          [ 0.5           ,  0.833333333333,  0.            ],
                          [ 0.5           ,  0.833333333333,  0.333333333333],
                          [ 0.5           ,  0.833333333333,  0.666666666667],
                          [ 0.833333333333,  0.166666666667,  0.            ],
                          [ 0.833333333333,  0.166666666667,  0.333333333333],
                          [ 0.833333333333,  0.166666666667,  0.666666666667],
                          [ 0.833333333333,  0.5           ,  0.            ],
                          [ 0.833333333333,  0.5           ,  0.333333333333],
                          [ 0.833333333333,  0.5           ,  0.666666666667],
                          [ 0.833333333333,  0.833333333333,  0.            ],
                          [ 0.833333333333,  0.833333333333,  0.333333333333],
                          [ 0.833333333333,  0.833333333333,  0.666666666667],
                          [ 0.166666666667,  0.            ,  0.166666666667],
                          [ 0.166666666667,  0.            ,  0.5           ],
                          [ 0.166666666667,  0.            ,  0.833333333333],
                          [ 0.166666666667,  0.333333333333,  0.166666666667],
                          [ 0.166666666667,  0.333333333333,  0.5           ],
                          [ 0.166666666667,  0.333333333333,  0.833333333333],
                          [ 0.166666666667,  0.666666666667,  0.166666666667],
                          [ 0.166666666667,  0.666666666667,  0.5           ],
                          [ 0.166666666667,  0.666666666667,  0.833333333333],
                          [ 0.5           ,  0.            ,  0.166666666667],
                          [ 0.5           ,  0.            ,  0.5           ],
                          [ 0.5           ,  0.            ,  0.833333333333],
                          [ 0.5           ,  0.333333333333,  0.166666666667],
                          [ 0.5           ,  0.333333333333,  0.5           ],
                          [ 0.5           ,  0.333333333333,  0.833333333333],
                          [ 0.5           ,  0.666666666667,  0.166666666667],
                          [ 0.5           ,  0.666666666667,  0.5           ],
                          [ 0.5           ,  0.666666666667,  0.833333333333],
                          [ 0.833333333333,  0.            ,  0.166666666667],
                          [ 0.833333333333,  0.            ,  0.5           ],
                          [ 0.833333333333,  0.            ,  0.833333333333],
                          [ 0.833333333333,  0.333333333333,  0.166666666667],
                          [ 0.833333333333,  0.333333333333,  0.5           ],
                          [ 0.833333333333,  0.333333333333,  0.833333333333],
                          [ 0.833333333333,  0.666666666667,  0.166666666667],
                          [ 0.833333333333,  0.666666666667,  0.5           ],
                          [ 0.833333333333,  0.666666666667,  0.833333333333],
                          [ 0.            ,  0.166666666667,  0.166666666667],
                          [ 0.            ,  0.166666666667,  0.5           ],
                          [ 0.            ,  0.166666666667,  0.833333333333],
                          [ 0.            ,  0.5           ,  0.166666666667],
                          [ 0.            ,  0.5           ,  0.5           ],
                          [ 0.            ,  0.5           ,  0.833333333333],
                          [ 0.            ,  0.833333333333,  0.166666666667],
                          [ 0.            ,  0.833333333333,  0.5           ],
                          [ 0.            ,  0.833333333333,  0.833333333333],
                          [ 0.333333333333,  0.166666666667,  0.166666666667],
                          [ 0.333333333333,  0.166666666667,  0.5           ],
                          [ 0.333333333333,  0.166666666667,  0.833333333333],
                          [ 0.333333333333,  0.5           ,  0.166666666667],
                          [ 0.333333333333,  0.5           ,  0.5           ],
                          [ 0.333333333333,  0.5           ,  0.833333333333],
                          [ 0.333333333333,  0.833333333333,  0.166666666667],
                          [ 0.333333333333,  0.833333333333,  0.5           ],
                          [ 0.333333333333,  0.833333333333,  0.833333333333],
                          [ 0.666666666667,  0.166666666667,  0.166666666667],
                          [ 0.666666666667,  0.166666666667,  0.5           ],
                          [ 0.666666666667,  0.166666666667,  0.833333333333],
                          [ 0.666666666667,  0.5           ,  0.166666666667],
                          [ 0.666666666667,  0.5           ,  0.5           ],
                          [ 0.666666666667,  0.5           ,  0.833333333333],
                          [ 0.666666666667,  0.833333333333,  0.166666666667],
                          [ 0.666666666667,  0.833333333333,  0.5           ],
                          [ 0.666666666667,  0.833333333333,  0.833333333333]]

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

gold_name = "gold"


# %% Set ForceFieldCalculator

# %% ForceFieldCalculator

potentialSet = MEAM_AgAu_2022()
calculator = TremoloXCalculator(parameters=potentialSet)


# %% Set Calculator

gold.setCalculator(calculator)

nlsave('Au_Ag_MonteCarlo_results.hdf5', gold)


# %% NVT Equilibration

method = NVTNoseHoover(
    initial_velocity=MaxwellBoltzmannDistribution(temperature=temperature),
    reservoir_temperature=temperature,
)

md_trajectory = MolecularDynamics(
    configuration=gold,
    trajectory_filename='Au_Ag_MonteCarlo_results.hdf5',
    steps=md_steps,
    log_interval=auto_log_interval,
    method=method,
    trajectory_interval=auto_trajectory_interval,
    trajectory_object_id='md',
)
last_image = md_trajectory.lastImage()

nlsave('Au_Ag_MonteCarlo_results.hdf5', last_image, object_id='md_1')


# %% NPT Equilibration

method = NPTMartynaTobiasKlein(
    initial_velocity=ConfigurationVelocities(),
    reservoir_temperature=temperature,
    reservoir_pressure=pressure,
)

md_trajectory_1 = MolecularDynamics(
    configuration=last_image,
    trajectory_filename='Au_Ag_MonteCarlo_results.hdf5',
    steps=md_steps,
    log_interval=auto_log_interval,
    method=method,
    trajectory_interval=auto_trajectory_interval,
    trajectory_object_id='md_2',
)
last_image_1 = md_trajectory_1.lastImage()

nlsave('Au_Ag_MonteCarlo_results.hdf5', last_image_1, object_id='md_3')


# %% Semii-grand canonical method


def semiigrand_canonical_method():
    # Defined variables.
    mc_temperature = temperature
    call_interval = 1
    swap_ratio = 0.01
    first_element_symbol = 'Au'
    second_element_symbol = 'Ag'
    chemical_potential = -1.06 * eV

    # Script.
    # This custom script supports atkpython syntax
    # and can perform almost any procedure.
    from NL.CommonConcepts.PeriodicTable import SYMBOL_TO_ELEMENT

    first_element = SYMBOL_TO_ELEMENT[first_element_symbol]
    second_element = SYMBOL_TO_ELEMENT[second_element_symbol]

    monte_carlo_hook = SemiGrandCanonicalMonteCarloHook(
        elements=(first_element, second_element),
        chemical_potential=chemical_potential,
        temperature=mc_temperature,
        swap_ratio=swap_ratio,
        call_interval=call_interval,
    )
    return monte_carlo_hook


monte_carlo_hook = semiigrand_canonical_method()

nlsave('Au_Ag_MonteCarlo_results.hdf5', monte_carlo_hook)


# %% Hybrid Monte Carlo


def hybrid_monte_carlo(configuration, monte_carlo_hook):
    # Defined variables.
    mc_temperature = temperature
    mc_pressure = pressure
    steps = md_steps
    log_interval = auto_log_interval
    trajectory_interval = auto_trajectory_interval

    # Script.
    # This custom script supports atkpython syntax
    # and can perform almost any procedure.
    dynamics_method = NPTMartynaTobiasKlein(
        initial_velocity=ConfigurationVelocities(),
        reservoir_temperature=mc_temperature,
        reservoir_pressure=mc_pressure,
    )
    hooks = HookFunctions(pre_step_hooks=monte_carlo_hook)
    trajectory = HybridMonteCarlo(
        configuration,
        steps=steps,
        log_interval=log_interval,
        trajectory_interval=trajectory_interval,
        method=dynamics_method,
        hook_functions=hooks,
    )

    # Get the composition from the proguction part of the calculation
    indices = numpy.arange(steps // 2, steps)
    composition = trajectory.measurement('composition')[1][indices].mean()

    return trajectory, composition


trajectory, composition = hybrid_monte_carlo(last_image_1, monte_carlo_hook)

nlsave('Au_Ag_MonteCarlo_results.hdf5', trajectory)

nlsave('Au_Ag_MonteCarlo_results.hdf5', composition)
