import argparse

from ase.build import bulk

from onescience.utils.mattersim import molecular_dynamics


parser = argparse.ArgumentParser(description="MatterSim molecular dynamics")
parser.add_argument(
    "--checkpoint",
    default="../weight/mattersim-v1.0.0-1M.pth",
    help="Path to MatterSim checkpoint",
)
parser.add_argument("--device", choices=("cpu", "cuda"), default="cuda")
args = parser.parse_args()

atoms = bulk("Si", "diamond", a=5.43).repeat((2, 2, 2))
dynamics = molecular_dynamics(
    atoms,
    checkpoint=args.checkpoint,
    ensemble="nvt_berendsen",
    temperature=300,
    timestep=1.0,
    logfile="-",
    device=args.device,
)
dynamics.run(10)