import os

import matplotlib.pyplot as plt
import numpy as np
import cv2
from scipy.ndimage import gaussian_filter
import scipy

from . import utils, io

def show_segmentation(fig, img, maski, flowi, channels=[0,0], file_name=None, labeli=None, api=None):
    """ plot segmentation results (like on website)
    
    Can save each panel of figure with file_name option. Use channels option if
    img input is not an RGB image with 3 channels.
    
    Parameters
    -------------

    fig: matplotlib.pyplot.figure
        figure in which to make plot

    img: 2D or 3D array
        image input into scellseg

    maski: int, 2D array
        for image k, masks[k] output from Cellpose.eval, where 0=NO masks; 1,2,...=mask labels

    flowi: int, 2D array 
        for image k, flows[k][0] output from Cellpose.eval (RGB of flows)

    channels: list of int (optional, default [0,0])
        channels used to run Cellpose, no need to use if image is RGB

    file_name: str (optional, default None)
        file name of image, if file_name is not None, figure panels are saved

    """
    ax = fig.add_subplot(1,4,1)
    img0 = img.copy()
    if img0.shape[0] < 4:
        img0 = np.transpose(img0, (1,2,0))
    if img0.shape[-1] < 3 or img0.ndim < 3:
        img0 = image_to_rgb(img0, channels=channels)
    else:
        if img0.max()<=50.0:
            img0 = np.uint8(np.clip(img0*255, 0, 1))
    ax.imshow(img0)
    ax.set_title('original image')
    ax.axis('off')

    ax = fig.add_subplot(1,4,2)
    ax.imshow(img0)
    # outlines = utils.masks_to_outlines(maski)
    # outX, outY = np.nonzero(outlines)
    # imgout= img0.copy()
    # imgout[outX, outY] = np.array([255,75,75])
    # ax.imshow(imgout)

    if labeli is not None:
        outpix_label = utils.outlines_list(labeli)
        for o_label in outpix_label:
            ax.plot(o_label[:, 0], o_label[:, 1], color='y', ls='-', lw=0.5)

    outpix = utils.outlines_list(maski)

    for o in outpix:
       ax.plot(o[:,0], o[:,1], color='r', ls='--', lw=0.5)

    if api is not None:
        ax.text(0.655, 0.033, 'ap@0.5=%.2f'%api[0], bbox=dict(facecolor='black', alpha=0.6, edgecolor='None', boxstyle='round'),
                transform=ax.transAxes, fontsize=7, color='white')
    ax.set_title('predicted outlines')
    ax.axis('off')

    overlay = mask_overlay(img0, maski)
    ax = fig.add_subplot(1,4,3)
    ax.imshow(overlay)
    ax.set_title('predicted masks')
    ax.axis('off')

    ax = fig.add_subplot(1,4,4)
    ax.imshow(flowi)
    ax.set_title('predicted cell pose')
    ax.axis('off')

    if file_name is not None:
        save_path = os.path.splitext(file_name)[0]
        io.imsave(save_path + '_overlay.jpg', overlay)
        # io.imsave(save_path + '_outlines.jpg', imgout)
        io.imsave(save_path + '_flows.jpg', flowi)


def show_segmentation2(fig, img, maski, flowi, channels=[0, 0], file_name=None, labeli=None, api=None):
    """ plot segmentation for scellseg paper
    """
    # ax = fig.add_subplot(1,3,1)

    img0 = img.copy()
    if img0.shape[0] < 4:
        img0 = np.transpose(img0, (1, 2, 0))
    if img0.shape[-1] < 3 or img0.ndim < 3:
        img0 = image_to_rgb(img0, channels=channels)
    else:
        if img0.max() <= 50.0:
            img0 = np.uint8(np.clip(img0 * 255, 0, 1))

    # ax.imshow(img0)
    # ax.set_title('original image')
    # ax.axis('off')

    ax = fig.add_subplot(1, 2, 1)
    ax.imshow(img0)
    # outlines = utils.masks_to_outlines(maski)
    # outX, outY = np.nonzero(outlines)
    # imgout= img0.copy()
    # imgout[outX, outY] = np.array([255,75,75])
    # ax.imshow(imgout)

    if labeli is not None:
        outpix_label = utils.outlines_list(labeli)
        for o_label in outpix_label:
            ax.plot(o_label[:, 0], o_label[:, 1], color='y', ls='-', lw=0.5)

    outpix = utils.outlines_list(maski)

    for o in outpix:
        ax.plot(o[:, 0], o[:, 1], color='r', ls='--', lw=0.5)

    # if api is not None:
    #     ax.text(0.655, 0.033, 'ap@0.5=%.2f' % api[0],
    #             bbox=dict(facecolor='black', alpha=0.6, edgecolor='None', boxstyle='round'),
    #             transform=ax.transAxes, fontsize=7, color='white')
    # ax.set_title('predicted outlines')
    ax.axis('off')

    overlay = mask_overlay(img0, maski)

    ax = fig.add_subplot(1,2,2)
    ax.imshow(flowi)
    # ax.set_title('predicted cell pose')
    ax.axis('off')


    if file_name is not None:
        save_path = os.path.splitext(file_name)[0]
        io.imsave(save_path + '_showimg.png', img0)
        io.imsave(save_path + '_overlay.jpg', overlay)
        # io.imsave(save_path + '_outlines.jpg', outpix)
        io.imsave(save_path + '_flows.jpg', flowi)


