from __future__ import print_function, unicode_literals, absolute_import, division

import numpy as np
import warnings
import os
import datetime
from tqdm import tqdm
from collections import defaultdict
from zipfile import ZipFile, ZIP_DEFLATED
from scipy.ndimage.morphology import distance_transform_edt, binary_fill_holes
from scipy.ndimage.measurements import find_objects
from scipy.optimize import minimize_scalar
from skimage.measure import regionprops
from csbdeep.utils import _raise
from csbdeep.utils.six import Path
from collections.abc import Iterable

from .matching import matching_dataset, _check_label_array


try:
    from edt import edt
    _edt_available = True
    try:    _edt_parallel_max = len(os.sched_getaffinity(0))
    except: _edt_parallel_max = 128
    _edt_parallel_default = 4
    _edt_parallel = os.environ.get('STARDIST_EDT_NUM_THREADS', _edt_parallel_default)
    try:
        _edt_parallel = min(_edt_parallel_max, int(_edt_parallel))
    except ValueError as e:
        warnings.warn(f"Invalid value ({_edt_parallel}) for STARDIST_EDT_NUM_THREADS. Using default value ({_edt_parallel_default}) instead.")
        _edt_parallel = _edt_parallel_default
    del _edt_parallel_default, _edt_parallel_max
except ImportError:
    _edt_available = False
    # warnings.warn("Could not find package edt... \nConsider installing it with \n  pip install edt\nto improve training data generation performance.")
    pass


def gputools_available():
    try:
        import gputools
    except:
        return False
    return True


def path_absolute(path_relative):
    """ Get absolute path to resource"""
    base_path = os.path.abspath(os.path.dirname(__file__))
    return os.path.join(base_path, path_relative)


def _is_power_of_2(i):
    assert i > 0
    e = np.log2(i)
    return e == int(e)


def _normalize_grid(grid,n):
    try:
        grid = tuple(grid)
        (len(grid) == n and
         all(map(np.isscalar,grid)) and
         all(map(_is_power_of_2,grid))) or _raise(TypeError())
        return tuple(int(g) for g in grid)
    except (TypeError, AssertionError):
        raise ValueError("grid = {grid} must be a list/tuple of length {n} with values that are power of 2".format(grid=grid, n=n))


def edt_prob(lbl_img, anisotropy=None):
    if _edt_available:
        return _edt_prob_edt(lbl_img, anisotropy=anisotropy)
    else:
        # warnings.warn("Could not find package edt... \nConsider installing it with \n  pip install edt\nto improve training data generation performance.")
        return _edt_prob_scipy(lbl_img, anisotropy=anisotropy)

def _edt_prob_edt(lbl_img, anisotropy=None):
    """Perform EDT on each labeled object and normalize.
    Internally uses https://github.com/seung-lab/euclidean-distance-transform-3d
    that can handle multiple labels at once
    """
    lbl_img = np.ascontiguousarray(lbl_img)
    constant_img = lbl_img.min() == lbl_img.max() and lbl_img.flat[0] > 0
    if constant_img:
        warnings.warn("EDT of constant label image is ill-defined. (Assuming background around it.)")
    # we just need to compute the edt once but then normalize it for each object
    prob = edt(lbl_img, anisotropy=anisotropy, black_border=constant_img, parallel=_edt_parallel)
    objects = find_objects(lbl_img)
    for i,sl in enumerate(objects,1):
        # i: object label id, sl: slices of object in lbl_img
        if sl is None: continue
        _mask = lbl_img[sl]==i
        # normalize it
        prob[sl][_mask] /= np.max(prob[sl][_mask]+1e-10)
    return prob

