import urllib
import urllib2

            
def simple(stream):
    
    return stream.read()

simple.ENCODING = "identity"
simple.ACCEPT = set( [ "", simple.ENCODING ] )


# HTTP functionality
AGENT = "HTTPInterface/0.0.2 +http://www.phenix-online.org/"
OPENER = urllib2.build_opener()

class Download(object):

    def __init__(self, url, data = {}, headers = {}, method = simple):
        
        request = urllib2.Request(
            url = url,
            data = urllib.urlencode( data ),
            headers = dict( headers.items()
                + [
                    ( "User-Agent", AGENT ),
                    ( "Accept-encoding", method.ENCODING ),
                    ]
                )
            )
        stream = OPENER.open( request )
        
        if stream.headers.get( "Content-Encoding", "" ).lower() not in method.ACCEPT:
            raise RuntimeError, "Service does not support encoding %s" % method.ENCODING
        
        self.data = method( stream = stream )
         

def is_online(url):
    
    try:
        stream = OPENER.open( url )

    except urllib2.URLError:
        return False

    return True
    

class RCSBHTTP(object):
    """
    Interface to RCSB service using HTTP
    """

    def __init__(self, identifier):

        if format not in [ "pdb", "fasta" ]:
            raise ValueError, "Unknown format: %s" % format
        
        if not is_online( url = "http://www.rcsb.org/pdb/home/home.do" ):
            raise RuntimeError, "Cannot contact RCSB service"
        
        self.format = format


    def fetch_single(self, identifier):

        try: 
            result = Download(
                url = "http://www.rcsb.org/pdb/download/downloadFile.do",
                data = {
                    "fileFormat": self.format,
                    "compression": "NO",
                    "structureId": identifier
                    },
                method = simple
                )
        
        except urllib2.URLError, e:
            raise ValueError, "Unknown identifier: %s" % identifier
        
        return result
    
    
    def __str__(self):
        
        return "RCSB HTTP service"
        

class DbfetchHTTP(object):
    """
    Interface to EBI Dbfetch service using HTTP
    """
    
    FORMATS_FOR = {
        "pdb": set( [ "pdb", "pdbml", "mmcif", "fasta" ] )
        }

    def __init__(self, database, format):

        if database not in self.FORMATS_FOR:
            raise ValueError, "Unknown database: %s" % database
        
        if format not in self.FORMATS_FOR[ database ]:
            raise ValueError, "Unknown format: %s" % format
        
        if not is_online( url = "http://www.ebi.ac.uk/cgi-bin/dbfetch" ):
            raise RuntimeError, "Cannot contact %s" % str( self )
        
        self.database = database
        self.format = format


    def fetch_single(self, identifier):

        return self.fetch_multiple( identifiers = [ identifier ] )
    
    
    def fetch_multiple(self, identifiers):

        result = Download(
            url = "http://www.ebi.ac.uk/Tools/dbfetch/dbfetch",
            data = {
                "db": self.database,
                "id": ",".join( set( identifiers ) ),
                "format": self.format,
                "style": "raw",
                }
            )
        
        if result.data.startswith( "ERROR" ):
            result.data = ""
        
        return result
    
    
    def __str__(self):
        
        return "EBI DbFetch service"
    
    
class SequenceFetch(object):
    """
    Uses a proxy to access sequence data
    """

    def __init__(self, proxy, format, header):
        
        self.proxy = proxy
        self.format = format
        self.header = header
        self.cache = {}
        
        
    def prefetch(self, identifiers):
        
        self.fetch_and_cache( identifiers = [ i.upper() for i in identifiers ] )


    def fetch_chain(self, identifier, chain):
        
        return self.fetch_entry( identifier = identifier ).get( chain )        
            
    
    def fetch_entry(self, identifier):
        
        std = identifier.upper()
        
        if std not in self.cache:
            self.fetch_and_cache( identifiers = [ std ] )

        return self.cache[ std ]
    
    
    def fetch_and_cache(self, identifiers):

        import iotbx.bioinformatics
        
        sequences_for = dict( [ ( i, {} ) for i in identifiers ] )
        result = self.proxy.fetch_multiple( identifiers = identifiers )
        
        try:
          sequences = iotbx.bioinformatics.parse_sequence_str(
            data = result.data,
            format = self.format
            )
          
        except ValueError, e:
          raise ValueError, "Incorrect format: expected %s, found %s" % e.args[:2]
        
        for s in sequences:
          try:
            s.reinterpret_header( header = self.header )
          
          except ValueError, e:
            raise ValueError, "Incorrect header format: expected %s, found %s" % e.args
          
          assert s.header.identifier in sequences_for
          sequences_for[ s.header.identifier ][ s.header.chain ] = s
          
        self.cache.update( sequences_for )


class PDBFetch(object):
    """
    Uses a proxy to get PDB data
    """
    
    def __init__(self, proxy):
        
        self.proxy = proxy
        
    
    def fetch_entry(self, identifier):
        
        import iotbx.pdb
        response = self.proxy.fetch_single( identifier = identifier )
        input = iotbx.pdb.input( lines = response.data, source_info = "PDB" )
        return input.construct_hierarchy()
                    
                    
    def fetch_chain(self, identifier, chain):
                
        root = self.fetch_entry( identifier = identifier )
        asc = root.atom_selection_cache().selection( "chain '%s'" % chain )
        return root.select( asc, True )
    
    
def parse_pdb_redirections(data):
    
    import re
    regex = re.compile(
        r"""
        ^ OBSLTE \s+
        [\w-]+ [ ]+
        ( \w{4} )
        (?: [ ]+ ( \w{4} ) )?
        \s* $
        """,
        re.VERBOSE | re.MULTILINE
        )
    
    replacement_for = {}
    
    for hit in regex.finditer( data ):
        ( old, new ) = hit.groups()
        replacement_for[ old.upper() ] = new.upper() if new else None
        
    return replacement_for
    
    
class PDBRedirections(object):
    """
    Downloads a list of redirections from the PDB
    """
    
    def __init__(self):
        
        try:
            download = Download(
                url = "ftp://ftp.wwpdb.org/pub/pdb/data/status/obsolete.dat"
                )
            self.redirection_for = parse_pdb_redirections( data = download.data )
            
        except urllib2.URLError:
            self.redirection_for = {}
        
        
    def obsoleted(self, identifier):
        
        std = identifier.upper()
        return ( std in self.redirection_for and self.redirection_for[ std ] )
        
        
    def replacement_for(self, identifier):
        
        std = identifier.upper()
        return self.redirection_for[ std ]
    
    
    def retracted(self, identifier):
        
        std = identifier.upper()
        return ( std in self.redirection_for and not self.redirection_for[ std ] )
    
    
    def count(self):
        
        return len( self.redirection_for )
    
    
REDIRECTIONS = None

def get_pdb_redirections():
    
    global REDIRECTIONS
    
    if not REDIRECTIONS:
        REDIRECTIONS = PDBRedirections()
        
    return REDIRECTIONS