#!/usr/bin/python
import os, sys
from string import lstrip,strip, rstrip, rsplit
from re import split

class MIO:
    log = []

    def prnLines(self, lines, index=None, **rest  ):
        #print "# prnlines: ", rest
        if not rest.get('prn', 1): return
        num = len(lines)
        chkrow = rest.get("chkrow", 1)
        line_thresh = rest.get("thresh", 400)
        _omit = False
        if num > line_thresh  and chkrow:
            num = 100
            _omit = True
        if index:
            for i in range(num): print '# %d %s ' % ( i, rstrip(lines[i]))
        else:
            for i in range(num): print rstrip(lines[i])
        if _omit: print "# lines no >%s, full txt, use < chkrow=0 >" % line_thresh


    def prnLns(self, lns, **rest):
        self.prnLines(lns, **rest)


    def readTxtFile(self, fp, vb=0):
        fh = open(fp,'r')
        txt = fh.readlines()
        fh.close()
        if vb:
            print "# reading fn: %s " % fp
            print "just to check"
            self.prnLines(txt)
        return txt

    def err_exit(self, err_msg):
        print "\n# **********   Error  Message  ************\n#"
        print err_msg
        print "#"
        print "# **********   Error  Exit  ************\n"
        sys.exit(0)

    def prnDict(self, _dict):
        dkeys = _dict.keys()
        for _key in dkeys:
            print "# %s - > %s " % (_key, _dict[_key])

    def add2log(self, lns, vb=0, prn=0):
        """ lns must be a list """
        if vb or prn: self.prnLns(lns)
        self.log += lns

