#! /usr/bin/python3 
# Last edited on 2025-12-18 22:37:34 by stolfi

# Functions to work on color images for VMS hacking.

import sys, re, os, numpy
from sys import stderr as err
from math import floor, inf, sqrt, comb, isfinite, pow
from PIL import Image, PngImagePlugin
from error_funcs import file_line_error, prog_error, arg_error
from process_funcs import bash, run_command

def read_color_png_image(cfile, maxval):
  # Reads a PNG image from file {cfile}. Converts the pixels from
  # integers to floats by dividing them by the specified {maxval}, without
  # any gamma correction, and clipping the result to {[0_1]}. The result
  # has shape {(ny,nx,nc)}, where {ny} and {nx} are the height and width of
  # the image, and {nc} is 1 for greyscale, 2 for greyscale+alpha,
  # 3 for RGB, and 4 for rgb+alpha.
  #
  # Pixels greater than {maxval} are counted and a report is written to
  # {stderr} at the end.
  
  err.write(f"reading {cfile} ...\n")
  P_PIL = Image.open(cfile)
  err.write(f"numpy.asarray ...\n")
  P = numpy.asarray(P_PIL)
  # debug("P.shape", numpy.shape(P))
  if len(numpy.shape(P)) == 2:
    err.write(f"upgrading from 2 to 3 axes ...\n")
    # Change shape to 3 axes where the third is trivial:
    ny, nx = numpy.shape(P)
    Q = numpy.zeros((nx,ny,1))
    Q[:,:,0] = P
    P = Q
    # debug("P.shape", numpy.shape(P))
  assert len(numpy.shape(P)) == 3
  ny, nx, nc = numpy.shape(P)
  
  # Convert pixels to floats and count outliers:
  err.write(f"allocating float image ...\n")
  C = numpy.zeros((ny, nx, nc,))

  err.write(f"converting to float ...\n")
  for ic in range(nc):
    nbig = 0; true_vlo = 999999999; true_vhi = 0
    for iy in range(ny):
      for ix in range(nx):
        ismp = P[iy,ix,ic]
        assert ismp >= 0
        true_vlo = min(true_vlo, ismp);
        true_vhi = max(true_vhi, ismp);
        if ismp > maxval: nbig += 1; ismp = maxval
        fval = float(ismp)/float(maxval)
        C[iy,ix,ic] = max(0.0, min(1.0, fval))
    if nbig > 0:
      err.write(f"channel {ic}: true sample range = {{{true_vlo:5d}..{true_vhi:5d}}}\n")
      fbig = float(nbig) / (nx*ny)
      err.write(f" {nbig:5d} samples over {maxval:5d} ({100*fbig:6.2f}%)\n")

  return C
  # ----------------------------------------------------------------------

def undo_SRGB_gamma_encoding(in_img):
  # Applies the sample decoding function ("gamma") of the SRGB standard.
  split = 0.04045
  gamma = 2.4
  ny, nx, nc = numpy.shape(in_img)
  for ic in range(nc):
    for iy in range(ny):
      for ix in range(nx):
        val = in_img[iy,ix,ic]
        sgn = +1 if val >= 0 else -1
        val = abs(val)
        if val != 0 and val != 1:
          if val <= 0.04045:
            val = val/12.92
          else:
            val = pow((val + 0.055)/1.055, gamma)
          in_img[iy,ix,ic] = sgn*val
  return 
  # ----------------------------------------------------------------------

def get_channel_quantiles(G, ic, M, frin):
  # Assumes that {G} is a {numpy} array with shape {(ny,nx,nc)}
  # representing an image with {nc} channels, and {ic} is 
  # a channel index in {0..nc-1}.
  #
  # Returns the tightest values {smin,smax} so that the fraction of the
  # samples of {G} that fall in the {[smin _
  # smax]}, is at least {frin}, equal numbers falling below {smin} and
  # above {smax}. 
  #
  # If {M} is not {None}, it must be a {numpy} array with shape
  # {(ny,nx)}. The quantiles above will be computed considering only
  # those pixels of {G} where {M} is nonzero.
  #
  ny, nx, nc = numpy.shape(G)
  assert numpy.shape(M) == (ny, nx), "bad mask shape"
  assert 0 <= ic and ic < nc, "bad {ic}"
  cvals = [ \
    G[iy,ix,ic] \
    for ix in range(nx) \
      for iy in range(ny) \
        if M is None or M[iy,ix] != 0 \
  ]

  # Normalize to {95%} of samples spanning {[0_1]}, with clipping:
  assert 0.0 <= frin and frin <= 1.0, "bad {frin}"
  smin = numpy.quantile(cvals, 0.5 - frin/2)
  smax = numpy.quantile(cvals, 0.5 + frin/2)
  smid = (smin + smax)/2
  srad = (smax - smin)/2
  return smin, smax
  # ----------------------------------------------------------------------

