#!/usr/bin/python3

import os, sys
import glob
import json
import pickle

# This will be a dictionary of all known products
# where the key is the MIBiG accesion and the value is a tuple ([SMILES], source)
allknownproducts = {}

##################################
# Add structures from JSON files #
##################################

mibigpath = './data/mibig/raw/mibig_json_2.0'
mibigjson = glob.glob(os.path.join(mibigpath, '*.json'))

# Some of the compounds in MIBiG are the predicted final product for multiple clusters
# but are only recorded in a subset of the JSON files for each cluster
# We therefore generate a dictionary mapping compound names to SMILES structures so that
# we can use these known structures later on.
mibigcompounds = {}

# This is a dictionary that associates all of the MIBiG accessio numbers to compound names
mibignames = {}

#print(mibigjson)

for jsonfile in mibigjson:
    with open(jsonfile) as f:
        data = json.loads(f.read())
    try:
        # Get names
        accession = os.path.basename(jsonfile).split('.')[0]
        #names = [x['compound'] for x in data['general_params']['compounds']] #old json files
        names = [x['compound'] for x in data['cluster']['compounds']]
        shortname = names[0].split()[0].lower()
        # Enter compound name into name dictionary
        # This must be done first, since structures may not be present resulting in KeyError
        mibignames[accession] = shortname
        # Get structures
        structures = [x['chem_struct'] for x in data['cluster']['compounds']]
        # Enter known product dictionary
        allknownproducts[accession] = (structures, 'mibig')
        # Enter known product into compound dictionary
        # Only use first of the structures for this compound dictionary
        mibigcompounds[shortname] = structures[0]
    except KeyError:
        continue

################################
# Get structures from ChemAxon #
################################

chemaxonsmi = './data/compounds/chemaxon.smi'

# This dictionary will also have name of compounds as keys and SMILES as values
chemaxoncompounds = {}

# Import structures from ChemAxon predictions
with open(chemaxonsmi) as f:
    for compound in [line.split() for line in f.readlines()]:
        name = compound[1].lower()
        structure = compound[0]
        # Some duplicate structures were generated by ChemAxon
        # so we ignore all compounds with '_' which are duplicates
        # some legitimate compounds maybe removed, so check manually if needed
        if len(name.split('_')) == 1:
            chemaxoncompounds[name] = structure

##################################################################
# Add duplicated structures from MIBiG, structures from ChemAxon #
##################################################################

# Iterate over all accession numbers
mibigaccessions = [os.path.basename(jsonfile).split('.')[0] for jsonfile in mibigjson]

for accession,compound in mibignames.items():
    # Check if we already have a product for the cluster
    if accession not in allknownproducts.keys():
        # Prioritize structures from MIBiG
        if compound in mibigcompounds.keys():
            structure = mibigcompounds[compound]
            allknownproducts[accession] = ([structure], 'mibig')
        # Otherwise try ChemAxon
        elif compound in chemaxoncompounds.keys():
            structure = chemaxoncompounds[compound]
            allknownproducts[accession] = ([structure], 'chemaxon')
        else:
            continue

#######################################
# Add structures from manual curation #
#######################################

manualsmi = './data/compounds/manual.smi'

with open(manualsmi) as f:
    for compound in [line.split() for line in f.readlines()]:
        accession = compound[0].strip()
        structure = compound[1].strip()
        allknownproducts[accession] = ([structure], 'manual')

#######################################
# Write aggregated structures to file #
#######################################

pickle.dump(allknownproducts, open('./data/compounds/all_known_products.p', 'wb'))
print("Product structures aggregated. ")
