import os
import numpy as np
import matplotlib.pyplot as plt
from pspylib.tiff.reader import TiffReader
from matplotlib.colors import BoundaryNorm
from matplotlib.ticker import NullLocator
from matplotlib.colors import LinearSegmentedColormap


def __auto_palette_range(gray,cmap):
    # 95.44 % 
    vmin = np.mean(gray) - 1.96 * np.std(gray)
    vmax = np.mean(gray) + 1.96 * np.std(gray)
    hist, edges = np.histogram(gray, bins=cmap.N, range=(vmin, vmax))
    boundary_norm = BoundaryNorm(edges, cmap.N)
    return boundary_norm

def plot_tiff(gray, cmap, size:tuple=(8,8)):
    fig = plt.figure(figsize=size, facecolor='white')
    boundary_norm = __auto_palette_range(gray, cmap)
    plt.imshow(gray,cmap=cmap,norm=boundary_norm)
    cbar = plt.colorbar()
    cbar.ax.yaxis.set_minor_locator(NullLocator())
    return fig

if __name__ == "__main__":
    samples_path = r"C:\Park Systems\SmartScan\samples"
    tiff_path = os.path.join(samples_path, "Image", "Cheese.tiff")
    tiff = TiffReader(tiff_path)
    Zdata = tiff.data.scanData.ZData
    header = tiff.data.scanHeader.scanHeader
    dshape = (int(header['height'][0]), int(header['width'][0]))
    tiff_image = np.reshape(Zdata,dshape)
    tiff_image = np.flipud(tiff_image)

    ori_cm = tiff.data.metaData.colorMap['colorMap'][0]
    r = np.array(ori_cm)[:256]
    g = np.array(ori_cm)[256:512]
    b = np.array(ori_cm)[512:768]
    rgb_cm = np.dstack([r,g,b])[0] / 65536
    cmap_name = 'meta_cmap'
    meta_cmap = LinearSegmentedColormap.from_list(cmap_name, rgb_cm,N=128)
    plot_tiff(tiff_image, meta_cmap)
    plt.xlabel('X (pixels)')
    plt.ylabel('Y (pixels)')
    plt.show()