def _edt_prob_scipy(lbl_img, anisotropy=None):
    """Perform EDT on each labeled object and normalize."""
    def grow(sl,interior):
        return tuple(slice(s.start-int(w[0]),s.stop+int(w[1])) for s,w in zip(sl,interior))
    def shrink(interior):
        return tuple(slice(int(w[0]),(-1 if w[1] else None)) for w in interior)
    constant_img = lbl_img.min() == lbl_img.max() and lbl_img.flat[0] > 0
    if constant_img:
        lbl_img = np.pad(lbl_img, ((1,1),)*lbl_img.ndim, mode='constant')
        warnings.warn("EDT of constant label image is ill-defined. (Assuming background around it.)")
    objects = find_objects(lbl_img)
    prob = np.zeros(lbl_img.shape,np.float32)
    for i,sl in enumerate(objects,1):
        # i: object label id, sl: slices of object in lbl_img
        if sl is None: continue
        interior = [(s.start>0,s.stop<sz) for s,sz in zip(sl,lbl_img.shape)]
        # 1. grow object slice by 1 for all interior object bounding boxes
        # 2. perform (correct) EDT for object with label id i
        # 3. extract EDT for object of original slice and normalize
        # 4. store edt for object only for pixels of given label id i
        shrink_slice = shrink(interior)
        grown_mask = lbl_img[grow(sl,interior)]==i
        mask = grown_mask[shrink_slice]
        edt = distance_transform_edt(grown_mask, sampling=anisotropy)[shrink_slice][mask]
        prob[sl][mask] = edt/(np.max(edt)+1e-10)
    if constant_img:
        prob = prob[(slice(1,-1),)*lbl_img.ndim].copy()
    return prob


def _fill_label_holes(lbl_img, **kwargs):
    lbl_img_filled = np.zeros_like(lbl_img)
    for l in (set(np.unique(lbl_img)) - set([0])):
        mask = lbl_img==l
        mask_filled = binary_fill_holes(mask,**kwargs)
        lbl_img_filled[mask_filled] = l
    return lbl_img_filled


def fill_label_holes(lbl_img, **kwargs):
    """Fill small holes in label image."""
    # TODO: refactor 'fill_label_holes' and 'edt_prob' to share code
    def grow(sl,interior):
        return tuple(slice(s.start-int(w[0]),s.stop+int(w[1])) for s,w in zip(sl,interior))
    def shrink(interior):
        return tuple(slice(int(w[0]),(-1 if w[1] else None)) for w in interior)
    objects = find_objects(lbl_img)
    lbl_img_filled = np.zeros_like(lbl_img)
    for i,sl in enumerate(objects,1):
        if sl is None: continue
        interior = [(s.start>0,s.stop<sz) for s,sz in zip(sl,lbl_img.shape)]
        shrink_slice = shrink(interior)
        grown_mask = lbl_img[grow(sl,interior)]==i
        mask_filled = binary_fill_holes(grown_mask,**kwargs)[shrink_slice]
        lbl_img_filled[sl][mask_filled] = i
    return lbl_img_filled


def sample_points(n_samples, mask, prob=None, b=2):
    """sample points to draw some of the associated polygons"""
    if b is not None and b > 0:
        # ignore image boundary, since predictions may not be reliable
        mask_b = np.zeros_like(mask)
        mask_b[b:-b,b:-b] = True
    else:
        mask_b = True

    points = np.nonzero(mask & mask_b)

    if prob is not None:
        # weighted sampling via prob
        w = prob[points[0],points[1]].astype(np.float64)
        w /= np.sum(w)
        ind = np.random.choice(len(points[0]), n_samples, replace=True, p=w)
    else:
        ind = np.random.choice(len(points[0]), n_samples, replace=True)

    points = points[0][ind], points[1][ind]
    points = np.stack(points,axis=-1)
    return points


def calculate_extents(lbl, func=np.median):
    """ Aggregate bounding box sizes of objects in label images. """
    if (isinstance(lbl,np.ndarray) and lbl.ndim==4) or (not isinstance(lbl,np.ndarray) and  isinstance(lbl,Iterable)):
        return func(np.stack([calculate_extents(_lbl,func) for _lbl in lbl], axis=0), axis=0)

    n = lbl.ndim
    n in (2,3) or _raise(ValueError("label image should be 2- or 3-dimensional (or pass a list of these)"))

    regs = regionprops(lbl)
    if len(regs) == 0:
        return np.zeros(n)
    else:
        extents = np.array([np.array(r.bbox[n:])-np.array(r.bbox[:n]) for r in regs])
        return func(extents, axis=0)


