#! /usr/bin/python3 
# Last edited on 2025-09-01 05:20:01 by stolfi

# Functions to work on linear grayscale images for VMS hacking.
  
import sys, re, os, numpy
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_mask_image(mfile):
  # Reads a mask image from file {mfile}. The result has shape
  # {(ny,nx)}, where {ny} and {nx} are the height and width of the
  # image, and pixel values either 0 or 255.
  
  M = numpy.asarray(Image.open(mfile))
  Msh = numpy.shape(M)
  assert len(Msh) == 2 or (len(Msh) == 3 and Msh[2] == 1), f"mask {mfile} is not grayscale"
  # debug("M.shape", numpy.shape(M))
  if len(Msh) == 3:
    # Reduce image from 3 to2 channels:
    M = M[:,:,0]; Msh = numpy.shape(M)
    # debug("M.shape", numpy.shape(M))
  assert len(numpy.shape(M)) == 2
  ny, nx = numpy.shape(M)

  # Print statistics:
  for tag, func in (("min", numpy.min), ("max", numpy.max)):
    val = func(M, axis=(0,1))
    ct = (M == val).sum()
    sys.stderr.write(f"  {tag} = {val:5d} ({ct:8d} elems, {ct/(ny*nx):10.8f})\n")  
  
  # debug("M[0,0]", M[0,0])
  return M
  # ----------------------------------------------------------------------

def read_png_image(cfile, maxval):
  # Reads a clip 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)}, where {ny} and {nx} are the height and width of
  # the image.
  #
  # Pixels greater than {maxval} are counted and a report is written to
  # {stderr} at the end.
  
  sys.stderr.write(f"reading {cfile} ...\n")
    
  P = numpy.asarray(Image.open(cfile))
  # debug("P.shape", numpy.shape(P))
  if len(numpy.shape(P)) == 3:
    # Reduce image from 3 to2 channels:
    nc = numpy.size(P,2)
    P = P[:,:,nc-1]
    # debug("P.shape", numpy.shape(P))
  assert len(numpy.shape(P)) == 2
  ny, nx = numpy.shape(P)
  
  # Convert pixels to floats and count outliers:
  C = numpy.zeros((ny, nx))
  nbig = 0; true_vlo = 999999999; true_vhi = 0
  for iy in range(ny):
    for ix in range(nx):
      ismp = P[iy,ix]
      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] = max(0.0, min(1.0, fval))
  if nbig > 0:
    sys.stderr.write(f"float range = {{{true_vlo:5d}..{true_vhi:5d}}}")
    fbig = float(nbig) / (nx*ny)
    sys.stderr.write(f" {nbig:5d} samples over {maxval:5d} ({100*fbig:6.2f}%)\n")

  return C
  # ----------------------------------------------------------------------
    
def blur_image(C):
  # Given a grayscale image {C} with shape {(ny,nx,)} and values in {[0_1]},
  # returns another image with same shape that is 
  # a blurred version of the same.
  #
  # Currently uses ImageMagick's "convert".  There should be
  # a better way...
  
  ny, nx = numpy.shape(C)

  nred = 10
  assert 100%nred == 0, "bug bad nred"
  assert ny%nred == 0 and nx%nred == 0, f"image dims must be multiples of {nred}"
  
  tpref = f"/tmp/{os.getpid()}"
  ifile = f"{tpref}-inp.png"
  rfile = f"{tpref}-red.png"
  ofile = f"{tpref}-out.png"
  write_image_as_gray_png(ifile, C)
  pred = 100//nred
  bash(f"convert {ifile} -set colorspace LinearGray -resize '{pred}%' -colorspace LinearGray -depth 16 {rfile}")
  pmag = 100*nred
  bash(f"convert {rfile} -set colorspace LinearGray -resize '{pmag}%' -colorspace LinearGray -depth 16 {ofile}")
  maxval = 65535
  B = read_png_image(ofile, maxval)
  assert numpy.shape(B) == (ny,nx,)
  bash(f"rm {ifile} {rfile} {ofile}")
  return B
  # ----------------------------------------------------------------------
  
