# Set minimal log verbosity
setVerbosity(MinimalLog)

# %% FCC Cobalt

# %% BulkConfiguration

# Set up lattice
lattice = SimpleCubic(3.5447*Angstrom)

# Define elements
elements = [Cobalt, Cobalt, Cobalt, Cobalt]

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

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

cobalt_fcc_name = "cobalt_fcc"


# %% Set ForceFieldCalculator

# %% ForceFieldCalculator

potentialSet = EAM_Co_2012()
calculator = TremoloXCalculator(parameters=potentialSet)


# %% Set Calculator

cobalt_fcc.setCalculator(calculator)

nlsave('Cobalt_FCC_HCP_Example_results.hdf5', cobalt_fcc)


# %% FCC_OptimizeGeometry

optimized_configuration = OptimizeGeometry(
    configuration=cobalt_fcc,
    max_forces=0.0005 * eV / Angstrom,
    max_stress=0.001 * GPa,
    constraints=[BravaisLatticeConstraint()],
    enable_optimization_stop_file=False,
)

nlsave(
    'Cobalt_FCC_HCP_Example_results.hdf5', optimized_configuration, object_id='optgeom'
)


# %% HCP Cobalt

# %% BulkConfiguration

# Set up lattice
lattice = Hexagonal(2.5071*Angstrom, 4.0686*Angstrom)

# Define elements
elements = [Cobalt, Cobalt]

# Define coordinates
fractional_coordinates = [[ 0.333333333333,  0.666666666667,  0.25          ],
                          [ 0.666666666667,  0.333333333333,  0.75          ]]

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

cobalt_hcp_name = "cobalt_hcp"


# %% Set ForceFieldCalculator

# %% ForceFieldCalculator

potentialSet = EAM_Co_2012()
calculator_1 = TremoloXCalculator(parameters=potentialSet)


# %% Set Calculator

cobalt_hcp.setCalculator(calculator_1)

nlsave('Cobalt_FCC_HCP_Example_results.hdf5', cobalt_hcp)


# %% HCP_OptimizeGeometry

restart_strategy = RestartFromTrajectory(
    trajectory_filename='Cobalt_FCC_HCP_Example_results.hdf5',
    object_id='optimize_trajectory',
)

optimized_configuration_1 = OptimizeGeometry(
    configuration=cobalt_hcp,
    max_forces=0.0005 * eV / Angstrom,
    max_stress=0.001 * GPa,
    constraints=[BravaisLatticeConstraint()],
    trajectory_filename='Cobalt_FCC_HCP_Example_results.hdf5',
    trajectory_object_id='optimize_trajectory',
    restart_strategy=restart_strategy,
)

nlsave(
    'Cobalt_FCC_HCP_Example_results.hdf5',
    optimized_configuration_1,
    object_id='optgeom_1',
)


# %% Iteration

# Iteration(preparation)
############################################################
#                 Iteration (temperature)                  #
############################################################


# Create iterator.
def temperaturesIterator():
    for i in range(8):
        v = (300.0 + i * 100.0) * Kelvin
        yield v


temperatures = temperaturesIterator()


