import imageio.v3 as iio
import numpy as np
import scipy.interpolate as si

# A X B x 3
def process(img, nx = 100, ny = 100, lowq = 0.02, highq = 0.98):
    img_bw = np.sum(img, axis = 2).astype(float)
    dx, dy = img_bw.shape

    xs = np.linspace(0, dx, nx + 1).astype(int)
    xs[0] = 0
    xs[-1] = dx
    ys = np.linspace(0, dy, ny + 1).astype(int)
    ys[0] = 0
    ys[-1] = dy

    low = np.zeros((nx, ny))
    high = np.zeros((nx, ny))

    for i in range(nx):
        for j in range(ny):
            chunk = img_bw[xs[i] : xs[i + 1], ys[j] : ys[j + 1]]
            low[i][j], high[i][j] = np.quantile(chunk, (lowq, highq))

    low = np.clip(low, 0, 255 * 3 / 8)

    xs_mid = (xs[1:] + xs[:-1] - 1) / 2
    ys_mid = (ys[1:] + ys[:-1] - 1) / 2
    XY = np.moveaxis(np.meshgrid(np.arange(dx), np.arange(dy), indexing = 'ij'), 0, 2)

    m = 'linear'

    low_fine = si.RegularGridInterpolator((xs_mid, ys_mid), low, bounds_error = False, fill_value = None, method = m)((XY))
    high_fine = si.RegularGridInterpolator((xs_mid, ys_mid), high, bounds_error = False, fill_value = None, method = m)((XY))

    adj_lum = (img_bw - low_fine) / (high_fine - low_fine)

    adj_lum = np.clip(adj_lum, 0, 1)
    adj = img * ((adj_lum * (255 * 3) / (img_bw + 1))[:, :, None])
    adj = np.clip(adj, 0, 255)
    adj = adj.astype(np.uint8)

    return adj

def run():
    import sys
    file_in = sys.argv[1]
    file_out = sys.argv[2]

    data = iio.imread(file_in)
    adj = process(data, 20, 20, lowq = 0.01, highq = 0.60)
    iio.imwrite(file_out, adj)

if __name__ == '__main__':
    run()