def normalize_RGB_image(RGB, M_avg, M_dev, Ymin, Ymax):
  # Assumes that {RGB} is a {numpy} array representng an image with
  # positive and negative samples and arbitrary mean, with shape
  # {(ny,nx,nc)}
  # 
  # The procedure scales and shifts all samples of {RGB} so that those
  # samples where the mask image {M_avg} is nonzero have average 0.50, 
  # and those samples where {M_dev} is nonzero have deviation 0.25. 
  #
  # Then hard-clips the brightness {Y} of every pixel, preserving its UV
  # components, to the range {[Ymin_Ymax]}, which must be a proper
  # subset of {[0.05_0.95]}.
  #
  # Then adjusts the UV components of all pixels, preserving their
  # {Y} value, so that most of those selected by {M_dev} are within the RGB
  # cube.  
  #
  # Finally soft-clips all pixels that fall outside the unit cube to that
  # cube, while preserving their {Y} and hue.
  #
  # Either mask image {M_avg} or {M_dev} is {None}, all pixels are
  # considered in the corresponding statistic. Othherwise, each image
  # must be an array with shape {(ny,nx)}. The selected pixels will be
  # those where the mask is nonzero.

  ny, nx, nc = numpy.shape(RGB)
  assert nc == 3, "image {RGB} must have three channels"
  if M_avg is not None: assert numpy.shape(M_avg) == (ny, nx)
  if M_dev is not None: assert numpy.shape(M_dev) == (ny, nx)
  assert 0.05 <= Ymin and Ymax - Ymin > 0.05 and Ymax < 0.95, \
    f"invalid {Ymin = }, {Ymax = }"
      
  err.write("normalizing channel means and deviations ...\n")
  for ic in range(nc):
    samples_for_avg = [ \
        RGB[iy,ix,ic] \
        for ix in range(nx) \
          for iy in range(ny) \
            if M_avg is None or M_avg[iy,ix] != 0 \
      ]
    savg = numpy.mean(samples_for_avg)
    samples_for_dev = [ \
        RGB[iy,ix,ic] - savg \
        for ix in range(nx) \
          for iy in range(ny) \
            if M_dev is None or M_dev[iy,ix] != 0 \
      ]
    sdev = numpy.std(samples_for_dev)
    err.write(f"  channel {ic} avg = {savg:.4f} dev = {sdev:.4f}\n") 
    # Normalize average and deviation:
    vavg = 0.50
    vdev = 0.25
    for iy in range(ny):
      for ix in range(nx):
        fsmp = vavg + vdev*(RGB[iy,ix,ic] - savg)/sdev
        RGB[iy,ix,ic] = fsmp
    
  err.write("converting to YUV ...\n")
  # Will hard-clip {Y} to {[Ymin_Ymax]}:
  YUV = YUV_image_from_RGB_image(RGB)
  
  err.write("clipping Y value ...\n")
  clip_Y_in_YUV_image(YUV, Ymin, Ymax)

  err.write("applying global UV scaling factor ...\n")
  sUV = choose_UV_scale(YUV, M_dev)
  err.write(f"  factor = {sUV:.6f}\n") 
  YUV[:,:,1] = sUV*YUV[:,:,1]
  YUV[:,:,2] = sUV*YUV[:,:,2]
    
  err.write("converting back to RGB ...\n")
  # Will clip to unit cube preserving {Y} and hue:
  RGB = RGB_image_from_YUV_image(YUV)

  return
  # ----------------------------------------------------------------------