class INPUT(MIO):
    """ read the input RNA secondary structure """

    def load2str(self, fn_inp=None, vb=0):
        """ read the input RNA secondary structure """
        if fn_inp is not None:
            self.setRNA2ndStrFn(fn_inp)
            self.setFnPre_by_inp()
        lns = self.readTxtFile(self.fn_rna_2str_inp)
        seg_lns = self.getSegLns(lns, vb=vb)
        self.parseRNA2ndStr(seg_lns, vb=0)
        #self.setOutPDBPath(vb=1)

    def writeG2GInput(self):
        """write out the input format for duplex generation and allatom 
         conversion by G2G"""
        ##infile = open(self.project_name+".inp", 'w')
        infile = open("g2g_input.inp", 'w')
        print >> infile, "### Do not change the first comment line of each block text, "
        print >> infile, "### otherwise, the program could not parse the input file, "
        print >> infile, "### Each block should be after one blank line \n"
        print >> infile, "# project name \nproject = %s\n" % self.project_name
        print >> infile, "# num_nt:%d \nstart_resid= 1 \nend_resid= %d\n" % (self.resnum, self.resnum)
        nt_seq = ""
        for nt in self.seq:
            if nt == " ": continue 
            else:
                nt_seq += nt
        print >> infile, "# Seqence \n%s\n\n" % nt_seq
        #print(self.duplex_dict)
        dup_count = 1
        for duplex in self.duplex_dict:
            if len(self.duplex_dict[duplex]) > 1:
              #print >> infile, "# helix: %d" %duplex
              print >> infile, "# helix: %d" %dup_count
              dup_count += 1
              for bp in self.duplex_dict[duplex]:
                  print >> infile, "%d  %s%s  %d" % (bp[0],self.seq_dict[bp[0]],self.seq_dict[bp[1]], bp[1]) 
              print >> infile , ""

    def setOutPDBPath(self, **rest):
        cur_dir = os.path.dirname(self.fn_pre)
        pdb_dir = os.path.join(cur_dir, 'pdb')
        if not os.path.exists(pdb_dir):
            os.makedirs(pdb_dir)
        self.out_pdb_pre = os.path.join(pdb_dir, os.path.basename(self.fn_pre))
        lns = [ "# output pdb directory: %s " % pdb_dir ]
        self.add2log(lns, **rest)

    def setRNA2ndStrFn(self, fn_2str):
        """ set secondary structure input filename """
        self.fn_rna_2str_inp = fn_2str
        self.fn_dir = os.path.dirname(self.fn_rna_2str_inp)

    def setFnPre_by_inp(self):
        self.fn_pre = self.fn_rna_2str_inp[:-4]

    def getSegLns(self, lns, vb=0):
        """ parse input 2nd structre file into section lines
        section starts with a blank line with a "#" in the first of the next line 
        """
        #self.prnLines(lns)
        if vb: self.prnLines(["### input lines for segment: " ] + lns )
        num_lns = 0
        all_seg_lns = []
        seg_lns = []
        pre_ln_blank = 1
        for ln in lns:
            ln = ln.strip()
            if ln:
                if (ln[0] == ">" or ln[0] == "$") and pre_ln_blank :
                    if len(seg_lns) > 0:
                        all_seg_lns.append(seg_lns)
                        seg_lns = [ln]
                    else:
                        seg_lns.append(ln)
                else:
                    seg_lns.append(ln)
                pre_ln_blank = 0
            else:
                pre_ln_blank = 1

        if len(seg_lns) > 0: all_seg_lns.append(seg_lns)
        if vb:
            for _segLns in all_seg_lns:
                self.prnLines(["######  seg #####"] + _segLns)
        #print all_seg_lns
        return all_seg_lns

    def parseRNA2ndStr(self, all_seg_lns, vb=0):
        """ parse RNA 2nd structure """
        self.helices_resids = []
        self.duplexs = {}
        for seg_lns in all_seg_lns:
            if seg_lns:
                first_ln = seg_lns[0]

                if  '> ' in first_ln:
                    self.parseRNASeq(seg_lns, vb=vb)
                elif 'bplist' in first_ln:
                    self.parseDuplex(seg_lns, vb=vb)
                else:
                    self.prnLines(["#******       None Formated text in the input file ****"]+ seg_lns)
            else:
                self.err_exit("# input block text format error, please check input file format ")

    def parseRNASeq(self, seq_lns, vb=0):
        """ set RNA seqence from file, sequence is the first section in all the input section lines """
        #seq_lns = seg_lns[0]
        if vb: print "# input seqence text: ", seq_lns
        if '> ' not in seq_lns[0]:
            raise """ Seqence not in the first section of input 2nd structure or the input file format is not compatible """
        title = split("\s+",seq_lns[0])
        self.project_name= title[1][:-4]
        log_lns = ["# The project name is %s " % self.project_name ]
        self.add2log(log_lns, vb=0)
        seq = ""
        for ln in seq_lns[1:]:
            seq += ln
        self.seq = seq
        seq_dict = {}
        resid = 1 # first resid
        for nt in seq:
            if nt == " ": continue
            elif nt in ['A',"G","C","U"]:
                seq_dict[resid] = nt
                resid += 1
            else:
                raise "# residue name %s  Not in ACGU, nonregular nucleotide " % nt
        self.seq_dict = seq_dict
        self.resids = self.seq_dict.keys()
        self.resids.sort()
        self.resnum = len(self.resids)
        log_lns = [ "# Num of nucletides in sequence= %3d "  % self.resnum  ]
        self.add2log(log_lns, vb=vb)
        if vb: self.prnDict(self.seq_dict)

    def parseDuplex(self, dup_lns, vb=0):
        if 'bplist' not in dup_lns[0]:
            #print dup_lns[0]
            self.err_exit( "# could not find duplex in the input file")
        duplex = []
        duplex_resids = []
        lns = []
        for ln in dup_lns[1:]:
            ln = ln.strip()
            if len(ln) < 4 or ln.startswith("##"): continue
            lns.append(ln)
            bp_left_resid, bp_right_resid = split("\s+",ln)
            bp_left_resid = int(bp_left_resid)
            bp_right_resid = int(bp_right_resid)
            duplex.append((bp_left_resid, bp_right_resid))
        duplex_dict = {}
        duplex_part = []
        duplex_id = 1
        for i in range(len(duplex)):
            if (i < len(duplex)-1) and abs(duplex[i][0] - duplex[i+1][0]) < 3 and abs(duplex[i][1] - duplex[i+1][1]) < 3 :
                duplex_part.append((duplex[i][0], duplex[i][1]))
            else:
                duplex_part.append((duplex[i][0], duplex[i][1]))
                branch_duplex = duplex_part[:]
                duplex_dict[duplex_id] = branch_duplex
                del duplex_part[:]
                duplex_id += 1
        self.duplex_dict = duplex_dict
        #print self.duplex_dict


###--Execution Part--##########
if (__name__ == "__main__"):
    import sys
    if len(sys.argv) != 2:
        Usage = """Usage:
                converts the *.bpl file (input to RNA2D3D) to the input for G2G program
                two sections (sequence & basepairs) in the *.bpl file should be separated by whitespace

                $ python %s input.bpl""" % __file__
        print Usage
    else:
        fp_inp = sys.argv[1]
        r_obj1 = INPUT() 
        r_obj1.load2str(fp_inp, vb=0) 
        r_obj1.writeG2GInput()     

  
