from phaser.pipeline import homology

from iotbx import bioinformatics
import iotbx.pdb

import StringIO
import gzip
import re
import xml.sax.saxutils
import os

# Exceptions
class ServiceException(Exception):
    pass

class NotFoundException(Exception):
    pass

class UnexpectedResponse(Exception):
    pass

class JobException(Exception):
    pass


# Job status objects
class JobBeingProcessed(object):
    
    def __str__(self):
        
        return "Job is being processed"
    

class JobPending(JobBeingProcessed):
    
    def __str__(self):
        
        return "Job is waiting in the queue"
    
    
class JobRunning(JobBeingProcessed):
    
    def __str__(self):
        
        return "Job is running"
     

class JobProcessingCompleted(object):
    
    def __str__(self):
        
        return "Job processing completed"


class JobFailure(JobProcessingCompleted):
    
    def __str__(self):
        
        return "Job failed"
    

class JobSuccess(JobProcessingCompleted):
    
    def __str__(self):
        
        return "Job finished successfully"
    

class JobNotFound(object):
    
    def __str__(self):
        
        return "Job not found"
    

# Common constants
IDENTIFIER = "identifier"
JOBID = "jobid"
FORMAT = "format"
PROGRAM = "program"
DATABASE = "database"
SEQUENCE = "sequence"
SEQUENCES = "sequences"


# Processor classes - simple text modification
class Validator(object):
    """
    Raises exception if check fails, else returns data passed
    """
    
    def __init__(self, check, exception, message):
        
        self.check = check
        self.exception = exception
        self.message = message
        
    
    def process(self, data):
        
        if not self.check( data ):
            raise self.exception, self.message
        
        else:
            return data
    
    
class FilterProcessor(object):
    """
    Convert data by a simple functional modification
    """
    
    def __init__(self, conversion):
        
        self.conversion = conversion
        
        
    def process(self, data):
        
        return self.conversion( data )
    
    
class ExtractProcessor(object):
    """
    Extract data from some text using a regular expression 
    """
    
    def __init__(self, regex):
        
        self.regex = regex
        
        
    def process(self, data):
        
        result = self.regex.search( data )
        
        if not result:
            raise UnexpectedResponse, "Unexpected response: %s" % data
        
        return result.group( 1 )
    
    
class GunzipProcessor(object):
    """
    Unzips data
    """
    
    def process(self, data):
        
        try:
            uncompressed_data = gzip.GzipFile(
                fileobj = StringIO.StringIO( data )
                ).read()
            
        except IOError, e:
            raise UnexpectedResponse, "Bad compression: error '%s'" % e 
            
        return uncompressed_data
    
    
class DictKeyTranslator(object):
    """
    Renames keys in dictionaries
    """
    
    def __init__(self, translation):
        
        self.translation = translation
        
        
    def process(self, data):
        
        return dict( [ ( self.translation.get( name, name ), value )
            for ( name, value ) in data.items() ] )
    
    
# Service objects
class ForwardedServiceMethod(object):
    """
    A function object that simply forwards the call to something else
    """
    
    def __init__(
        self,
        function,
        preprocessors = [],
        postprocessors = [],
        arguments = [],
        defaults = {}
        ):
        
        self.function = function
        self.preprocessors = preprocessors
        self.postprocessors = postprocessors
        self.arguments = arguments
        self.defaults = defaults
        
        
    def add_default(self, argument, value):
        
        self.defaults[ argument ] = value
        
        
    def __call__(self, **kwargs):
        
        request = self.defaults.copy()
        request.update(
            dict(
                [ ( name, value ) for ( name, value ) in kwargs.items()
                    if name in self.arguments ]
                )
            )
        
        for processor in self.preprocessors:
            request = processor.process( request )
            
        try:
            response = self.function( **request )
            
        except KeyboardInterrupt, e:
            raise
            
        except Exception, e:
            raise ServiceException, e 
        
        for processor in self.postprocessors:
            response = processor.process( response )
            
        return response
    
    