def clip_Y_in_YUV_image(YUV, Ymin, Ymax):
  # Assumes that {YUV} a {numpy} array representing an image in YUV
  # color space, with shape {(ny,nx,3)}.
  #
  # Hard-clipes the {Y} channel of the image the range {[Ymin_Ymax]},
  # which must be a proper subset of {[0.05_0.95]}
  
  ny,nx,nc = numpy.shape(YUV)
  assert nc == 3, "bad nc"
  assert 0.05 <= Ymin and Ymax - Ymin > 0.05 and Ymax < 0.95, \
    f"invalid {Ymin = }, {Ymax = }"
  ndark = 0
  nlite = 0
  Ymin_true = +inf; Ymax_true = -inf;
  for iy in range(ny):
    for ix in range(nx):
      Y = YUV[iy,ix,0]
      Ymin_true = min(Ymin_true, Y)
      Ymax_true = min(Ymax_true, Y)
      # Clip the {Y} value to {[Ymin _ Ymax]}:
      if Y < Ymin:
        Y = Ymin; ndark += 1
      elif Y > Ymax:
        Y = Ymax; nlite += 1
      YUV[iy,ix,0] = Y

  if ndark > 0 or nlite > 0:
    err.write(f"  {ndark:5d} pixels too dark, {nlite:5d} too bright (clipped)\n")
  return
  # ----------------------------------------------------------------------

def normalize_YUV_image(YUV, M_avg, M_dev, Yavg, Ydev, UVdev):
  # Assumes that {YUV} is a {numpy} array representing image in the
  # linear YUV color space, with positive and negative samples and
  # arbitrary mean, with shape {(ny,nx,3)}.
  # 
  # The procedure scales and shifts the samples in each channel, so that
  # they are normalized. The normalization leaves the Y channel (channel
  # 0) with mean {Yavg} over the subsect selected by the mask {M_avg} and
  # deviation {Ydev} on the subset selected by the mask image {M_dev};
  # while the U and V channels (1 and 2) are left with zero mean over
  # {M_avg} and deviation {UVdev} over {M_dev}.
  #
  # Either mask image {M_avg} or {M_dev} is {None}, all pixels are
  # considered in the corresponding statistic. Otherwise, each mask must
  # be a {numpy} array with shape {(ny,nx)}. The selected pixels will be
  # those where the mask is nonzero.

  ny, nx, nc = numpy.shape(YUV)
  assert nc == 3, "image {YUV} must have three channels"
  if M_avg is not None: assert numpy.shape(M_avg) == (ny, nx)
  if M_dev is not None: assert numpy.shape(M_dev) == (ny, nx)
      
  err.write("normalizing channel means and deviations ...\n")
  for ic in range(nc):
    samples_for_avg = [ \
        YUV[iy,ix,ic] \
        for ix in range(nx) \
          for iy in range(ny) \
            if M_avg is None or M_avg[iy,ix] != 0 \
      ]
    savg = numpy.mean(samples_for_avg)
    samples_for_dev = [ \
        YUV[iy,ix,ic] - savg \
        for ix in range(nx) \
          for iy in range(ny) \
            if M_dev is None or M_dev[iy,ix] != 0 \
      ]
    sdev = numpy.std(samples_for_dev)
    err.write(f"  channel {ic} avg = {savg:.4f} dev = {sdev:.4f}\n") 
    vavg = Yavg if ic == 0 else 0
    vdev = Ydev if ic == 0 else UVdev
    for iy in range(ny):
      for ix in range(nx):
        fsmp = vavg + vdev*(YUV[iy,ix,ic] - savg)/sdev
        YUV[iy,ix,ic] = fsmp
  return
  # ----------------------------------------------------------------------

def choose_UV_scale(YUV, M):
  # Assumes that {YUV} is a {numpy} array representing an image in the
  # linear YUV color space, with positive and negative samples and
  # arbitrary mean, with shape {(ny,nx,3)}.
  # 
  # Chooses a scaling factor for the U and V coordinates of {YUV} so
  # that most of the pixels will be inside the {[0 _ 1]} cube when the
  # image is converted to RGB.
  #
  # Requires that the Y values in {YUV} be in the range {[0.2 _ 0.8]}.
  #
  # If {M} is not {None}, it must be a {numpy} array with shape
  # {(ny,nx)}. In this case, the scaing factor will be computed
  # considering only those pixels of {YUV} where {M} is nonzero.
  
  ny, nx, nc = numpy.shape(YUV)
  assert nc == 3, "image {YUV} msut have three channels"
  if M is not None: assert numpy.shape(M) == (ny, nx)
      
  frmin = 0.05
  svals = []  
  for ix in range(nx):
    for iy in range(ny):
      if M is None or M[iy,ix] != 0:
        Y, U, V = YUV[iy,ix,:]
        assert Y >= 0.19999 and Y <= 0.80001, "bug: Y range" 
        RGB_UV = RGB_pixel_from_YUV_pixel((0,U,V))
        sUV = 1.0e100
        for kc in range(3):
          if RGB_UV[kc] > 1.0e-100:
            sUV = min (sUV, (1.0-Y)/RGB_UV[kc])
          elif RGB_UV[kc] < -1.0e-100:
            sUV = min(sUV, (0.0-Y)/RGB_UV[kc])
        svals.append(sUV)
  sUV_max = numpy.quantile(svals, frmin)
  return sUV_max
  # ----------------------------------------------------------------------

