import logging
import multiprocessing
import sys
import time
import types
from enum import Enum
from multiprocessing import Semaphore, Process
from multiprocessing.managers import SharedMemoryManager
from typing import Optional

import numpy as np
from osgeo import gdal
from dill import dumps, loads

from . import constants
from .rounding import are_pixel_scales_equal_enough, round_up_pixels, round_down_pixels
from .window import Area, PixelScale, Window
from .backends import backend
from .backends.enumeration import operators as op

logger = logging.getLogger(__name__)
logger.setLevel(logging.WARNING)

class WindowOperation(Enum):
    NONE = 1
    UNION = 2
    INTERSECTION = 3
    LEFT = 4
    RIGHT = 5

class LayerConstant:
    def __init__(self, val):
        self.val = val

    def __str__(self):
        return str(self.val)

    def _eval(self, _area, _index, _step, _target_window):
        return self.val


class LayerMathMixin:

    def __add__(self, other):
        return LayerOperation(self, op.ADD, other, window_op=WindowOperation.UNION)

    def __sub__(self, other):
        return LayerOperation(self, op.SUB, other, window_op=WindowOperation.UNION)

    def __mul__(self, other):
        return LayerOperation(self, op.MUL, other, window_op=WindowOperation.INTERSECTION)

    def __truediv__(self, other):
        return LayerOperation(self, op.TRUEDIV, other, window_op=WindowOperation.INTERSECTION)

    def __floordiv__(self, other):
        return LayerOperation(self, op.FLOORDIV, other, window_op=WindowOperation.INTERSECTION)

    def __mod__(self, other):
        return LayerOperation(self, op.REMAINDER, other, window_op=WindowOperation.INTERSECTION)

    def __pow__(self, other):
        return LayerOperation(self, op.POW, other, window_op=WindowOperation.UNION)

    def __eq__(self, other):
        return LayerOperation(self, op.EQ, other, window_op=WindowOperation.INTERSECTION)

    def __ne__(self, other):
        return LayerOperation(self, op.NE, other, window_op=WindowOperation.UNION)

    def __lt__(self, other):
        return LayerOperation(self, op.LT, other, window_op=WindowOperation.UNION)

    def __le__(self, other):
        return LayerOperation(self, op.LE, other, window_op=WindowOperation.UNION)

    def __gt__(self, other):
        return LayerOperation(self, op.GT, other, window_op=WindowOperation.UNION)

    def __ge__(self, other):
        return LayerOperation(self, op.GE, other, window_op=WindowOperation.UNION)

    def __and__(self, other):
        return LayerOperation(self, op.AND, other, window_op=WindowOperation.INTERSECTION)

    def __or__(self, other):
        return LayerOperation(self, op.OR, other, window_op=WindowOperation.UNION)

    def _eval(self, area, index, step, target_window=None):
        try:
            window = self.window if target_window is None else target_window
            return self.read_array_for_area(area, 0, index, window.xsize, step)
        except AttributeError:
            return self.read_array_for_area(area, 0, index, target_window.xsize if target_window else 1, step)

    def nan_to_num(self, nan=0, posinf=None, neginf=None):
        return LayerOperation(
            self,
            op.NAN_TO_NUM,
            window_op=WindowOperation.NONE,
            copy=False,
            nan=nan,
            posinf=posinf,
            neginf=neginf,
        )

    def isin(self, test_elements):
        return LayerOperation(
            self,
            op.ISIN,
            window_op=WindowOperation.NONE,
            test_elements=test_elements,
        )

    def abs(self):
        return LayerOperation(
            self,
            op.ABS,
            window_op=WindowOperation.NONE,
        )

    def log(self):
        return LayerOperation(
            self,
            op.LOG,
            window_op=WindowOperation.NONE,
        )

    def log2(self):
        return LayerOperation(
            self,
            op.LOG2,
            window_op=WindowOperation.NONE,
        )

    def log10(self):
        return LayerOperation(
            self,
            op.LOG10,
            window_op=WindowOperation.NONE,
        )

    def exp(self):
        return LayerOperation(
            self,
            op.EXP,
            window_op=WindowOperation.NONE,
        )

    def exp2(self):
        return LayerOperation(
            self,
            op.EXP2,
            window_op=WindowOperation.NONE,
        )

    def clip(self, min=None, max=None): # pylint: disable=W0622
        # In the numpy 1 API np.clip(array) used a_max, a_min arguments and array.clip() used max and min as arguments
        # In numpy 2 they moved so that max and min worked on both, but still support a_max, and a_min on np.clip.
        # For now I'm only going to support the newer max/min everywhere notion, but I have to internally call
        # a_max, a_min so that yirgacheffe can work on older numpy installs.
        return LayerOperation(
            self,
            op.CLIP,
            window_op=WindowOperation.NONE,
            a_min=min,
            a_max=max,
        )

    def conv2d(self, weights):
        # A set of limitations that are just down to implementation time restrictions
        weights_width, weights_height = weights.shape
        if weights_width != weights_height:
            raise ValueError("Currently only square matrixes are supported for weights")
        padding = (weights_width - 1) / 2
        if padding != int(padding):
            raise ValueError("Currently weights dimensions must be odd")

        return LayerOperation(
            self,
            op.CONV2D,
            window_op=WindowOperation.NONE,
            buffer_padding=padding,
            weights=weights.astype(np.float32),
        )

    def numpy_apply(self, func, other=None):
        return LayerOperation(self, func, other)

    def shader_apply(self, func, other=None):
        return ShaderStyleOperation(self, func, other)

    def save(self, destination_layer, and_sum=False, callback=None, band=1):
        return LayerOperation(self).save(destination_layer, and_sum, callback, band)

    def parallel_save(self, destination_layer, and_sum=False, callback=None, parallelism=None, band=1):
        return LayerOperation(self).parallel_save(destination_layer, and_sum, callback, parallelism, band)

    def parallel_sum(self, callback=None, parallelism=None, band=1):
        return LayerOperation(self).parallel_sum(callback, parallelism, band)

    def sum(self):
        return LayerOperation(self).sum()

    def min(self):
        return LayerOperation(self).min()

    def max(self):
        return LayerOperation(self).max()


