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

# Functions to work on multispectral 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
 
import vms_linear_gray_image_funcs as gfn

def read_clip_npy_images(page, clip, bands):
  # Reads the images of the clip {clip} in the spectral bands listed
  # in {bands}. The result is a {numpy} array {C} of floats with shape
  # {(ny,nx,nb)} with float samples in {[0 _ 1]}.

  C = None # For now. Later, the clips in array form.
  nb = len(bands)
  for ib in range(nb):
    band = bands[ib]
    cfile = f"pages/{page}/clips/{clip}/npy/{band}.npy"
    Ci = numpy.load(cfile)
    assert len(numpy.shape(Ci)) == 2, "clip image should be monochrome"
    if C is None: 
      ny, nx = numpy.shape(Ci)
      C = numpy.zeros((ny,nx,nb))
    else:
      assert numpy.shape(Ci) == (ny, nx), "bug inconsistent clip shape"
    C[:,:,ib] = Ci
  return C
  # ----------------------------------------------------------------------
 
def map_images_to_unit_range(G, shift, M, bands):
  # Given a set of monochrome images {G} as an array with shape {(ny,nx,nb,)}
  # returns a copy of it with values mapped to span
  # {[0_1]}
  #
  # If {shift} is true, the mapping will be such that most samples in
  # {G} 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 {G} are expected to
  # be non-negative. The mapping will be such that most samples in
  # {G} 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} of the {G} images where {M} is nonzero
  # (but will still be applied to every whole image).
  #
  
  ny, nx, nb = numpy.shape(G)
  
  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:
        for ib in range(nb):
          samples.append(G[iy,ix,ib])

  frac_val = 0.001
  gmin, gmax, fmin, fmax = gfn.choose_sample_map(samples, frac_val, shift)

  sys.stderr.write(f"  scaling images to [0_1] ...\n")
  F = numpy.zeros((ny,nx,nb,))
  for ib in range(nb):
    F[:,:,ib] = gfn.clip_image_squashed(G[:,:,ib], bands[ib], gmin, gmax, fmin,fmax, 0.0, 1.0)
  return F
  # ----------------------------------------------------------------------

def map_images_to_unit_deviation(G, shift, M, bands):
  # Given a set of monochrome images {G} as an array with shape
  # {(ny,nx,nb,)} 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 {G} images where {M} is nonzero
  # (but will still be applied to every whole image).
  #

  samples = [ \
      G[iy,ix,ib] \
      for ib in range(nb) \
        if M is None or M[iy,ix] != 0 \
          for iy in range(ny) \
            for ix in range(nx) \
    ]
  
  avg = numpy.mean(samples)
  dev = numpy.std(samples, mean = avg)
  
  F = (G - avg)/dev
  if not shift: F += avg
  return F
  # ----------------------------------------------------------------------

def scale_and_shift_to_fit_unit_cube(B, shift, M):
  # Receives a set {B} of multispectral images with shape {(ny,nx,nb,)},
  # a spectrum {shift} with shape {(nb,)}, and a two-level mask {M} with
  # shape {(ny,nx,)}.  The components of {B} may be negative or greater than 1.
  # 
  # Scales very pixel spectrum in {B} by a single factor {scale} and
  # adds the spctrum {shift}, where {scale} is the largest possible so
  # that every sample of most of the pixels selected by {M} lies
  # comfortably in the {[0 _ 1]} range. Then clips every sample of every
  # image to that range.
  #
  # Returns the scaled and shifted image {B}, and the {scale} factor used.
  
  ny, nx, nb = numpy.shape(B)
  assert numpy.shape(shift) == (nb,)
  assert numpy.shape(M) == (ny,nx,)
  
  # Find max scaling factor {scale} to fit in unit cube:
  sys.stderr.write("computing max scalings for relevant pixels ...\n")
  scales = []
  for iy in range(ny):
    for ix in range(nx):
      if M[iy,ix] != 0:
        bpix = B[iy,ix,:]
        for ib in range(nb):
          # Compute chroma scale {scale_ib} to fit this component in {[0 _ 1]}: 
          if bpix[ib] > +1.0e-100:
            scale_ib = (1.0 - shift[ib])/bpix[ib]
          elif bpix[ib] < -1.0e-100:
            scale_ib = (0.0 - shift[ib])/bpix[ib]
          else:
            scale_ib = +inf
          if scale_ib != +inf: scales.append(scale_ib)
  ns = len(scales)
  sys.stderr.write(f"collected {ns} max scales\n")
  
  frac_scale = 0.005
  scale = numpy.quantile(numpy.array(scales), frac_scale, axis=0)
  sys.stderr.write(f"chosen scaling = {scale:10.6}\n")

  sys.stderr.write("applying scale and shift ...\n")
  F = scale_shift_and_clip_images(B, scale, shift)
  return F, scale
  # ----------------------------------------------------------------------
   
def scale_shift_and_clip_images(D, scale, shift):
  # Receives a set {D} of multispectral images with shape {(ny,nx,nb,)}
  # and a spectrum {bpix} with shape {(nb,)}.
  # 
  # Returns a copy of {D} with every pixel scaled by the float {scale},
  # shifted by the spectrum {bpix}, and clipped to the {[0_1]^nd} unit
  # hypercube.
  
  ny, nx, nb = numpy.shape(D)
  assert numpy.shape(shift) == (nb,)
  
  sys.stderr.write("applying scale, shift, and clipping ...\n")
  B = numpy.zeros((ny,nx,nb,))
  vlo = 0.05; vhi = 0.95
  for ib in range(nb):
    C_ib = shift[ib] + scale*D[:,:,ib] 
    tag_ib = f"channel {ib:2d}"
    B[:,:,ib] = clip_image_squashed(C_ib, tag_ib, vlo, vhi, vlo, vhi, vmin, vmax)
  return B
  # ----------------------------------------------------------------------

def spectrum_to_color_matrix(bands, bands_tb):
  nb = len(bands)
  RfromS = numpy.zeros((nb,3))
  wlens = [ bands_tb[bands[ib]]['wlen'] for ib in range(nb) ]
  wmin = min(wlens)
  wctr = 535
  wmax = max(wlens)
  for ib in range(nb):
    band = bands[ib]
    wlen = bands_tb[band]['wlen']
    RfromS[ib,:] = spectral_weights(wlen, wmin, wctr, wmax)

  # Normalize to total weight per channel = 1:
  for ic in range(3):
    Wsum = numpy.sum(RfromS[:,ic])
    RfromS[:,ic] /= Wsum
  return RfromS
  # ----------------------------------------------------------------------
  
def spectral_weights(wlen, wmin, wctr, wmax):
  CC = wmin
  BB = 4*wctr - wmax - 3*wmin
  AA = 2*wmax + 2*wmin - 4*wctr
  t = (- BB + sqrt(BB**2 - 4*AA*(CC-wlen)))/(2*AA)
  W = numpy.zeros((3,))
  for ic in range(3):
    W[ic] = (1-t)**ic * t**(2-ic) * comb(3,ic)
  return W
  # ----------------------------------------------------------------------