#!/usr/bin/python
# Build A form RNA duplex
# Convert glob model to an allatom explicit form
import sys
import numpy as np 
from lib.libblock import BlockDuplex
from lib.libSAXSPM import SAXSPM, read_pdb_xyz
from lib.librna import ConfRMSD
import os

class Block2Duplex(BlockDuplex): 
    def start(self):
        self.load2str(vb=0)
        self.buildDuplexes()

class GetAllatom(BlockDuplex):
    def start(self):
        self.load2str(vb=0)
        #let us first generate the single strands
        # and align their center of geometry (cog) the target glob structure
        cog = {"A":(-10.077, 6.409, 4.698),"U":(-16.378, -5.194, 3.223),\
               "G":(-3.249, 3.149, -0.868),"C":(-17.204, 2.798, 9.777)}
        xyz = read_pdb_xyz(target_pdb)
        pqr = read_pdb_xyz(templ_pdb)
        mod = []
        for nt in self.sstrands:
            for resid in nt:
                #print cog[self.seq_dict[resid]]
                for m in range(len(xyz)):
                    if (xyz[m][1] == self.seq_dict[resid] and xyz[m][2] == resid):
                        transx = xyz[m][3] - cog[self.seq_dict[resid]][0]
                        transy = xyz[m][4] - cog[self.seq_dict[resid]][1]
                        transz = xyz[m][5] - cog[self.seq_dict[resid]][2]
                        #print transx, transy, transz
                        mod.append([xyz[m][1], xyz[m][2], transx, transy, transz])
        #print len(mod)
        counter = 1
        allatom_ss = []
        fmt_pdb = '%-6s%5d %-4s %-4s%1s%4d%1s%3s%8.3f%8.3f%8.3f%18s%-4s'
        #f3 = open("sstrands-translate.pdb", 'w')
        for ss in mod:
            for temp in pqr:
                if (ss[0] == temp[1]):
                    allatom_ss.append([temp[0], temp[1], ss[1],temp[3]+ss[2], temp[4]+ss[3], temp[5]+ss[4]]) 
                    #print >> f3, fmt_pdb % ('ATOM', counter,temp[0],temp[1],'X',ss[1],\
                    #                        "","",temp[3]+ss[2],temp[4]+ss[3],temp[5]+ss[4],'0.00','')
                    counter += 1
        #print len(self.duplexs)
        #print self.project_name
        duplex = []
        for index in range(len(self.duplexs)):
             duplex_pdb = self.project_name + "_duplex" + str(index+1) + ".pdb"
             duplex += align_pdb(duplex_pdb, target_pdb)
        full_pdb = allatom_ss + duplex
        serial = 1
        base_pdb_name=os.path.basename(target_pdb) 
        final_pdb =  base_pdb_name[:-4] + "-allatom.pdb"
        f4 = open(final_pdb, 'w')
        for resid in range(self.rna_start_resid, self.rna_end_resid+1):
            for line in range(len(full_pdb)):
                if (resid == full_pdb[line][2]):
                    print >> f4, fmt_pdb % ('ATOM', serial,full_pdb[line][0],full_pdb[line][1],'X',full_pdb[line][2], \
                                            "","",full_pdb[line][3],full_pdb[line][4],full_pdb[line][5],'0.00','') 
                    serial += 1
        #print self.sstrands





def check_oddatoms(xyz):
    allatoms = ["P", "OP1", "OP2", "O5'", "OP3", "C5'", "C4'", "O4'", "C3'", "O3'",\
                "C2'", "O2'", "C1'","N9", "C8", "N7", "C5", "C6", "N6", "O6", "N1",\
                "C2", "O2", "N2", "N3", "C4", "O4", "N4"]
    for i in range(len(xyz)):
        if (SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][0] not in allatoms):
            print "atom %s not in list: check list" % SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][0]
            sys.exit(1)

