import json
from sys import argv
"""
Parse table output from HMMer and save it as a JSON dictionary
"""

# Pfam tbl path
pfam_in = argv[1]
# output path
dict_out = open(argv[2], "w")
# Initialize pfam dictionary
pfam_db = {}

for line in open(pfam_in):
    # Skip comment lines
    if not line.startswith("#"):
        # Process each line
        data_line = [x for x in line.split(' ') if x != ""]
        acc = data_line[0].split("|")[2]
        ## This was added because for some datasets, gene accessions can be non-unique
        ## so the first 3 letters of the genus/species name was added to the accession
        # sp_list = data_line[0].split("|")[0].split("_")
        # if len(sp_list) == 4:
        #     gen = sp_list[1][:3]
        #     spe = sp_list[2][:3]
        # else:
        #     gen = sp_list[0][:3]
        #     spe = sp_list[1][:3]
        # seq = "{}_{}_{}".format(gen, spe, acc)
        seq = acc
        pfam = data_line[3]
        if not pfam in pfam_db:
            pfam_db[pfam] = [seq]
        else:
            pfam_db[pfam].append(seq)

# Save to file
dict_out.write(json.dumps(pfam_db))