import time

class Queue(object):
    """
    Process queue
    """
    
    def __init__(self, nproc, proxy):
        
        assert 0 < nproc
        self.proxy = proxy
        self.nproc = nproc
        self.waiting = []
        self.running = set()
        self.complete = set()
        self.interval = 0.01
        
        
    def submit(self, owner, calculation, extra):
        
        self.waiting.append( ( owner, calculation, extra ) )
        self.poll()
        
        
    def empty(self):
        
        return not self.running
    
    
    def full(self):
        
        return self.nproc <= len( self.running )
    
    
    def wait(self):
        
        while not self.complete:
            self.poll()
            time.sleep( self.interval )
            
            
    def join(self):
        
        while self.running:
            self.poll()
            time.sleep( self.interval )
    
    
    def poll(self):
        
        # Process finished jobs
        for job in list( self.running ):
            try:
                job.poll()
                
            except RuntimeError, e:
                from libtbx.utils import Sorry
                raise Sorry, e
            
            if job.finished():
                self.running.remove( job )
                self.complete.add( job )
            
        # Submit new jobs
        while len( self.running ) < self.nproc and self.waiting:
            ( owner, calculation, extra ) = self.waiting.pop( 0 )
            
            try:
                job = self.proxy( owner = owner, calculation = calculation, extra = extra )
                
            except RuntimeError, e:
                from libtbx.utils import Sorry
                raise Sorry, e
            
            self.running.add( job )
        
    
    def completed(self):
        
        return list( self.complete )
    
        
    def remove(self, job):
        
        self.complete.remove( job )


class Job(object):
    """
    Stores related job info
    """
    
    def __init__(self, owner, calculation, extra):
        
        self.owner = owner
        self.calculation = calculation
        self.extra = extra
        self.result = self.calculation.run()
        
        
    def poll(self):
        
        pass
    
    
    def finished(self):
        
        return True
    
    
    def __str__(self):
        
        return "(main thread) %s" % self.calculation


class MPJob(object):
    """
    A job being run on a multiprocessing process
    """
    
    MANAGER = None
    
    def __init__(self, owner, calculation, extra):
        
        self.owner = owner
        self.calculation = calculation
        self.extra = extra
        self.result = None
        
        if not self.is_initialized():
            self.initialize()
            
        import multiprocessing
        self.queue = self.MANAGER.Queue()
        self.process = multiprocessing.Process(
            target = mpjob_run,
            args = ( self.calculation, self.queue )
            )
        self.process.start()
        
        
    def poll(self):
        
        if self.process and not self.process.is_alive():
            self.result = self.queue.get()
            self.process.join()
            self.process = None
            
            
    def finished(self):
        
        return self.process is None
    
    
    def __str__(self):
        
        return "(multiprocessing) %s" % self.calculation
    
    
    @classmethod
    def initialize(cls):
        
        import multiprocessing
        cls.MANAGER = multiprocessing.Manager()
        
        
    @classmethod
    def is_initialized(cls):
        
        return cls.MANAGER is not None
        
        
def mpjob_run(calculation, queue):
    
    result = calculation.run()
    queue.put( result )


class TJob(object):
    """
    A job being run on a threading thread
    """
    
    def __init__(self, owner, calculation, extra):
        
        self.owner = owner
        self.calculation = calculation
        self.extra = extra
        self.result = None
        
        import threading
        import Queue
        self.queue = Queue.Queue()
        self.thread = threading.Thread( target = self.run )
        self.thread.start()
        
        
    def run(self):
        
        result = self.calculation.run()
        self.queue.put( result )
        
        
    def poll(self):
        
        if self.thread and not self.thread.is_alive():
            self.result = self.queue.get()
            self.thread.join()
            self.thread = None
            self.queue = None
            
            
    def finished(self):
        
        return self.thread is None
    
    
    def __str__(self):
        
        return "(threading) %s" % self.calculation
    
    
class QueueTarget(object):
    """
    An executable class
    """
    
    def __init__(self, calculation, queue, cpus):
        
        self.calculation = calculation
        self.queue = queue
        self.cpus = cpus
        
    
    def __call__(self):
    
        result = self.calculation.run( cpus = self.cpus )
        self.queue.put( result )
        

class QueueJob(object):
    """
    A job being run through a queuing system
    """
    
    TIMEOUT = 60
    
    def __init__(self, owner, calculation, extra, queue, job):
        
        self.owner = owner
        self.calculation = calculation
        self.extra = extra
        self.result = None
        
        self.queue = queue
        self.job = job
        self.job.start()
        
        
    def poll(self):
        
        from libtbx.queuing_system_utils import generic
        
        if self.job and not self.job.is_alive():
            try:
                self.result = self.queue.get( block = True, timeout = self.TIMEOUT )
                
            except generic.QueueEmptyException, e:
                raise RuntimeError, "Job %s: %s" % ( self.job.name, e )
              
            self.job.join()
            self.job = None
            self.queue = None
            
            
    def finished(self):
        
        return self.job is None
    
    
    def __str__(self):
        
        return "(queue) %s" % self.calculation
        
        
    @classmethod
    def Generic(cls, owner, calculation, extra, qinterface, factory, qslot_cpus):
        
        from libtbx.queuing_system_utils import generic
        queue = generic.Queue( identifier = "brt" )
        job = factory(
            name = "brt",
            target = QueueTarget(
                calculation = calculation,
                queue = queue,
                cpus = qslot_cpus
                ),
            qinterface = qinterface
            )
        return cls(
            owner = owner,
            calculation = calculation,
            extra = extra,
            queue = queue,
            job = job
            )