def polyroi_bytearray(x,y,pos=None,subpixel=True):
    """ Byte array of polygon roi with provided x and y coordinates
        See https://github.com/imagej/imagej1/blob/master/ij/io/RoiDecoder.java
    """
    import struct
    def _int16(x):
        return int(x).to_bytes(2, byteorder='big', signed=True)
    def _uint16(x):
        return int(x).to_bytes(2, byteorder='big', signed=False)
    def _int32(x):
        return int(x).to_bytes(4, byteorder='big', signed=True)
    def _float(x):
        return struct.pack(">f", x)

    subpixel = bool(subpixel)
    # add offset since pixel center is at (0.5,0.5) in ImageJ
    x_raw = np.asarray(x).ravel() + 0.5
    y_raw = np.asarray(y).ravel() + 0.5
    x = np.round(x_raw)
    y = np.round(y_raw)
    assert len(x) == len(y)
    top, left, bottom, right = y.min(), x.min(), y.max(), x.max() # bbox

    n_coords = len(x)
    bytes_header = 64
    bytes_total = bytes_header + n_coords*2*2 + subpixel*n_coords*2*4
    B = [0] * bytes_total
    B[ 0: 4] = map(ord,'Iout')   # magic start
    B[ 4: 6] = _int16(227)       # version
    B[ 6: 8] = _int16(0)         # roi type (0 = polygon)
    B[ 8:10] = _int16(top)       # bbox top
    B[10:12] = _int16(left)      # bbox left
    B[12:14] = _int16(bottom)    # bbox bottom
    B[14:16] = _int16(right)     # bbox right
    B[16:18] = _uint16(n_coords) # number of coordinates
    if subpixel:
        B[50:52] = _int16(128)   # subpixel resolution (option flag)
    if pos is not None:
        B[56:60] = _int32(pos)   # position (C, Z, or T)

    for i,(_x,_y) in enumerate(zip(x,y)):
        xs = bytes_header + 2*i
        ys = xs + 2*n_coords
        B[xs:xs+2] = _int16(_x - left)
        B[ys:ys+2] = _int16(_y - top)

    if subpixel:
        base1 = bytes_header + n_coords*2*2
        base2 = base1 + n_coords*4
        for i,(_x,_y) in enumerate(zip(x_raw,y_raw)):
            xs = base1 + 4*i
            ys = base2 + 4*i
            B[xs:xs+4] = _float(_x)
            B[ys:ys+4] = _float(_y)

    return bytearray(B)


def export_imagej_rois(fname, polygons, set_position=True, subpixel=True, compression=ZIP_DEFLATED):
    """ polygons assumed to be a list of arrays with shape (id,2,c) """

    if isinstance(polygons,np.ndarray):
        polygons = (polygons,)

    fname = Path(fname)
    if fname.suffix == '.zip':
        fname = fname.with_suffix('')

    with ZipFile(str(fname)+'.zip', mode='w', compression=compression) as roizip:
        for pos,polygroup in enumerate(polygons,start=1):
            for i,poly in enumerate(polygroup,start=1):
                roi = polyroi_bytearray(poly[1],poly[0], pos=(pos if set_position else None), subpixel=subpixel)
                roizip.writestr('{pos:03d}_{i:03d}.roi'.format(pos=pos,i=i), roi)


