#!/usr/bin/env amspython
# coding: utf-8

# The ASE .xyz format is a convenient format for storing single point properties of a collection of structures. It is commonly used and supported by academic machine learning potential projects, so you may find it useful.
#
# Here, we will initialize/save a ResultsImporter object from/to ASE .xyz format.
#
# For advanced usage you can also work directly with lists of ASE Atoms, and the ParAMS ``JobCollection`` and ``DataSet`` classes. See the documentation for ``params_to_ase``, and ``ase_to_params``.

from scm.params import ResultsImporter
import os


# ## One-liner conversion from ParAMS .yaml to ASE .xyz
#
# Here we use the examples directory that already contains files in the ParAMS .yaml format:

yaml_dir = os.path.expandvars("$AMSHOME/scripting/scm/params/examples/DataSets/LiquidAr_32Atoms_100Frames")

# put training_set.xyz, validation_set.xyz in current working directory
xyz_target_dir = os.getcwd()

xyz_files = ResultsImporter.from_yaml(yaml_dir).store_ase(xyz_target_dir, format="extxyz")

print(xyz_files)


# ## One-liner conversion from ASE .xyz to ParAMS .yaml

xyz_dir = os.getcwd()

# put job_collection.yaml etc. in a folder yaml_ref_data
yaml_target_dir = "yaml_ref_data"

yaml_files = ResultsImporter.from_ase(f"{xyz_dir}/training_set.xyz", f"{xyz_dir}/validation_set.xyz").store(
    yaml_target_dir, backup=False
)

print(yaml_files)


# ## More about ASE .xyz format
#
# Let's look at the first few lines of the ASE .xyz files:

for file in xyz_files:
    print(f"--- first 5 lines of {file} ---")
    with open(file) as f:
        print("".join(f.readlines()[:5]))


# To read the ASE format, use ``ase.io.read("filename.xyz", ":")`` to get a list of ASE Atoms. For more details, see the ASE documentation.

import ase.io

list_of_ase_atoms = ase.io.read(xyz_files[0], ":")  # files[0] == "training_set.xyz"
atoms = list_of_ase_atoms[0]
print("First structure:")
print(atoms)
print(f"Energy: {atoms.get_potential_energy()}")


# ### More about ParAMS ResultsImporter initialized from ASE .xyz

ri = ResultsImporter.from_ase(f"{xyz_dir}/training_set.xyz", f"{xyz_dir}/validation_set.xyz")


# The structure is stored in the job collection:

for name in ri.job_collection:
    print(f"ID: {name}")
    print(ri.job_collection[name])
    break


# The reference values are stored in the data sets:

# print the first training set entry
for ds_entry in ri.get_data_set("training_set"):
    print(ds_entry)
    break


# print the first validation set entry
for ds_entry in ri.get_data_set("validation_set"):
    print(ds_entry)
    break


# When you initialize a ResultsImporter using ``from_ase``, it will automatically set the ResultsImporter units to the ASE units. If you use the results importer to import new data, the data will be stored in the ASE units (eV, eV/angstrom, etc.).

print(f"{ri.settings['units']['energy']=}")
print(f"{ri.settings['units']['forces']=}")