#
# title:   torsion.py
# summary: a class to make access to torsion files easier
# author:  Nick Fitzkee (nickfitzkee at jhu dot edu)
# date:    February 1, 2005
#

class TorResidue:
    def __init__(self, line=None):
        self.res = None
        self.idx = None
        self.icd = None
        self.phi = None
        self.psi = None
        self.ome = None
        self.chi = []
        self.rcs = []
        self.ssc = []

        if line:
            ridx = line[0:8]
            ic   = line[9]
            aa   = line[11:14]
            data = line[15:]
            
            phi, psi, ome, chi1, chi2, chi3, chi4, om, om_ss, \
                 nm, nm_ss = data.split()
            ridx, aa, phi, psi, ome, chi1, chi2, chi3, chi4 = \
                 int(ridx), aa, float(phi), float(psi), float(ome), \
                 float(chi1), float(chi2), float(chi3), float(chi4)

            self.res, self.idx, self.icd = aa, ridx, ic
            
            if phi < 900.0: self.phi = phi
            if psi < 900.0: self.psi = psi
            if ome < 900.0: self.ome = ome
            
            for c in (chi1, chi2, chi3, chi4):
                if c < 900.0: self.chi.append(c)

            self.rcs = [om, nm]
            self.ssc = [om_ss, nm_ss]

    def index(self):
        return self.idx, self.icd

class TorChain:
    def __init__(self, fstream):
        self.data = {}
        self.didx = {}
        self.ordr = []
        self.chid = None
        self.nres = None

        self._make(fstream)

    def _make(self, f):
        line = f.readline()
        self.chid = line[0:1]
        self.nres = int(line[2:])
        n = 0
        
        for i in xrange(self.nres):
            tr = TorResidue(f.readline())
            id = tr.index()
            self.data[id] = tr
            self.ordr.append(id)
            self.didx[id] = n
            n = n+1

    def chain(self): return self.chid

    def __len__(self): return len(self.data)

    def __getitem__(self, i): return self.data[i]

class TorFile:
    def __init__(self, fn):
        fh = fn
        tc = 1
        
        if type(fn) == type('hi'):
            if fn[-2:] == 'gz':
                import gzip
                fh = gzip.open(fn)
            elif fn[-3:] == 'bz2':
                import bz2
                fh = bz2.BZ2File(fn)
            else:
                fh = open(fn)
        else:
            tc = 0
            
        self.data = {}
        self.didx = {}
        self.ordr = []
        self.nchn = None

        self._make(fh)

        if tc: fh.close()

    def _make(self, f):
        self.nchn = int(f.readline())
        n = 0
        
        for i in xrange(self.nchn):
            tc = TorChain(f)
            id = tc.chain()
            self.data[id] = tc
            self.ordr.append(id)
            self.didx[id] = n
            n = n+1

    def __len__(self): return len(self.data)

    def __getitem__(self, i): return self.data[i]