def optimize_threshold(Y, Yhat, model, nms_thresh, measure='accuracy', iou_threshs=[0.3,0.5,0.7], bracket=None, tol=1e-2, maxiter=20, verbose=1):
    """ Tune prob_thresh for provided (fixed) nms_thresh to maximize matching score (for given measure and averaged over iou_threshs). """
    np.isscalar(nms_thresh) or _raise(ValueError("nms_thresh must be a scalar"))
    iou_threshs = [iou_threshs] if np.isscalar(iou_threshs) else iou_threshs
    values = dict()

    if bracket is None:
        max_prob = max([np.max(prob) for prob, dist in Yhat])
        bracket = max_prob/2, max_prob
    # print("bracket =", bracket)

    with tqdm(total=maxiter, disable=(verbose!=1), desc="NMS threshold = %g" % nms_thresh) as progress:

        def fn(thr):
            prob_thresh = np.clip(thr, *bracket)
            value = values.get(prob_thresh)
            if value is None:
                Y_instances = [model._instances_from_prediction(y.shape, *prob_dist, prob_thresh=prob_thresh, nms_thresh=nms_thresh)[0] for y,prob_dist in zip(Y,Yhat)]
                stats = matching_dataset(Y, Y_instances, thresh=iou_threshs, show_progress=False, parallel=True)
                values[prob_thresh] = value = np.mean([s._asdict()[measure] for s in stats])
            if verbose > 1:
                print("{now}   thresh: {prob_thresh:f}   {measure}: {value:f}".format(
                    now = datetime.datetime.now().strftime('%H:%M:%S'),
                    prob_thresh = prob_thresh,
                    measure = measure,
                    value = value,
                ), flush=True)
            else:
                progress.update()
                progress.set_postfix_str("{prob_thresh:.3f} -> {value:.3f}".format(prob_thresh=prob_thresh, value=value))
                progress.refresh()
            return -value

        opt = minimize_scalar(fn, method='golden', bracket=bracket, tol=tol, options={'maxiter': maxiter})

    verbose > 1 and print('\n',opt, flush=True)
    return opt.x, -opt.fun


def _invert_dict(d):
    """ return  v-> [k_1,k_2,k_3....] for k,v in d"""
    res = defaultdict(list)
    for k,v in d.items():
        res[v].append(k)
    return res


def mask_to_categorical(y, n_classes, classes, return_cls_dict=False):
    """generates a multi-channel categorical class map

    Parameters
    ----------
    y : n-dimensional ndarray
        integer label array
    n_classes : int
        Number of different classes (without background)
    classes: dict, integer, or None
        the label to class assignment
        can be
        - dict {label -> class_id}
           the value of class_id can be
                             0   -> background class
                  1...n_classes  -> the respective object class (1 ... n_classes)
                           None  -> ignore object (prob is set to -1 for the pixels of the object, except for background class)
        - single integer value or None -> broadcast value to all labels

    Returns
    -------
    probability map of shape y.shape+(n_classes+1,) (first channel is background)

    """

    _check_label_array(y, 'y')
    if not (np.issubdtype(type(n_classes), np.integer) and n_classes>=1):
        raise ValueError(f"n_classes is '{n_classes}' but should be a positive integer")

    y_labels = np.unique(y[y>0]).tolist()

    # build dict class_id -> labels (inverse of classes)
    if np.issubdtype(type(classes), np.integer) or classes is None:
        classes = dict((k,classes) for k in y_labels)
    elif isinstance(classes, dict):
        pass
    else:
        raise ValueError("classes should be dict, single scalar, or None!")

    if not set(y_labels).issubset(set(classes.keys())):
        raise ValueError(f"all gt labels should be present in class dict provided \ngt_labels found\n{set(y_labels)}\nclass dict labels provided\n{set(classes.keys())}")

    cls_dict = _invert_dict(classes)

    # prob map
    y_mask = np.zeros(y.shape+(n_classes+1,), np.float32)

    for cls, labels in cls_dict.items():
        if cls is None:
            # prob == -1 will be used in the loss to ignore object
            y_mask[np.isin(y, labels)] = -1
        elif np.issubdtype(type(cls), np.integer) and 0 <= cls <= n_classes:
            y_mask[...,cls] = np.isin(y, labels)
        else:
            raise ValueError(f"Wrong class id '{cls}' (for n_classes={n_classes})")

    # set 0/1 background prob (unaffected by None values for class ids)
    y_mask[...,0] = (y==0)

    if return_cls_dict:
        return y_mask, cls_dict
    else:
        return y_mask


def _is_floatarray(x):
    return isinstance(x.dtype.type(0),np.floating)


def abspath(root, relpath):
    from pathlib import Path
    root = Path(root)
    if root.is_dir():
        path = root/relpath
    else:
        path = root.parent/relpath
    return str(path.absolute())