class ProxyService(object):
    """
    Class to interact with potentially remote resources
    """
    
    def __init__(self, proxy, name):
        
        self.proxy = proxy
        self.name = name
        
        
    def register(
        self,
        name,
        proxy_method_name,
        request_constructors = [],
        response_evaluators = [],
        arguments = [],
        defaults = {}
        ):
        
        if hasattr( self, name ):
            raise AttributeError, "Cannot overwrite existing '%s'" % name
        
        method = ForwardedServiceMethod(
            function = getattr( self.proxy, proxy_method_name ),
            preprocessors = request_constructors,
            postprocessors = response_evaluators,
            arguments = arguments,
            defaults = defaults,
            exception_type = exception_type 
            )
        
        setattr( self, name, method ) 
        
        
    def unregister(self, name):
        
        if not hasattr( self, name ):
            raise AttributeError, "No such method: %s" % name
        
        if hasattr( self.__class__, name ):
            raise AttributeError, "Cannot unregister class method"
        
        delattr( self, name )
        
        
    def __str__(self):
        
        return "%s (provided by %s)" % ( self.name, str( self.proxy ) )
    
    
# Caching
class MemoizedFunction(object):
    """
    Holds information about a memoized function
    """
    
    def __init__(self, function, key_generator, memoizer):
        
        self.function = function
        self.key_generator = key_generator
        self.memoizer = memoizer
        
        
    def __call__(self, *args, **kwargs):
        
        key = self.key_generator( *args, **kwargs )
        
        if self.memoizer.contains( key ):
            return self.memoizer.retrieve( key = key )
        
        else:
            data = self.function( *args, **kwargs )
            self.memoizer.insert( key = key, data = data )
            return data
        
        
    @staticmethod
    def simple_key(*args, **kwargs):
        
        all = args + tuple( kwargs.values() )
        
        return all[0] if all else all
        
        
class Memoizer(object):
    """
    General caching-related functionality
    """ 
        
    def attach_method(self, method, key_generator):
        
        memoized_function = MemoizedFunction(
            function = method,
            key_generator = key_generator,
            memoizer = self
            )
        
        setattr( method.im_self, method.__name__, memoized_function ) 
    
    
    def detach_method(self, memoized_function):
        
        if not isinstance( memoized_function, MemoizedFunction ):
            raise ValueError, "Cannot detach non-memoized function"
        
        setattr(
            memoized_function.function.im_self,
            memoized_function.function.__name__,
            memoized_function.function
            )
        

# Cache objects
class MemoryMemoizer(Memoizer):
    """
    Caches data in the memory
    """
    
    def __init__(self):
        
        self.cache = {}
        
        
    def insert(self, key, data):
        
        self.cache[ key ] = data
        
        
    def retrieve(self, key):
        
        return self.cache[ key ]
    
    
    def contains(self, key):
        
        return key in self.cache
    
    
    def purge(self):
        
        self.cache = {}
        
        
class PersistentMemoizer(Memoizer):
    """
    Caches data on the file system
    """
    
    def __init__(self, tmp_dir, template):
        
        assert os.path.exists( tmp_dir ) and os.path.isdir( tmp_dir )
        self._tmp_dir = tmp_dir
        self._template = template
        
        
    def file_name(self, key):
        
        return self._template % key
    
    
    def file_path(self, key):
        
        return os.path.join( self._tmp_dir, self.file_name( key = key ) )
        
        
    def insert(self, key, data):
        
        try:
            ofile = file( self.file_path( key = key ), "w" )
            ofile.write( data )
        
        finally:
            ofile.close()
        
        
    def retrieve(self, key):
        
        return file( self.file_path( key = key ) ).read()
    
    
    def contains(self, key):
        
        return self.file_name( key = key ) in os.listdir( self._tmp_dir )
    
    
    def purge(self):
        
        pass
        
    
# Predefined proxy service method parameters
OCA_HTTP_FETCH_PARAMS = {
    "proxy_method_name": "fetch",
    "request_constructors": [],
    "response_evaluators": [
        ExtractProcessor(
            regex = re.compile( r"<pre>(.*)</pre>", re.DOTALL )
            ),
        Validator(
            check = ( lambda s: s ),
            exception = NotFoundException,
            message = "Not found" 
            ),
        FilterProcessor(
            conversion = lambda s: bioinformatics.seq_sequence_parse( s )
            ),
        ],
    "arguments": [ IDENTIFIER ],
    "defaults": {
        FORMAT: "seq"
        },
    }