def show_classic_segmentation(fig, img, maski, channels=[0, 0], file_name=None, labeli=None, api=None):
    """ plot segmentation for scellseg paper
    """
    ax = fig.add_subplot(1, 3, 1)
    img0 = img.copy()
    if img0.shape[0] < 4:
        img0 = np.transpose(img0, (1, 2, 0))
    if img0.shape[-1] < 3 or img0.ndim < 3:
        img0 = image_to_rgb(img0, channels=channels)
    else:
        if img0.max() <= 50.0:
            img0 = np.uint8(np.clip(img0 * 255, 0, 1))
    ax.imshow(img0)
    ax.set_title('original image')
    ax.axis('off')

    ax = fig.add_subplot(1, 3, 2)
    ax.imshow(img0)
    # outlines = utils.masks_to_outlines(maski)
    # outX, outY = np.nonzero(outlines)
    # imgout= img0.copy()
    # imgout[outX, outY] = np.array([255,75,75])
    # ax.imshow(imgout)

    if labeli is not None:
        outpix_label = utils.outlines_list(labeli)
        for o_label in outpix_label:
            ax.plot(o_label[:, 0], o_label[:, 1], color='y', ls='-', lw=0.5)

    outpix = utils.outlines_list(maski)

    for o in outpix:
        ax.plot(o[:, 0], o[:, 1], color='r', ls='--', lw=0.5)

    if api is not None:
        ax.text(0.655, 0.033, 'ap@0.5=%.2f' % api[0],
                bbox=dict(facecolor='black', alpha=0.6, edgecolor='None', boxstyle='round'),
                transform=ax.transAxes, fontsize=7, color='white')
    ax.set_title('predicted outlines')
    ax.axis('off')

    overlay = mask_overlay(img0, maski)
    ax = fig.add_subplot(1, 3, 3)
    ax.imshow(overlay)
    ax.set_title('predicted masks')
    ax.axis('off')

    if file_name is not None:
        save_path = os.path.splitext(file_name)[0]
        io.imsave(save_path + '_overlay.jpg', overlay)
        # io.imsave(save_path + '_outlines.jpg', imgout)


def show_classic_segmentation2(fig, img, maski, channels=[0, 0], file_name=None, labeli=None, api=None):
    """ plot segmentation results (like on website)

    Can save each panel of figure with file_name option. Use channels option if
    img input is not an RGB image with 3 channels.

    Parameters
    -------------

    fig: matplotlib.pyplot.figure
        figure in which to make plot

    img: 2D or 3D array
        image input into scellseg

    maski: int, 2D array
        for image k, masks[k] output from Cellpose.eval, where 0=NO masks; 1,2,...=mask labels

    flowi: int, 2D array
        for image k, flows[k][0] output from Cellpose.eval (RGB of flows)

    channels: list of int (optional, default [0,0])
        channels used to run Cellpose, no need to use if image is RGB

    file_name: str (optional, default None)
        file name of image, if file_name is not None, figure panels are saved

    """
    # ax = fig.add_subplot(1, 3, 1)
    img0 = img.copy()
    if img0.shape[0] < 4:
        img0 = np.transpose(img0, (1, 2, 0))
    if img0.shape[-1] < 3 or img0.ndim < 3:
        img0 = image_to_rgb(img0, channels=channels)
    else:
        if img0.max() <= 50.0:
            img0 = np.uint8(np.clip(img0 * 255, 0, 1))
    # ax.imshow(img0)
    # ax.set_title('original image')
    # ax.axis('off')

    ax = fig.add_subplot(1, 1, 1)
    ax.imshow(img0)
    # outlines = utils.masks_to_outlines(maski)
    # outX, outY = np.nonzero(outlines)
    # imgout= img0.copy()
    # imgout[outX, outY] = np.array([255,75,75])
    # ax.imshow(imgout)

    if labeli is not None:
        outpix_label = utils.outlines_list(labeli)
        for o_label in outpix_label:
            ax.plot(o_label[:, 0], o_label[:, 1], color='y', ls='-', lw=0.5)

    outpix = utils.outlines_list(maski)

    for o in outpix:
        ax.plot(o[:, 0], o[:, 1], color='r', ls='--', lw=0.5)

    # if api is not None:
    #     ax.text(0.655, 0.033, 'ap@0.5=%.2f' % api[0],
    #             bbox=dict(facecolor='black', alpha=0.6, edgecolor='None', boxstyle='round'),
    #             transform=ax.transAxes, fontsize=7, color='white')
    # ax.set_title('predicted outlines')
    ax.axis('off')

    overlay = mask_overlay(img0, maski)
    # ax = fig.add_subplot(1, 2, 2)
    # ax.imshow(overlay)
    # ax.set_title('predicted masks')
    # ax.axis('off')

    if file_name is not None:
        save_path = os.path.splitext(file_name)[0]
        io.imsave(save_path + '_showimg.png', img0)
        io.imsave(save_path + '_overlay.jpg', overlay)
        # io.imsave(save_path + '_outlines.jpg', imgout)


