#
#
#

__author__ = "Magdalena A. Jonikas"
__version__ = "3.0"
__date__ = "04/08/09"
__doc__ = """Find full atomic matches to a coarse-grain template structure"""

import sys, math, time, pickle
import c2a


def findMatches(templateTrace,
                centerAtoms,
                fragmentFile,
                conversionData,
                searchStructureFile,
                nhradius=0.1,
                hradius=0.1):
    t0=time.time()
# Define names:
    templateName = templateTrace[:-4] # Remove the .pdb part of the template name
    molName = fragmentFile[0:-7] # Use the first 4 characters as the name of the molecule
    libName = searchStructureFile[:-4] # Name of the full atomic source molecule

# Process the fragments and primary sequence
    try:
        primaryData = open('%s-primary.txt' % molName).readlines() # Primary sequence
    except IOError:
        try:
            primaryData = open('%s.seq' % molName).readlines()
        except IOError:
            print "need a primary sequence filed named either:"
            print "%s-primary.txt" % molName
            print "or"
            print "%s.seq" % molName
            sys.exit()
    fragmentData = open(fragmentFile).readlines() # Fragment definitions
    pieces = c2a.Pieces(fragmentData)

# Load the search structure
    print "Loading data from the full atomic structure %s" % libName
    searchStructureData = open(searchStructureFile).readlines()
    searchStructure = c2a.Coarse(searchStructureData, typelist=centerAtoms)
    searchStructure.makeSearchLibs()
    searchStructure.dfudge = nhradius
    searchStructure.dhfudge = hradius
    searchStructure.fullAtomic = c2a.FullAtomic(searchStructureData)
    t1 = time.time()
    print "  This took %.2f seconds" % (t1-t0)

# Setup the nucleotide conversion data
    conversion = pickle.load(conversionData)
    convStructure = c2a.ConvStructure(primaryData, pieces, conversion)

# Parse the frames in the template trace file
    print "Using file: ", templateTrace
    frameList = c2a.parseFrames(templateTrace)

    coordsLib = {}
    countData = {}
    frameIndex=0
    t3 = time.time()

    for frame in frameList:
        print "Processing frame", frameIndex
        libOut = open('%s-%s-%i-lib.pkl' % (templateName, libName, frameIndex), 'w')
        finalLib = {}
        modelStruct = c2a.Coarse(frame, typelist=centerAtoms)

        for end in pieces.ends.keys():
            e = pieces.ends[end]
            resList = [e.o1ResList]
            finalLib[end] = c2a.getOptions(end, resList, modelStruct, searchStructure, templateName, frameIndex, convStructure)
            t4 = time.time()
            print "It took %.2f sec. to process end" % (t4-t3), end
            t3=t4

        for loop in pieces.loops.keys():
            l = pieces.loops[loop]
            resList = [l.o1ResList]
            finalLib[loop] = c2a.getOptions(loop, resList, modelStruct, searchStructure, templateName, frameIndex, convStructure)
            t4 = time.time()
            print "It took %.2f sec. to process loop" % (t4-t3), loop
            t3=t4
            

        for junction in pieces.junctions.keys():
            j = pieces.junctions[junction]
            if len(j.resList)==0: continue
            resList = [j.o1ResList]
            finalLib[junction] = c2a.getOptions(junction, resList, modelStruct, searchStructure, templateName, frameIndex, convStructure)
            t4 = time.time()
            print "It took %.2f sec. to process junction" % (t4-t3), junction
            t3=t4

        for helix in pieces.helices.keys():
            h = pieces.helices[helix]
            resList = [h.H1resList, h.H2resList]
            finalLib[helix]= c2a.getOptions(helix, resList, modelStruct, searchStructure, templateName, frameIndex, convStructure)
            t4 = time.time()
            print "It took %.2f sec. to process helix" % (t4-t3), helix
            t3=t4
        pickle.dump(finalLib, libOut)
        frameIndex+=1
        libOut.close()


def main():
    t0 = time.time() # Start timing

# Inputs:
    findMatches(templateTrace = sys.argv[1], # Coarse grain model e.g. 6TNA-C3.pdb
                fragmentFile = sys.argv[2], # Fragment definition file e.g. 6TNA-PS.txt
                conversionData = open(sys.argv[3]), # Conversion data between residues
                searchStructureFile = sys.argv[4])# PDB file containing a full atomic structure e.g. 1N32.pdb
    return (time.time()-t0)


if __name__=='__main__':
    print "\nTotal time: %.2f\n" % (main())