import dpdata
from deepmd.infer.deep_pot import DeepPot
import numpy as np
import os
import gc
import glob
import logging

datadir = "./data-clean-v2-7-20873-npy"
modelpath = "./FeCHO-dpa231-v2-7-3heads-150w.pt"
savedir = "descriptors"

omp = 16
proc = 4
os.environ['OMP_NUM_THREADS'] = f'{omp}'

def descriptor_from_model(sys: dpdata.LabeledSystem, model:DeepPot):
    coords = sys.data["coords"]
    cells = sys.data["cells"]
    model_type_map = model.get_type_map()
    type_trans = np.array([model_type_map.index(i) for i in sys.data['atom_names']])
    atypes = list(type_trans[sys.data['atom_types']])
    predict = model.eval_descriptor(coords, cells, atypes)
    return predict
#alldata = dpdata.MultiSystems.from_dir(datadir,datakey,fmt="deepmd/npy")
all_set_directories = glob.glob(os.path.join(
    datadir, '**', 'set.*'), recursive=True)
all_directories = set()
for directory in all_set_directories:
    coord_path = os.path.join(directory, 'coord.npy')
    if os.path.exists(coord_path):
        all_directories.add(os.path.dirname(directory))
all_directories = list(all_directories)

model = DeepPot(modelpath, head="Target_FTS")

logging.basicConfig(
    level=logging.INFO, 
    format='%(asctime)s - %(levelname)s - %(message)s',  
    datefmt='%Y-%m-%d %H:%M:%S'  
)

logging.info("Start Generating Descriptors")

if not os.path.exists(savedir):
    os.mkdir(savedir)

with open("running", "w") as fo:
    for onedir in all_directories:
        onedata = dpdata.LabeledSystem(onedir, fmt="deepmd/npy")
        key = onedata.short_name
        save_key = f"{savedir}/{key}"
        logging.info(f"Generating descriptors for {key}")
        if os.path.exists(save_key):
            if os.path.exists(f"{save_key}/desc.npy"):
                logging.info(f"Descriptors for {key} already exist, skip")
                continue
        else:
            os.mkdir(save_key)
        desc = descriptor_from_model(onedata, model)
        logging.info(f"Descriptors for {key} generated")
        
        np.save(f"{savedir}/{key}/desc.npy", desc)
        logging.info(f"Descriptors for {key} saved")

logging.info("All Done !!!")
os.system("mv running done")

2025-01-06 15:58:57 - INFO - Start Generating Descriptors
2025-01-06 15:58:57 - INFO - Generating descriptors for O0H6Fe48C8
2025-01-06 15:59:00 - INFO - Descriptors for O0H6Fe48C8 generated
2025-01-06 15:59:00 - INFO - Descriptors for O0H6Fe48C8 saved
2025-01-06 15:59:00 - INFO - Generating descriptors for O3H4Fe0C6
Traceback (most recent call last):
  File "/home/mps/liuzq/FeCHO-dpa2/300rc0/v2-7-3h-100w/desc-gen/gen_desc.py", line 64, in <module>
    desc = descriptor_from_model(onedata, model)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mps/liuzq/FeCHO-dpa2/300rc0/v2-7-3h-100w/desc-gen/gen_desc.py", line 25, in descriptor_from_model
    predict = model.eval_descriptor(coords, cells, atypes)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mps/miniconda3/envs/deepmd-3rc0/lib/python3.11/site-packages/deepmd/infer/deep_eval.py", line 445, in eval_descriptor
    descriptor = self.deep_eval.eval_descriptor(
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mps/miniconda3/envs/deepmd-3rc0/lib/python3.11/site-packages/deepmd/pt/infer/deep_eval.py", line 658, in eval_descriptor
    descriptor = model.eval_descriptor()
                 ^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
  File "/home/mps/miniconda3/envs/deepmd-3rc0/lib/python3.11/site-packages/deepmd/pt/model/model/dp_model.py", line 66, in eval_descriptor
    def eval_descriptor(self) -> torch.Tensor:
        """Evaluate the descriptor."""
        return self.atomic_model.eval_descriptor()
               ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
  File "/home/mps/miniconda3/envs/deepmd-3rc0/lib/python3.11/site-packages/deepmd/pt/model/atomic_model/dp_atomic_model.py", line 76, in eval_descriptor
    def eval_descriptor(self) -> torch.Tensor:
        """Evaluate the descriptor."""
        return torch.concat(self.eval_descriptor_list)
               ~~~~~~~~~~~~ <--- HERE
RuntimeError: Sizes of tensors must match except in dimension 0. Expected size 62 but got size 13 for tensor number 1 in the list.