class LayerOperation(LayerMathMixin):

    @staticmethod
    def where(cond, a, b):
        return LayerOperation(
            cond,
            op.WHERE,
            rhs=a,
            other=b
        )

    @staticmethod
    def maximum(a, b):
        return LayerOperation(
            a,
            op.MAXIMUM,
            b,
            window_op=WindowOperation.UNION,
        )

    @staticmethod
    def minimum(a, b):
        return LayerOperation(
            a,
            op.MINIMUM,
            rhs=b,
            window_op=WindowOperation.UNION,
        )

    def __init__(
        self,
        lhs,
        operator=None,
        rhs=None,
        other=None,
        window_op=WindowOperation.NONE,
        buffer_padding=0,
        **kwargs
    ):
        self.ystep = constants.YSTEP
        self.kwargs = kwargs
        self.window_op = window_op
        self.buffer_padding = buffer_padding

        if lhs is None:
            raise ValueError("LHS on operation should not be none")
        self.lhs = lhs

        self.operator = operator

        if rhs is not None:
            if backend.isscalar(rhs):
                self.rhs = LayerConstant(rhs)
            elif isinstance(rhs, (backend.array_t)):
                if rhs.shape == ():
                    self.rhs = LayerConstant(rhs.item())
                else:
                    raise ValueError("Numpy arrays are no allowed")
            else:
                if not are_pixel_scales_equal_enough([lhs.pixel_scale, rhs.pixel_scale]):
                    raise ValueError("Not all layers are at the same pixel scale")
                self.rhs = rhs
        else:
            self.rhs = None

        if other is not None:
            if backend.isscalar(other):
                self.other = LayerConstant(other)
            elif isinstance(other, (backend.array_t)):
                if other.shape == ():
                    self.rhs = LayerConstant(other.item())
                else:
                    raise ValueError("Numpy arrays are no allowed")
            else:
                if not are_pixel_scales_equal_enough([lhs.pixel_scale, other.pixel_scale]):
                    raise ValueError("Not all layers are at the same pixel scale")
                self.other = other
        else:
            self.other = None

    def __str__(self):
        try:
            return f"({self.lhs} {self.operator} {self.rhs})"
        except AttributeError:
            try:
                return f"({self.operator} {self.lhs})"
            except AttributeError:
                return str(self.lhs)

    def __len__(self):
        return len(self.lhs)

    def __getstate__(self) -> object:
        odict = self.__dict__.copy()
        if isinstance(self.operator, types.LambdaType):
            odict['operator_dill'] = dumps(self.operator)
            del odict['operator']
        return odict

    def __setstate__(self, state):
        if 'operator_dill' in state:
            state['operator'] = loads(state['operator_dill'])
            del state['operator_dill']
        self.__dict__.update(state)

    @property
    def area(self) -> Area:
        # The type().__name__ here is to avoid a circular import dependancy
        lhs_area = self.lhs.area if not type(self.lhs).__name__ == "ConstantLayer" else None
        try:
            rhs_area = self.rhs.area if not type(self.rhs).__name__ == "ConstantLayer" else None
        except AttributeError:
            rhs_area = None
        try:
            other_area = self.other.area if not type(self.other).__name__ == "ConstantLayer" else None
        except AttributeError:
            other_area = None

        all_areas = []
        if lhs_area is not None:
            all_areas.append(lhs_area)
        if rhs_area is not None:
            all_areas.append(rhs_area)
        if other_area is not None:
            all_areas.append(other_area)

        match self.window_op:
            case WindowOperation.NONE:
                return all_areas[0]
            case WindowOperation.LEFT:
                return lhs_area
            case WindowOperation.RIGHT:
                assert rhs_area is not None
                return rhs_area
            case WindowOperation.INTERSECTION:
                intersection = Area(
                    left=max(x.left for x in all_areas),
                    top=min(x.top for x in all_areas),
                    right=min(x.right for x in all_areas),
                    bottom=max(x.bottom for x in all_areas)
                )
                if (intersection.left >= intersection.right) or (intersection.bottom >= intersection.top):
                    raise ValueError('No intersection possible')
                return intersection
            case WindowOperation.UNION:
                return Area(
                    left=min(x.left for x in all_areas),
                    top=max(x.top for x in all_areas),
                    right=max(x.right for x in all_areas),
                    bottom=min(x.bottom for x in all_areas)
                )

    @property
    def pixel_scale(self) -> PixelScale:
        # Because we test at construction that pixel scales for RHS/other are roughly equal,
        # I believe this should be sufficient...
        try:
            pixel_scale = self.lhs.pixel_scale
        except AttributeError:
            pixel_scale = None

        if pixel_scale is None:
            return self.rhs.pixel_scale
        return pixel_scale

    @property
    def window(self) -> Window:
        pixel_scale = self.pixel_scale
        area = self.area

        return Window(
            xoff=round_down_pixels(area.left / pixel_scale.xstep, pixel_scale.xstep),
            yoff=round_down_pixels(area.top / (pixel_scale.ystep * -1.0), pixel_scale.ystep * -1.0),
            xsize=round_up_pixels(
                (area.right - area.left) / pixel_scale.xstep, pixel_scale.xstep
            ),
            ysize=round_up_pixels(
                (area.top - area.bottom) / (pixel_scale.ystep * -1.0),
                (pixel_scale.ystep * -1.0)
            ),
        )

    @property
    def datatype(self):
        # TODO: Work out how to indicate type promotion via numpy
        return self.lhs.datatype

    @property
    def projection(self):
        try:
            return self.lhs.projection
        except AttributeError:
            return self.rhs.projection

    def _eval(self, area: Area, index: int, step: int, target_window:Optional[Window]=None):

        if self.buffer_padding:
            if target_window:
                target_window = target_window.grow(self.buffer_padding)
            pixel_scale = self.pixel_scale
            area = area.grow(self.buffer_padding * pixel_scale.xstep)
            # The index doesn't need updating because we updated area/window
            step += (2 * self.buffer_padding)

        lhs_data = self.lhs._eval(area, index, step, target_window)

        if self.operator is None:
            return lhs_data

        try:
            operator = backend.operator_map[self.operator]
        except KeyError:
            # Handles things like `numpy_apply` where a custom operator is provided
            operator = self.operator

        if self.other is not None:
            assert self.rhs is not None
            rhs_data = self.rhs._eval(area, index, step, target_window)
            other_data = self.other._eval(area, index, step, target_window)
            return operator(lhs_data, rhs_data, other_data, **self.kwargs)

        if self.rhs is not None:
            rhs_data = self.rhs._eval(area, index, step, target_window)
            return operator(lhs_data, rhs_data, **self.kwargs)

        return operator(lhs_data, **self.kwargs)

    def sum(self):
        # The result accumulator is float64, and for precision reasons
        # we force the sum to be done in float64 also. Otherwise we
        # see variable results depending on chunk size, as different parts
        # of the sum are done in different types.
        res = 0.0
        computation_window = self.window
        for yoffset in range(0, computation_window.ysize, self.ystep):
            step=self.ystep
            if yoffset+step > computation_window.ysize:
                step = computation_window.ysize - yoffset
            chunk = self._eval(self.area, yoffset, step, computation_window)
            res += backend.sum_op(chunk)
        return res

    def min(self):
        res = None
        computation_window = self.window
        for yoffset in range(0, computation_window.ysize, self.ystep):
            step=self.ystep
            if yoffset+step > computation_window.ysize:
                step = computation_window.ysize - yoffset
            chunk = self._eval(self.area, yoffset, step, computation_window)
            chunk_min = backend.min_op(chunk)
            if (res is None) or (res > chunk_min):
                res = chunk_min
        return res

    def max(self):
        res = None
        computation_window = self.window
        for yoffset in range(0, computation_window.ysize, self.ystep):
            step=self.ystep
            if yoffset+step > computation_window.ysize:
                step = computation_window.ysize - yoffset
            chunk = self._eval(self.area, yoffset, step, computation_window)
            chunk_max = backend.max_op(chunk)
            if (res is None) or (chunk_max > res):
                res = chunk_max
        return res

    def save(self, destination_layer, and_sum=False, callback=None, band=1):
        """
        Calling save will write the output of the operation to the provied layer.
        If you provide sum as true it will additionall compute the sum and return that.
        """

        if destination_layer is None:
            raise ValueError("Layer is required")
        try:
            band = destination_layer._dataset.GetRasterBand(band)
        except AttributeError as exc:
            raise ValueError("Layer must be a raster backed layer") from exc

        computation_window = self.window
        destination_window = destination_layer.window

        if (computation_window.xsize != destination_window.xsize) \
                or (computation_window.ysize != destination_window.ysize):
            raise ValueError("Destination raster window size does not match input raster window size.")

        total = 0.0

        for yoffset in range(0, computation_window.ysize, self.ystep):
            if callback:
                callback(yoffset / computation_window.ysize)
            step=self.ystep
            if yoffset+step > computation_window.ysize:
                step = computation_window.ysize - yoffset
            chunk = self._eval(self.area, yoffset, step, computation_window)
            if isinstance(chunk, (float, int)):
                chunk = backend.full((step, destination_window.xsize), chunk)
            band.WriteArray(
                backend.demote_array(chunk),
                destination_window.xoff,
                yoffset + destination_window.yoff,
            )
            if and_sum:
                total += backend.sum_op(chunk)
        if callback:
            callback(1.0)

        return total if and_sum else None

    def _parallel_worker(self, index, shared_mem, sem, np_dtype, width, input_queue, output_queue):
        arr = np.ndarray((self.ystep, width), dtype=np_dtype, buffer=shared_mem.buf)

        try:
            while True:
                # We acquire the lock so we know we have somewhere to put the
                # result before we take work. This is because in practice
                # it seems the writing to GeoTIFF is the bottleneck, and
                # we had workers taking a task, then waiting for somewhere to
                # write to for ages when other workers were exiting because there
                # was nothing to do.
                sem.acquire()

                task = input_queue.get()
                if task is None:
                    sem.release()
                    output_queue.put(None)
                    break
                yoffset, step = task

                result = self._eval(self.area, yoffset, step)
                backend.eval_op(result)

                arr[:step] = backend.demote_array(result)

                output_queue.put((index, yoffset, step))

        except Exception as e: # pylint: disable=W0718
            logger.exception(e)
            sem.release()
            output_queue.put(None)

    def _park(self):
        try:
            self.lhs._park()
        except AttributeError:
            pass
        try:
            self.rhs._park()
        except AttributeError:
            pass
        try:
            self.other._park()
        except AttributeError:
            pass

    def _parallel_save(self, destination_layer, and_sum=False, callback=None, parallelism=None, band=1):
        assert (destination_layer is not None) or and_sum
        computation_window = self.window
        if destination_layer is not None:
            try:
                band = destination_layer._dataset.GetRasterBand(band)
            except AttributeError as exc:
                raise ValueError("Layer must be a raster backed layer") from exc

            destination_window = destination_layer.window

            if (computation_window.xsize != destination_window.xsize) \
                    or (computation_window.ysize != destination_window.ysize):
                raise ValueError("Destination raster window size does not match input raster window size.")

            np_dtype = {
                gdal.GDT_Byte:    np.dtype('byte'),
                gdal.GDT_Float32: np.dtype('float32'),
                gdal.GDT_Float64: np.dtype('float64'),
                gdal.GDT_Int8:    np.dtype('int8'),
                gdal.GDT_Int16:   np.dtype('int16'),
                gdal.GDT_Int32:   np.dtype('int32'),
                gdal.GDT_Int64:   np.dtype('int64'),
                gdal.GDT_UInt16:  np.dtype('uint16'),
                gdal.GDT_UInt32:  np.dtype('uint32'),
                gdal.GDT_UInt64:  np.dtype('uint64'),
            }[band.DataType]
        else:
            band = None
            np_dtype = np.dtype('float64')

        # The parallel save will cause a fork on linux, so we need to
        # remove all SWIG references
        self._park()

        total = 0.0

        with multiprocessing.Manager() as manager:
            with SharedMemoryManager() as smm:

                worker_count = parallelism or multiprocessing.cpu_count()
                work_blocks = len(range(0, computation_window.ysize, self.ystep))
                worker_count = min(work_blocks, worker_count)

                mem_sem_cast = []
                for i in range(worker_count):
                    shared_buf = smm.SharedMemory(size=np_dtype.itemsize * self.ystep * computation_window.xsize)
                    cast_buf = np.ndarray((self.ystep, computation_window.xsize), dtype=np_dtype, buffer=shared_buf.buf)
                    cast_buf[:] = np.zeros((self.ystep, computation_window.xsize), np_dtype)
                    mem_sem_cast.append((shared_buf, Semaphore(), cast_buf))

                source_queue = manager.Queue()
                result_queue = manager.Queue()

                for yoffset in range(0, computation_window.ysize, self.ystep):
                    step = ((computation_window.ysize - yoffset)
                        if yoffset+self.ystep > computation_window.ysize
                        else self.ystep)
                    source_queue.put((
                        yoffset,
                        step
                    ))
                for _ in range(worker_count):
                    source_queue.put(None)

                if callback:
                    callback(0.0)

                workers = [Process(target=self._parallel_worker, args=(
                    i,
                    mem_sem_cast[i][0],
                    mem_sem_cast[i][1],
                    np_dtype,
                    computation_window.xsize,
                    source_queue,
                    result_queue
                )) for i in range(worker_count)]
                for worker in workers:
                    worker.start()

                sentinal_count = len(workers)
                retired_blocks = 0
                while sentinal_count > 0:
                    res = result_queue.get()
                    if res is None:
                        sentinal_count -= 1
                        continue
                    index, yoffset, step = res
                    _, sem, arr = mem_sem_cast[index]
                    if band:
                        band.WriteArray(
                            arr[0:step],
                            destination_window.xoff,
                            yoffset + destination_window.yoff,
                        )
                    if and_sum:
                        total += np.sum(np.array(arr[0:step]).astype(np.float64))
                    sem.release()
                    retired_blocks += 1
                    if callback:
                        callback(retired_blocks / work_blocks)

                processes = workers
                while processes:
                    candidates = [x for x in processes if not x.is_alive()]
                    for candidate in candidates:
                        candidate.join()
                        if candidate.exitcode:
                            for victim in processes:
                                victim.kill()
                            sys.exit(candidate.exitcode)
                        processes.remove(candidate)
                    time.sleep(0.01)

        return total if and_sum else None

    def parallel_save(self, destination_layer, and_sum=False, callback=None, parallelism=None, band=1):
        if destination_layer is None:
            raise ValueError("Layer is required")
        return self._parallel_save(destination_layer, and_sum, callback, parallelism, band)

    def parallel_sum(self, callback=None, parallelism=None, band=1):
        return self._parallel_save(None, True, callback, parallelism, band)

