# -*- coding: utf-8 -*-

from abc import ABCMeta, abstractmethod
from copy import copy

from joblib import Parallel, delayed
from shapely.affinity import translate

from .timing import WorkflowTiming
from .image import Image
from .information import ChainInformation, WorkflowInformation
from .logging import Loggable, SilentLogger
from .util import batch_split

__author__ = "Romain Mormont <romainmormont@hotmail.com>"
__version__ = "0.1"


class ImageProvider(Loggable):
    """An image provider is a class which generates a set of images to be processed by a workflow chein
    """
    __metaclass__ = ABCMeta

    def __init__(self, silent_fail=False, logger=SilentLogger()):
        """Constructs instances of ImageProvider

        Parameters
        ----------
        silent_fail: bool
            True for putting the image provider in silent fail mode. In this situation, when an image cannot be
            extracted, the provider simply ignore the error and skip the image. Otherwise, when set to False,
            the provider raises an error when an image extraction fails.
        logger: Logger (optional, default: SilentLogger)
            A logger object
        """
        Loggable.__init__(self, logger)
        self._silent_fail = silent_fail

    @abstractmethod
    def get_images(self):
        """
        Return the images to be processed by instances of the workflow

        Returns
        -------
        images: iterable (subtype: Image)
            An iterable of images

        Exceptions
        ----------
        ImageExtractionException:
            Raised when an image cannot be extracted. This error is never raised when the image provider is in
            silent_fail mode. In this situation, the provider fetches as many images as possible and returns only the
            successfully fetched images in the array.
        """
        pass

    def __iter__(self):
        """Make ImageProvider iterable"""
        for image in self.get_images():
            yield image


def _translate(window, workflow_information):
    """Translate the polygons contained in the workflow information object so that their reference system
    is the base image and not the window.
    """
    offset_x, offset_y = window.abs_offset
    polygons = workflow_information.polygons
    for i, polygon in enumerate(polygons):
        polygons[i] = translate(polygon, offset_x, offset_y)


def _execute_workflow(workflow, windows, do_translate, logger=SilentLogger()):
    """Execute a workflow on the given windows

    Parameters
    ----------
    workflow: Workflow
        THe workflow to use for processing the images
    windows: iterable (subtype: Image|ImageWindow)
        The windows to be processed
    do_translate: bool
        True if the polygons must be translated in the reference system of the base image, False otherwise.
    logger: Logger (optional, default: SilentLogger)
        The logger to use for displaying the results
    """
    workflow_information = None
    window_count = len(windows)
    for i, window in enumerate(windows):
        logger.info("WorkflowExecutor : process image {}/{}.".format(i + 1, window_count))
        returned_info = workflow.process(window)
        if do_translate:
            _translate(window, returned_info)
        workflow_information = returned_info if workflow_information is None else workflow_information.merge(returned_info)
    return workflow_information


class WorkflowExecutor(Loggable):
    """A class for encapsulating the execution of a workflow on a set of windows. This class can handle the translation
     of polygon in the full image reference system (optional, see translate parameter at construction).
    """
    __metaclass__ = ABCMeta

    def __init__(self, workflow, translate=True, logger=SilentLogger(), n_jobs=1):
        """Constructs a WorkflowExecutor for the given workflow object

        Parameters
        ----------
        workflow: Workflow
            The workflow object for processing windows
        translate: boolean
            True for translating polygons in the full coordinate system.
        logger: Logger (optional, default: a SilentLogger instance)
            A logger object
        n_jobs: int (optional, default: 1)
            Number of jobs to execute the workflow on the underlying image windows
        """
        Loggable.__init__(self, logger)
        self._workflow = workflow
        self._do_translate = translate
        self._n_jobs = n_jobs
        self._pool = None

    def execute(self, windows):
        """Execute the workflow with the images generated by get_images

        Parameters
        ----------
        windows: iterable (subtype: ImageWindow)
            The image windows to be processed by the executor

        Returns
        -------
        workflow_information: WorkflowInformation
            The workflow information object containing the data generated by all the executions of the stored workflow
        """
        if self._pool is None:
            self._pool = Parallel(n_jobs=self._n_jobs)
        window_count = len(windows)
        self.logger.info("WorkflowExecutor: {} window(s) to process.".format(window_count))
        self.logger.info("WorkflowExecutor: translation {}.".format("enabled" if self._do_translate else "disabled"))

        # prepare parallel processing
        window_batches = batch_split(self._n_jobs, windows)

        results = self._pool(
            delayed(_execute_workflow)(
                self._workflow, batch,
                self._do_translate,
                logger=self.logger)
                for batch in window_batches
        )

        # merge produced workflow information
        if len(results) == 0:
            return WorkflowInformation([], [], WorkflowTiming())

        workflow_information = results[0]
        for workflow_info in results[1:]:
            workflow_information = workflow_information.merge(workflow_info)

        return workflow_information

    def __getstate__(self):
        """Make workflow executor serializable"""
        d = copy(self.__dict__)
        d._pool = None
        return d


