import os, sys  
from string import lstrip,strip, rstrip, rsplit
from re import split

import numpy as Npy
from math import  log, sin, cos, sqrt, pi, fabs, radians,degrees,ceil

from copy import copy 
from cPickle import dump as _pdump
from cPickle import load as _pload 
from time import strftime, localtime
from  libdisclaimer import version 

from libntatoms import nt_atoms_woh
_array = Npy.array
pi2 = pi * 2.0



def prnUT(t1, t2, cmt=""):
    used_time = t2 - t1 
    print "# %s used time: %8.4f " %  (cmt, used_time)
    
def dumpObj(obj, fp, **rest):
    """ dump obj into file fn on disk"""
    fmode = rest.get("fmode", 'wb')
    if 'dirpath' in rest: 
        fp = os.path.join(rest['dirpath'],fp)
    elif 'dpath' in rest:
        fp = os.path.join(rest['dpath'],fp)
    cmt = rest.get('cmt', "")
    msg = '# dump < %s > to file: %s' % (cmt, fp )
    fh = open(fp,fmode)                
    _pdump(obj,fh,1)
    fh.close()
    vb = rest.get('vb',0)
    if vb:  print msg
    return

def loadObj(fp, **rest):
    #print "# cwd ", os.getcwd()
    if 'dpath' in rest: 
        fp = os.path.join(rest['dpath'],fp)
    elif 'dirpath' in rest:
        fp = os.path.join(rest['dirpath'],fp)
    vb=rest.get('vb', 1)
    if vb: print '# load file: %s' % fp 
    fh = open(fp, 'rb')
    _obj =  _pload(fh)
    fh.close()
    return _obj