class ShaderStyleOperation(LayerOperation):

    def _eval(self, area, index, step, target_window=None):
        if target_window is None:
            target_window = self.window
        lhs_data = self.lhs._eval(area, index, step, target_window)
        if self.rhs is not None:
            rhs_data = self.rhs._eval(area, index, step, target_window)
        else:
            rhs_data = None

        # Constant results make this a bit messier. Might in future
        # be nicer to promote them to arrays sooner?
        if isinstance(lhs_data, (int, float)):
            if rhs_data is None:
                return self.operator(lhs_data, **self.kwargs)
            if isinstance(rhs_data, (int, float)):
                return self.operator(lhs_data, rhs_data, **self.kwargs)
            else:
                result = np.empty_like(rhs_data)
        else:
            result = np.empty_like(lhs_data)

        window = self.window
        for yoffset in range(step):
            for xoffset in range(window.xsize):
                try:
                    lhs_val = lhs_data[yoffset][xoffset]
                except TypeError:
                    lhs_val = lhs_data
                if rhs_data is not None:
                    try:
                        rhs_val = rhs_data[yoffset][xoffset]
                    except TypeError:
                        rhs_val = rhs_data
                    result[yoffset][xoffset] = self.operator(lhs_val, rhs_val, **self.kwargs)
                else:
                    result[yoffset][xoffset] = self.operator(lhs_val, **self.kwargs)

        return result

# We provide these module level accessors as it's often nicer to write `log(x/y)` rather than `(x/y).log()`
where = LayerOperation.where
minumum = LayerOperation.minimum
maximum = LayerOperation.maximum
clip = LayerOperation.clip
log = LayerOperation.log
log2 = LayerOperation.log2
log10 = LayerOperation.log10
exp = LayerOperation.exp
exp2 = LayerOperation.exp2
nan_to_num = LayerOperation.nan_to_num
isin = LayerOperation.isin
abs = LayerOperation.abs # pylint: disable=W0622