def write_uint_array_as_png(pfile, P):
  # Expects {P} to be a {numpy} array with type {numpy.uint8} or
  # {numpy.uint16} with shape {(ny,nx,nc)} or {(ny,nx)}. In the first
  # case, {nc} must be in {1..4}. Writes the image to file "{pfile}" as
  # a linear (grayscale or RGB) PNG image.
  #
  sh = numpy.shape(P); dim = len(sh)
  if dim == 2:
    ny, nx = sh; nc = 1
  elif dim == 3:
    ny, nx, nc = sh
  else:
    assert False, f"invalid image shape {sh}"

  if nc == 1:
    if P.dtype == numpy.uint8:
      pil_mode = "I;8"
    elif P.dtype == numpy.uint16:
      pil_mode = "I;16"
    else:
      assert False, "invalid matrix dtype for grayscale"
  elif nc == 2:
    assert P.dtype == numpy.uint8, "invalid matrix dtype for gray+alpha"
    pil_mode = "LA"
  elif nc == 3 or nc == 4:
    assert P.dtype == numpy.uint8, "invalid matrix dtype for RGB"
    pil_mode = "RGB" if nc == 3 else "RGBA"
  else: 
    assert False, f"invalid channel count {nc}"

  img = Image.fromarray(P, mode=pil_mode)
  img.save(pfile, pnginfo=PngImagePlugin.PngInfo())
  return
  # ----------------------------------------------------------------------

def write_image_as_gray_png(pfile, C):
  # Assumes that {C} is a grayscale image represented as a float 
  # matrix with samples in {[0_1]}.  Writes {C} to file
  # "{pfile}" as a linear grayscale PNG image file.
  
  sh = numpy.shape(C); dim = len(sh)
  if dim == 2:
    ny, nx = sh
  elif dim == 3:
    ny, nx, nc = sh
    assert nc == 1, f"invalid channel count {nc}"
  else:
    assert False, f"invalid image shape {sh}"
  
  maxval = 65535
  sys.stderr.write(f"{pfile}: ")
  P = quantize_image(C, maxval)
  write_uint_array_as_png(pfile, P)
  return
  # ----------------------------------------------------------------------

def quantize_image(C, maxval):
  # Expects {C} to be a float array with samples in {[0_1]} and shape {(ny,nx,nc,)}
  # or {(ny,nx,)}.  
  # 
  # Returns {C} converted to a matrix of integers in {0..maxval} by linear scale.
  #
  # The data type of the resutlt will be {numpy.uint8} if {maxval}
  # is in {1..255}, and {numpy.uint16} if it is in {256..65535}
  
  assert isinstance(maxval,int) and maxval >= 1 and maxval <= 65535, \
    f"bad {{maxval}} = {maxval}"

  sh = numpy.shape(C); dim = len(sh)
  if dim == 2:
    ny, nx = sh; nc = 1
    fval_min = numpy.min(C, axis=(0,1))
    fval_max = numpy.max(C, axis=(0,1))
  elif dim == 3:
    ny, nx, nc = sh
    fval_min = numpy.min(C, axis=(0,1,2))
    fval_max = numpy.max(C, axis=(0,1,2))
  else:
    assert False, f"invalid image shape {sh}"

  dtype = numpy.uint8 if maxval <= 255 else numpy.uint16
  
  P = numpy.zeros(sh, dtype=dtype)
  fval_min = +inf;    fval_max = -inf
  ival_min = 9999999; ival_max = -1
  nsma = 0; nbig = 0
  for iy in range(ny):
    for ix in range(nx):
      for ic in range(nc):
        fval = C[iy,ix] if dim == 2 else C[iy,ix,ic]
        assert fval >= 0, f"invalid sample {fval:24.16e}"
        fval_min = min(fval_min, fval)
        fval_max = max(fval_max, fval)
        if fval < 0:
          nsma += 1; fval = 0.0
        elif fval > 1.0:
          nbig += 1; fval = 1.0
        ival = int(floor(maxval*fval + 0.5))
        assert 0 <= ival and ival <= maxval, "bug quantization"
        ival_min = min(ival_min, ival)
        ival_max = max(ival_max, ival)
        if dim == 2: 
          P[iy,ix] = ival
        else:
          P[iy,ix,ic] = ival
  sys.stderr.write(f"quantize_image: float range [{fval_min:8.6f} _ {fval_max:8.6f}]")
  sys.stderr.write(f" int range {{{ival_min:5d} .. {ival_max:5d}}}\n")
  if nsma > 0 or nbig > 0:
    sys.stderr.write(f" {nsma:6d} input samples below 0.00")
    sys.stderr.write(f" {nbig:6d} input samples above 1.00")
  return P
  # ----------------------------------------------------------------------
  