class MIO:
    log = []
    version = version 
    def setG2G_Wdir(self):
        """ set G2G directory """
        cwd = os.getcwd()
        if "examples" in cwd:
            wdir, tmp  = cwd.split('examples')
            g2g_wdir = os.path.dirname(wdir)    
        else:
            g2g_wdir = cwd
        print "# g2g_wdir: %s " % g2g_wdir
        self.g2g_wdir = g2g_wdir 
        
            
    def add2log(self, lns, vb=0, prn=0):
        """ lns must be a list """
        if vb or prn: self.prnLns(lns)
        self.log += lns 
    
    def saveLog(self, fn_pre=None, fn_mid=None, vb=0, prn=0):
        if fn_pre is None:
            fn_pre = self.fn_pre 
        fn_ver = "v%s" % self.getVersion()
        if fn_mid is None:
            fn_segs = [fn_pre, fn_ver]
        else:
            fn_segs = [fn_pre, fn_mid, fn_ver]
        fn_log = "%s_%s.log" %  ('_'.join(fn_segs), strftime("%m_%d_%Y_%H%M%S", localtime()))
        if prn: self.prnLines(self.log)
        self.saveLines(self.log, fn_log, vb=vb)
        
     
    def getVersion(self):
        return self.version
    
    
    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 cpTxtFile(self, src_fn, dst_fn, vb=0):
        if vb: print "# copying  %s --> %s " % (src_fn, dst_fn)
        lns = self.readTxtFile(src_fn)
        self.saveLines(lns, dst_fn,vb=0)
        
    
    def saveLines(self, lines, fp,  fmode='w', vb=0, **rest):
        """ save lines list into a disk file """
        ln = "# save to file: %s " % fp 
        self.add2log([ln], vb=vb)
        fh = open(fp, fmode)
        for line in lines:
            line = rstrip(line)
            fh.write(line + '\n')
        fh.close()

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

            
    def txt2col(self, ccs, dat_type="arr"):    
        """ only works for block data """
        all_col_strs = [] 
        for line in ccs:
            line = strip(line)
            if line and line[0] != "#":
                col_strs = split(r'\s+',line)
                all_col_strs.append(col_strs)
        return all_col_strs 
    
                
    def txt2colDat(self, ccs, dat_type="arr"):    
        """ only works for block data, dat_type: list or array(default) """
        bdata = []  # block data (2 dimensions's list
        for line in ccs:
            line = strip(line)
            if line and line[0] != "#":
                col_strs = split(r'\s+',line)
                col_ln = [ float(_item) for _item in col_strs]
                bdata.append(col_ln)
        if dat_type == 'list':
            return bdata
        elif dat_type == 'arr':
            return npy.array(bdata)
        else:
            print "#  output data type: %s " % dat_type 
            raise "# no matched data type(list, arr) "
  

    def getFnPre(self, fn, vb=0):
        """ return return filename without extension"""
        tmp = fn.rsplit(".")
        if vb:
            print "# fn:     %s " % fn 
            print "# fn_pre: %s" % tmp[0]
            print "# fn_ext: %s" % tmp[1]
        return tmp[0]
    
    def prnList(self, var_list, index=0):
        if index:
            for i in range(len(var_list)):
                print "i= %2d %s " % (i, var_list[i])
        else:
            for i in range(len(var_list)):
                print "%s " % (var_list[i])    
    
    def err_exit(self, err_msg):
        print "\n# **********   Error  Message  ************\n#"
        print err_msg 
        print "#"
        print "# **********   Error  Exit  ************\n"
        sys.exit(0)        

    def getDateStr(self):
        """ return date in fomrat of Mon-Day-Year"""
        return strftime("%m_%d_%Y")
    
    
    def getWorkInf(self, vb=0):
        lns = ["# ---------  Job Run informaton  -------------------------------"]
        lns +=["#    date :  %s " % self.getDateStr() ]
        lns +=["# diretory:  %s " % os.getcwd()]
        lns +=["#  command:  python %s " % " ".join(self.argv)]
        lns +=["# -----------------------"]
        if vb: self.prnLines(lns)
        return lns
    
    def loadObj(self, fp, **rest):
        self.add2log(['# load file: %s' % fp ])
        return loadObj(fp, **rest)
        
    def dumpObj(self, obj, fp, **rest):
        dumpObj(obj, fp, **rest)
    

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=vb)
        self.setOutPDBPath(vb=1)
    
    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.startswith('##'): 
                    continue 
                elif ln[0] == "#" and pre_ln_blank and ln[1] != "#":
                    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)
        return all_seg_lns
    
    def parseProjectName(self, seg_lns, vb=0):
        #self.prnLines(seg_lns)
        for ln in seg_lns:
            ln = ln.strip()
            if ln and ln[0] != '#':
                tmp = split('\s+=\s+', ln)
                if tmp[0] == 'project': 
                    self.project_name = tmp[1]
                    self.fn_pre = self.project_name
                    self.add2log(["# project name: %s " % self.project_name])
                
                
                
    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  'num_nt' in first_ln:
                    self.parseResid(seg_lns, vb=vb)
                elif 'project' in first_ln:
                    self.parseProjectName(seg_lns, vb=vb)
                elif  'Seqence' in first_ln:
                    self.parseRNASeq(seg_lns, vb=vb)
                elif 'filename' in first_ln and 'rdc' in first_ln:
                    self.parseFnRDC(seg_lns, vb=vb)
                elif 'Angles' in first_ln:
                    self.parseAng_Dups(seg_lns, vb=vb)
                elif 'Orientations' in first_ln:
                    self.parseDupOrients(seg_lns, vb=vb)
                elif 'helix' in first_ln:
                    self.parseDuplex(seg_lns, vb=vb)
                elif 'control' in first_ln:
                    self.parseControl(seg_lns)
                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 ")
        self.duplex_ids = self.duplexs.keys()
        self.duplex_ids.sort()
        self.parseSingleStrand(vb=vb)
        
    def getInpNameProp(self, ln, vb=0):
        if vb: print "# input line: ", ln 
        if '=' in ln:
            tmp = split("\s*=\s*", ln)
            if vb: print "# return: ", tmp 
            return tmp 
        
                
    def parseResid(self, seg_lns, vb=0):
        """ set RNA  start resid or/and end resid """
        self.rna_end_resid = "NA" # -1
        self.rna_start_resid = 0 
        if vb: print "# vb: input seg_lns for resid: ", seg_lns 
        for ln in seg_lns:
            ln = ln.strip()
            if not ln or ln[0] == "#": continue 
            if 'start' in ln:
                tmp = split("=", ln)
                self.rna_start_resid = int(tmp[1])
            elif 'end' in ln:
                tmp = split("=", ln)
                self.rna_end_resid = int(tmp[1])
        lns = [ "# RNA start= %3d  end= %s" % (self.rna_start_resid, self.rna_end_resid) ]
        if vb: self.prnLines(lns)
        self.add2log(lns)
    
    
    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 'seq' not in seq_lns[0].lower():
            raise """ Seqence not in the first section of input 2nd structure or the input file format is not compatible """
        seq = ""
        for ln in seq_lns[1:]:
            seq += ln 
        seq_dict = {}
        resid = self.rna_start_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 'helix' in dup_lns[0]:
            dup_id = self.getDuplexID(dup_lns[0])
        else:
            self.err_exit( "# could not find duplex ID in the input file: eg.   # helix 1 " )
        duplex = []
        duplex_resids = []
        lns = []
        for ln in dup_lns[1:]:
            ln = ln.strip()
            if len(ln) < 5 or ln.startswith("##"): continue 
            lns.append(ln)
            bp_left_resid, bp_name, bp_right_resid = split("\s+",ln)
            bp_left_resid = int(bp_left_resid)
            bp_right_resid = int(bp_right_resid)
            #lns += [ "# %3d --%3d :  %s " % (bp_left_resid, bp_right_resid, bp_name) ]
            duplex_resids += [bp_left_resid, bp_right_resid]
            duplex.append((bp_left_resid, bp_right_resid, bp_name))
        self.helices_resids += duplex_resids
        self.duplexs[dup_id] = duplex
        lns +=  ["# self.duplexs:  %s " %   self.duplexs ]
        lns += [ "# self.helices_resids: %s " %  self.helices_resids ]
        self.add2log(lns, vb=vb)
        #return helix_id, duplex_resids, duplex
        
    
    def parseFnRDC(self, lns, vb=0):
        """ parse RDC file name """
        self.fn_exp_rdc = None 
        #self.prnLines(lns)
        log_lns = []
        for _ln in lns:
            _ln.strip()
            if _ln[0] != '#':
                if '=' in _ln:
                    if 'fn_rdc' in _ln:
                        tmp =  split("\s*=\s*", _ln)
                        self.fn_exp_rdc = tmp[1]
                        log_lns = ["# reading  RDC filename(self.fn_exp_rdc): %s " % self.fn_exp_rdc ]
                        #self.add2log(log_lns)
                        break 
                    if 'ex_rdcs' in _ln:
                        log_lns = ["# excluded rdcs in duplex periodicty fitting "]
                        log_lns += [_ln]
                        tmp =  split("\s*=\s*", _ln)
                        ex_rdcs = tmp[1].split()
                        ex_rdcs = [ int(_rdc) for _rdc in ex_rdcs]
                        self.ex_rdcs = ex_rdcs
        if self.fn_exp_rdc is None:
            self.err_exit("No rdc filename in the ipnut file\n, fn_rdc= rdc_example.tbl")
        self.fn_exp_rdc = os.path.join(self.fn_dir, self.fn_exp_rdc)
        log_lns += [ "#*  fn_exp_rdc: %s " % self.fn_exp_rdc    ]
        self.add2log(log_lns, vb=vb)
    
    def __parseAngeRange(self, ln):
        tmp = ln.split('+-')
        med = float(tmp[0])
        _da = float(tmp[1])
        return  (med - _da, med + _da)
    
        
    def parseAng_Dups(self, lns, vb=0):
        """ parse angles between duplexes """
        dup_angles = []
        for ln in lns:
            ln = ln.strip()
            if not ln or ln[0] == "#" or ln.startswith("##"): continue 
            tmp = split("\s+", ln)
            #print tmp 
            dup1, dup2 = tmp[0].split('-')
            dup1 = dup1.strip()
            dup2 = dup2.strip()
            ang_ranges = []
            da1 = self.__parseAngeRange(tmp[1])
            ang_ranges.append(da1)
            if len(tmp) == 4 and "+-" in tmp[3]:
                da2 = self.__parseAngeRange(tmp[3])
                ang_ranges.append(da2)
            dup_angles.append(((dup1, dup2), ang_ranges))
        self.dup_angles = dup_angles
        if vb: print "# self.dup_angles: ", dup_angles 
        self.add2log(lns)
        
        
    def getSecondStr(self):
        return self.sstrands, self.duplexs
    def prnSeconStr(self):
        print "# duplex: self.duplexs: ", self.duplexs
        print "# single strand: self.sstrands", self.sstrands 
        
    
    def getDuplexID(self, ln):
        tmp = None
        for sep in ["=", ":", "helix", "stem"]:
            if sep in ln:
                tmp = split(sep,ln)
                #print "sep: ", sep, "tmp: ", tmp     
                break 
        if tmp is None: raise "# could not find duplex id seperator: %s "  % ln 
        #duplex_num = int(tmp[1])    
        return tmp[1].strip()
    
    
    def prnHelixInp(self, seg_lns, index=1):
        """ print input helices """
        for helix_lns in seg_lns:
            if 'helix' in helix_lns[0]:
                print "\n# number of basepairs= %3d " % (len(helix_lns)-1)
                self.prnLns(helix_lns, index=index)
                    
    def parseSingleStrand(self, vb=0):
        """ locate single strand in RNA molecules """
        sstrands = []
        ss_resids = []
        for resid in self.resids:
            if resid not in self.helices_resids:
                ss_resids.append(resid)
        _ss = []
        pre_resid = -1 
        for resid in ss_resids:
            if resid == pre_resid+1:
                _ss.append(resid)
            else:
                if _ss: sstrands.append(_ss)
                _ss = [resid]
            pre_resid = resid 
        sstrands.append(_ss)
        if vb: print  "# single strands:", sstrands
        self.sstrands = sstrands
    
        
    def test(self):
        print "# fn_inp: ", self.fn_2str_input
        self.load2str()
        

        
