# Build RNA second strucure database from all X-ray structures 

import os
import numpy
from librna import nt_atoms_woh, RNA2D  
from libio  import INPUT 
from libaction import HomoMove,Mutation,Sequence
array = numpy.array 
from lib.libdisclaimer import DISCLAIMER

class BlockBase(HomoMove, INPUT, RNA2D, Mutation,Sequence):
    """ public functions for block subclasses """
    ref_atom_ids = {"A":(0, 5, 7, 8), "G":(0, 5, 7, 8),"U":(0, 5, 7, 8), "C":(0, 5, 7, 8)}
    def chkRefAtomIds(self):
        """ chekc reference frame atom ids """
        for nt_name in ['A','G','U','C']:
            print "# nt: %s " % nt_name
            res_ref_aids = self.ref_atom_ids[nt_name]
            res_atoms = nt_atoms_woh[nt_name]
            for aid in res_ref_aids:
                print "# %s -->  %s " % (aid, res_atoms[aid])
            

class BlockDuplex(BlockBase, DISCLAIMER):
    def __init__(self, fn):
        #self.prnDis()
        if not hasattr(self, 'g2g_wdir'):
            self.setG2G_Wdir()
        self.bp_bk_db = self.loadObj(os.path.join(self.g2g_wdir, 'db/bp_block.db'))
        self.bk_keys = self.bp_bk_db.keys()
        self.setRNA2ndStrFn(fn)
        self.setFnPre_by_inp()
        
        
    def getBpRefXYZ(self, bp_xyz):
        """ return basepair reference frame """
        ref_xyz = []
        ref_xyz += bp_xyz[0].tolist() #self.getResRefXYZ(bp_xyz[0])
        ref_xyz += bp_xyz[1].tolist() #self.getResRefXYZ(bp_xyz[1])
        return array(ref_xyz)
    

    def duplex2blocks(self, duplex, vb=0):
        """ break duplex into the continuous two basepair  building blocks """
        blocks = []
        for i in range(len(duplex)-1):
            bp1 = duplex[i]
            bp2 = duplex[i+1]
            bkey = "%s%s" % (bp1[2], bp2[2])
            if vb: print "#  %s - %s  --> %s " % (bp1[2], bp2[2], bkey)
            blocks.append((bkey, (bp1[0], bp1[1]), (bp2[0], bp2[1])))
        return blocks
    
    
    def addBP2Duplex3D(self, bp, bp_xyz):
        """ add basepair to duplex """
        for i in [0,1]:
            self.new_duplex[bp[i]] = bp_xyz[i].copy()
    
    def getCloseBP(self, bkey):
        num_nt = len(bkey)
        tmp = [ bkey[i] for i in range(num_nt) ]
        # one NT mutation 
        for i in range(num_nt):
            tmp[i]  = self.comp_seq[bkey[i]]
            tpt_bk = "".join(tmp)
            if tpt_bk in self.bk_keys:
                return (tpt_bk, [i]) 
            else:
                tmp[i] = bkey[i]
        for i in range(num_nt):
            for j in range(num_nt):
                if i== j: continue 
                tmp[i] = self.comp_seq[bkey[i]]
                tmp[j] = self.comp_seq[bkey[j]]
                tpt_bk = "".join(tmp)
                print "# tpt_bk: ", tpt_bk 
                print "i=", i, 'J', j 
                if tpt_bk in self.bk_keys:
                    return (tpt_bk, [i,j]) 
                else:
                    tmp[i] = bkey[i]
                    tmp[j] = bkey[j]
        print "# lookfor : %s " % bkey 
        self.err_exit("no match")
        
    
    def getBestMatchBk(self, bkey):
        scores = []
        for _bk in self.bk_keys:
            scores.append(self.getMatchSeqScore(bkey, _bk))
        sort_ids = self.bestBySeqScore(scores)
        #print bkey, "-->", self.bk_keys[sort_ids[0]]
        return self.bk_keys[sort_ids[0]]
    
    
    def getBlock3D(self, bkey):
        """ get building block by build block key in database"""
        if bkey in self.bk_keys:
            return self.bp_bk_db[bkey]
        else: # mutation
            #mbk, pos = self.getCloseBP(bkey)
            #print mbk, pos 
            #tpt_bk = self.bp_bk_db[mbk]
            #print tpt_bk
            #return self.mutateBpBk(bkey, mbk, pos)
            mkey = self.getBestMatchBk(bkey)
            pos = []
            for i in range(len(mkey)):
                if mkey[i] != bkey[i]:
                    pos.append(i)
            self.add2log([ "# mkey --> bkey: %s --> %s  " % (mkey, bkey)])
            return self.mutateBpBk(bkey, mkey, pos)        
            #print "# mut pos: ", pos 
            #self.err_exit("mutaion")
    
            
    def mutateBpBk(self, bkey,  mbk, pos):
        """ mutate bp blocks to generate new 2bp blocks
        block3d: [(xyz1_bp1, xyz2_bp1), (xyz1_bp2, xyz2_bp2)]
        """
        tpt_bk = self.bp_bk_db[mbk]
        new_bk = [[None,None],[None,None]]
        for i in pos:
            if i < 2:
                bp_id = 0 
                res_id = i 
            else:
                bp_id = 1
                res_id = i - 2 
            new_res = self.mutate(mbk[i], bkey[i], tpt_bk[bp_id][res_id])    
            new_bk[bp_id][res_id] = new_res 
        for i in range(2):
            for j in range(2):
                if new_bk[i][j] is None:
                    new_bk[i][j] = tpt_bk[i][j]
        return new_bk 
            
            
    def genDuplex3D(self, blocks):
        """ generate A-form like single strand:
        block: ('41', (1, 20), (2, 19))
        block3d: [(xyz1_bp1, xyz2_bp1), (xyz1_bp2, xyz2_bp2)]
        """
        self.new_duplex = {}
        if type(blocks) != type([]): blocks = [blocks]
        first_block3D = self.getBlock3D(blocks[0][0])
        self.addBP2Duplex3D(blocks[0][1], first_block3D[0])
        self.addBP2Duplex3D(blocks[0][2], first_block3D[1])
        anchor_ref = self.getBpRefXYZ(first_block3D[1])
        for _block in blocks[1:]:
            b3d = self.getBlock3D(_block[0])
            tpt_ref = self.getBpRefXYZ(b3d[0])
            HM = self.getHM(anchor_ref, tpt_ref, vb=0)
            tgt_bp = (self.homoMove(b3d[1][0],HM), self.homoMove(b3d[1][1], HM))
            self.addBP2Duplex3D(_block[2], tgt_bp)
            anchor_ref = self.getBpRefXYZ(tgt_bp)
            
    
    def resDict2pdb(self, pdb_xyz, seq_dt, norm=0):
        """ save residue xyz array into pdb file """
        fmt_pdb = '%-6s%5d %-4s %-4s%1s%4d%1s%3s%8.3f%8.3f%8.3f%18s%-4s'
        resids = pdb_xyz.keys()
        resids.sort()
        lns = []
        atom_id = 0
        anchor_P = pdb_xyz[resids[0]][0]
        for resid in resids:
            nt_name = seq_dt[resid] 
            nt_atom_names = nt_atoms_woh[nt_name]
            res_xyz = pdb_xyz[resid] 
            for i in range(len(nt_atom_names)):
                atom_name = nt_atom_names[i]
                if len(atom_name) <= 3: atom_name= " %s" % atom_name
                if norm:
                    atom_xyz = res_xyz[i,:] - anchor_P 
                else:
                    atom_xyz = res_xyz[i,:]
                atom_id += 1
                ln =  fmt_pdb  % ('ATOM', atom_id, atom_name, nt_name," ",  resid, "","", atom_xyz[0],atom_xyz[1], atom_xyz[2], "", "") #_atom.x, _atom.y, _atom.z, _atom.misc, _atom.segid)
                lns.append(ln)
        return lns 
    
    
    def getDuplexSeq(self, duplex, vb=0):
        seq_dt = {}
        for _bp in duplex:
            seq_dt[_bp[0]] = _bp[2][0]
            seq_dt[_bp[1]] = _bp[2][1]
        if vb: print "# duplex seq_dt: ", seq_dt
        return seq_dt 

    
    def buildDuplexes(self):
        duplex_ids = self.duplexs.keys()
        duplex_ids.sort()
        for dup_id in duplex_ids:
            duplex = self.duplexs[ dup_id]
            dup_seq_dt = self.getDuplexSeq(duplex)
            blocks = self.duplex2blocks(duplex, vb=0)
            self.genDuplex3D(blocks)
            lns = self.resDict2pdb(self.new_duplex, dup_seq_dt, norm=1)
            fp_pdb = "%s_duplex%s.pdb" % (self.fn_pre, dup_id)
            if os.path.exists(fp_pdb):
                fp_pdb =  "%s_ndup%s.pdb" % (self.fn_pre, dup_id)
            self.saveLines(lns, fp_pdb , vb=1)
            
        
