from phaser.pipeline import services

import iotbx.pdb

import unittest
import urllib2
import urllib
import StringIO
import gzip
import re
import time
import os
import tempfile
        
        
class TestValidator(unittest.TestCase):
    
    def testNotFoundIsEmpty(self):
        
        validator = services.Validator(
            check = ( lambda s: not ( not s ) ),
            exception = services.NotFoundException,
            message = "Not found" 
            )
        
        self.assertEqual( validator.process( "Foo" ), "Foo" )
        
        self.assertRaises(
            services.NotFoundException,
            validator.process,
            ""
            )
        
        self.assertRaises(
            services.NotFoundException,
            validator.process,
            None
            )
        
        
class TestFilterProcessor(unittest.TestCase):
    
    def testProcess(self):
        
        processor = services.FilterProcessor(
            conversion = lambda s: s.lower() if s else s 
            )
        
        self.assertEqual( processor.process( "3adk" ), "3adk" )
        self.assertEqual( processor.process( "3ADK" ), "3adk" )
        self.assertEqual( processor.process( "3ADk" ), "3adk" )
        self.assertEqual( processor.process( "" ), "" )
        self.assertEqual( processor.process( None ), None )
        
        
class TestExtractProcessor(unittest.TestCase):
    
    def testProcess(self):
        
        processor = services.ExtractProcessor(
            regex = re.compile( r"<pre>(.*)</pre>", re.DOTALL ) 
            )
        
        self.assertEqual(
            processor.process( "<pre>Some random data</pre>" ),
            "Some random data"
            )
        self.assertEqual(
            processor.process( "<title><pre>Some random data</pre></title>" ),
            "Some random data"
            )
        self.assertEqual(
            processor.process( "<pre></pre>" ),
            ""
            )
        self.assertRaises(
            services.UnexpectedResponse,
            processor.process,
            "<pre>Some random data</pre"
            )
        
        
class TestGunzipProcessor(unittest.TestCase):
        
    def testProcess(self):
        
        data = "Some random data"
        compressed_data_file = StringIO.StringIO()
        compressor_stream = gzip.GzipFile( fileobj = compressed_data_file, mode = "wb" )
        compressor_stream.write( data )
        compressor_stream.close()
        compressed_data = compressed_data_file.getvalue()
        
        processor = services.GunzipProcessor()
        
        self.assertEqual( processor.process( compressed_data ), data )
        self.assertRaises(
            services.UnexpectedResponse,
            processor.process,
            data
            )
        
        
class TestMemoryMemoizer(unittest.TestCase):
    
    DATA1 = "Some data"
    DATA2 = "Some other data"
    DATA3 = "Some distinct data"
    DATA4 = "Some default data"
    
    def slow_method(self, identifier):
        
        time.sleep( 1.0 )
        
        if identifier == 1:
            return self.DATA1
        
        elif identifier == 2:
            return self.DATA2
        
        elif identifier == 3:
            return self.DATA3
        
        else:
            return self.DATA4
        
        
    def testCaching(self):
        
        memoizer = services.MemoryMemoizer()
        
        self.assertEqual( memoizer.cache, {} )
        
        memoizer.insert( key = 1, data = self.DATA1 )
        memoizer.insert( key = 2, data = self.DATA2 )
        
        self.assert_( memoizer.contains( key = 1 ) )
        self.assert_( memoizer.contains( key = 2 ) )
        self.assert_( not memoizer.contains( key = 3 ) )
        self.assertEqual(
            memoizer.cache,
            {
                1: self.DATA1,
                2: self.DATA2,
                }
            )
        
        self.assertEqual( memoizer.retrieve( key = 1 ), self.DATA1 )
        self.assertEqual( memoizer.retrieve( key = 2 ), self.DATA2 )
        self.assertRaises( KeyError, memoizer.retrieve, 3 )
        
        memoizer.purge()
        self.assertEqual( memoizer.cache, {} )
        
    
    def testMemoize(self):
        
        start = time.time()
        self.assertEqual( self.slow_method( 1 ), self.DATA1 )
        end = time.time()
        self.assertAlmostEqual( end - start, 1.0, 1 )
        
        memoizer = services.MemoryMemoizer()
        original_method = self.slow_method
        memoizer.attach_method(
            self.slow_method,
            key_generator = services.MemoizedFunction.simple_key
            )
        self.assert_( original_method != self.slow_method )
        self.assert_( hasattr( self.slow_method, "function" ) )
        self.assert_( isinstance( self.slow_method, services.MemoizedFunction ) )
        
        start = time.time()
        self.assertEqual( self.slow_method( 1 ), self.DATA1 )
        end = time.time()
        self.assertAlmostEqual( end - start, 1.0, 1 )
        start = time.time()
        self.assertEqual( self.slow_method( 1 ), self.DATA1 )
        self.assertEqual( self.slow_method( 1 ), self.DATA1 )
        self.assertEqual( self.slow_method( 1 ), self.DATA1 )
        self.assertEqual( self.slow_method( 1 ), self.DATA1 )
        end = time.time()
        self.assertAlmostEqual( end - start, 0.0, 1 )
        
        start = time.time()
        self.assertEqual( self.slow_method( 4 ), self.DATA4 )
        end = time.time()
        self.assertAlmostEqual( end - start, 1.0, 1 )
        start = time.time()
        self.assertEqual( self.slow_method( 4 ), self.DATA4 )
        self.assertEqual( self.slow_method( 4 ), self.DATA4 )
        self.assertEqual( self.slow_method( 4 ), self.DATA4 )
        self.assertEqual( self.slow_method( 4 ), self.DATA4 )
        end = time.time()
        self.assertAlmostEqual( end - start, 0.0, 1 )
        
        self.assertEqual(
            memoizer.cache,
            {
                1: self.DATA1,
                4: self.DATA4,
                }
            )
        memoizer.purge()
        start = time.time()
        self.assertEqual( self.slow_method( 3 ), self.DATA3 )
        end = time.time()
        self.assertEqual(
            memoizer.cache,
            {
                3: self.DATA3,
                }
            )
        
        
        memoizer.detach_method( self.slow_method )
        self.assert_( original_method == self.slow_method )
        self.assert_( not hasattr( self.slow_method, "function" ) )
        start = time.time()
        self.assertEqual( self.slow_method( 3 ), self.DATA3 )
        end = time.time()
        self.assertAlmostEqual( end - start, 1.0, 1 )
        