DBFETCH_HTTP_FETCH_PARAMS = {
    "proxy_method_name": "fetch",
    "request_constructors": [],
    "response_evaluators": [
        Validator(
            check = ( lambda s: s != "No entries found. \n" ),
            exception = NotFoundException,
            message = "Not found" 
            ),
        FilterProcessor(
            conversion = lambda s: iotbx.pdb.input(
                lines = s,
                source_info = "EBI Dbfetch HTTP"
                )
            ),
        ],
    "arguments": [ IDENTIFIER, "database" ],
    "defaults": {
        "database": "pdb",
        FORMAT: "pdb",
        },
    }

DBFETCH_SOAP_FETCH_PARAMS = {
    "proxy_method_name": "fetchData",
    "request_constructors": [
        Validator(
            check = lambda s: IDENTIFIER in s,
            exception = ValueError,
            message = "Incomplete query" 
            ),
        FilterProcessor(
            conversion = lambda s: {
                "query": "%s:%s" % ( s[ "database" ].lower(), s[ IDENTIFIER ].lower() ),
                "format": s[ FORMAT ],
                "style": s[ "style" ],
                }
            ),
        ],
    "response_evaluators": [
        Validator(
            check = lambda s: "fetchDataReturn" in s,
            exception = UnexpectedResponse,
            message = "Badly formed response from server" 
            ),
        FilterProcessor(
            conversion = lambda s: xml.sax.saxutils.unescape(
                s[ "fetchDataReturn" ].encode( "ascii", "ignore" )
                )
            ),
        FilterProcessor(
            conversion = lambda s: iotbx.pdb.input(
                lines = s,
                source_info = "EBI Dbfetch SOAP"
                )
            ),
        ],
    "arguments": [ IDENTIFIER, "database", "style" ],
    "defaults": {
        "database": "pdb",
        FORMAT: "pdb",
        "style": "raw",
        },
    }
    

RCSB_HTTP_FETCH_UNCOMPRESSED_PARAMS = {
    "proxy_method_name": "fetch",
    "request_constructors": [],
    "response_evaluators": [
        FilterProcessor(
            conversion = lambda s: bioinformatics.seq_sequence_parse( s )
            ),
        ],
    "arguments": [ IDENTIFIER ],
    "defaults": {
        "compression": False,
        FORMAT: "fasta",
        },
    }

RCSB_HTTP_FETCH_COMPRESSED_PARAMS = {
    "proxy_method_name": "fetch",
    "request_constructors": [],
    "response_evaluators": [
        GunzipProcessor(),
        FilterProcessor(
            conversion = lambda s: iotbx.pdb.input(
                lines = s,
                source_info = "RCSB PDB"
                )
            ),
        ],
    "arguments": [ IDENTIFIER ],
    "defaults": {
        "compression": True,
        FORMAT: "pdb",
        },
    }

# EBI generic proxy handlers
EBI_GENERIC_SOAP_STATUS_REQUEST_CONSTRUCTORS = [
    Validator(
        check = lambda s: JOBID in s,
        exception = ValueError,
        message = "Incomplete call"
        ),
    FilterProcessor(
        conversion = lambda s: { "jobid": s[ JOBID ], }
        ),
    ]
EBI_GENERIC_SOAP_RETRIEVE_REQUEST_CONSTRUCTORS = [
    Validator(
        check = lambda s: JOBID in s and FORMAT in s,
        exception = ValueError,
        message = "Incomplete call"
        ),
    FilterProcessor(
        conversion = lambda s: {
            "jobid": s[ JOBID ],
            "type": s[ FORMAT ],
            }
        ),
    ]
EBI_GENERIC_SOAP_SUBMIT_RESPONSE_EVALUATORS = [
    Validator(
        check = lambda s: "jobid" in s,
        exception = UnexpectedResponse,
        message = "Invalid response from service"
        ),
    FilterProcessor(
        conversion = lambda s: s[ "jobid" ]
        ),
    ]
