"""Generate a tiny extxyz smoke dataset for NequIP demo training."""

from __future__ import annotations

import os
from pathlib import Path

import numpy as np
from ase import Atoms
from ase.build import bulk
from ase.io import write


def main() -> None:
    out_dir = Path(__file__).parent / "reference_data"
    out_dir.mkdir(parents=True, exist_ok=True)
    out_file = out_dir / "smoke.xyz"

    rng = np.random.default_rng(123)
    structures = []
    # A few small Cu clusters with random displacements.
    base = bulk("Cu", "fcc", a=3.6) * (2, 2, 2)
    for i in range(8):
        atoms = base.copy()
        atoms.positions += rng.normal(scale=0.05, size=atoms.positions.shape)
        atoms.info["energy"] = float(-len(atoms) * 3.5 + rng.normal(scale=0.5))
        atoms.arrays["forces"] = rng.normal(scale=0.1, size=atoms.positions.shape)
        structures.append(atoms)

    write(out_file, structures, format="extxyz")
    print(f"Wrote {len(structures)} structures to {out_file}")


if __name__ == "__main__":
    main()