"""      
class TestOCAService(unittest.TestCase):
        
    oca = None
    sequence = ( ">3adk_  mol:protein-het length:195     Adenylate Kinase (E.C. 2.7.4.3) - Chain \n" 
        + "XMEEKLKKSKIIFVVGGPGSGKGTQCEKIVQKYGYTHLSTGDLLRAEVSSGSARGKMLSEIMEKGQLVPLETVLD\n"
        + "MLRDAMVAKVDTSKGFLIDGYPREVKQGEEFERKIGQPTLLLYVDAGPETMTKRLLKRGETSGRVDDNEETIKKR\n"
        + "LETYYKATEPVIAFYEKRGIVRKVNAEGSVDDVFSQVCTHLDTLK\n" )
    
    
    def setUp(self):
        
        if not self.oca:
            self.oca = services.OCAService()
        
        
    def testServiceOnline(self):
        
        self.assert_( self.oca.service_online() )
        
        
    def testName(self):
        
        self.assertEqual( str( self.oca ), "OCA service at EBI" )
        
        
    def testRetrieve(self):
        
        self.assertEqual( self.oca.retrieve( "3ADK" ), self.sequence )
        self.assertEqual( self.oca.retrieve( "3adK" ), self.sequence )
        
        self.assertRaises(
            services.NotFoundException,
            self.oca.retrieve,
            "4BZQ"
            )
        
        
class TestDbfetchService(unittest.TestCase):
    
    http = None
    soap = None
    
    def setUp(self):
        
        self.cwd = os.getcwd()
        
        if not self.http:
            self.http = services.DbfetchService()
            
        if not self.soap and self.cwd != tempfile.gettempdir():
            os.chdir( tempfile.gettempdir() )
            self.soap = services.DbfetchService( interface = "soap" )
            
            
    def tearDown(self):
        
        os.chdir( self.cwd )
        
        
    def testServiceOnline(self):
        
        self.assert_( self.http.service_online() )
        
        if self.soap:
            self.assert_( self.soap.service_online() )
            
        else:
            self.fail( "Cannot perform test" )
        
        
    def testName(self):
        
        self.assertEqual( str( self.http ), "Dbfetch service at EBI" )
        
        if self.soap:
            self.assertEqual( str( self.soap ), "Dbfetch service at EBI" )
            
        else:
            self.fail( "Cannot perform test" )
        
        
    def testRetrieve(self):
        
        data = self.http.retrieve( "3ADK" )
        self.checkData( data )
        
        if self.soap:
            data = self.soap.retrieve( "3ADK" )
            self.checkData( data )
            
        else:
            self.fail( "Cannot perform test" )
        
    
    def testNotFound(self):
        
        self.assertRaises(
            services.NotFoundException,
            self.http.retrieve,
            "4BZQ"
            )
        
        if self.soap:
            self.assertRaises(
                services.NotFoundException,
                self.soap.retrieve,
                "4BZQ"
                )
        
        else:
            self.fail( "Cannot perform test" )
            
    
    def checkData(self, data):
        
        hierarchy = iotbx.pdb.input( lines = data, source_info = "PDB" ).construct_hierarchy()
        
        self.assertEqual( len( hierarchy.models() ), 1 )
        self.assertEqual( len( hierarchy.models()[0].chains() ), 2 )
        self.assertEqual(
            len( hierarchy.models()[0].chains()[0].residue_groups() ),
            195
            )
        self.assertEqual(
            len( hierarchy.models()[0].chains()[1].residue_groups() ),
            2
            )
        self.assertEqual( len( hierarchy.atoms() ), 1529 )
        
        
class TestRCSBService(unittest.TestCase):
    
    pdb = None
    seq = None
    
    def setUp(self):
        
        if not self.pdb:
            self.pdb = services.RCSBService()
            
        if not self.seq:
            self.seq = services.RCSBService( format = "fasta", compression = False )
        
        
    def testServiceOnline(self):
        
        self.assert_( self.pdb.service_online() )
        self.assert_( self.seq.service_online() )
        
        
    def testName(self):
        
        self.assertEqual( str( self.pdb ), "PDB service at RCSB" )
        self.assertEqual( str( self.seq ), "PDB service at RCSB" )
        
        
    def testRetrieve(self):
        
        data = self.pdb.retrieve( "3ADK" )
        
        hierarchy = iotbx.pdb.input( lines = data, source_info = "PDB" ).construct_hierarchy()
        
        self.assertEqual( len( hierarchy.models() ), 1 )
        self.assertEqual( len( hierarchy.models()[0].chains() ), 2 )
        self.assertEqual(
            len( hierarchy.models()[0].chains()[0].residue_groups() ),
            195
            )
        self.assertEqual(
            len( hierarchy.models()[0].chains()[1].residue_groups() ),
            2
            )
        self.assertEqual( len( hierarchy.atoms() ), 1529 )
        
        sequence = ( ">3ADK:A|PDBID|CHAIN|SEQUENCE\n"
        + "XMEEKLKKSKIIFVVGGPGSGKGTQCEKIVQKYGYTHLSTGDLLRAEVSSGSARGKMLSEIMEKGQLVPLETVLDMLRDA\n"
        + "MVAKVDTSKGFLIDGYPREVKQGEEFERKIGQPTLLLYVDAGPETMTKRLLKRGETSGRVDDNEETIKKRLETYYKATEP\n"
        + "VIAFYEKRGIVRKVNAEGSVDDVFSQVCTHLDTLK\n" )
        
        self.assertEqual( self.seq.retrieve( "3ADK" ), sequence )
        self.assertEqual( self.seq.retrieve( "3adK" ), sequence )
    
    
    def testNotFound(self):
        
        self.assertRaises(
            services.NotFoundException,
            self.pdb.retrieve,
            "4BZQ"
            )
        self.assertRaises(
            services.NotFoundException,
            self.seq.retrieve,
            "4BZQ"
            )
        
        
class TestEBIWUBlastProxy(unittest.TestCase):
    
    proxy = None
    
    def setUp(self):
        
        self.cwd = os.getcwd()
        
        if not self.proxy and self.cwd != tempfile.gettempdir():
            os.chdir( tempfile.gettempdir() )
            self.proxy = services.EBIWUBlastProxy()
            
            
    def tearDown(self):
        
        os.chdir( self.cwd )
        
    
    def testCompulsorySubmitParameters(self):
        
        if self.proxy:
            self.assertEqual(
                self.proxy.compulsory_submit_parameters(),
                [ "program", "database", "email", "sequence" ]
                )
            
        else:
            self.fail( "Cannot perform test" )
        
    
    def testCompulsoryStatusParameters(self):
        
        if self.proxy:
            self.assertEqual(
                self.proxy.compulsory_status_parameters(),
                []
                )
        
        else:
            self.fail( "Cannot perform test" )
        
        
    def testCompulsoryRetrieveParameters(self):
        
        if self.proxy:
            self.assertEqual(
                self.proxy.compulsory_retrieve_parameters(),
                [ "format" ]
                )
        
        else:
            self.fail( "Cannot perform test" )
            
        
    def testAcceptedSubmitParameters(self):
        
        if self.proxy:
            self.assertEqual(
                self.proxy.accepted_submit_parameters(),
                [ "program", "database", "email", "sequence", "sort", "filter",
                    "matrix", "numal", "scores", "exp", "echofilter", "stats",
                    "strand", "sensitivity", "async" ]
                )
        
        else:
            self.fail( "Cannot perform test" )
            
    
    def testAcceptedStatusParameters(self):
        
        if self.proxy:
            self.assertEqual(
                self.proxy.accepted_status_parameters(),
                []
                )
        
        else:
            self.fail( "Cannot perform test" )
        
        
    def testAcceptedRetrieveParameters(self):
        
        if self.proxy:
            self.assertEqual(
                self.proxy.accepted_retrieve_parameters(),
                [ "format" ]
                )
        
        else:
            self.fail( "Cannot perform test" )
        
        
    def testConstructSubmitRequest(self):
        
        if self.proxy:
            result = {
                "params": {
                    "program": "blastp",
                    "database": "pdb",
                    "email": "gb360@cam.ac.uk",
                    "sort": "pvalue",
                    },
                "content": [ {
                    "type": "sequence",
                    "content": ">3ADK\nABCD",
                    } ]
                }
            self.assertEqual(
                self.proxy.construct_submit_request(
                    params = {
                        "program": "blastp",
                        "database": "pdb",
                        "email": "gb360@cam.ac.uk",
                        "sequence": ">3ADK\nABCD",
                        "sort": "pvalue",
                        }
                    ),
                result
                )
            self.assertRaises(
                ValueError,
                self.proxy.construct_submit_request,
                {
                    "program": "blastp",
                    "database": "pdb",
                    "sequence": ">3ADK\nABCD",
                    "sort": "pvalue",
                    }
                )
            self.assertEqual(
                self.proxy.construct_submit_request(
                    params = {
                        "program": "blastp",
                        "database": "pdb",
                        "email": "gb360@cam.ac.uk",
                        "sequence": ">3ADK\nABCD",
                        "sort": "pvalue",
                        "unknown": "unknown",
                        }
                    ),
                result  
                )
        
        else:
            self.fail( "Cannot perform test" )
        
        
    def testConstructStatusRequest(self):
        
        if self.proxy:
            self.assertEqual(
                self.proxy.construct_status_request(
                    jobid = "some jobid"
                    ),
                {
                    "jobid": "some jobid",
                    }
                )
            self.assertEqual(
                self.proxy.construct_status_request(
                    jobid = "some jobid",
                    params = {
                        "unknown": "unknown",
                        }
                    ),
                {
                    "jobid": "some jobid",
                    }
                )
        
        else:
            self.fail( "Cannot perform test" )
        
        
    def testConstructRetrieveRequest(self):
        
        if self.proxy:
            self.assertEqual(
                self.proxy.construct_retrieve_request(
                    jobid = "some jobid",
                    params = {
                        "format": "tooloutput",
                        }
                    ),
                {
                    "jobid": "some jobid",
                    "type": "tooloutput",
                    }
                )
            self.assertRaises(
                ValueError,
                self.proxy.construct_retrieve_request,
                "some jobid",
                {
                    }
                )
            self.assertEqual(
                self.proxy.construct_retrieve_request(
                    jobid = "some jobid",
                    params = {
                        "format": "tooloutput",
                        "unknown": "unknown",
                        }
                    ),
                {
                    "jobid": "some jobid",
                    "type": "tooloutput",
                    }
                )
            
        else:
            self.fail( "Cannot perform test" )
        
        
    def testEvaluateSubmitResponse(self):
        
        if self.proxy:
            self.assertEqual(
                self.proxy.evaluate_submit_response(
                    response = {
                        "jobid": "some jobid",
                        }
                    ),
                "some jobid"
                )
            self.assertEqual(
                self.proxy.evaluate_submit_response(
                    response = {
                        "jobid": "some jobid",
                        "unknown": "unknown",
                        }
                    ),
                "some jobid"
                )
            self.assertRaises(
                services.UnexpectedResponse,
                self.proxy.evaluate_submit_response,
                {
                    "unknown": "some jobid",
                    }
                )
            
        else:
            self.fail( "Cannot perform test" )
        
        
    def testEvaluateStatusResponse(self):
        
        if self.proxy:
            self.assert_(
                isinstance(
                    self.proxy.evaluate_status_response( { "status": "DONE", } ),
                    services.JobSuccess
                    )
                )
            self.assertRaises(
                services.UnexpectedResponse,
                self.proxy.evaluate_status_response,
                { "unknown": "DONE", }
                )
            self.assertRaises(
                services.UnexpectedResponse,
                self.proxy.evaluate_status_response,
                { "status": "finished", }
                )
            
        else:
            self.fail( "Cannot perform test" )
        
        
    def testEvaluateRetrieveResponse(self):
        
        if self.proxy:
            self.assertEqual(
                self.proxy.evaluate_retrieve_response(
                    response = {
                        "result": "some result",
                        }
                    ),
                "some result"
                )
            self.assertEqual(
                self.proxy.evaluate_retrieve_response(
                    response = {
                        "result": "some result",
                        "unknown": "unknown",
                        }
                    ),
                "some result"
                )
            self.assertRaises(
                services.UnexpectedResponse,
                self.proxy.evaluate_retrieve_response,
                {
                    "unknown": "some result",
                    }
                )
            
        else:
            self.fail( "Cannot perform test" )
        
        
class TestWUBlastEBI(unittest.TestCase):
    
    service = None
    sequence = ( ">3ADK:A|PDBID|CHAIN|SEQUENCE\n"
        + "XMEEKLKKSKIIFVVGGPGSGKGTQCEKIVQKYGYTHLSTGDLLRAEVSSGSARGKMLSEIMEKGQLVPLETVLDMLRDA\n"
        + "MVAKVDTSKGFLIDGYPREVKQGEEFERKIGQPTLLLYVDAGPETMTKRLLKRGETSGRVDDNEETIKKRLETYYKATEP\n"
        + "VIAFYEKRGIVRKVNAEGSVDDVFSQVCTHLDTLK\n" )
    
    def setUp(self):
        
        if not self.service:
            self.service = services.WUBlastEBI()
            
            
    def testExercise(self):
        
        jobid = self.service.submit(
            sequence = self.sequence
            )
        self.assert_( jobid )
        
        status = self.service.status( jobid = jobid )
        
        self.assert_( isinstance( status, services.JobBeingProcessed )
            or isinstance( status, services.JobProcessingCompleted ) )
        
        try:
            result = self.service.poll( jobid = jobid )
            
        except services.JobException, e:
            return
        
        print result
        
        self.assert_(
            isinstance( self.service.status( jobid = jobid ), services.JobSuccess )
            )
        self.assertEqual( self.service.retrieve( jobid = jobid ), result )
        
        
    def testConstruct(self):
        
        self.assertRaises(
            ValueError,
            services.WUBlastEBI,
            "blastm"
            )
        
        self.assertRaises(
            ValueError,
            services.WUBlastEBI,
            "blastp",
            "pdb",
            "lvalue"
            )
        
        
    def testConstructQuery(self):
        
        request = self.service.construct_submit_request(
            { 
                "sequence": self.sequence,
                "numal": 75,
                }
            )
        self.assertEqual( len( request[ "content" ] ), 1 )
        self.assertEqual( request[ "content" ][0]["type"], "sequence" )
        self.assertEqual( request[ "content" ][0]["content"], self.sequence )
        
        self.assert_( "numal" in request[ "params" ] )
        self.assertEqual( request[ "params" ][ "numal" ], 75 )
        
        self.assertEqual(
            request,
            self.service.construct_submit_request(
                {
                    "sequence": self.sequence,
                    "numal": 75,
                    "unknown": 10,
                    }
                )
            )
        
        request = self.service.construct_submit_request(
            {
                "dbfetch": "gi765",
                "numal": 100,
                }
            )
        self.assertEqual( len( request[ "content" ] ), 1 )
        self.assertEqual( request[ "content" ][0]["type"], "dbfetch" )
        self.assertEqual( request[ "content" ][0]["content"], "gi765" )
        self.assert_( "numal" in request[ "params" ] )
        self.assertEqual( request[ "params" ][ "numal" ], 100 )
        
        self.assertRaises(
            ValueError,
            self.service.construct_submit_request,
            { "numal": 50 }
            )
"""
            

suite_validator = unittest.TestLoader().loadTestsFromTestCase(
    TestValidator
    )
suite_filter_processor = unittest.TestLoader().loadTestsFromTestCase(
    TestFilterProcessor
    )
suite_extract_processor = unittest.TestLoader().loadTestsFromTestCase(
    TestExtractProcessor
    )
suite_gunzip_processor = unittest.TestLoader().loadTestsFromTestCase(
    TestGunzipProcessor
    )
suite_memory_memoizer = unittest.TestLoader().loadTestsFromTestCase(
    TestMemoryMemoizer
    )

alltests = unittest.TestSuite(
    [
        suite_validator,
        suite_filter_processor,
        suite_extract_processor,
        suite_gunzip_processor,
        suite_memory_memoizer,
        ]
    )
unittest.TextTestRunner( verbosity = 2 ).run( alltests )