EBI_GENERIC_SOAP_STATUS_RESPONSE_EVALUATORS = [
    Validator(
        check = lambda s: "status" in s,
        exception = UnexpectedResponse,
        message = "Invalid response from service"
        ),
    FilterProcessor( conversion = lambda s: s[ "status" ] ),
    Validator(
        check = lambda s: s in [ "DONE", "ERROR", "NOT_FOUND", "PENDING", "RUNNING" ],
        exception = UnexpectedResponse,
        message = "Invalid response from service"
        ),
    FilterProcessor(
        conversion = lambda s: {
            "DONE": JobSuccess(),
            "ERROR": JobFailure(),
            "NOT_FOUND": JobNotFound(),
            "PENDING": JobPending(),
            "RUNNING": JobRunning()
            }[ s ]
        ),
    ]
EBI_GENERIC_SOAP_RETRIEVE_RESPONSE_EVALUATORS = [
    Validator(
        check = lambda s: "result" in s,
        exception = UnexpectedResponse,
        message = "Invalid response from service"
        ),
    FilterProcessor( conversion = lambda s: s[ "result" ] ),
    ]

# EBI service specific proxy handlers
def EBICustomisedSOAPSubmitParams(
    proxy_method_name,
    parameter_conversions = [],
    compulsory_arguments = [],
    additional_arguments = [],
    defaults = {}
    ):
    
    return {
        "proxy_method_name": proxy_method_name,
        "request_constructors": ( [
            Validator(
                check = lambda s: all( [ key in s for key in compulsory_arguments ] ),
                exception = ValueError,
                message = "Incomplete call"
                ),
            ]
            + parameter_conversions ),
        "response_evaluators": EBI_GENERIC_SOAP_SUBMIT_RESPONSE_EVALUATORS,
        "arguments": compulsory_arguments + additional_arguments,
        "defaults": defaults,
    }
    
    
def EBICustomisedSOAPRetrieveParams(parser, format):
    
    return {
        "proxy_method_name": "poll",
        "request_constructors": EBI_GENERIC_SOAP_RETRIEVE_REQUEST_CONSTRUCTORS,
        "response_evaluators": EBI_GENERIC_SOAP_RETRIEVE_RESPONSE_EVALUATORS + [ parser ],
        "arguments": [ JOBID ],
        "defaults": {
            FORMAT: format,
            },
    }
    
# EBI generic service proxy methods
EBI_GENERIC_SOAP_STATUS_PARAMS = {
    "proxy_method_name": "checkStatus",
    "request_constructors": EBI_GENERIC_SOAP_STATUS_REQUEST_CONSTRUCTORS,
    "response_evaluators": EBI_GENERIC_SOAP_STATUS_RESPONSE_EVALUATORS,
    "arguments": [ JOBID ],
    "defaults": {},
    }

EBI_GENERIC_SOAP_RETRIEVE_PARAMS = {
    "proxy_method_name": "poll",
    "request_constructors": EBI_GENERIC_SOAP_RETRIEVE_REQUEST_CONSTRUCTORS,
    "response_evaluators": EBI_GENERIC_SOAP_RETRIEVE_RESPONSE_EVALUATORS,
    "arguments": [ JOBID ],
    "defaults": {},
    }

# EBI service specific proxy method parameters
EBI_WUBLAST_SOAP_SUBMIT_PARAMS = EBICustomisedSOAPSubmitParams(
    proxy_method_name = "runWUBlast",
    parameter_conversions = [
        FilterProcessor(
            conversion = lambda s: {
                "params": dict(
                    [ ( name, s[ name ] ) for name in s if name != SEQUENCE ]
                    ),
                "content": [
                    {
                        "type": "sequence",
                        "content": str( s[ SEQUENCE ] ),
                        },
                    ],
                }
            ),
        ],
    compulsory_arguments = [ PROGRAM, DATABASE, SEQUENCE, "email" ],
    additional_arguments = [ "sort", "filter", "matrix", "numal", "scores",
        "exp", "echofilter", "stats", "strand", "sensitivity", "async" ],
    defaults = {
        DATABASE: "pdb",
        PROGRAM: "blastp",
        "async": True,
        "email": "gb360@cam.ac.uk",
        },
    )
