"""
Defines the base API for the property estimator task calculation backend.
"""
import re
from enum import Enum

from propertyestimator import unit


class ComputeResources:
    """An object which stores how many of each type of computational resource
    (threads or gpu's) is available to a calculation worker.

    TODO: The use of the terminology here is questionable, and is used interchangable
          with process which may lead to some confusion.
    """

    class GPUToolkit(Enum):
        """An enumeration of the different GPU toolkits to
        make available to different calculations.
        """
        CUDA = 'CUDA'
        OpenCL = 'OpenCL'

    @property
    def number_of_threads(self):
        """int: The number of threads available to a calculation worker."""
        return self._number_of_threads

    @property
    def number_of_gpus(self):
        """int: The number of GPUs available to a calculation worker."""
        return self._number_of_gpus

    @property
    def preferred_gpu_toolkit(self):
        """ComputeResources.GPUToolkit: The preferred toolkit to use when running on GPUs."""
        return self._preferred_gpu_toolkit

    @property
    def gpu_device_indices(self):
        """str: The indices of the GPUs to run on. This is purely an internal
        implementation detail and should not be relied upon externally."""
        return self._gpu_device_indices

    def __init__(self, number_of_threads=1, number_of_gpus=0, preferred_gpu_toolkit=None):
        """Constructs a new ComputeResources object.

        Parameters
        ----------
        number_of_threads: int
            The number of threads available to a calculation worker.
        number_of_gpus: int
            The number of GPUs available to a calculation worker.
        preferred_gpu_toolkit: ComputeResources.GPUToolkit, optional
            The preferred toolkit to use when running on GPUs.
        """

        self._number_of_threads = number_of_threads
        self._number_of_gpus = number_of_gpus

        self._preferred_gpu_toolkit = preferred_gpu_toolkit
        # A workaround for when using a local cluster backend which is
        # strictly for internal purposes only for now.
        self._gpu_device_indices = None
        
        assert self._number_of_threads >= 0
        assert self._number_of_gpus >= 0

        assert self._number_of_threads > 0 or self._number_of_gpus > 0

        if self._number_of_gpus > 0:
            assert self._preferred_gpu_toolkit is not None

    def __getstate__(self):
        return {
            'number_of_threads': self.number_of_threads,
            'number_of_gpus': self.number_of_gpus,
            'preferred_gpu_toolkit': self.preferred_gpu_toolkit,
            '_gpu_device_indices': self._gpu_device_indices
        }

    def __setstate__(self, state):

        self._number_of_threads = state['number_of_threads']
        self._number_of_gpus = state['number_of_gpus']
        self._preferred_gpu_toolkit = state['preferred_gpu_toolkit']
        self._gpu_device_indices = state['_gpu_device_indices']

    def __eq__(self, other):
        return self.number_of_threads == other.number_of_threads and \
               self.number_of_gpus == other.number_of_gpus and \
               self.preferred_gpu_toolkit == other.preferred_gpu_toolkit

    def __ne__(self, other):
        return not self.__eq__(other)


class QueueWorkerResources(ComputeResources):
    """An extended resource object with properties specific
    to calculations which will run on queue based resources,
    such as LSF, PBS or SLURM.
    """

    @property
    def per_thread_memory_limit(self):
        """simtk.Quantity: The maximum amount of memory available to each thread, such that
        the total memory limit will be `per_cpu_memory_limit * number_of_threads`."""
        return self._per_thread_memory_limit

    @property
    def wallclock_time_limit(self):
        """str: The maximum amount of wall clock time that a worker can run for. This should
        be a string of the form `HH:MM` where HH is the number of hours and MM the number of minutes"""
        return self._wallclock_time_limit

    def __init__(self, number_of_threads=1, number_of_gpus=0, preferred_gpu_toolkit=None,
                 per_thread_memory_limit=1*unit.gigabytes, wallclock_time_limit="01:00"):
        """Constructs a new ComputeResources object.

        Notes
        -----
        Both the requested `number_of_threads` and the `number_of_gpus` must be less than
        or equal to the number of threads (/cpus/cores) and GPUs available to each compute
        node in the cluster respectively, such that a single worker is able to be accommodated
        by a single compute node.

        Parameters
        ----------
        per_thread_memory_limit: simtk.Quantity
            The maximum amount of memory available to each thread.
        wallclock_time_limit: str
            The maximum amount of wall clock time that a worker can run for. This should
            be a string of the form `HH:MM` where HH is the number of hours and MM the number
            of minutes
        """

        super().__init__(number_of_threads, number_of_gpus, preferred_gpu_toolkit)

        self._per_thread_memory_limit = per_thread_memory_limit
        self._wallclock_time_limit = wallclock_time_limit

        assert self._per_thread_memory_limit is not None

        assert (isinstance(self._per_thread_memory_limit, unit.Quantity) and
                unit.get_base_units(unit.byte)[-1] == unit.get_base_units(self._per_thread_memory_limit.units)[-1])

        assert self._per_thread_memory_limit > 0 * unit.byte

        assert wallclock_time_limit is not None
        assert isinstance(wallclock_time_limit, str) and len(wallclock_time_limit) > 0

        wallclock_pattern = re.compile(r'\d\d:\d\d')
        assert wallclock_pattern.match(wallclock_time_limit) is not None

    def __getstate__(self):

        base_dict = super(QueueWorkerResources, self).__getstate__()

        base_dict.update({
            'per_thread_memory_limit': self.number_of_threads,
            'wallclock_time_limit': self.number_of_threads,
        })

        return base_dict

    def __setstate__(self, state):
        super(QueueWorkerResources, self).__setstate__(state)

        self._per_thread_memory_limit = state['per_thread_memory_limit']
        self._wallclock_time_limit = state['wallclock_time_limit']

    def __eq__(self, other):
        return super(QueueWorkerResources, self).__eq__(other) and \
               self.per_thread_memory_limit == other.per_thread_memory_limit and \
               self.wallclock_time_limit == other.wallclock_time_limit

    def __ne__(self, other):
        return not self.__eq__(other)


class PropertyEstimatorBackend:
    """An abstract base representation of a property estimator backend. A backend is
    responsible for coordinating, distributing and running calculations on the
    available hardware. This may range from a single machine to a multinode cluster,
    but *not* accross multiple cluster or physical locations.

    Notes
    -----
    All estimator backend classes must inherit from this class, and must implement the
    `start`, `stop`, and `submit_task` method.
    """

    def __init__(self, number_of_workers=1, resources_per_worker=ComputeResources()):

        """Constructs a new PropertyEstimatorBackend object.

        Parameters
        ----------
        number_of_workers : int
            The number of works to run the calculations on. One worker
            can perform a single task (e.g run a simulation) at once.
        resources_per_worker: ComputeResources
            The number of resources to request per worker.
        """

        self._number_of_workers = number_of_workers
        self._resources_per_worker = resources_per_worker

    def _get_worker_resources_dict(self):
        """Get dict representation of the resources requested
        by a worker.

        Returns
        -------
        dict of str and int
        """
        return {
            'number_of_threads': self._resources_per_worker.number_of_threads,
            'number_of_gpus': self._resources_per_worker.number_of_gpus,
        }

    def start(self):
        """Start the calculation backend."""
        pass

    def stop(self):
        """Stop the calculation backend."""
        pass

    def submit_task(self, function, *args, **kwargs):
        """Submit a task to the compute resources
        managed by this backend.

        Parameters
        ----------
        function: function
            The function to run.

        Returns
        -------
        Future
            Returns a future object which will eventually point to the results
            of the submitted task.
        """
        pass