def get_globs(pdb_name):
    #f1 = open('thf_ribo_deform.pdb','r')
    #f2 = open("glob_CM.pdb", 'w')
    xyz = read_pdb_xyz(pdb_name)
    check_oddatoms(xyz)
    allatoms = ["P", "OP1", "OP2", "O5'", "OP3", "C5'", "C4'", "O4'", "C3'", "O3'",\
                "C2'", "O2'", "C1'","N9", "C8", "N7", "C5", "C6", "N6", "O6", "N1",\
                "C2", "O2", "N2", "N3", "C4", "O4", "N4"]
    #print "%d  atoms found in %s" %(len(xyz), pdb_name)
    num_lines = len(xyz)
    serial = 1
    fmt_pdb = '%-6s%5d %-4s %-4s%1s%4d%1s%3s%8.3f%8.3f%8.3f%18s%-4s'
    nx,ny,nz,n_tot=0,0,0,0  # initiatlize
    glob = []
    for i in range(0,num_lines):
        count_begin = 0;
        if i != num_lines -1 and xyz[i][2] == xyz[i+1][2] : #check the resids are the same
             nx += SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][2]*xyz[i][3]
             ny += SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][2]*xyz[i][4]
             nz += SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][2]*xyz[i][5]
             n_tot += SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][2]
        elif i == num_lines:
             nx += SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][2]*xyz[i][3]
             ny += SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][2]*xyz[i][4]
             nz += SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][2]*xyz[i][5]
             n_tot += SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][2]
             #print >> f2, fmt_pdb % ('ATOM', serial,'CA',str(xyz[i][1]),'X',xyz[i][2],\
             #                        "","",nx/n_tot,ny/n_tot,nz/n_tot,'0.00','')
             glob.append([str(xyz[i][1]), xyz[i][2], nx/n_tot,ny/n_tot,nz/n_tot])
             #print xyz[i][2], nx/n_tot, ny/n_tot, nz/n_tot
        else:
             nx += SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][2]*xyz[i][3]
             ny += SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][2]*xyz[i][4]
             nz += SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][2]*xyz[i][5]
             n_tot += SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][2]
             #print >> f2, fmt_pdb % ('ATOM', serial,'CA',str(xyz[i][1]),'X',xyz[i][2],"",\
             #                        "",nx/n_tot,ny/n_tot,nz/n_tot,'0.00','')
             glob.append([str(xyz[i][1]), xyz[i][2], nx/n_tot,ny/n_tot,nz/n_tot])
             #print xyz[i][2], nx/n_tot, ny/n_tot, nz/n_tot
             nx,ny,nz,n_tot=0,0,0,0  # reinitiatlize
             serial += 1
    return glob

def align_pdb(allatom_pdb, target_pdb):
    glob = get_globs(allatom_pdb)
    #print glob
    xyz = read_pdb_xyz(target_pdb)
    pqr = read_pdb_xyz(allatom_pdb)
    targetglobs = []
    for k in range(len(glob)):
        #print glob[k][0], glob[k][1]
        for m in range(len(xyz)):
            if (glob[k][0] == xyz[m][1] and glob[k][1] == xyz[m][2]):
                #print glob[k][0],xyz[m][1],glob[k][1],xyz[m][2]
                targetglobs.append([xyz[m][1], xyz[m][2], xyz[m][3],xyz[m][4],xyz[m][5]])
    #print(targetglobs)
    #print(len(glob))
    #print targetglobs[25][4]
    conf1 = []
    conf2 = []
    conf3 = []
    for i in range(len(glob)):
        #print(i)
        conf2.append([glob[i][2],glob[i][3],glob[i][4]])
        conf1.append([targetglobs[i][2],targetglobs[i][3],targetglobs[i][4]])
    #print conf1
    for i in range(len(pqr)):
        conf3.append([pqr[i][3],pqr[i][4],pqr[i][5]])
    array1 = np.asarray(conf1)
    array2 = np.asarray(conf2)
    array3 = np.asarray(conf3)
    mycal = ConfRMSD()
    mycal.calculate_rmsd(array1, array2,array3, is_align=True, vb=True)
    #f2 = open("aligned-glob.pdb", 'w')
    #f3 = open("aligned-allatom.pdb", 'w')
    fmt_pdb = '%-6s%5d %-4s %-4s%1s%4d%1s%3s%8.3f%8.3f%8.3f%18s%-4s'
    duplex = []
    for p in range(0,len(glob)):
        #print >> f2, fmt_pdb % ('ATOM', p,'CA',str(glob[p][0]),'X',glob[p][1],"","",array2[p][0],array2[p][1],array2[p][2],'0.00','')
        pass
    for q in range(0,len(pqr)):
        #print >> f3, fmt_pdb % ('ATOM', q,str(pqr[q][0]),str(pqr[q][1]),'X',(pqr[q][2]),"","",array3[q][0],array3[q][1],array3[q][2],'0.00','')
        duplex.append([str(pqr[q][0]), str(pqr[q][1]), pqr[q][2], array3[q][0], array3[q][1], array3[q][2]])
    return duplex


############### Program Execution ###################################

if (__name__ == "__main__"):
    import sys
    if len(sys.argv) != 4:
        Usage = """Usage: 
                        python   %s   fn_inp target.pdb nt-tmpl.pdb 
                        
                        fn_inp:   input secondary structure file (sequence and duplexes)
                   
                        target.pdb: target glob model

                        nt-tmpl.pdb: template pdb file used to write out single strands 
                    """ % __file__
        print Usage
    else:
        fp_inp = sys.argv[1]
        target_pdb = sys.argv[2]
        templ_pdb = sys.argv[3]
        r_obj1 = Block2Duplex(fp_inp)
        r_obj1.start()
        r_obj2 = GetAllatom(fp_inp)
        r_obj2.start()
    
