#!/usr/bin/env python3

"""
Convert an IFF EquilibriumCrystalStructure_Unconstrained_TypeLabels test from CHARMM-GUI atom types to PCFF
atom types. 

Read from STDIN:
    CHARMM-GUI-IFF test generators

Printed to STDOUT:
    PCFF-IFF test generators
"""
import sys, os, pathlib
import random
import json

CHARMM_GUI_TO_PCFF_MAPPING={
    "IAY1": "ay1",
    "IAY2": "ay2",
    "ISY1": "sy1",
    "ISY2": "sy2",
    "IOY7": "oy7",
    "IOY8": "oy8",
    "IOY1": "oy1",
    "IOY2": "oy2",
    "IOY3": "oy3",
    "IOY9": "oy9",
    "IOY5": "oy5",
    "IOY4": "oy4",
    "IOY6": "oy6",
    "IHOY": "hoy",
    "IHOK": "hok",
    "IOC24": "oc24",
    "IOAP1": "oap1",
    "IPAP": "pap",
    "IHOP": "hop",
    "IOC11": "oc11",
    "ISC2": "sc2",
    "ISC3": "sc3",
    "IOC12": "oc12",
    "IOC13": "oc13",
    "IOC14": "oc14",
    "IAC2": "ac2",
    "IOC10": "oc10",
    "IOC7": "oc7",
    "IHOC": "hoc",
    "IOC9": "oc9",
    "IOC6": "oc6",
    "ISC4": "sc4",
    "IOC8": "oc8",
    "IOC5": "oc5",
    "IOAP2": "oap2",
    "ISC1": "sc1",
    "IOC2": "oc2",
    "IOC1": "oc1",
    "IOC23": "oc23",
    "IAC1": "ac1",
    "IOC3": "oc3",
    "IOC4": "oc4",
    "IAL": "Al",
    "IAG": "Ag",
    "IPD": "Pd",
    "IPT": "Pt",
    "IAU": "Au",
    "INI": "Ni",
    "ICU": "Cu",
    "IPB": "Pb",
    "IK_CM": "k+",
    "UNACCOUNTED1": "na+",
    "UNACCOUNTED2": "s_m",
    "UNACCOUNTED3": "oc15",
    "IO_SC": "o*",
    "ICA_T": "ca+t",
    "IH_SC": "h*",
    "UNACCOUNTED7": "ac3",
    "UNACCOUNTED8": "ca+m",
    "UNACCOUNTED9": "oc21",
    "UNACCOUNTED10": "oc18",
    "UNACCOUNTED11": "oc16",
    "UNACCOUNTED12": "oc20",
    "ICA_G": "ca+g",
    "UNACCOUNTED14": "oc19",
    "UNACCOUNTED15": "oc17",
    "IS_AN": "s'",
    "UNACCOUNTED17": "oc22",
    "ICA_H": "ca+h",
    "ICA_E": "ca+e",
    "ICA_S": "ca++",
    "ICPEO": "ce1",
    "IOPEO": "oe1",
    "IHPEO": "he1",
    "ICA_A": "ca+a"
}

MODELTYPE = "PCFF-INTERFACE"
KIM_USER_ID = "4ad03136-ed7f-4316-b586-1e94ccceb311"
MATCHING_MODELS = "[\"class2\"]"

# relative to script file location, will be converted to absolute
TEST_DIR_PATH = "../Tests/PCFF-IFF"
SYM_FILE_PATH="symfiles/pcff/sym.txt"

def get_random_kim_id() -> str:
    """
    Returns random 12-digit numerical string

    Returns:
        Random 12-digit numerical string
    """
    return "".join(["{}".format(random.randint(0, 9)) for num in range(12)])        

if __name__ == "__main__":
    for line in sys.stdin:
        test_generator=json.loads(line)
        material = test_generator["material"]
        #print("Processing %s" % material.replace("-"," "))

        # overwrite old values
        test_generator["modeltype"]=MODELTYPE
        test_generator["kimnum"]=get_random_kim_id()
        test_generator["matching_models"]=MATCHING_MODELS

        # get charmm-gui data file
        charmm_files_to_copy=test_generator["FILES_TO_COPY"][1:] # first entry is symbol file
        assert(len(charmm_files_to_copy)==1)
        assert os.path.exists(charmm_files_to_copy[0]),"Can't find charmm-gui data file to convert, are you running from the top-level directory inside the TD?"
 
        # get new paths
        test_dir = os.path.join(pathlib.Path(__file__).parent.resolve(),TEST_DIR_PATH,material)
        pcff_data_path = os.path.join(test_dir,"input0.dat")
        test_generator["FILES_TO_COPY"]=[os.path.join(pathlib.Path(__file__).parent.resolve(),SYM_FILE_PATH),pcff_data_path]

        # get pcff type labels, check if all are supported
        pcff_type_labels_list = []
        unsupported_type_label_found = False
        for charmm_gui_type_label in test_generator["atom_type_labels"].split():
            if charmm_gui_type_label not in CHARMM_GUI_TO_PCFF_MAPPING:
                #print("Type label %s not supported by PCFF model, skipping this test" % charmm_gui_type_label)
                unsupported_type_label_found = True
                break
            pcff_type_labels_list.append(CHARMM_GUI_TO_PCFF_MAPPING[charmm_gui_type_label])
        if unsupported_type_label_found:
            continue
        test_generator["atom_type_labels"]=" ".join(pcff_type_labels_list)

        # create material dir
        pathlib.Path(test_dir).mkdir(parents=True, exist_ok=True)  

        with open(charmm_files_to_copy[0]) as charmm_dat, open(pcff_data_path,"w") as pcff_dat:
            for charmm_line in charmm_dat:
                pcff_line = charmm_line
                for charmm_gui_type_label in CHARMM_GUI_TO_PCFF_MAPPING:
                    pcff_line = pcff_line.replace(charmm_gui_type_label,CHARMM_GUI_TO_PCFF_MAPPING[charmm_gui_type_label])
                pcff_dat.write(pcff_line)
        
        print(json.dumps(test_generator))