import os
import numpy as np
import matplotlib.pyplot as plt
from pspylib.tiff.reader import TiffReader

def one_d_flatten(array_tiff, scope='line_x', region=None):
    assert scope in ['line_x', 'whole_x', 'line_y', 'whole_y']
    if region is not None:
        assert (np.array(array_tiff.shape) == np.array(region.shape)).all() # shape should match
    else:
        region = np.ones_like(array_tiff, dtype=np.bool_)

    image = array_tiff.copy()
    H, W = image.shape
    match scope:
        case 'line_x' | 'whole_x':
            A = np.ones((W, 2))
            A[:, 0] = (np.arange(1, W + 1) * 1.0 / W)
            results = []
            for i in range(H):
                img_row, reg_row = image[i, :], region[i, :]
                valid_n = np.sum(reg_row)
                if valid_n <= 2: # no region selected
                    bias = np.mean(img_row)
                    result_row = img_row - bias
                else:
                    coef, resids, rank, s = np.linalg.lstsq(A[reg_row, :], img_row[reg_row], rcond=None)
                    result_row = img_row - np.dot(A, coef)
                results.append(result_row)
            result_image = np.stack(results, axis=0)
            line_image = image - result_image
            avg_line = np.mean(line_image, axis=0, keepdims=True)
    
            if scope == 'whole_x':
                result_image = image - avg_line

        case 'line_y' | 'whole_y':
            A = np.ones((H, 2))
            A[:, 0] = (np.arange(1, H + 1) * 1.0 / H)
            results = []
            for i in range(W):
                img_column, reg_column = image[:, i], region[:, i]
                valid_n = np.sum(reg_column)
                if valid_n <= 2: # no region selected
                    bias = np.mean(img_column)
                    result_column = img_column - bias
                else:
                    coef, resids, rank, s = np.linalg.lstsq(A[reg_column, :], img_column[reg_column], rcond=None)
                    result_column = img_column - np.dot(A, coef)
                results.append(result_column)
            result_image = np.stack(results, axis=1)
            line_image = image - result_image
            avg_line = np.mean(line_image, axis=1, keepdims=True)
    
            if scope == 'whole_y':
                result_image = image - avg_line

    return result_image

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)


    flattened_image = one_d_flatten(tiff_image, scope='line_x')

    plt.subplot(1, 2, 1)
    plt.title('Original Image')
    plt.imshow(tiff_image, cmap='gray')
    plt.axis('off')

    plt.subplot(1, 2, 2)
    plt.title('1D Flattened Image')
    plt.imshow(flattened_image, cmap='gray')
    plt.axis('off')
    plt.tight_layout()
    plt.show()

import os
import numpy as np
import matplotlib.pyplot as plt
from pspylib.tiff.reader import TiffReader

def one_d_flatten(array_tiff, scope='line_x', region=None):
    assert scope in ['line_x', 'whole_x', 'line_y', 'whole_y']
    if region is not None:
        assert (np.array(array_tiff.shape) == np.array(region.shape)).all() # shape should match
    else:
        region = np.ones_like(array_tiff, dtype=np.bool_)

    image = array_tiff.copy()
    H, W = image.shape
    match scope:
        case 'line_x' | 'whole_x':
            A = np.ones((W, 2))
            A[:, 0] = (np.arange(1, W + 1) * 1.0 / W)
            results = []
            for i in range(H):
                img_row, reg_row = image[i, :], region[i, :]
                valid_n = np.sum(reg_row)
                if valid_n <= 2: # no region selected
                    bias = np.mean(img_row)
                    result_row = img_row - bias
                else:
                    coef, resids, rank, s = np.linalg.lstsq(A[reg_row, :], img_row[reg_row], rcond=None)
                    result_row = img_row - np.dot(A, coef)
                results.append(result_row)
            result_image = np.stack(results, axis=0)
            line_image = image - result_image
            avg_line = np.mean(line_image, axis=0, keepdims=True)
    
            if scope == 'whole_x':
                result_image = image - avg_line

        case 'line_y' | 'whole_y':
            A = np.ones((H, 2))
            A[:, 0] = (np.arange(1, H + 1) * 1.0 / H)
            results = []
            for i in range(W):
                img_column, reg_column = image[:, i], region[:, i]
                valid_n = np.sum(reg_column)
                if valid_n <= 2: # no region selected
                    bias = np.mean(img_column)
                    result_column = img_column - bias
                else:
                    coef, resids, rank, s = np.linalg.lstsq(A[reg_column, :], img_column[reg_column], rcond=None)
                    result_column = img_column - np.dot(A, coef)
                results.append(result_column)
            result_image = np.stack(results, axis=1)
            line_image = image - result_image
            avg_line = np.mean(line_image, axis=1, keepdims=True)
    
            if scope == 'whole_y':
                result_image = image - avg_line

    return result_image

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)


    flattened_image = one_d_flatten(tiff_image, scope='line_x')

    plt.subplot(1, 2, 1)
    plt.title('Original Image')
    plt.imshow(tiff_image, cmap='gray')
    plt.axis('off')

    plt.subplot(1, 2, 2)
    plt.title('1D Flattened Image')
    plt.imshow(flattened_image, cmap='gray')
    plt.axis('off')
    plt.tight_layout()
    plt.show()