# Copyright (c) 1996, 1997, The Regents of the University of California.
# All rights reserved.  See Legal.htm for full text and disclaimer.

"""
PDB basic read-access class PR.py by Paul Dubois, LLNL
Version 1.1 July 1997
"""
import pprint
import pypdb
import sys

class PR:
    "PDB file read-access class."

    no_file_message = '(PR object not open on any file)'

    Error = 'PR error'

    def __del__(self):
        self.close()

    def __getattr__(self, name):
        w = self.read(self._registry[name])
        if self._cache:
            self.__dict__[name] = w
        return w

    def __init__(self, filename='', path='', pdbtype='', 
                 cache=1, verbose = 1):
        """PR(filename='', path='', pdbtype='',cache=1,verbose=1)
        opens filename if given, registers names specified by 
        PDB-name and type specifiers. 
        Object will use a cache by default. """

        self.__dict__['_handle'] = None
        self.__dict__['_registry'] = {}
        self.__dict__['_cache'] = cache
        self.set_verbosity(verbose)
        if filename:
            self.open(filename, path, pdbtype)

    def __repr__(self):
        if self.is_open():
            current_mode = 'opened for ' + self.inquire_handle().mode()
        else:
            current_mode = PR.no_file_message
        return 'PDB file %s %s.' % (self.inquire_filename(), current_mode)

    __str__ = __repr__

    def check_open(self):
        "check_open(): raise exception if not open for read."
        if not self.is_open():
            raise PR.Error, 'PR object not open for read.'

    def close(self):
        "close(): close the file."
        if self.is_open():
            if self.inquire_verbosity():
                print "Closing PDB file", \
                      self.inquire_filename()
            self.inquire_handle().close()
            r = self._registry
            c = self._cache
            d = self.__dict__
            v = self.inquire_verbosity()

            for n in d.keys():
                del d[n]

            d['_handle'] = None
            d['_registry'] = {}
            d['_cache'] = c
            self.set_verbosity(v)


    def inquire_filename(self):
        "inquire_filename() = name of this file."
        self.check_open()
        return self.inquire_handle().filename()

    def inquire_handle(self):
        "inquire_handle() = PDBfile object open on this file."
        return self._handle

    def inquire_high(self, name):
        "inquire_high(name) = high indices of name."
        self.check_open()
        pname = self.inquire_pdbname(name)
        return self.inquire_handle().high(pname)

    def inquire_low(self, name):
        "inquire_low(name) = low indices of the name."
        self.check_open()
        pname = self.inquire_pdbname(name)
        return self.inquire_handle().low(pname)

    def inquire_ls(self,path='',pdbtype=''):
        """inquire_ls(path='',pdbtype='') = list of those pdb names 
        in the file which represent objects with names matching the
        path or type equal to the given type. Blank matches all."""

        self.check_open()
        return self.inquire_handle().ls(path, pdbtype)

    def inquire_mode(self):
        "inquire_mode() = mode ('r', 'w', or 'a') of this file."
        self.check_open()
        return self.inquire_handle().mode()

    def inquire_names(self):
        "inquire_names() = sorted list of registered names"
        self.check_open()
        r = self._registry.keys()
        r.sort()
        return r

    def inquire_pdbname(self, name):
        """inquire_pdbname(name) = registered pdbname corresponding to
        given name."""
        self.check_open()
        return self._registry[name] 

    def inquire_pwd(self):
        "inquire_pwd() = present PDB directory"
        self.check_open()
        return self.inquire_handle().pwd()

    def inquire_shape(self,name):
        "inquire_shape(name) = shape of name."
        self.check_open()
        pname = self.inquire_pdbname(name)
        low = self.inquire_low(pname)
        high = self.inquire_high(pname)
        n = len(low)
        result = []
        for i in range(n):
            result.append(high[i] - low[i] + 1)
        return tuple(result)

    def inquire_type(self, name):
        "inquire_type(name) = type of name."
        self.check_open()
        pname = self.inquire_pdbname(name)
        return self.inquire_handle().type(pname)

    def inquire_verbosity(self):
        "inquire_verbosity() = current value of verbose flag."
        return self._verbose

    def is_cached(self):
        "is_cached() = value of cache flag."
        return self._cache

    def is_column_major(self):
        "is_column_major() = true if file was written by Fortran."
        self.check_open()
        return self.inquire_handle().is_column_major()

    def is_open(self):
        "is_open() = true if file is open for read"
        if self.inquire_handle() == None:
            return 0
        return 1

    def print_names(self, file=sys.stdout):
        """print_list(file=sys.stdout): pretty print the list of
        registered names to file."""

        self.check_open()
        pprint.pprint(self.inquire_names(), file)


    def open(self, filename, path='',pdbtype=''):
        """open(name, path='', pdbtype=''): open file, register names
        using register_pattern(path, pdbtype). The default is to 
        register all the names."""

        self.close()
        if filename:
            self.__dict__['_handle'] =  pypdb.open(filename)
            self.register_pattern(path, pdbtype)

    def read(self,pdbname):
        "read(pdbname) = the value of pdbname as a Python object."
        self.check_open()
# Note: cannot handle partial reads via the full read by naming
        if pdbname[-1] == ')':
            raise PR.Error, pdbname + ' -- must use read_part'
        return self.inquire_handle().read(pdbname)

    def read_part(self, name, triples):
        """read_part(name, triples) =  part of the value of name 
        as a Python object. Triples must be an array of 3 integers
        per subscript giving the low/high/stride values desired. 
        Note that these subscripts are relative to the values stored
        in the PDB file, and do not follow Python conventions."""

        self.check_open()
        pname = self.inquire_pdbname(name)
        low = self.inquire_low(name)
        high = self.inquire_high(name)
        n = len(low)
        if (3*n) != len(triples):
            raise PR.Error, \
                  name + ', wrong number of subscripts, need '+`n`
        for i in range(n):
            if triples[3*i] > triples[3*i+1] or \
               low[i] > triples[3*i] or \
               high[i] < triples[3*i+1] or \
               triples[3*i + 2] < 1:
                raise PR.Error, name + ', bad triple =' \
                      +`triples[3*i:3*i+3]`
        return self.inquire_handle().read_part(pname, triples)

    def register_name(self, name, pdbname):
        """register_name(name, pdbname) registers name as a key to 
        get the variable whose name in the file is pdbname"""
        self._registry[name] = pdbname

    def register_pattern(self, path, pdbtype):
        """register_pattern(path, pdbtype) -- 
        Register names found using inquire_ls via register_name"""

        self.check_open()
        for name in self.inquire_handle().ls(path,pdbtype):
            self.register_name(name, name)

    def set_directory(self, name):
        """set_directory(name) 
        -- change PDB directory to name, return status"""
        self.check_open()
        return self.inquire_handle().cd(name)

    def set_verbosity(self, flag):
        """verbose(flag) sets verbosity level to flag.
        0 for quiet operation, 
        1 to report closing only, 
        2 to report access to data."""
        if 0 <= flag <= 2:
            self.__dict__['_verbose'] = flag
        else:
            raise PR.Error, 'Illegal value for verbosity: '+`flag`