class OUTPUT:
    def tmp_resDict2pdb(self, pdb_xyz, vb=0): # deprecated for consistence
        """ 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
        #print "# resid= %3d nt_name= %s " % (resid,nt_name)
        for resid in resids:
            nt_name = self.seq_dict[resid] 
            if vb: print "# resid= %3d %s " % (resid,nt_name)
            nt_atom_names = nt_atoms_woh[nt_name]
            res_xyz = pdb_xyz[resid] 
            #res_xyz = array(res_xyz)
            for i in range(len(nt_atom_names)):
                atom_name = nt_atom_names[i]
                #print "# i= %3d nt_name=%s "  % (i, nt_name)
                if len(atom_name) <= 3: atom_name= " %s" % atom_name
                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 resPdbDt2pdb(self, pdb_xyz, pdb_resids=None, vb=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'
        #print "# pdb_xyz: ", pdb_xyz
        resids = pdb_xyz.keys()
        resids.sort()
        lns = []
        atom_id = 0
        #print "# resid= %3d nt_name= %s " % (resid,nt_name)
        if pdb_resids is None: pdb_resids = resids
        for i in range(len(resids)):
            resid = resids[i]
            pdb_resid = pdb_resids[i]
            nt_name = self.seq_dict[resid] 
            if vb: print "# resid= %3d %s " % (resid,nt_name)
            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]
                #print "# i= %3d nt_name=%s "  % (i, nt_name)
                if len(atom_name) <= 3: atom_name= " %s" % atom_name
                atom_xyz = res_xyz[i,:]
                atom_id += 1
                ln =  fmt_pdb  % ('ATOM', atom_id, atom_name, nt_name," ",  pdb_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 prnRelXYZ(self):
        """ print atom relative positions """
        apx, apy, apz = 41.535, -30.076,  -0.7872 # riboA 67 P atom position in PDB
        for i in range(4):
            rot_xyz = self.rotXYZ_by_refTPR(h2_o4[0], ref_h2_o4[i], hex_xyz[67])
            pivot = rot_xyz[0] 
            dx = pivot[0] - apx 
            dy = pivot[1] - apy 
            dz = pivot[2] - apz 
            for i in range(len(rot_xyz)):
                _xyz = rot_xyz[i]
                rx = _xyz[0] - dx 
                ry = _xyz[1] - dy 
                rz = _xyz[2] - dz 
                print "#  %3d  %8.3f %8.3f %8.3f :R: %8.3f %8.3f %8.3f  " % (i, _xyz[0], _xyz[1], _xyz[2], rx, ry, rz )

    def prnXYZ(self, xyz):
        for i in range(len(xyz)):
            _xyz = xyz[i]
            print "#  %3d  %8.3f %8.3f %8.3f " % (i, _xyz[0], _xyz[1], _xyz[2])
    

    def resDict2pdb(self):
        """ 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 = self.new_pdb_xyz.keys()
        resids.sort()
        lns = []
        atom_id = 0
        for resid in resids:
            nt_name = self.seq_dict[resid] 
            print "# resid= %3d nt_name= %s " % (resid,nt_name)
            nt_atom_names = nt_atoms_woh[nt_name]
            res_xyz = self.new_pdb_xyz[resid] 
            for i in range(len(nt_atom_names)):
                atom_name = nt_atom_names[i]
                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)
                #print ln 
        return lns 



            
        

if (__name__ == "__main__"):
    pass 