def mask_rgb(masks, colors=None):
    """ masks in random rgb colors

    Parameters
    ----------------

    masks: int, 2D array
        masks where 0=NO masks; 1,2,...=mask labels

    colors: int, 2D array (optional, default None)
        size [nmasks x 3], each entry is a color in 0-255 range

    Returns
    ----------------

    RGB: uint8, 3D array
        array of masks overlaid on grayscale image

    """
    if colors is not None:
        if colors.max()>1:
            colors = np.float32(colors)
            colors /= 255
        colors = utils.rgb_to_hsv(colors)
    
    HSV = np.zeros((masks.shape[0], masks.shape[1], 3), np.float32)
    HSV[:,:,2] = 1.0
    for n in range(int(masks.max())):
        ipix = (masks==n+1).nonzero()
        if colors is None:
            HSV[ipix[0],ipix[1],0] = np.random.rand()
        else:
            HSV[ipix[0],ipix[1],0] = colors[n,0]
        HSV[ipix[0],ipix[1],1] = np.random.rand()*0.5+0.5
        HSV[ipix[0],ipix[1],2] = np.random.rand()*0.5+0.5
    RGB = (utils.hsv_to_rgb(HSV) * 255).astype(np.uint8)
    return RGB

def mask_overlay(img, masks, colors=None):
    """ overlay masks on image (set image to grayscale)

    Parameters
    ----------------

    img: int or float, 2D or 3D array
        img is of size [Ly x Lx (x nchan)]

    masks: int, 2D array
        masks where 0=NO masks; 1,2,...=mask labels

    colors: int, 2D array (optional, default None)
        size [nmasks x 3], each entry is a color in 0-255 range

    Returns
    ----------------

    RGB: uint8, 3D array
        array of masks overlaid on grayscale image

    """
    if colors is not None:
        if colors.max()>1:
            colors = np.float32(colors)
            colors /= 255
        colors = utils.rgb_to_hsv(colors)
    if img.ndim>2:
        img = img.astype(np.float32).mean(axis=-1)
    else:
        img = img.astype(np.float32)
    img = utils.normalize99(img)
    img -= img.min()
    img /= img.max()
    HSV = np.zeros((img.shape[0], img.shape[1], 3), np.float32)
    HSV[:,:,2] = np.clip(img*1.5, 0, 1.0)
    for n in range(int(masks.max())):
        ipix = (masks==n+1).nonzero()
        if colors is None:
            HSV[ipix[0],ipix[1],0] = np.random.rand()
        else:
            HSV[ipix[0],ipix[1],0] = colors[n,0]
        HSV[ipix[0],ipix[1],1] = 1.0
    RGB = (utils.hsv_to_rgb(HSV) * 255).astype(np.uint8)
    return RGB

def image_to_rgb(img0, channels=[0,0]):
    """ image is 2 x Ly x Lx or Ly x Lx x 2 - change to RGB Ly x Lx x 3 """
    img = img0.copy()
    img = img.astype(np.float32)
    if img.ndim<3:
        img = img[:,:,np.newaxis]
    if img.shape[0]<5:
        img = np.transpose(img, (1,2,0))
    if channels[0]==0:
        img = img.mean(axis=-1)[:,:,np.newaxis]
    for i in range(img.shape[-1]):
        if np.ptp(img[:,:,i])>0:
            img[:,:,i] = utils.normalize99(img[:,:,i])
            img[:,:,i] = np.clip(img[:,:,i], 0, 1)
    img *= 255
    img = np.uint8(img)
    RGB = np.zeros((img.shape[0], img.shape[1], 3), np.uint8)
    if img.shape[-1]==1:
        RGB = np.tile(img,(1,1,3))
    else:
        RGB[:,:,channels[0]-1] = img[:,:,0]
        if channels[1] > 0:
            RGB[:,:,channels[1]-1] = img[:,:,1]
    return RGB

def interesting_patch(mask, bsize=130):
    """ get patch of size bsize x bsize with most masks """
    Ly,Lx = mask.shape
    m = np.float32(mask>0)
    m = gaussian_filter(m, bsize/2)
    y,x = np.unravel_index(np.argmax(m), m.shape)
    ycent = max(bsize//2, min(y, Ly-bsize//2))
    xcent = max(bsize//2, min(x, Lx-bsize//2))
    patch = [np.arange(ycent-bsize//2, ycent+bsize//2, 1, int),
             np.arange(xcent-bsize//2, xcent+bsize//2, 1, int)]
    return patch

def disk(med, r, Ly, Lx):
    """ returns pixels of disk with radius r and center med """
    yy, xx = np.meshgrid(np.arange(0,Ly,1,int), np.arange(0,Lx,1,int),
                         indexing='ij')
    inds = ((yy-med[0])**2 + (xx-med[1])**2)**0.5 <= r
    y = yy[inds].flatten()
    x = xx[inds].flatten()
    return y,x
