"""
    RSA <= 0.2 Buried (Inaccessible)
    RSA > 0.2 Exposed (Accessible)

    SST = [H,I,G] - Helix
    SST = [B,E] - Beta
    SST = [T] - Turn
    SST = [S,-] - Coil
"""
from Bio.PDB import PDBParser, DSSP
import pickle
import os
import sys
import warnings

warnings.filterwarnings("ignore")

#CURRENT_FOLDER = '/home/local/BHRI/sportelli/Desktop/Important_Code/structural/aaindex'
CURRENT_FOLDER = '/home/tanu/git/LSHTM_analysis/scripts/aa_index_scripts/aaindex'
DATA_FOLDER = os.path.join(CURRENT_FOLDER,'data')

RSA_SST_DEPENDENT = {
    'exposed_helix' : 'KOSJ950101',
    'exposed_beta'  : 'KOSJ950102',
    'exposed_turn'  : 'KOSJ950103',
    'exposed_coil'  : 'KOSJ950104',
    'buried_helix'  : 'KOSJ950105',
    'buried_beta'   : 'KOSJ950106',
    'buried_turn'   : 'KOSJ950107',
    'buried_coil'   : 'KOSJ950108',
}

SST_DEPENDENT = {
    'helix'  : 'KOSJ950109',
    'beta'   : 'KOSJ950110',
    'turn'   : 'KOSJ950111',
    'coil'   : 'KOSJ950112',
}

RSA_DEPENDENT1 = {
    'exposed' : 'KOSJ950113',
    'buried'  : 'KOSJ950114',
}

RSA_DEPENDENT2 = {
    'exposed' : 'OVEJ920104',
    'buried'  : 'OVEJ920105',
}

def get_environment(pdb_file, chain, position, insertion_code=' '):
    parser = PDBParser()
    structure = parser.get_structure(pdb_file, pdb_file)
    model = structure[0]

    dssp = DSSP(model, pdb_file, dssp='mkdssp')
    dssp_key = [item for item in dssp.keys() if item[0] == chain and item[1][1] == int(position) and item[1][2] == insertion_code]

    dssp_key = dssp_key[0]
    sst = dssp[dssp_key][2]
    rsa = float(dssp[dssp_key][3])

    return{'sst':sst, 'rsa':rsa}

def main():
    """
        READ INPUT
    """
    pdb_file = sys.argv[1]
    chain_id = sys.argv[2]
    mutation_code = sys.argv[3]

    aa_from = mutation_code[0]
    aa_to = mutation_code[-1]
    position = mutation_code[1:-1]
    insertion_code = ' '
    if not position[-1].isdigit():
        insertion_code = position[-1]
        position = position[:-1]

    """
        READ DATABASES
        index2 - Amino acid substitution indexes
        index3 - Statistical protein contact potentials
    """
    index2 = pickle.load(open('{}/aaindex2.p'.format(DATA_FOLDER),'rb'))
    index3 = pickle.load(open('{}/aaindex3.p'.format(DATA_FOLDER),'rb'))

    """
        LOOP THROUGH TABLES AND EXTRACT VALUES
    """
    results_index2 = dict()
    results_index3 = dict()
    for key in index2.keys():
        if index2[key][aa_from][aa_to] != None:
            results_index2[key] = index2[key][aa_from][aa_to]
        else:
            results_index2[key] = index2[key][aa_to][aa_from]

    for key in index3.keys():
        if index3[key][aa_from][aa_to] != None:
            results_index3[key] = index3[key][aa_from][aa_to]
        else:
            results_index3[key] = index3[key][aa_to][aa_from]

    """
        GET ENVIRONMENT CHARACTERISTICS
    """
    environment = get_environment(pdb_file, chain_id, position, insertion_code)

    buried = 'buried'
    sst = str()
    if environment['rsa'] <= 0.2:
        buried = 'exposed'

    if environment['sst'] in ['H','I','G']:
        sst = 'helix'
    elif environment['sst'] in ['B','E']:
        sst = 'beta'
    elif environment['sst'] in ['T']:
        sst = 'turn'
    else:
        sst = 'coil'

    results_index2['KOSJ950100_RSA_SST'] = results_index2[RSA_SST_DEPENDENT['{}_{}'.format(buried,sst)]]
    results_index2['KOSJ950100_SST'] = results_index2[SST_DEPENDENT[sst]]
    results_index2['KOSJ950110_RSA'] = results_index2[RSA_DEPENDENT1[buried]]
    results_index2['OVEJ920100_RSA'] = results_index2[RSA_DEPENDENT2[buried]]

    for value in RSA_SST_DEPENDENT.values():
        results_index2.pop(value)
    for value in SST_DEPENDENT.values():
        results_index2.pop(value)
    for value in RSA_DEPENDENT1.values():
        results_index2.pop(value)
    for value in RSA_DEPENDENT2.values():
        results_index2.pop(value)

    """
        PRINT RESULTS
    """
    output_dict = dict()
    output_dict.update(results_index2)
    output_dict.update(results_index3)

    keys = list(output_dict.keys())
    keys.sort()
    values = [str(output_dict[item]) for item in keys]

    # print(",".join(keys))
    print(",".join(values))

    return True


if __name__ == "__main__":

    if len(sys.argv) != 4:
        print("Error on parsing argument list")
        print("Please provide a one letter code for wild-type and mutant residues")
        print("Eg.: python get_scores.py pdb_file chain_id mutation_code")
        sys.exit(1)
    main()
