Module library.static.topologies

POINTER_TYPES = ["[moleculetype]", "[atoms]",
                 "[bonds]", "[angles]", "[pairs]"]

DOPC_CG_NAME_TO_TYPE_MAP = {
    'NC3': 'Q0',
    'PO4': 'Qa',
    'GL1': 'Na',
    'GL2': 'Na',
    'C1A': 'C1',
    'D2A': 'C3',
    'C3A': 'C1',
    'C4A': 'C1',
    'C1B': 'C1',
    'D2B': 'C3',
    'C3B': 'C1',
    'C4B': 'C1',
}

DOPC_BEAD_TYPE_NAME_IDS = {
    'Q0': 0,
    'Qa': 1,
    'Na': 2,
    'C1': 3,
    'C3': 4,
}

DOPC_ELEMENT_TYPE_NAME_IDS = {
    'C': 0,
    'N': 1,
    'O': 2,
    'P': 3,
    'H': 4,
}


DOPC_CG_BOND_MAP = {
    1: 2,
    2: 3,
    3: [4, 5],
    4: 9,
    5: 6,
    6: 7,
    7: 8,
    9: 10,
    10: 11,
    11: 12
}


DOPC_AT_NAMES = [
    "N",
    "C12",
    "H12A",
    "H12B",
    "C13",
    "H13A",
    "H13B",
    "H13C",
    "C14",
    "H14A",
    "H14B",
    "H14C",
    "C15",
    "H15A",
    "H15B",
    "H15C",
    "C11",
    "H11A",
    "H11B",
    "P",
    "O13",
    "O14",
    "O12",
    "O11",
    "C1",
    "HA",
    "HB",
    "C2",
    "HS",
    "O21",
    "C21",
    "O22",
    "C22",
    "H2R",
    "H2S",
    "C3",
    "HX",
    "HY",
    "O31",
    "C31",
    "O32",
    "C32",
    "H2X",
    "H2Y",
    "C23",
    "H3R",
    "H3S",
    "C24",
    "H4R",
    "H4S",
    "C25",
    "H5R",
    "H5S",
    "C26",
    "H6R",
    "H6S",
    "C27",
    "H7R",
    "H7S",
    "C28",
    "H8R",
    "H8S",
    "C29",
    "H9R",
    "C210",
    "H10R",
    "C211",
    "H11R",
    "H11S",
    "C212",
    "H12R",
    "H12S",
    "C213",
    "H13R",
    "H13S",
    "C214",
    "H14R",
    "H14S",
    "C215",
    "H15R",
    "H15S",
    "C216",
    "H16R",
    "H16S",
    "C217",
    "H17R",
    "H17S",
    "C218",
    "H18R",
    "H18S",
    "H18T",
    "C33",
    "H3X",
    "H3Y",
    "C34",
    "H4X",
    "H4Y",
    "C35",
    "H5X",
    "H5Y",
    "C36",
    "H6X",
    "H6Y",
    "C37",
    "H7X",
    "H7Y",
    "C38",
    "H8X",
    "H8Y",
    "C39",
    "H9X",
    "C310",
    "H10X",
    "C311",
    "H11X",
    "H11Y",
    "C312",
    "H12X",
    "H12Y",
    "C313",
    "H13X",
    "H13Y",
    "C314",
    "H14X",
    "H14Y",
    "C315",
    "H15X",
    "H15Y",
    "C316",
    "H16X",
    "H16Y",
    "C317",
    "H17X",
    "H17Y",
    "C318",
    "H18X",
    "H18Y",
    "H18Z",
]

def cg_name_to_type_dict(residue_map_path):
    """
        Returns a dict of CG residue names to CG residue types from a given residue map file
    """
    residue_map = {}
    type_pointer = None

    with open(residue_map_path, 'r') as f:
        for line in f.readlines():
            # Skip comments and empty lines
            if line.startswith(';') or len(line.split()) == 0:
                continue
        
            # Check if the line is a pointer
            if line.replace(" ", "").split()[0] in POINTER_TYPES:
                type_pointer = line.replace(" ", "").split()[0]
                continue

            # Check if the line is an atom type
            if type_pointer == POINTER_TYPES[1]:
                residue_map[line.split()[4]] = line.split()[1]

    return residue_map


def cg_bond_map_dict(residue_map_path):
    """
        Returns a dict of CG residue names to CG residue types from a given residue map file
    """
    bond_map = dict()
    type_pointer = None

    with open(residue_map_path, 'r') as f:
        for line in f.readlines():
            # Skip comments and empty lines
            if line.startswith(';') or len(line.split()) == 0:
                continue

            # Check if the line is a pointer
            if line.replace(" ", "").split()[0] in POINTER_TYPES:
                type_pointer = line.replace(" ", "").split()[0]
                continue

            # Check if the line is a bond type
            if type_pointer == POINTER_TYPES[2]:
                if int(line.split()[0]) not in bond_map.keys():
                    bond_map[int(line.split()[0])] = [ int(line.split()[1]) ]
                else:
                    bond_map[int(line.split()[0])].append(int(line.split()[1]))

    return bond_map