# read RNA sequence and base pairs from data file # from paper with secondary structure
# output RNA 2nd  noe distance, torsion angle, hydrogen bond restraints for xplor 

from re import split 
from string import strip,rstrip
import numpy as npy 
import os 
from math import degrees
from libio import INPUT
#from ibduplex import Duplex
from librna import  RNA2D, RNAPDB 


SS_SEQ_CS = {"AA":[("H8", "H1'", 7.5), ("H2", "H8", 6.5), ("H2", "H1'", 6.5), ("H1'", "H8", 6.5)], # 1
             "AC":[("H1'", "H6", 6.5), ("H2", "H1'", 3.5), ("H2", "H5", 6.5), ("H2", "H6", 5.0), ("H8", "H6", 6.5)], # 2
             "AG":[("H1'","H8", 5.0), ("H8", "H8", 5.0), ("H2", "H1'", 3.5), ("H2","H1'", 3.5), ("H2", "H8", 5.0)], # 3
             "GA":[("H8", "H8", 6.5), ("H1", "H8", 6.5), ("H1", "H2", 7.5), ("H1", "H1'", 6.5), ("H1'", "H8", 6.5)], #4
             'AU':[("H8", "H6", 6.5), ("H8","H1'", 7.5), ("H1'","H5", 6.5), ("H1'", "H6", 6.5)], #5
             "UA":[("H6", "H8",6.5), ("H5","H8",7.5), ("H1'","H8",5.0), ("H6", "H1'", 7.5)], # 6
             "CC":[("H1'", "H5", 5.0), ("H1'", "H6", 5.0), ("H5", "H6", 7.5), ("H6", "H1'", 7.5)],  # 7
             
             "GC":[("H1'", "H5", 6.5), ("H1'", "H6", 5.0), ("H1","H1'", 5.0), ("H1", "H5", 6.5), ("H1", "H6", 6.5), ("H8", "H1'", 7.5), ("H8", "H6", 5.0)], #8
             "CG":[("H5", "H1'", 6.5), ("H6", "H1'", 5.0), ("H1'","H1", 5.0), ("H5", "H1", 6.5), ("H6", "H1", 6.5), ("H1'", "H8", 7.5), ("H6", "H8", 5.0)], #8
             
             "CU":[("H5", "H1'", 7.5), ("H5", "H5", 5.0), ("H5", "H6", 5.0), ("H6","H1'", 7.5), ("H6", "H5", 5.0), ("H6", "H6", 5.0)], # 9
             "UC":[("H1'", "H5", 5.0), ("H5", "H5", 5.0), ("H5", "H6", 6.5), ("H6", "H5", 5.0), ("H6", "H6", 5.0)], # 10
             
             "GG":[("H8", "H8", 5.0), ("H8", "H1", 7.5), ("H1", "H8", 7.5), ("H1", "H1", 3.5), ("H8", "H1'", 7.5), ("H1'", "H8", 5.0)], # 11
             "GU":[("H8", "H6", 6.5), ("H1'", "H5", 6.5), ("H1'", "H6", 6.5), ("H8", "H1'", 7.5)], # 12
             "UG":[("H1'","H8", 6.5), ("H5", "H8", 7.5), ("H6", "H8", 6.5)], # 13
             "UU":[("H6", "H1'", 7.5), ("H1'", "H6", 6.5)] # 14 
             }
              
              

