"""
    Optimal Image Subtraction (OIS) module
    --------------------------------------

    A collection of tools to perform optimal image differencing
    for the Transient Optical Robotic Observatory of the South (TOROS).

    ### Usage example (from python):

        >>> import ois
        >>> difference, optimalImage, optimalKernel, background =
                                   ois.optimal_system(image, referenceImage)

    (conv_image is the least square optimal approximation to image)

    See optimalkernelandbkg docstring for more options.

    ### Command line arguments:
    * -h, --help: Prints this help and exits.
    * -v, --version: Prints version information and exits.

    (c) Martin Beroiz

    email: <martinberoiz@gmail.com>

    University of Texas at San Antonio
"""


__version__ = '0.1.4'

import numpy as np
from scipy import signal
from scipy import ndimage


class EvenSideKernelError(ValueError):
    pass


def _has_mask(image):
    is_masked_array = isinstance(image, np.ma.MaskedArray)
    if is_masked_array and isinstance(image.mask, np.ndarray):
        return True
    return False


class SubtractionStrategy(object):

    def __init__(self, image, refimage, kernelshape, bkgdegree):
        self.k_shape = kernelshape

        # Check here for dimensions
        if image.ndim != 2:
            raise ValueError("Wrong dimensions for image")
        if refimage.ndim != 2:
            raise ValueError("Wrong dimensions for refimage")
        if image.shape != refimage.shape:
            raise ValueError("Images have different shapes")
        self.h, self.w = image.shape
        self.image, self.refimage, self.badpixmask =\
            self.separate_data_mask(image, refimage)

        self.coeffs = None
        self.bkgdegree = bkgdegree
        self.optimal_image = None
        self.background = None
        self.kernel = None
        self.difference = None

    def separate_data_mask(self, image, refimage):
        def ret_data(image):
            if isinstance(image, np.ma.MaskedArray):
                image_data = image.data
            else:
                image_data = image
            return image_data
        badpixmask = None
        if _has_mask(refimage):
            badpixmask = ndimage.binary_dilation(
                refimage.mask.astype('uint8'),
                structure=np.ones(self.k_shape)).astype('bool')
            if _has_mask(image):
                badpixmask += image.mask
        elif _has_mask(image):
            badpixmask = image.mask
        return ret_data(image), ret_data(refimage), badpixmask

    def coeffstobackground(self, coeffs):
        "Given a list of coefficients, return an array with the polynomial background"
        bkgdeg = int(-1.5 + 0.5 * np.sqrt(9 + 8 * (len(coeffs) - 1)))
        h, w = self.h, self.w
        y, x = np.mgrid[:h, :w]
        allxs = [pow(x, i) for i in range(bkgdeg + 1)]
        allys = [pow(y, i) for i in range(bkgdeg + 1)]
        mybkg = np.zeros(self.image.shape)
        ind = 0
        for i, anX in enumerate(allxs):
            for aY in allys[:bkgdeg + 1 - i]:
                mybkg += coeffs[ind] * anX * aY
                ind += 1
        return mybkg

    def get_cmatrices_background(self):
        h, w = self.refimage.shape
        y, x = np.mgrid[:h, :w]
        allxs = [pow(x, i) for i in range(self.bkgdegree + 1)]
        allys = [pow(y, i) for i in range(self.bkgdegree + 1)]
        bkg_c = [anX * aY for i, anX in enumerate(allxs)
                 for aY in allys[:self.bkgdegree + 1 - i]]
        return bkg_c

    def get_coeffs():
        "Override this function to solve the matrix minimization system"
        return self.coeffs

    def get_optimal_image(self):
        if self.optimal_image is not None:
            return self.optimal_image
        opt_image = signal.convolve2d(
            self.refimage, self.get_kernel(), mode='same')
        if self.bkgdegree is not None:
            opt_image += self.get_background()
        if self.badpixmask is not None:
            self.optimal_image = np.ma.array(opt_image, mask=self.badpixmask)
        else:
            self.optimal_image = opt_image
        return self.optimal_image

    def get_background(self):
        if self.background is not None:
            return self.background
        if self.bkgdegree is not None:
            bkgdof = (self.bkgdegree + 1) * (self.bkgdegree + 2) // 2
            coeffs = self.get_coeffs()
            self.background = self.coeffstobackground(coeffs[-bkgdof:])
        else:
            self.background = np.zeros(self.image.shape)
        return self.background

    def get_kernel(self):
        "Override this method to return the kernel"
        return self.kernel

    def get_difference(self):
        if self.difference is not None:
            return self.difference
        opt_image = self.get_optimal_image()
        if self.badpixmask is not None:
            self.difference = np.ma.array(
                self.image - opt_image, mask=self.badpixmask)
        else:
            self.difference = self.image - opt_image
        return self.difference


