import cv2
import sys
import fire
import histomicstk
import histomicstk as htk
import openslide
import dask
import tqdm
import numpy as np
from dask.diagnostics import ProgressBar
from pathflowai.utils import generate_tissue_mask
from histomicstk.preprocessing.color_normalization.\
    deconvolution_based_normalization import deconvolution_based_normalization

W_target = np.array([
            [0.6185391,  0.1576997,  -0.01119131],
            [0.7012888,  0.8638838,  0.45586256],
            [0.3493163,  0.4657428, -0.85597752]
        ])

def return_norm_image(img,mask,W_source=None,W_target=None):
    img=deconvolution_based_normalization(
        img, W_source=W_source, W_target=W_target, im_target=None,
        stains=['hematoxylin', 'eosin'], mask_out=~mask,
        stain_unmixing_routine_params={"I_0":215})
    return img

def check_ext(image_file):
    return any([image_file.endswith(ext) for ext in ['.svs','.png','.jpg','.jpeg','.tiff','.tif']])

def stain_norm(image_file,compression=10,patch_size=1024):
    if check_ext(image_file):
        img = openslide.open_slide(image_file)
        image = np.array(img.read_region((0,0), 0, img.level_dimensions[0]))[...,:3]
    elif image_file.endswith(".npy"):
        image=np.load(image_file)
    else: raise NotImplementedError
    mask=generate_tissue_mask(image,compression=compression,keep_holes=False)
    img_small=cv2.resize(image,None,fx=1/compression,fy=1/compression)
    mask_small=cv2.resize(mask.astype(int),None,fx=1/compression,fy=1/compression,interpolation=cv2.INTER_NEAREST).astype(bool)
    W_source = htk.preprocessing.color_deconvolution.rgb_separate_stains_macenko_pca(img_small, 215)
    W_source = htk.preprocessing.color_deconvolution._reorder_stains(W_source)
    res=[]
    coords=[]
    for i in np.arange(0,image.shape[0]-patch_size,patch_size):
        for j in np.arange(0,image.shape[1]-patch_size,patch_size):
            if mask[i:i+patch_size,j:j+patch_size].mean():
                coords.append((i,j))
                res.append(dask.delayed(return_norm_image)(image[i:i+patch_size,j:j+patch_size],mask[i:i+patch_size,j:j+patch_size],W_source,W_target))
    with ProgressBar():
        res_returned=dask.compute(*res,scheduler="processes")
    img_new=np.ones(image.shape).astype(np.uint8)*255
    for k in tqdm.trange(len(coords)):
        i,j=coords[k]
        img_new[i:i+patch_size,j:j+patch_size]=res_returned[k]
    return img_new

def stain_norm_pipeline(image_file="stain_in.svs",
                        npy_out='stain_out.npy',
                        compression=10,
                        patch_size=1024):
    np.save(npy_out,stain_norm(image_file,compression,patch_size))

if __name__=="__main__":
    fire.Fire(stain_norm_pipeline)