def YUV_image_from_RGB_image(RGB):
  # Assumes that {RGB} is a {numpy} array representing an image in linear RGB
  # color space, with shape {(ny,nx,3)}. 
  #
  # Converts each pixel by linear transformation from RGB space to
  # linear YUV, still with positive and negative samples.

  ny, nx, nc = numpy.shape(RGB)
  assert nc == 3, "bug {nc}"

  YUV = numpy.zeros((ny, nx, 3))
  for iy in range(ny):
    for ix in range(nx):
      YUV[iy,ix,:] = YUV_pixel_from_RGB_pixel(RGB[iy,ix,:])
  return YUV
  # ----------------------------------------------------------------------

def RGB_image_from_YUV_image(YUV):
  # Assumes that {YUV} a {numpy} array representing an image in YUV
  # color space, with shape {(ny,nx,3)}.
  #
  # Converts each pixel by linear transformation from YUV space to
  # linear RGB, after down-scaling the UV components of each pixel, if
  # necessary, so that the result lies in the unit cube.
  #
  # If the {Y} component of a pixel is negative, the result is
  # black.  If it is greater than 1, the result is white.
  #
  ny, nx, nc = numpy.shape(YUV)
  assert nc == 3, "bug {nc}"
  ndark = 0; nlite = 0; nclip = 0
  RGB_min = [+inf]*3; RGB_max = [-inf]*3
  RGB = numpy.zeros((ny,nx,nc))
  for iy in range(ny):
    for ix in range(nx):
      yuv = YUV[iy,ix,:]
      rgb = RGB_pixel_from_YUV_pixel(yuv)
      ok = yuv[0] >= 0 and yuv[0] <= 1.0
      for ic in range(3): 
        RGB_min[ic] = min(RGB_min[ic], rgb[ic]);
        RGB_max[ic] = max(RGB_max[ic], rgb[ic]);
        ok = ok and (rgb[ic] >= 0 and rgb[ic] <= 1.0)
      if yuv[0] < 0:
        rgb = numpy.full((3,), 0.0); ndark += 1
      elif yuv[0] > 1:
        rgb = numpy.full((3,), 1.0); nlite += 1
      elif not ok:
        rgb = clip_RGB_pixel_to_unit_cube(rgb)
        nclip += 1
      RGB[iy,ix,:] = rgb
  if ndark > 0 or nlite > 9 or nclip > 0:
    err.write(f"  raw ranges:\n")
    for ic in range(3):
      err.write(f"   channel {ic}: [{RGB_min[ic]:7.3f} _ {RGB_max[ic]:7.3f}]\n")
    err.write(f"  {nclip:5d} pixels outside cube (clipped)\n")
    err.write(f"  {ndark:5d} Y values were below zero, {nlite:5d} above 1\n")
  return RGB
  # ----------------------------------------------------------------------

def clip_RGB_image_to_unit_cube(RGB):
  # If any pixel of {RGB} is outside the unit cube, 
  # brings it into the unit cube with {clip_RGB_pixel_to_unit_cube(rgb)}
  #
  ny, nx, nc = numpy.shape(RGB)
  assert nc == 3, "bug {nc}"

  nout = 0; RGB_min = +inf; RGB_max = -inf
  for iy in range(ny):
    for ix in range(nx):
      rgb = RGB[iy,ix,:]
      ok = True
      for ic in range(3): 
        if rgb[ic] < 0:
          nout += 1; RGB_min = min(RGB_min, rgb[ic]); ok = False
        elif rgb[ic] > 1:
          nout += 1; RGB_max = max(RGB_max, rgb[ic]); ok = False
      if not ok:
        rgb = clip_RGB_pixel_to_unit_cube(rgb)
      RGB[iy,ix,:] = rgb
  if nout > 0:
    err.write(f"clip_RGB_image_to_unit_cube: {nout:5d} pixels outside cube\n")
    err.write(f"Actual range [{RGB_min:7.3f} _ {RGB_max:7.3f}]\n")
  return
  # ----------------------------------------------------------------------

