# NOEtools.py: A python module for predicting NOE coordinates from
#                assignment data.
#
#    The input and output are modelled on nmrview peaklists.
#    This modules is suitable for directly generating an nmrview
#    peaklist with predicted crosspeaks directly from the
#    input assignment peaklist.

from . import xpktools


def predictNOE(peaklist, originNuc, detectedNuc, originResNum, toResNum):
# Predict the i->j NOE position based on self peak (diagonal) assignments
#
# example predictNOE(peaklist,"N15","H1",10,12)
#    where peaklist is of the type xpktools.peaklist
#    would generate a .xpk file entry for a crosspeak
#    that originated on N15 of residue 10 and ended up
#    as magnetization detected on the H1 nucleus of
#    residue 12.
# CAVEAT: The initial peaklist is assumed to be diagonal (self peaks only)
#       and currently there is not checking done to insure that this
#       assumption holds true.  Check your peaklist for errors and
#       off diagonal peaks before attempting to use predictNOE.

    returnLine = ""  # The modified line to be returned to the caller

    datamap = _data_map(peaklist.datalabels)

    # Construct labels for keying into dictionary
    originAssCol = datamap[originNuc + ".L"] + 1
    originPPMCol = datamap[originNuc + ".P"] + 1
    detectedPPMCol = datamap[detectedNuc + ".P"] + 1

    # Make a list of the data lines involving the detected
    if str(toResNum) in peaklist.residue_dict(detectedNuc) \
    and str(originResNum) in peaklist.residue_dict(detectedNuc):
        detectedList = peaklist.residue_dict(detectedNuc)[str(toResNum)]
        originList = peaklist.residue_dict(detectedNuc)[str(originResNum)]
        returnLine = detectedList[0]

        for line in detectedList:
            aveDetectedPPM = _col_ave(detectedList, detectedPPMCol)
            aveOriginPPM = _col_ave(originList, originPPMCol)
            originAss = originList[0].split()[originAssCol]

        returnLine = xpktools.replace_entry(returnLine, originAssCol + 1, originAss)
        returnLine = xpktools.replace_entry(returnLine, originPPMCol + 1, aveOriginPPM)

    return returnLine


def _data_map(labelline):
    # Generate a map between datalabels and column number
    #   based on a labelline
    i = 0  # A counter
    datamap = {}  # The data map dictionary
    labelList = labelline.split()  # Get the label line

    # Get the column number for each label
    for i in range(len(labelList)):
        datamap[labelList[i]] = i

    return datamap


def _col_ave(list, col):
    # Compute average values from a particular column in a string list
    total = 0
    n = 0
    for element in list:
        total += float(element.split()[col])
        n += 1
    return total / n
