from ase.build import bulk, molecule, fcc100, add_adsorbate

# 自动定位 UMA 旋转基文件 Jd.pt
import os
_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
_JD_PATH = os.path.join(_REPO_ROOT, "weight", "Jd.pt")
if os.path.isfile(_JD_PATH):
    os.environ.setdefault("ONESCIENCE_UMA_JD_PATH", _JD_PATH)

from onescience.utils.uma.units.mlip_unit import load_predict_unit
from onescience.datapipes.materials.custom_stack.core.atomic_data import AtomicData, atomicdata_list_to_batch

# 1. 创建异构结构
h2o = molecule("H2O")
h2o.info.update({"charge": 0, "spin": 1})

pt = bulk("Pt")

slab = fcc100("Cu", (3, 3, 3), vacuum=8, periodic=True)
adsorbate = molecule("CO")
add_adsorbate(slab, adsorbate, 2.0, "bridge")

# 2. 结构转为 AtomicData，并指定不同 task_name
atomic_data_list = [
    AtomicData.from_ase(
        h2o, task_name="omol", r_data_keys=["spin", "charge"], molecule_cell_size=12
    ),
    AtomicData.from_ase(pt, task_name="omat"),
    AtomicData.from_ase(slab, task_name="oc20"),
]

# 3. 合成 batch
batch = atomicdata_list_to_batch(atomic_data_list)

# 4. 加载 UMA 模型
predictor = load_predict_unit(
    "../weight/uma-s-1p1_converted.pt", device="cuda"#替换为你的检查点路径
)

# 5. 执行联合推理
preds = predictor.predict(batch)

# 6. 输出每个结构结果
for i in range(len(preds["energy"])):
    energy = preds["energy"][i].item()
    forces = preds["forces"][batch.batch == i].cpu().numpy()

    print(f"\n[Structure {i}]")
    print("Predicted energy:", energy)
    print("Predicted forces:\n", forces)