class PolygonFilter(object):
    """An interface to be implemented by classes which filter polygons to be passed to workflow executors"""
    __metaclass__ = ABCMeta

    @abstractmethod
    def filter(self, chain_information):
        """Filter the polygons from in the workflow information collection according to some rules.

        Parameters
        ----------
        chain_information: ChainInformation
            The chain information object containing the previously generated polygons

        Returns
        -------
        polygons: iterable (subtype: shapely.geometry.Polygon)
            The filtered polygons
        """
        pass

    def windows(self, image, chain_information):
        """Return the windows boxing the polygons filtered by self.filter()

        Parameters
        ----------
        image: Image
            The image from which the windows can be extracted
        chain_information: ChainInformation
            The chain information object containing the previously generated polygons

        Returns
        -------
        windows: iterable (subtype: ImageWindow)
            The generated windows
        """
        polygons = self.filter(chain_information)
        return [image.window_from_polygon(polygon, mask=True) for polygon in polygons]


class DefaultFilter(PolygonFilter):
    """A polygon filter which filters no polygon"""
    def filter(self, chain_information):
        return [p for p in chain_information.polygons]


class WorkflowChain(Loggable):
    """A class for chaining the execution of several workflows on images
    """

    def __init__(self, workflow, executors, filters, workflow_labels=None, logger=SilentLogger()):
        """Constructor for WorkflowChain objects

        Parameters
        ----------
        workflow: Workflow
            The first workflow to apply to the full image
        executors: iterable (subtype: WorkflowExecutor, size: N)
            The executors encapsulating the execution of the workflows to apply after the first one;
        filters: iterable (subtype: PolygonFilter, size: N)
            The filters to use for filtering polygons before passing their windows to a workflow executor.
            Especially, the ith filter corresponds to the ith executor in 'executors'
        workflow_labels: iterable (subtype: hashable, size: N+1, optional, default: range(1, N+2))
            The labels identifying the various recorded workflows. The first one corresponds to the first workflow.
        logger: Logger (optional, default: a SilentLogger instance)
            A logger object
        """
        Loggable.__init__(self, logger)
        self._first_workflow = workflow
        self._executors = executors
        self._filters = filters
        self._labels = workflow_labels if workflow_labels is not None else range(1, len(filters)+2)
        self._chain_information = ChainInformation()

    def process(self, image):
        """Apply the processing chain on a given image.

        Parameters
        ----------
        image: Image
            The image to process

        Returns
        -------
        chain_information: ChainInformation
            A chain information object containing the generated objects
        """
        # Run 1st workflow and save its results in the collection
        self.logger.info("WorkflowChain : start chain on image ({} x {} x {}).".format(image.height, image.width,
                                                                                       image.channels))
        workflow_information = self._first_workflow.process(image)
        chain_info = ChainInformation()
        chain_info.append(self._labels[0], workflow_information)

        # Run subsequent workflows
        for i, executor in enumerate(self._executors):
            self.logger.info("WorkflowChain : run workflow executor {}.".format(i + 1))
            filtered = self._filters[i].windows(image, chain_info)
            workflow_information = executor.execute(filtered)
            chain_info.append(self._labels[i+1], workflow_information)

        return chain_info

    def process_batch(self, images):
        """Apply the processing chain to a bunch of images

        Parameters
        ----------
        images: iterable (subtype: Image, size: N)
            The image to process (can be a list, or an ImageProvider)

        Returns
        -------
        chain_infos: iterable (subtype: ChainInformation, size: N)
            The chain information objects generated by processing the given images
        """
        return [self.process(image) for image in images]