class AlardLuptonStrategy(SubtractionStrategy):

    def __init__(self, image, refimage, kernelshape, bkgdegree, gausslist):
        super(AlardLuptonStrategy, self).\
            __init__(image, refimage, kernelshape, bkgdegree)
        if gausslist is None:
            self.gausslist = [{}]
        else:
            self.gausslist = gausslist
        self.clean_gausslist()

    def gauss(self, center, sx, sy):
        h, w = self.k_shape
        x0, y0 = center
        x, y = np.meshgrid(list(range(w)), list(range(h)))
        k = np.exp(-0.5 * ((x - x0) ** 2 / sx ** 2 + (y - y0) ** 2 / sy ** 2))
        norm = k.sum()
        return k / norm

    def clean_gausslist(self):
        for agauss in self.gausslist:
            if 'center' not in agauss:
                h, w = self.k_shape
                agauss['center'] = ((h - 1) // 2., (w - 1) // 2.)
            if 'modPolyDeg' not in agauss:
                agauss['modPolyDeg'] = 2
            if 'sx' not in agauss:
                agauss['sx'] = 2.
            if 'sy' not in agauss:
                agauss['sy'] = 2.

    def get_cmatrices(self):
        kh, kw = self.k_shape
        v, u = np.mgrid[:kh, :kw]
        c = []
        for aGauss in self.gausslist:
            n = aGauss['modPolyDeg'] + 1
            allus = [pow(u, i) for i in range(n)]
            allvs = [pow(v, i) for i in range(n)]
            gaussk = self.gauss(center=aGauss['center'],
                                sx=aGauss['sx'], sy=aGauss['sy'])
            newc = [signal.convolve2d(self.refimage, gaussk * aU * aV,
                                      mode='same')
                    for i, aU in enumerate(allus)
                    for aV in allvs[:n - i]
                    ]
            c.extend(newc)
        return c

    def get_kernel(self):
        if self.kernel is not None:
            return self.kernel
        nkcoeffs = 0
        for aGauss in self.gausslist:
            n = aGauss['modPolyDeg'] + 1
            nkcoeffs += n * (n + 1) // 2
        coeffs = self.get_coeffs()
        kcoeffs = coeffs[:nkcoeffs]
        kh, kw = self.k_shape
        v, u = np.mgrid[:kh, :kw]
        kernel = np.zeros((kh, kw))
        for aGauss in self.gausslist:
            n = aGauss['modPolyDeg'] + 1
            allus = [pow(u, i) for i in range(n)]
            allvs = [pow(v, i) for i in range(n)]
            gaussk = self.gauss(center=aGauss['center'],
                                sx=aGauss['sx'], sy=aGauss['sy'])
            ind = 0
            for i, aU in enumerate(allus):
                for aV in allvs[:n - i]:
                    kernel += kcoeffs[ind] * aU * aV
                    ind += 1
            kernel *= gaussk
        self.kernel = kernel
        return self.kernel        

    def get_coeffs(self):
        if self.coeffs is not None:
            return self.coeffs
        c = self.get_cmatrices()
        if self.bkgdegree is not None:
            c_bkg = self.get_cmatrices_background()
            c.extend(c_bkg)

        if self.badpixmask is None:
            m = np.array([[(ci * cj).sum() for ci in c] for cj in c])
            b = np.array([(self.image * ci).sum() for ci in c])
        else:
            # These next two lines take most of the computation time
            m = np.array([[(ci * cj)[~self.badpixmask].sum()
                         for ci in c] for cj in c])
            b = np.array([(self.image * ci)[~self.badpixmask].sum()
                         for ci in c])
        self.coeffs = np.linalg.solve(m, b)
        return self.coeffs


class BramichStrategy(SubtractionStrategy):

    def get_cmatrices(self):
        kh, kw = self.k_shape
        h, w = self.refimage.shape
        c = []
        for i in range(kh):
            for j in range(kw):
                cij = np.zeros(self.refimage.shape)
                max_r = min(h, h - kh // 2 + i)
                min_r = max(0, i - kh // 2)
                max_c = min(w, w - kw // 2 + j)
                min_c = max(0, j - kw // 2)
                max_r_ref = min(h, h - i + kh // 2)
                min_r_ref = max(0, kh // 2 - i)
                max_c_ref = min(w, w - j + kw // 2)
                min_c_ref = max(0, kw // 2 - j)
                cij[min_r:max_r, min_c:max_c] = \
                    self.refimage[min_r_ref:max_r_ref, min_c_ref:max_c_ref]
                c.extend([cij])

        # This is more pythonic but much slower (50 times slower)
        # canonBasis = np.identity(kw*kh).reshape(kh*kw,kh,kw)
        # c.extend([signal.convolve2d(refimage, kij, mode='same')
        #                 for kij in canonBasis])
        # canonBasis = None
        return c

    def get_kernel(self):
        if self.kernel is not None:
            return self.kernel
        coeffs = self.get_coeffs()
        kh, kw = self.k_shape
        self.kernel = coeffs[:(kh * kw)].reshape(self.k_shape)
        return self.kernel

    def get_coeffs(self):
        if self.coeffs is not None:
            return self.coeffs
        c = self.get_cmatrices()
        if self.bkgdegree is not None:
            c_bkg = self.get_cmatrices_background()
            c.extend(c_bkg)

        if self.badpixmask is None:
            m = np.array([[(ci * cj).sum() for ci in c] for cj in c])
            b = np.array([(self.image * ci).sum() for ci in c])
        else:
            # These next two lines take most of the computation time
            m = np.array([[(ci * cj)[~self.badpixmask].sum()
                         for ci in c] for cj in c])
            b = np.array([(self.image * ci)[~self.badpixmask].sum()
                         for ci in c])
        self.coeffs = np.linalg.solve(m, b)
        return self.coeffs


class AdaptiveBramichStrategy(SubtractionStrategy):
    def __init__(self, image, refimage, kernelshape, bkgdegree, poly_degree=2):
        self.poly_deg = poly_degree
        self.poly_dof = (poly_degree + 1) * (poly_degree + 2) // 2
        self.k_side = kernelshape[0]

        # Check here for types
        if image.dtype != np.float64:
            self.img64 = image.astype('float64')
        else:
            self.img64 = image
        if refimage.dtype != np.float64:
            self.ref64 = refimage.astype('float64')
        else:
            self.ref64 = refimage

        super(AdaptiveBramichStrategy, self).\
            __init__(image, refimage, kernelshape, bkgdegree)

    def get_optimal_image(self):
        # AdaptiveBramich has to override this function because it uses a
        # special type of convolution for optimal_image
        if self.optimal_image is not None:
            return self.optimal_image
        import varconv
        opt_image = varconv.convolve2d_adaptive(
            self.ref64, self.get_kernel(), self.poly_deg)
        if self.bkgdegree is not None:
            opt_image += self.get_background()
        if self.badpixmask is not None:
            self.optimal_image = np.ma.array(opt_image, mask=self.badpixmask)
        else:
            self.optimal_image = opt_image
        return self.optimal_image

    def get_kernel(self):
        if self.kernel is not None:
            return self.kernel
        poly_dof = (self.poly_deg + 1) * (self.poly_deg + 2) // 2
        k_dof = self.k_side * self.k_side * poly_dof
        ks = self.k_side
        coeffs = self.get_coeffs()
        self.kernel = coeffs[:k_dof].reshape((ks, ks, self.poly_dof))
        return self.kernel

    def get_coeffs(self):
        if self.coeffs is not None:
            return self.coeffs
        import varconv
        c_bkgdegree = -1 if self.bkgdegree is None else self.bkgdegree
        m, b, conv = varconv.gen_matrix_system(self.img64, self.ref64,
                                               self.badpixmask is not None,
                                               self.badpixmask,
                                               self.k_side, self.poly_deg,
                                               c_bkgdegree)
        self.coeffs = np.linalg.solve(m, b)
        return self.coeffs


def convolve2d_adaptive(image, kernel, poly_degree):
    "Convolve image with the adaptive kernel of `poly_degree` degree."
    import varconv

    # Check here for dimensions
    if image.ndim != 2:
        raise ValueError("Wrong dimensions for image")
    if kernel.ndim != 3:
        raise ValueError("Wrong dimensions for kernel")

    # Check here for types
    if image.dtype != np.float64:
        img64 = image.astype('float64')
    else:
        img64 = image
    if kernel.dtype != np.float64:
        k64 = kernel.astype('float64')
    else:
        k64 = kernel

    conv = varconv.convolve2d_adaptive(img64, k64, poly_degree)
    return conv


def eval_adpative_kernel(kernel, x, y):
    "Return the adaptive kernel at position (x, y) = (col, row)."
    if kernel.ndim == 2:
        return kernel

    kh, kw, dof = kernel.shape
    # The conversion from degrees of freedom (dof) to the polynomial degree
    # The last 0.5 is to round to nearest integer
    deg = int(-1.5 + np.sqrt(1 + 8 * dof) / 2 + 0.5)
    k_rolled = np.rollaxis(kernel, 2, 0)
    k_xy = np.zeros((kh, kw))
    d = 0
    for powx in range(deg + 1):
        for powy in range(deg - powx + 1):
            k_xy += k_rolled[d] * np.power(y, powy) * np.power(x, powx)
            d += 1
    return k_xy


def optimal_system(image, refimage, kernelshape=(11, 11), bkgdegree=None,
                   method="Bramich", gridshape=None, **kwargs):
    """Do Optimal Image Subtraction and return optimal image, kernel
    and background.

    This is an implementation of a few Optimal Image Subtraction algorithms.
    They all (optionally) simultaneously fit a background.

    gridshape: a tuple containing the number of vertical and horizontal
    divisions of a grid. Subtraction will be performed on each grid element.
    `None` is equivalent to a (1, 1) grid (no grid).

    kernelshape: shape of the kernel to use. Must be of odd size.

    bkgdegree: degree of the polynomial to fit the background.
    To turn off background fitting set this to None.

    method: One of the following strings
    * Bramich: A Delta basis for the kernel (all pixels fit
      independently)
    * AdaptiveBramich: Same as Bramich, but with a polynomial variation across
      the image.
      It needs the parameter poly_degree, which is the polynomial degree of the
      variation.
    * Alard-Lupton: A modulated multi-Gaussian kernel.
      It needs the gausslist keyword.
      gausslist is a list of dictionaries containing data of the gaussians
      used in the decomposition of the kernel. Dictionary keywords are:
      center, sx, sy, modPolyDeg

    Extra parameters are passed to the individual methods.
    poly_degree: needed only for AdaptiveBramich. It is the degree
    of the polynomial for the kernel spatial variation.

    gausslist: needed only for Alard-Lupton. A list of dictionaries with info
    for the modulated multi-Gaussian.
        Dictionary keys are:
        center: a (row, column) tuple for the center of the Gaussian.
            Default: kernel center.
        modPolyDeg: the degree of the modulating polynomial. Default: 2
        sx: sigma in x direction. Default: 2.
        sy: sigma in y direction. Deafult: 2.
        All keys are optional.

    Return (difference, optimal_image, kernel, background)
    """

    kh, kw = kernelshape

    if (kw % 2 == 0) or (kh % 2 == 0):
        raise EvenSideKernelError("Kernel sides must be odd.")

    DefaultStrategy = BramichStrategy  # noqa
    all_strategies = {"AdaptiveBramich": AdaptiveBramichStrategy,
                      "Bramich": BramichStrategy,
                      "Alard-Lupton": AlardLuptonStrategy}
    try:
        DiffStrategy = all_strategies[method]  # noqa
    except KeyError:
        raise ValueError("No method named {}".format(method))

    if gridshape is None or gridshape == (1, 1):
        # If there's no grid, do without it
        subt_strat = DiffStrategy(image, refimage, kernelshape, bkgdegree,
                                  **kwargs)
        opt_image = subt_strat.get_optimal_image()
        kernel = subt_strat.get_kernel()
        background = subt_strat.get_background()
        difference = subt_strat.get_difference()
        return difference, opt_image, kernel, background

    else:
        ny, nx = gridshape
        h, w = image.shape

        # normal slices with no border
        stamps_y = [slice(h * i // ny, h * (i + 1) // ny, None) for i in range(ny)]
        stamps_x = [slice(w * i // nx, w * (i + 1) // nx, None) for i in range(nx)]

        # slices with borders where possible
        # Slices should be in (h * i // ny, h * (i + 1) // ny) but we add and
        # subtract the kernel spill k_spill and then we clip to keep it inside
        # image boundaries.
        k_spill = (kh - 1) // 2
        slc_wborder_y = [slice(np.clip(h * i // ny - k_spill, 0, h),
                               np.clip(h * (i + 1) // ny + k_spill, 0, h),
                               None) for i in range(ny)]
        slc_wborder_x = [slice(np.clip(w * i // nx - k_spill, 0, w),
                               np.clip(w * (i + 1) // nx + k_spill, 0, w),
                               None) for i in range(nx)]

        img_stamps = [image[sly, slx] for sly in slc_wborder_y
                      for slx in slc_wborder_x]
        ref_stamps = [refimage[sly, slx] for sly in slc_wborder_y
                      for slx in slc_wborder_x]

        # After we do the subtraction we need to crop the extra borders in the
        # stamps.
        # The recover_slices are the prescription for what to crop on each stamp.
        recover_slices = []
        for i in range(ny):
            start_border_y = slc_wborder_y[i].start
            stop_border_y = slc_wborder_y[i].stop
            # Slice should end at h * (i + 1) // ny, any other pixels should
            # be trimmed. sly_stop is either negative or 0.
            # In the special case where 0 pixels need to be trimmed
            # we use None so slice goes to the end.
            sly_stop = (h * (i + 1) // ny - stop_border_y) or None
            # Same with initial pixels, but sly_start is positive or 0.
            # Zero is not a special case now (0 is array initial pixel)
            sly_start = h * i // ny - start_border_y
            sly = slice(sly_start, sly_stop, None)
            for j in range(nx):
                start_border_x = slc_wborder_x[j].start
                stop_border_x = slc_wborder_x[j].stop
                slx_stop = (w * (j + 1) // nx - stop_border_x) or None
                slx_start = w * j // nx - start_border_x
                slx = slice(slx_start, slx_stop, None)
                recover_slices.append([sly, slx])

        # Here do the subtraction on each stamp
        if _has_mask(image) or _has_mask(refimage):
            optimal_collage = np.ma.empty(image.shape)
            subtract_collage = np.ma.empty(image.shape)
        else:
            optimal_collage = np.empty(image.shape)
            subtract_collage = np.empty(image.shape)
        bkg_collage = np.empty(image.shape)
        kernel_collage = []
        stamp_slices = [[asly, aslx] for asly in stamps_y for aslx in stamps_x]
        for ind, ((sly_out, slx_out), (sly_in, slx_in)) in \
                enumerate(zip(recover_slices, stamp_slices)):

            subt_strat = DiffStrategy(img_stamps[ind], ref_stamps[ind],
                                      kernelshape,
                                      bkgdegree,
                                      **kwargs)
            opti = subt_strat.get_optimal_image()
            ki = subt_strat.get_kernel()
            bgi = subt_strat.get_background()
            di = subt_strat.get_difference()
            optimal_collage[sly_in, slx_in] = opti[sly_out, slx_out]
            bkg_collage[sly_in, slx_in] = bgi[sly_out, slx_out]
            subtract_collage[sly_in, slx_in] = di[sly_out, slx_out]
            kernel_collage.append(ki)

        return subtract_collage, optimal_collage, kernel_collage, bkg_collage
