# Manuel A. Morales (moralesq@mit.edu)
# Harvard-MIT Department of Health Sciences & Technology  
# Athinoula A. Martinos Center for Biomedical Imaging

import numpy as np
from skimage import measure
import matplotlib.pylab as plt
from scipy.ndimage.morphology import binary_fill_holes

def Contours(mask, tissue_labels=[1,2,3]):
    contours = []
    for i in tissue_labels:
        mask_ = binary_fill_holes(mask==i)
        c = measure.find_contours(mask_,0.8)
        c = c[np.argmax([len(c) for c in c])]
        contours.append(c)
    return contours

def PlotContours(ax, mask, tissue_labels=[1,2,3], 
                 contour_colors=['lime','magenta','red'],
                 contour_labels=['RV','LVM','LV'],
                 tolerance=0.1,
                 alpha=1,
                 linewidth=2,
                 legend=False):
    
    contours = Contours(mask, tissue_labels=tissue_labels)
    for i, contour in enumerate(contours):
        ax.plot(contour[:, 1], contour[:, 0],alpha=alpha, linewidth=linewidth, color=contour_colors[i],label=contour_labels[i])
    if legend:
        ax.legend(fontsize=26)  
        
        
def Plot(image, mask=None, figsize=(7,7), crop=False):
    
    if crop:
        nx, ny = image.shape
        image = image[nx//2-64:nx//2+64,ny//2-64:ny//2+64]
        if mask is not None:
            mask = mask[nx//2-64:nx//2+64,ny//2-64:ny//2+64]
    
    fig, ax = plt.subplots(1,1, figsize=figsize)
    ax.imshow(image, cmap='gray')
    if mask is not None:
        try:
            PlotContours(ax, mask, alpha=1, linewidth=1.3)  
        except:
            print('No contours found!')