class BlockOther(BlockBase):
    """ functions for generating tloops and linkers """
    def __init__(self, fn):
        self.setRNA2ndStrFn(fn)
        self.setFnPre_by_inp() 
        self.loadNonDUBDB()
        
    def loadNonDUBDB(self, **rest):
        """ load non-duplex motifs' database """
        if not hasattr(self, 'g2g_wdir'):
            self.setG2G_Wdir()
        self.tloop_db = self.loadObj(os.path.join(self.g2g_wdir, 'db/db_tloop.db'), **rest)
        self.ss_db = self.loadObj(os.path.join(self.g2g_wdir, 'db/db_ss.db'), **rest)
        
     
    def matchOBKSeq(self, obk_tpt, obk_inf, vb=0):
        """ obk_tpt: ( 'GA', [46, 50], 'UAC', [47, 48, 49], xyz(dict) )
            obk_inf:  ['tloop', 7, 'UUGAUAUGG', [31, 39, 32, 33, 34, 35, 36, 37, 38]]

        """
        if vb: print "# obk_inf: ", obk_inf 
        tpt_seq = obk_tpt[0] + obk_tpt[2]
        tpt_xyz = obk_tpt[-1]
        tpt_resids = obk_tpt[1] + obk_tpt[3]
        tgt_seq = obk_inf[2]  # target sequence 
        tgt_resids = obk_inf[-1]
        assert len(tpt_resids) == len(tgt_resids)
        tgt_xyz = {}
        for i in range(len(tpt_resids)):
            tpt_resid = tpt_resids[i]
            tpt_nt = tpt_seq[i]
            tgt_nt = tgt_seq[i]
            tgt_resid = tgt_resids[i]
            tgt_xyz[tgt_resid] = self.mutate(tpt_nt, tgt_nt, tpt_xyz[tpt_resid])
            if vb: print "# %s %s  -- >  %s %s " % (tpt_resid, tpt_nt, tgt_resid, tgt_nt)
        return tgt_resids[:2], tgt_xyz
        
    
    def bestMatchOBK(self, obk_db, obk_seq, lnk=1, stack=0, vb=0): # pass
        """ best sequence match can produce geometric clashes in structure.
        Therefore, use loops with C1'--N1/N9 vector pointing to duplex axle """
        match_score = []
        cnt = 0
        best_score = 10 * len(obk_seq)
        for _obk in obk_db:
            _obk_in_db_seq = _obk[0] + _obk[2]
            
            tscore = self.getMatchSeqScore(_obk_in_db_seq, obk_seq)
            #if tscore == best_score: 
            #    return _obk
            match_score.append(tscore)
            cnt += 1
        if not lnk: # tloop
            ms = max(match_score)
            if len(_obk[2]) == 7 and cnt > 20:
                best_id = 19
            else:
                best_id =  match_score.index(ms)
        elif stack:
            best_id = self.getStackedTpt(obk_db, match_score)
        else:
            ms = max(match_score)
            best_id =  match_score.index(ms)
        if vb: 
            print "# match_score: ", match_score
            print "# best tpt: ", obk_db[best_id]
        return obk_db[best_id]
    
        
    def getStackedTpt(self, obk_db, scores ):
        """  ('GU', [58, 61], 'AU', [59, 60], xyz) """
        sids = numpy.argsort(numpy.array(scores))
        sids = sids.tolist()
        sids.reverse()
        for i in range(len(sids)):
            obk_id = sids[i]
            obk_tpt = obk_db[obk_id]
            obk_resids = obk_tpt[1] + obk_tpt[3]
            obk_resids.sort()
            obk_seq = {}
            for i in range(2):
                obk_seq[obk_tpt[1][i]] = obk_tpt[0][i]
            for i in range(len(obk_tpt[3])):
                obk_seq[obk_tpt[3][i]] = obk_tpt[2][i]
            if self.isStackTpt(obk_resids, obk_seq, obk_tpt[4]):
                return obk_id
        #print "stack: best score: ",  scores[sids[0]]
        return sids[0]
    
    def isStackTpt(self, tpt_resids, tpt_seq, tpt_xyz):
        for i in range(len(tpt_resids)-1):
            resid = tpt_resids[i]
            res_xyz = tpt_xyz[resid]
            nt = tpt_seq[resid]
            tpt_ref, anchor = self.getResBaseRefXYZ_UV(nt, res_xyz)
            next_resid = tpt_resids[i+1]
            next_nt = tpt_seq[next_resid]
            next_res_xyz = tpt_xyz[next_resid]
            next_tpt_ref, next_anchor = self.getResBaseRefXYZ_UV(next_nt, next_res_xyz)
            _ang = self.getVecAngle(tpt_ref[2,:], next_tpt_ref[2,:], deg=1)
            _dist = self.getDist(anchor, next_anchor)
            if _ang > 30.0 or  _dist > 5.5:
                return False
            #print "# resid: %3d-%3d  ang= %7.1f _dist= %7.1f " % (resid, next_resid, _ang, _dist )
        return True                          
        
    def genTloop(self): # pass
        self.tpt_tloops = []
        for _tloop in self.sec_inf['tloop']:
            obk_tpt = self.obk2tpt(_tloop)
            self.tpt_tloops.append(obk_tpt)
            
    
    def obk2tpt(self, _obk):
        """   obk_inf:  ['tloop', 7, 'UUGAUAUGG', [31, 39, 32, 33, 34, 35, 36, 37, 38]]"""
        obk_inf = self.getOBInf(_obk, vb=0)
        obk_type = obk_inf[0]
        stack = 0
        lnk = 0
        if obk_type == 'tloop':
            obk_db = self.tloop_db[obk_inf[1]]
        elif obk_type == 'linker':
            lnk = 1
            obk_db = self.ss_db[obk_inf[1]]
            if hasattr(self, 'stack'):
                stack = self.stack
            else:
                stack = 1
        else:
            self.err_exit("%s database is not included in this version " % _obk_type)
        obk_tpt = self.bestMatchOBK(obk_db, obk_inf[2], lnk=lnk, stack=stack, vb=0)
        obk_tpt = self.matchOBKSeq(obk_tpt, obk_inf, vb=0)
        return obk_tpt 
    
        
    def genLinker(self, vb=0):
        self.tpt_linkers = []
        linkers = self.sec_inf['linker']
        num_linker = len(linkers)
        ln = "# num_linker_in_pdb: %3d " % num_linker 
        lns = [ln]    
        #print "# linkers: ", linkers 
        for i in range(num_linker):
            _linker = linkers[i] 
            lns += [  "# gen linker %s: %s " % (i, _linker) ]
            obk_tpt  = self.obk2tpt(_linker)
            self.tpt_linkers.append(obk_tpt)
        self.add2log(lns, vb=vb)    
            
    def saveLinkers2PDB(self, vb=0):
        for i in range(len(self.tpt_linkers)):
            lnk_resids, lnk_xyz = self.tpt_linkers[i]
            lns = self.resXYZ2pdb(pdb_xyz=lnk_xyz)
            self.saveLines(lns, "%s_linker%s.pdb" % (self.fn_pre, i), vb=vb)
    
            
    def saveTloops2PDB(self, vb=0):
        for i in range(len(self.tpt_tloops)):
            lnk_resids, lnk_xyz = self.tpt_tloops[i]
            lns = self.resXYZ2pdb(pdb_xyz=lnk_xyz)
            self.saveLines(lns, "%s_tloop%s.pdb" % (self.fn_pre, i), vb=vb)
            
    def getOBInf(self, other_bk, vb=0): # pass
        """ # resids =  """
        num_obk = len(other_bk)
        if num_obk == 3: # linker block
            obk_type = 'linker'
            resids = other_bk[-1]
            obk_len = len(resids)
            obk_resids = [resids[0]-1, resids[-1]+1 ] + resids
            #tbp = (resids[0]-1, resids[-1]+1)
            #tbp_seq = "%s%s" % (self.seq_dict[tbp[0]], self.seq_dict[tbp[1]])
            #obk_bp = (tbp[0], tbp[1], tbp_seq)
        elif num_obk == 2: # terminal loop 
            obk_type = 'tloop'
            resids = other_bk[-1]
            obk_len = len(resids)
            obk_resids = [other_bk[0][0], other_bk[0][1]] + resids  
        else:
            self.err_exit("abnormal other block ")
        obk_seq = self.resids2seq(obk_resids)
        obk_inf = [obk_type, obk_len, obk_seq, obk_resids]
        if vb: print "# obk_inf: ", obk_inf
        return obk_inf
    
        
        
if (__name__ == "__main__"):
    import sys 
    if len(sys.argv) >= 2:
        r_obj = BlockDuplex(sys.argv[1])
        r_obj.start()
    else:
        print """# Usage:   python   %s   input_file
        
        input_file:   secondary structure file (sequence and duplexes)
        """  % __file__
        
    