class GenXplorTBL(INPUT, RNA2D, RNAPDB):
    """ generate restraints file for xplor """
    nt_names = {"U":"Uri", "G":"Gua", "A":"Ade", "C":"Cyt"} 
    nt_hids = {"A":8, "G":8, "C":6,"U":6} # the index of the hydrogen in pyrimine(pyrimidines) next to sugar " 
    nt_Cnames = {"U":"URI", "G":"GUA", "A":"ADE", "C":"CYT"}
    def __init__(self, fn_pdb):
        RNAPDB.__init__(self, fn_pdb=fn_pdb)
        
    def sugarPucker(self, vb=0):
        """
        !  Torsion angle restraints for sugar pucker in the RNA hairpin duplex region  
        !  The following infomration come from Xplor for RNA 
        !  Pucker is defined with ribose torsion angles nu1 and nu2.
        !               nu1 = O4'-C1'-C2'-C3'
        !               nu2 = C1'-C2'-C3'-C4'
        !       "N" C3' endo:
        !               nu1 = -20 +/- 10
        !               nu2 = +35 +/-  5
        !       "S" C2' endo:     # usually in the loop  or single strand region
        !               nu1 = +35 +/-  5
        !               nu2 = -35 +/-  5
        
        !  C3' endo nucleotides:
        !
        !       Gua  1
        {
        assign (resid  1 and name O4')
               (resid  1 and name C1')
               (resid  1 and name C2')
               (resid  1 and name C3') 1.0 -20.0 10.0 2
        assign (resid  1 and name C1')
               (resid  1 and name C2')
               (resid  1 and name C3')
               (resid  1 and name C4') 1.0  35.0 5.0 2
        }

        """
        sugar_pucker_lns = ["!  C3' endo nucleotides:  "]
        sugar_ntids = []
        for nt_id in self.helices_resids:
            if nt_id not in sugar_ntids:
                nt_name = self.residues[nt_id].nt_name #self.nt_names[nt_sn]
                nt_lns = ["!"]
                nt_lns += ["!       %s  %3d " %  (nt_name, nt_id)]  
                nt_lns += ["{"]
                nt_lns += ["assign (resid  %2d and name O4')" % nt_id]
                nt_lns += ["       (resid  %2d and name C1')" % nt_id]
                nt_lns += ["       (resid  %2d and name C2')" % nt_id]
                nt_lns += ["       (resid  %2d and name C3') 1.0 %4.1f 10.0 2" % (nt_id, (self.residues[nt_id].getMu1()))]
                nt_lns += ["assign (resid  %2d and name C1')" % nt_id]
                nt_lns += ["       (resid  %2d and name C2')" % nt_id]
                nt_lns += ["       (resid  %2d and name C3')" % nt_id]
                nt_lns += ["       (resid  %2d and name C4') 1.0 %4.1f 5.0 2 " % (nt_id,degrees(self.residues[nt_id].getMu2()))]
                nt_lns += ["}"]
                nt_lns += ["assign (resid %2d and name c5') (resid %2d and name c4')" % (nt_id, nt_id)]      # delta angles 
                nt_lns += ["       (resid %2d and name c3') (resid %2d and name o3') 1.0 %4.1f 20.0 2 " % (nt_id, nt_id, self.residues[nt_id].getDelta())]
                sugar_ntids.append(nt_id)
                sugar_pucker_lns += nt_lns
        if vb: self.prnLns(sugar_pucker_lns, no_strip=1)            
        return sugar_pucker_lns
     
     
    def genXhi(self, vb=0):
        """
        !
        !
        !  Torsion angle restraints for the xhi torsion angle in the
        !       GZ33 - Theophylline RNA complex.
        !
        !  Xhi is defined as follows...
        !               Purines     = O4'-C1'-N9-C4
        !               Pyrimidines = O4'-C1'-N1-C2
        !       Anti:
        !               Xhi = 180.0 +/- 90.0
        !       including high-anti and excluding high-syn:
        !               Xhi = -150.0 +/- 90.0
        !       Syn:
        !               Xhi =   0.0 +/- 90.0
        !       See Saenger page 22.
        !
        !  There are no syn bases in the GZ33-Theophylline complex.
        !     No large H1'-Haro nOes were observed, and large H3'-Haro
        !     (or H2'-Haro) were observed for all nucleotides.
        !
        assign (resid  1 and name O4') (resid  1 and name C1')
               (resid  1 and name N9 )
               (resid  1 and name C4 ) 1.0 -150.0 90.0 2
        assign (resid  6 and name O4') (resid  6 and name C1')
        (resid  6 and name N1 )
       (resid  6 and name C2 ) 1.0 -150.0 90.0 2
       
               """
        xhi_lns = []
        for nt_id in self.resids: 
            nt_sn = self.residues[nt_id].nt_name
            if nt_sn in ['A','G']:
                xhi_lns += ["assign (resid %2d and name O4') (resid %2d and name C1')" % (nt_id, nt_id)]
                xhi_lns += ["       (resid %2d and name N9 )" % nt_id]
                xhi_lns += ["       (resid %2d and name C4 ) 1.0 %5.1f 90.0 2 "  % (nt_id, self.residues[nt_id].getXhi()) ]
            else:
                xhi_lns += ["assign (resid %2d and name O4') (resid %2d and name C1')" % (nt_id, nt_id)]
                xhi_lns += ["       (resid %2d and name N1 )" % nt_id]
                xhi_lns += ["       (resid %2d and name C2 ) 1.0 %5.1f 90.0 2 " % (nt_id, self.residues[nt_id].getXhi())]
                
        if vb: self.prnLns(xhi_lns, no_strip=1)
        return xhi_lns 
    
    
    def genEpsZeta(self, vb=0):
        """ generate Epsion and Zeta angle restraints only for duplex regions
        !!1
        # Epsilon
        assign  (resid  1 and name  c4') (resid  1 and name  c3')
                (resid  1 and name  o3') (resid  2 and name   p ) 1.0 -160.0 50.0 2
        # Zeta
        assign  (resid  1 and name  c3') (resid  1 and name  o3')
                (resid  2 and name   p ) (resid  2 and name  o5') 1.0 -70.0 50.0 2
        """
        eps_lns = []
        zeta_lns = []
        for nt_id in self.helices_resids:
            try:
                eps = self.getResEpsilon(nt_id)
                eps_lns += ["!!%d" % nt_id]
                eps_lns += ["assign  (resid %2d and name  c4') (resid %2d and name  c3')" % (nt_id, nt_id)]
                eps_lns += ["        (resid %2d and name  o3') (resid %2d and name   p ) 1.0 %5.1f 50.0 2 " % (nt_id, nt_id+1, eps)] 
            except:
                pass 
            try:
                zeta = self.getResZeta(nt_id)
                zeta_lns += ["!!%d" % nt_id]
                zeta_lns += ["assign  (resid %2d and name  c3') (resid %2d and name  o3')" % (nt_id, nt_id)]
                zeta_lns += ["        (resid %2d and name   p ) (resid %2d and name  o5') 1.0 %5.1f 50.0 2 " % (nt_id+1, nt_id+1, zeta)]
            except:
                pass 
        if vb: 
            self.prnLns(eps_lns, no_strip=1)
            self.prnLns(zeta_lns, no_strip=1)
        return eps_lns, zeta_lns
    
    
    def genAlphaBeta(self, vb=0):
        """
        !!2
        #alpha
        assign  (resid  1 and name  o3') (resid  2 and name   p )
                (resid  2 and name  o5') (resid  2 and name  c5') 1.0 -60.0 40.0 2
        #beta         
        assign  (resid  2 and name   p ) (resid  2 and name  o5')
                (resid  2 and name  c5') (resid  2 and name  c4') 1.0 180.0 50.0 2
        
        """ 
        ntids = self.helices_resids
        ntids.sort()
        pre_ntid = ntids[0]
        alpha_lns = []
        beta_lns = []
        for nt_id in ntids[1:]: 
            if pre_ntid  + 1 == nt_id: 
                alpha_lns += ["!!%d" % nt_id]
                alpha_lns += ["assign  (resid %2d and name  o3') (resid %2d and name   p )" % (nt_id-1, nt_id)]
                alpha_lns += ["        (resid %2d and name  o5') (resid %2d and name  c5') 1.0 %5.1f 40.0 2 " % (nt_id, nt_id, self.getResAlpha(nt_id))] 
                beta_lns += ["!!%d" % nt_id]
                beta_lns += ["assign  (resid %2d and name   p ) (resid %2d and name  o5')" % (nt_id, nt_id)]
                beta_lns += ["        (resid %2d and name  c5') (resid %2d and name  c4')  1.0 %5.1f 50.0 2 " % (nt_id, nt_id, self.residues[nt_id].getBeta())]
            pre_ntid = nt_id
        if vb: 
            self.prnLns(alpha_lns, no_strip=1)
            self.prnLns(beta_lns, no_strip=1)
        return alpha_lns, beta_lns
    
    def genGamma(self, vb=0):
        """
        assign (resid 1 and name o5') (resid 1 and name c5')
        (resid 1 and name c4') (resid 1 and name c3')   1.0 60.0 30.0 2
        """
        gamma_lns  =  ["!gamma"]
        for nt_id in self.helices_resids:
            gamma_lns += ["assign (resid %2d and name o5') (resid %2d and name c5') " % (nt_id, nt_id) ]
            gamma_lns += ["       (resid %2d and name c4') (resid %2d and name c3')   1.0 %5.1f 30.0 2 " % (nt_id, nt_id, self.residues[nt_id].getGamma())]
        if vb: self.prnLns(gamma_lns, no_strip=1)
        return gamma_lns
    
       
    def genHbond(self, helices, vb=0):
        """ generate Hydrogen bonds restraints for duplexes
        """ 
        hbond_lns = []
        for hex_id in helices.keys():
            for _bp in helices[hex_id]:  #_bp: (1, 20, 'GU')
                nt1_sn = _bp[2][0]
                nt2_sn = _bp[2][1]
                bp=[(nt1_sn, _bp[0]), (nt2_sn, _bp[1])]
                if nt1_sn == 'G' and nt2_sn == "C": # G-C pair 
                    hbond_lns += self.__getGCHbond(bp[0], bp[1])
                elif nt1_sn == "C" and nt2_sn == "G": 
                    hbond_lns += self.__getGCHbond(bp[1], bp[0])
                elif nt1_sn == "A" and nt2_sn == "U":
                    hbond_lns += self.__getAUHbond(bp[0], bp[1])
                elif nt1_sn == "U" and nt2_sn == "A":
                    hbond_lns += self.__getAUHbond(bp[1], bp[0])
                elif nt1_sn == "G" and nt2_sn == "U":
                    hbond_lns += self.__getGUHbond(bp[0], bp[1])
                elif nt1_sn == "U" and nt2_sn == "G":
                    hbond_lns += self.__getGUHbond(bp[1], bp[0])
                else:
                    print "#  no  %s - %s basepair hbond information " % (nt1_sn, nt2_sn)
        if vb: self.prnLns(hbond_lns, no_strip=1)
        return hbond_lns 
    
    
    def __getGCHbond(self, nt1, nt2):
        """  
        # bairpair format in program : [('G', 7), ('C', 24)]
        ! Gua 1 - Cyt 33
        assign (resid 1 and name n1) (resid 33 and name n3) 2.87 .2 .2
        assign (resid 1 and name h1) (resid 33 and name n3) 1.86 0.2 0.2
        assign (resid 1 and name o6) (resid 33 and name n4) 2.81 .2 .2
        assign (resid 1 and name n2) (resid 33 and name o2) 2.81 .2 .2
        assign (resid 1 and name n2) (resid 33 and name n3) 3.58 0.2 0.2
        assign (resid 1 and name o6) (resid 33 and name n3) 3.63 0.2 0.2
        """
        hbond_lns = [" ", "!", "! %s %3d - %s %3d " % (self.nt_names[nt1[0]], nt1[1], self.nt_names[nt2[0]], nt2[1])] 
        hbond_lns += ["assign (resid %2d and name H1 ) (resid %2d and name N3 )  2.00 0.2 0.2" % (nt1[1], nt2[1])]
        hbond_lns += ["assign (resid %2d and name N1 ) (resid %2d and name N3 )  2.90 0.3 0.3" % (nt1[1], nt2[1])]
        hbond_lns += ["assign (resid %2d and name O6 ) (resid %2d and name H41)  2.00 0.2 0.2" % (nt1[1], nt2[1])]
        hbond_lns += ["assign (resid %2d and name O6 ) (resid %2d and name N4 )  2.90 0.3 0.3" % (nt1[1], nt2[1])]
        hbond_lns += ["assign (resid %2d and name H22) (resid %2d and name O2 )  2.00 0.2 0.2" % (nt1[1], nt2[1])]
        hbond_lns += ["assign (resid %2d and name N2 ) (resid %2d and name O2 )  2.90 0.3 0.3" % (nt1[1], nt2[1])]
        return hbond_lns 
    
    def __getGUHbond(self, nt1, nt2):
        """
         ! Gua 2 - Uri 32   *standard anti-anti G-U wobble pair
        assign (resid 2 and name O6) (resid 32 and name H3)  1.9 0.1 0.30
        assign (resid 2 and name O6) (resid 32 and name N3)  3.0 0.28 0.08
        assign (resid 2 and name H1) (resid 32 and name O2)  1.9 0.1 0.30
        assign (resid 2 and name N1) (resid 32 and name O2)  3.0 0.28 0.08
        """
        hbond_lns = [" ", "!","! %s %3d - %s %3d *standard anti-anti G-U wobble pair " % (self.nt_names[nt1[0]], nt1[1], self.nt_names[nt2[0]], nt2[1])] 
        hbond_lns += ["assign (resid %2d and name O6 ) (resid %2d and name H3 )  2.00 0.2 0.2 " % (nt1[1], nt2[1])]
        hbond_lns += ["assign (resid %2d and name O6 ) (resid %2d and name N3 )  2.90 0.3 0.3 " % (nt1[1], nt2[1])]
        hbond_lns += ["assign (resid %2d and name H1 ) (resid %2d and name O2 )  2.00 0.2 0.2 " % (nt1[1], nt2[1])]
        hbond_lns += ["assign (resid %2d and name N1 ) (resid %2d and name O2 )  2.90 0.3 0.3 " % (nt1[1], nt2[1])]
        return hbond_lns 
    
     
    def __getAUHbond(self, nt1, nt2):
        """
        ! Ade2 - Uri 32   
        assign (resid 2 and name N6) (resid 32 and name O4)  3.0 0.28 0.08 # 2.95
        assign (resid 2 and name N1) (resid 32 and name N3)  2.9 0.28 0.08 # 2.82
        """
        hbond_lns = [" ", "!","! %s %3d - %s %3d " % (self.nt_names[nt1[0]], nt1[1], self.nt_names[nt2[0]], nt2[1])] 
        hbond_lns += ["assign (resid %2d and name H61) (resid %2d and name O4 )  2.00 0.2 0.2" % (nt1[1], nt2[1])]
        hbond_lns += ["assign (resid %2d and name N6 ) (resid %2d and name O4 )  2.90 0.3 0.3" % (nt1[1], nt2[1])]
        hbond_lns += ["assign (resid %2d and name N1 ) (resid %2d and name H3 )  2.00 0.2 0.2" % (nt1[1], nt2[1])]
        hbond_lns += ["assign (resid %2d and name N1 ) (resid %2d and name N3 )  2.90 0.3 0.3" % (nt1[1], nt2[1])]
        return hbond_lns 
    
        
    def genNoe(self, _hbp_ntids, vb=0):
        """
        generate noe distance restraints in duplex regions:
                 iH' -  i   H8/H6
                 iH' -  i+1 H8/H6
                 i H8/H6 - i+1  H8/H6
        """          
        noe_lns = []
        _hbp_ntids.sort()
        for i in range(len(_hbp_ntids)-1):
            nt1_id = _hbp_ntids[i]
            nt2_id = _hbp_ntids[i+1]  # next nucleotide id
            nt1_sn = self.residues[nt1_id].nt_name
            nt2_sn = self.residues[nt2_id].nt_name
            noe_lns += ["!----------------------------------------------------- "]
            noe_lns += ["! %s %d" % (self.nt_names[nt1_sn], nt1_id)]
            noe_lns += ["assign (resid %2d and name H1') (resid %2d and name H%d ) 3.2 1.2 1.3 " % (nt1_id, nt1_id, self.nt_hids[nt1_sn])]
            if nt1_id+1 == nt2_id:
                noe_lns += ["assign (resid %2d and name H1') (resid %2d and name H%d ) 3.2 1.2 1.3"  % (nt1_id, nt2_id, self.nt_hids[nt2_sn])]
                noe_lns += ["assign (resid %2d and name H%d ) (resid %2d and name H%d ) 4.1 1.9 1.9 " % (nt1_id, self.nt_hids[nt1_sn], nt2_id, self.nt_hids[nt2_sn])] 
            noe_lns += ["!----------------------------------------------------- "]
        nt1_id = _hbp_ntids[-1]       
        noe_lns += ["!----------------------------------------------------- "]
        noe_lns += ["! %s %d" % (self.nt_names[nt1_sn], nt1_id)]
        noe_lns += ["assign (resid %2d and name H1') (resid %2d and name H%d ) 3.2 1.2 1.3 " % (nt1_id, nt1_id, self.nt_hids[nt1_sn])]
        noe_lns += ["!----------------------------------------------------- "]
        if vb: self.prnLns(noe_lns, no_strip=1)
        return noe_lns 

    def genBaseStack_NOE(self, resid1, resid2):
        """ generate consecutive nucleotide NOE : 
        assign (resi 6  and name H8) (resi 5 and name H2') 5.000 3.200 1.000 
        """
        lns = []
        base_stack = {}
        stack_key = "%s%s" % (self.residues[resid1].nt_name, self.residues[resid2].nt_name)
        sps = SS_SEQ_CS[stack_key]
        for an1, an2, _dist in sps: 
            noe = self.getStackNoe(_dist, inter=1)
            lns += ["assign (resi %s  and name %s) (resi %s  and name %s) %s" %  ( resid1,  an1,  resid2, an2, noe) ]
        return lns             
    
    def getStackNoe(self, _dist, inter=0):
        if _dist < 3.6: return   "3.500 1.500 1.500"
        elif _dist < 5.1: return "5.000 1.500 1.500"
        elif _dist < 6.6: return "6.500 1.500 1.500"
        elif _dist < 7.6: return "7.500 1.500 1.500"
        else: return None 
        
    def genAllBaseStack(self, resids, vb=0):
        """ generate  base stacking for single strand regions or loop regions """
        self.all_base_stack = {}
        lns = []
        for i in range(len(resids)-1):
            resid1 = resids[i]
            resid2 = resids[i+1]
            if resid2 - resid1 != 1: continue 
            lns += self.genBaseStack_NOE(resid1, resid2)
        if vb: self.prnLns(lns) 
        #self.prnDict(self.all_base_stack)
        return lns 
    
    
    def prnXplorSeq(self, num_nt_row=10):
        seq_lns = ["       sequence   "]
        ln = "        "
        cnt = 0
        for resid in self.resids:
            ln += "%s " % self.nt_Cnames[self.residues[resid].nt_name]
            cnt += 1
            if cnt == num_nt_row:
                seq_lns.append(ln)
                ln = "        "
                cnt = 0
        if ln.strip(): seq_lns.append(ln)
        seq_lns += ["       end "]
        self.prnLns(seq_lns, no_strip=1)
        return seq_lns
    
    
    def getLinkerStackTBL(self, vb=0):
        """ return linker base stacking TBL without hairpin loop"""
        segs = self.sec_inf['linker']
        resids = []
        for _seg in segs:
            resids += _seg[-1]
            #if _seg[0] -1 >= self.resids[0]:
            #    resids += [_seg[0]-1] +  _seg
            #if _seg[-1] + 1 <= self.resids[-1]:
            #    resids.append(_seg[-1] + 1)
        #print resids
        stack_tbl = self.genAllBaseStack(resids, vb=vb)
        return stack_tbl
        
    def getDupHbondTBL(self, vb=0):
        #hbp_tbl = self.genHbond(self.hbps)
        return self.genHbond(self.duplexs, vb=vb)
        #pbp_tbl = self.genHbond(self.pbps)
        #hbond_tbl = hbp_tbl + pbp_tbl 
    
    
    def getDupNoeTBL(self, vb=0):
        """ return duplex NOE table """
        noe_tbl = self.genNoe(self.helices_resids, vb=vb)
        return noe_tbl      
    
    def getDupAngleTBL(self, vb=0):
        """ return duplex angle restraints """ 
        ang_tbl = []
        ang_tbl += self.sugarPucker()
        ang_tbl += self.genXhi()
        eps_tbl, zeta_tbl = self.genEpsZeta()
        ang_tbl += eps_tbl
        ang_tbl += zeta_tbl
        
        alpha_tbl, beta_tbl = self.genAlphaBeta()
        ang_tbl += alpha_tbl
        ang_tbl += beta_tbl
        ang_tbl += self.genGamma()
        if vb: self.prnLines(ang_tbl, chkrow=0)
        return ang_tbl 
    
        
    