for temperature in temperatures:

    # %% OptimizeGeometryParameters

    optimize_geometry_parameters = OptimizeGeometryParameters(
        constraints=[BravaisLatticeConstraint()]
    )

    # %% DynamicalMatrixParameters

    dynamical_matrix_parameters = DynamicalMatrixParameters()

    # %% FCC_OptimizeBulkFreeEnergy

    optimize_result_table = OptimizeBulkFreeEnergy(
        configuration=optimized_configuration,
        temperature=temperature,
        optimize_geometry_parameters=optimize_geometry_parameters,
        dynamical_matrix_parameters=dynamical_matrix_parameters,
    )
    nlsave(
        'Cobalt_FCC_HCP_Example_results.hdf5',
        optimize_result_table,
        object_id=f'optimize_result_table_optenergy_temperature_{temperature}',
    )

    optimize_configuration = optimize_result_table[0][0]
    nlsave(
        'Cobalt_FCC_HCP_Example_results.hdf5',
        optimize_configuration,
        object_id=f'optimize_configuration_optenergy_temperature_{temperature}',
    )

    optimize_free_energy = optimize_result_table[0][4]
    nlsave(
        'Cobalt_FCC_HCP_Example_results.hdf5',
        optimize_free_energy,
        object_id=f'optimize_free_energy_optenergy_temperature_{temperature}',
    )

    # %% HCP_OptimizeBulkFreeEnergy

    optimize_result_table_1 = OptimizeBulkFreeEnergy(
        configuration=optimized_configuration_1,
        temperature=temperature,
        bounding_factor=0.05,
        optimize_geometry_parameters=optimize_geometry_parameters,
        dynamical_matrix_parameters=dynamical_matrix_parameters,
    )
    nlsave(
        'Cobalt_FCC_HCP_Example_results.hdf5',
        optimize_result_table_1,
        object_id=f'optimize_result_table_1_optenergy_temperature_{temperature}',
    )

    optimize_configuration_1 = optimize_result_table_1[0][0]
    nlsave(
        'Cobalt_FCC_HCP_Example_results.hdf5',
        optimize_configuration_1,
        object_id=f'optimize_configuration_1_optenergy_temperature_{temperature}',
    )

    optimize_free_energy_1 = optimize_result_table_1[0][4]
    nlsave(
        'Cobalt_FCC_HCP_Example_results.hdf5',
        optimize_free_energy_1,
        object_id=f'optimize_free_energy_1_optenergy_temperature_{temperature}',
    )

    # %% Convert to energy per atom

    def convert_to_energy_per_atom(fcc_energy, hcp_energy):
        # This custom script supports atkpython syntax
        # and can perform almost any procedure.
        fcc_energy_per_atom = fcc_energy / 4
        hcp_energy_per_atom = hcp_energy / 2

        return fcc_energy_per_atom, hcp_energy_per_atom

    fcc_energy_per_atom, hcp_energy_per_atom = convert_to_energy_per_atom(
        optimize_free_energy, optimize_free_energy_1
    )

    nlsave(
        'Cobalt_FCC_HCP_Example_results.hdf5',
        fcc_energy_per_atom,
        object_id=f'fcc_energy_per_atom_Convert_to_energy_per_atom_temperature_{temperature}_1',
    )

    nlsave(
        'Cobalt_FCC_HCP_Example_results.hdf5',
        hcp_energy_per_atom,
        object_id=f'hcp_energy_per_atom_Convert_to_energy_per_atom_temperature_{temperature}_1',
    )

    # %% Results table

    if 'results_table' not in locals():
        results_table = Table('Cobalt_FCC_HCP_Example_results.hdf5', object_id='table')
        results_table.addQuantityColumn(key='temperature', unit=Kelvin)
        results_table.addQuantityColumn(key='fcc_energy', unit=eV)
        results_table.addQuantityColumn(key='hcp_energy', unit=eV)
        results_table.setMetatext('Results table')

    results_table.append(temperature, fcc_energy_per_atom, hcp_energy_per_atom)


# %% Plot Table

# Create PlotModel.
plot_model = Plot.PlotModel(x_unit=Kelvin, y_unit=eV)

plot_model.framing().setTitle('Plot_Table')

# Add line
line = Plot.Line(results_table[:, 'temperature'], results_table[:, 'fcc_energy'])
line.setLabel('')
line.setColor('#b82832')
line.setLineStyle('-')
line.setMarkerStyle('')
plot_model.addItem(line)

# Add line
line_1 = Plot.Line(results_table[:, 'temperature'], results_table[:, 'hcp_energy'])
line_1.setLabel('')
line_1.setColor('#0a3b76')
line_1.setLineStyle('-')
line_1.setMarkerStyle('')
plot_model.addItem(line_1)

# Auto-adjust axis limits.
plot_model.setLimits()

# Save plot to file.
Plot.save(plot_model, 'Cobalt_FCC_HCP_Example_results.hdf5')