import os
import numpy as np
from ase.io import read, write
from ase.build import make_supercell

# Input Parameters
input_cif = "./graphene.cif"
supercell_matrix = np.diag([4, 4, 1])  # 4x4x1
n_structures = 100 # You usually wants more!
perturb_range = 0.1  # Å
template_script = "./run.py" # template python script
dataset_root = "batch_calc"

# Read input structure and make supercell
atoms = read(input_cif)
supercell = make_supercell(atoms, supercell_matrix)

# Generate perturbed copies
for i in range(n_structures):
    # Create a perturbed copy
    perturbed = supercell.copy()
    perturbations = np.random.uniform(-perturb_range, perturb_range, size=perturbed.positions.shape)
    perturbed.positions += perturbations

    # Create output directory
    outdir = os.path.join(dataset_root, f"{i}")
    os.makedirs(outdir, exist_ok=True)

    # Write to geometry.in in FHI-aims format
    outfile = os.path.join(outdir, "geometry.in")
    write(outfile, perturbed, format="aims")
    os.system(f"cp {template_script} {outdir}/")

print(f"Generated {n_structures} perturbed structures under '{dataset_root}/'")