def clip_RGB_pixel_to_unit_cube(rgb):
  # If {rgb} is outside the unit cube, 
  # brings it into the unit cube by preserving its Y (luminosity)
  # component and downscaling the chroma component by 
  # a suitable factor.
  # 
  # If Y is negative, the result is black [0,0,0]. If Y is greater than 1,
  # the result is white [1,1,1].
  #
  Y = Y_from_RGB_pixel(rgb)
  
  if Y <= 0: return [0,0,0]
  if Y >= 1: return [1,1,1]
  
  # Find the scaling factor {cscale} for the chroma component that
  # will keep the result inside the RGB cube:
  cscale = 1;
  for ic in range(3):
    chri = rgb[ic] - Y
    if chri > 0:
      # We want {Y + cscale*chri <= 1}:
      cscale = min(cscale, (1-Y)/chri)
    elif chri < 0:
      # We want {Y + cscale*chri >= 0}:
      cscale = min(cscale, (0-Y)/chri)
  rgb_out = rgb.copy()
  if cscale < 1.0:
    for ic in range(3):
      rgb_out [ic] = Y + cscale*(rgb[ic] - Y)
  return rgb_out 
  # ----------------------------------------------------------------------

def Y_from_RGB_pixel(rgb):
  # Returns the Y (luminance) value of an RGB triple {rgb}.
  # 
  # A linear function from {\RR^3} to {\RR}. Does not care
  # if any input or output coord is negative or 
  # greater than 1.
  #
  return 0.298911*rgb[0] + 0.586611*rgb[1] + 0.114478*rgb[2]
  # ----------------------------------------------------------------------

def YUV_pixel_from_RGB_pixel(rgb):
  R, G, B = rgb
  Y = +0.298911*R + 0.586611*G +0.114478*B;
  U = -0.147000*R - 0.289000*G +0.436000*B;
  V = +0.615000*R - 0.515000*G -0.100000*B;
  return [Y, U, V]
  # ----------------------------------------------------------------------

def RGB_pixel_from_YUV_pixel(yuv):
  # Converts a pixel color from YUV coords {yuv} to RGB coords.
  #
  # The standard linear map from {\RR^3} tp {\RR^3}.
  # Does not care if any input or output coord is negative or 
  # greater than 1.
  #
  Y, U, V = yuv;
  R = +1.000000*Y -0.001164*U +1.139704*V;
  G = +1.000000*Y -0.395735*U -0.580624*V;
  B = +1.000000*Y +2.030875*U -0.000606*V;
  return [R, G, B]
  # ----------------------------------------------------------------------

def write_image_as_png(rfile, R):
  # Assumes {R} is a numpy array with shape {(ny,nx,nc)}
  # where {nc} is either 1 or 3. Writes it out to file "{rfile}".
  # 
  # ??? !!! Crock -- do better !!!
  # 
  dim = len(numpy.shape(R))
  if dim == 2:
    ny, nx = numpy.shape(R); nc = 1
  elif dim == 3:
    ny, nx, nc = numpy.shape(R)
  else:
    assert False, "wrong array shape"
  assert nc == 1 or nc == 3, "bug {nc}"
  ext = "pgm" if nc == 1 else "ppm" # File name extension.
  magic = "P2" if nc == 1 else "P3" # Magic code for PGM or PPM file.
  inSpace = "linearGray" if nc == 1 else "linearRGB"
  outSpace = "Gray" if nc == 1 else "sRGB"
  tfile = f"/tmp/{os.getpid()}-R.{ext}"
  # Write an ascii PGM or PPM:
  maxval = 1023
  with open(tfile, "w") as wr:
    wr.write(f"{magic}\n")
    wr.write(f"{nx} {ny}\n")
    wr.write(f"{maxval}\n")
    for iy in range(ny):
      for ix in range(nx):
        if ix % 20 == 0: wr.write("\n")
        for ic in range(nc):
          fsmp = R[iy,ix,ic] if dim == 3 else R[iy,ix]
          ival = int(floor(maxval*fsmp + 0.5))
          if ival < 0 or ival > maxval: prog_error("int sample out of bounds")
          wr.write(f" {ival}")
      wr.write("\n")
    wr.close()
  
  # Convert to PNG:
  bash\
    ( f"convert {tfile}" + \
      f"    -set colorspace {inSpace}" + \
      f"    -depth 8" + \
      f"    -colorspace {outSpace}" +
      f"  {rfile}"
    )
  return
  # ----------------------------------------------------------------------