from pds4_tools import pds4_read
import matplotlib.pyplot as plt
import matplotlib
%matplotlib inline

import numpy as np
from PIL import Image

from skimage import exposure
from skimage import data, img_as_float
import colour

from colour_demosaicing import (
    demosaicing_CFA_Bayer_bilinear,
    demosaicing_CFA_Bayer_Malvar2004,
    demosaicing_CFA_Bayer_Menon2007,
    mosaicing_CFA_Bayer)
cctf_encoding = colour.cctf_encoding
_ = colour.utilities.filter_warnings()

path = './input/文件名.2CL'
d = pds4_read(path, quiet=True)
fig, axes = plt.subplots(1,1,figsize=(10,10))
img = np.array(d[0].data)
axes.imshow(img, cmap='gray')

def read_pds(path):
    data = pds4_read(path, quiet=True)
    img = np.array(data[0].data)
    img = img_as_float(img)
    return img

def plot_img_and_hist(image, hist=True, bins=128):
    """Plot an image along with its histogram.
    """
    if hist:
        fig, axes = plt.subplots(2,1, figsize=(10,10), gridspec_kw={'height_ratios': [3, 1]})
        ax_img, ax_hist = axes
    else:
        fig, ax_img = plt.subplots(figsize=(10,10))

    # Display image
    ax_img.imshow(image, cmap='gray')
    ax_img.set_axis_off()

    if hist:
        # Display histogram
        ax_hist.hist(image[:,:,0].ravel(), bins=bins, histtype='step', color='red')
        ax_hist.hist(image[:,:,1].ravel(), bins=bins, histtype='step', color='green')
        ax_hist.hist(image[:,:,2].ravel(), bins=bins, histtype='step', color='blue')

        ax_hist.ticklabel_format(axis='y', style='scientific', scilimits=(0, 0))
        ax_hist.set_xlabel('Pixel intensity')
        ax_hist.set_xlim(0, 1)
        ax_hist.set_yticks([])

def export_img(name, img):
    pil_img = Image.fromarray(np.uint8(img*255))
    pil_img.save(name)

path = './input/文件名.2CL'
img = read_pds(path)
plot_img_and_hist(img, hist=True)
plt.show()
export_img(f"{path}.png", img)

def debayer_img(img, CFA='RGGB'):
    # Menon2007 yields better edges than bilinear
    debayered = cctf_encoding(demosaicing_CFA_Bayer_Menon2007(img, CFA))
    return debayered
def stretch_img(img):
    p2, p98 = np.percentile(img, (2, 98))
    img = exposure.rescale_intensity(img, in_range=(p2, p98))
    return img

path = './input/文件名.2BL'
img = read_pds(path)
print(img.shape)
debayered = debayer_img(img)
final = stretch_img(debayered)
plot_img_and_hist(final, hist=True)
plt.show()
export_img(f"{path}.png", final)