def write_image_as_gray_png_scaled(pfile, C, clo, chi, plo, phi):
  # Writes {C} as an grayscale linear PNG file,
  # mapping {[clo _ chi]} affinely to {[plo _ phi]} and clipping
  # the result to {[0 _ 1]}.
  #
  ny, nx = numpy.shape(C)
  
  ptag = re.sub(r"^.*/", "", pfile)
  P = clip_image_squashed(C, ptag, clo, chi, plo, phi, 0.0, 1.0)
  write_image_as_gray_png(pfile, P)
  return 
  # ----------------------------------------------------------------------
  
def map_value_squashed(val, vlo, vhi, clo, chi, cmin, cmax):
  # Maps value {val} in the range {[vlo _ vhi]} affinely to the range
  # {[clo _ chi]}.  
  #
  # Values below {vlo} are squashed non-linearly
  # to the range {[cmin _ clo]}.  Values above {vhi} are squashed
  # non-linearly to the range {[chi _ cmax]}. 
  #
  # If {cmin=clo}, values below {vlo} are simply clipped to {cmin}.
  # otherwise the mapping is C1 at {vlo}. Likewise, if {cmax=chi}
  # values above {vhi} are simly clipped do {cmax}, otherwise
  # the function is C1 at {vhi}.
  #
  assert vlo < vhi, "bad vlo,vhi"
  assert cmin <= clo and clo < chi and chi <= cmax, "bad cmin,clo,chi,cmax"
  
  def squash(tt, RR):
    # A function of {tt} that is 1 for {tt=1}, then grows monotonically
    # with initial slope 1 and asymptotic value {RR}. Valid only for {tt
    # >= 1} and {RR >= tt}

    assert tt >= 1
    assert RR > 1
    ss = (tt - 1)/(RR - 1)
    qq = ss*(ss + 4)
    uu = qq/(qq + 4)
    rr = 1 + uu*(RR - 1)
    return rr
    # ....................................................................
    
  # Map {val} so that {[vlo_vhi]} goes to {[-1 _ +1]}: 
  t = (2*val - (vlo + vhi))/(vhi - vlo)
  
  # Compute {r} that is same as {t} in {[-1 _ +1]} but
  # tends to the proper values beyond that range:
  if val < vlo:
    if cmin == clo:
      r = -1.0
    else:
      R = (2*cmin - (clo + chi))/(chi - clo)
      r = -squash(-t, -R)
  elif val > vhi:
    if cmax == chi:
      r = +1.0
    else:
      R = (2*cmax - (clo + chi))/(chi - clo)
      r = squash(t, R)
  else:
    r = t

  # TESTING:
  # for k in range(3):
  #   ttt = 1.0 + k*0.000001
  #   RRR = 2.01
  #   rrr = squash(ttt, RRR)
  #   sys.stderr.write(f"squash({ttt}, {RRR}) = {rrr}\n")
  # assert abs(squash(1.00001, RRR) - 1.00001) < 0.0001, "bug squash"

  c = 0.5*((clo + chi) + r*(chi - clo))
  assert c >= cmin and c <= cmax, "bug squash"
  return c
  # ----------------------------------------------------------------------
 
def map_image_to_unit_range(B, shift, M, bname):
  # Given a monochrome image {B} as an array with shape {(ny,nx,)}
  # returns a copy of it with values mapped to span {[0_1]}.
  #
  # If {shift} is true, the mapping will be such that most
  # samples will be mapped shifted and scaled by the same amounts,
  # with smooth clipping to {[0_1]} at both ends of the range.
  #
  # If {shift} is false, the samples in {B} are expected to
  # be non-negative. The mapping will be such that most samples will
  # be scaled (not shifted) by the same factor, with soft clipping to 1
  # at the upper end and hard clipping at 0 at the lower end.
  #
  # If {M} is not {None}, it should be a binary mask image with shape
  # {(ny,nx,)}. Then the mapping will be selected by considering only
  # those pixels {iy,ix} where {M} is nonzero (but will still be applied to the
  # whole image.
  #
  
  ny, nx = numpy.shape(B)
  if M is not None: assert numpy.shape(M) == (ny,nx,)
  
  sys.stderr.write(f"  finding approx range of sample values ...\n")
  samples = []
  for iy in range(ny):
    for ix in range(nx):
      if M is None or M[iy,ix] != 0:
        samples.append(B[iy,ix])

  frac_val = 0.001
  bmin, bmax, fmin, fmax = choose_sample_map(samples, frac_val, shift)

  sys.stderr.write(f"  shifting and scaling images to [0_1] ...\n")
  F = clip_image_squashed(B, bname, bmin, bmax, fmin, fmax, 0.0, 1.0)

  return F
  # ----------------------------------------------------------------------
 