EBI_WUBLAST_SOAP_RETRIEVE_PARAMS = dict(
    EBI_GENERIC_SOAP_RETRIEVE_PARAMS.items()
    + [ ( "defaults", { FORMAT: "toolxml" } ) ]
    )

EBI_CLUSTALW2_SOAP_SUBMIT_PARAMS = EBICustomisedSOAPSubmitParams(
    proxy_method_name = "runClustalW2",
    parameter_conversions = [
        FilterProcessor(
            conversion = lambda s: {
                "params": dict(
                    [ ( name, s[ name ] ) for name in s if name != SEQUENCES ]
                    ),
                "content": [
                    {
                        "type": "sequence",
                        "content": "\n".join(
                            [ str( seq ) for seq in s[ SEQUENCES ] ]
                            ),
                        },
                    ],
                }
            ),
        ],
    compulsory_arguments = [ SEQUENCES, "email" ],
    additional_arguments = [ "async", "align", "tree", "iteration",
        "numiter", "outorder", "alignment", "quicktree", "ktup", "window",
        "topdiags", "pairgap", "score", "gapopen", "endgaps", "gapext",
        "gapdist", "matrix", "outputtree", "kimura", "tossgaps", "clustering" ],
    defaults = {
        "async": True,
        "email": "gb360@cam.ac.uk",
        "output": "aln1",
        },
    )

EBI_CLUSTALW2_SOAP_RETRIEVE_PARAMS = dict(
    EBI_GENERIC_SOAP_RETRIEVE_PARAMS.items()
    + [ ( "defaults", { FORMAT: "toolaln" } ) ]
    )

# NCBI Blast proxy service method parameters
NCBI_BLAST_HTTP_SUBMIT_PARAMS = {
    "proxy_method_name": "put",
    "request_constructors": [
        DictKeyTranslator( translation = { "sequence": "query" } )
        ],
    "response_evaluators": [],
    "arguments": [ PROGRAM, DATABASE, SEQUENCE, "auto_format",
        "composition_based_statistics", "db_genetic_code", "endpoints",
        "entrez_query", "expect", "filter", "gapcosts", "genetic_code",
        "hitlist_size", "I_thresh", "layout", "lcase_mask", "matrix_name",
        "nucl_penalty", "nucl_reward", "other_advanced", "perc_ident",
        "phi_pattern", "query_file", "query_believe_defline", "query_from",
        "query_to", "searchsp_eff", "service", "threshold",
        "ungapped_alignment", "word_size" ],
    "defaults": {
        DATABASE: "pdb",
        PROGRAM: "blastp",
        },
    }

NCBI_BLAST_HTTP_STATUS_PARAMS = {
    "proxy_method_name": "status",
    "request_constructors": [
        DictKeyTranslator( translation = { JOBID: "rid" } )
        ],
    "response_evaluators": [
        FilterProcessor(
        conversion = lambda s: s.strip().upper()
        ),
    FilterProcessor(
        conversion = lambda s: {
            "READY": JobSuccess(),
            "WAITING": JobRunning(),
            "UNKNOWN": JobNotFound(),
            }[ s ]
        )
        ],
    "arguments": [ JOBID ],
    "defaults": {},
    }

NCBI_BLAST_HTTP_RETRIEVE_PARAMS = {
    "proxy_method_name": "get",
    "request_constructors": [
        DictKeyTranslator( translation = {
            JOBID: "rid",
            FORMAT: "format_type",
            } )
        ],
    "response_evaluators": [
        FilterProcessor(
            conversion = lambda s: homology.NCBIBlastXMLParse.parse( s )
            ),
        ],
    "arguments": [ JOBID, "alignments", "alignment_view", "descriptions",
        "entrez_links_new_window", "expect_low", "expect_high",
        "format_entrez_query", "format_object", "ncbi_gi", "results_file",
        "service", "show_overview" ],
    "defaults": {
        FORMAT: "XML",
        },
    }