def map_image_to_unit_deviation(B, shift, M, bname):
  # Given amonochrome image {B} as an array with shape
  # {(ny,nx,)} returns a copy of it with values mapped so that most
  # of them have unit deviation relative to their mean.
  #
  # If {shift} is false, the mean value of the samples will be
  # preserved, and the scaling will be applied to the deviations from
  # that mean value. If {shift} is true, the values will also be shifted
  # so that their mean becomes zero.
  #
  # If {M} is not {None}, it should be a binary mask image with shape
  # {(ny,nx,)}. Then the mapping will be selected by considering only
  # those pixels {iy,ix} of the {B} images where {M} is nonzero (but
  # will still be applied to every whole image).
  #

  samples = []
  for iy in range(ny):
    for ix in range(nx):
      if M is None or M[iy,ix] != 0:
        samples.append(B[iy,ix])
  
  avg = numpy.mean(samples)
  dev = numpy.std(samples, mean = avg)
  
  F = (B - avg)/dev
  if not shift: F += avg
  return F
  # ----------------------------------------------------------------------
  
def choose_sample_map(samples, frac, shift):
  # Chooses the parameters for an affine or linear map 
  # that takes most of the given {samples} to span {[0_1]}.
  # 
  # If {shift} is true, the mapping will be affine (shift and scale)
  # for all samples between the lower and upper {frac/2} quantiles.
  #
  # If {shift} is false, the samples are expected to
  # be non-negative. The mapping will be linear (just scale, no shift)
  # for all samples between 0 and upper {frac} quantile.
  #
  if shift:
    fmin = frac/2;     bmin = numpy.quantile(samples, fmin)     
    fmax = 1 - frac/2; bmax = numpy.quantile(samples, fmax)
  else:
    fmin = 0.0;        bmin = 0.0
    fmax = 1 - frac;   bmax = numpy.quantile(samples, fmax)
    assert min(samples) >= 0.0, "negative samples"
  return bmin, bmax, fmin, fmax
  # ----------------------------------------------------------------------

def clip_image_squashed(C, cname, clo, chi, vlo, vhi, vmin, vmax):
  # Given a greyscale image {C} with shape {(ny,nx,)},
  # maps all samples from the range {[clo _ chi]} to {[vlo _ vhi]}.
  # clipping hard to {[vmin _ vmax]}. 
    
  ny, nx = numpy.shape(C)
  assert clo < chi, "bad input range"
  assert vmin <= vlo and vlo < vhi and vhi <= vmax, "bad output ranges"
  
  sys.stderr.write(f"  scaling {cname} from [{clo:9.6f} _ {chi:9.6f}] to [{vlo:9.6f} _ {vhi:9.6f}] clipped at [{vmin:9.6f} _ {vmax:9.6f}] ...\n")

  P = numpy.zeros((ny,nx,))
  cval_min = +inf; pval_min = +inf; nsma = 0
  cval_max = -inf; pval_max = -inf; nbig = 0
  nbad = 0; # Infinities or NANs.
  for iy in range(ny):
    for ix in range(nx):
      cval = C[iy,ix]
      if isfinite(cval):
        cval_min = min(cval_min, cval);
        cval_max = max(cval_max, cval);
        if cval < clo:
          nsma += 1; 
        elif cval > chi:
          nbig += 1
        # Soft clip to {[vmin _ vmax]}:
        pval = map_value_squashed(cval, clo, chi, vlo, vhi, vmin, vmax)
        assert pval >= vmin and pval <= vmax, f"bug map_value_squashed {pval}"
      else:
        pval = 0.5
        nbad += 1
      pval_min = min(pval_min, pval)
      pval_max = max(pval_max, pval)
      P[iy,ix] = pval
  sys.stderr.write("  %s: input range [%+8.5f _ %+8.5f]" % (cname, cval_min, cval_max))
  sys.stderr.write(" mapped to [%+8.5f _ %+8.5f]\n" % (pval_min, pval_max))
  if nsma > 0 or nbig > 0 or nbad > 0:
    sys.stderr.write(f"    {nsma:6d} ({nsma/(ny*nx):.5f}) input samples were below {clo:9.6f}\n")
    sys.stderr.write(f"    {nbig:6d} ({nbig/(ny*nx):.5f}) input samples were above {chi:9.6f}\n")
    sys.stderr.write(f"    {nbad:6d} ({nbad/(ny*nx):.5f}) input samples were INF or NAN\n")
  return P
  # ----------------------------------------------------------------------