# Last edited on 2025-08-30 22:50:14 by stolfi

      
def normalize_pure_chroma(D, M):
  # Takes a multispectral imageset {D} as an array with shape
  # {(ny,nx,nb,)} which has been projected on some linear subspace of
  # {\RR^{nd}} orthogonal to {(1,1,...,1)}, and a two-level mask {M}
  # with shape {(ny,nx,)}.
  #
  # Computes the principal axes of the cloud of points of
  # {\RR^{nd}} which are the spectra of the pixels selected by {M}.
  # Those are the eigenvectors of the covariance matrix of those spectra
  #
  # There should be at most {nd-1} eigenvectors with nonzero
  # eigenvalues, possibly fewer. Expands all pixels of {D} (ignoring
  # {M}) in the directions of those eigenvectors so that those {M}-selected
  # pixels have unit standard deviation on each of those directions.
  #
  # Then finds the approximate maximum norm {mchr} of the {M}-selected
  # expanded pixels. Smoothly shrinks any pixels ({M}-selected or not)
  # whose norm exceeds {mchr} so that its norm does not exceed {mchr}
  
  ny, nx, nb = numpy.shape(D);
  assert numpy.shape(M) == (ny,nx), "bug M shape"

  sys.stderr.write(f"extracting chromas of selected pixels ...\n")
  S = extract_masked_pixel_spectra(D, M)
  ns = numpy.size(S,0)
  debug("ns", ns)
  assert numpy.shape(S) == (ns,nb,)
  
  sys.stderr.write(f"computing their principal axes ...\n")
  sctr = numpy.zeros((nb,))
  evec, eval = get_pixel_cloud_axes(S, sctr)
  assert numpy.shape(evec) == (nb,nb,)
  assert numpy.shape(eval) == (nb,)
  ne = nb # Number of eigenvalues and eigenvectors.
    
  sys.stderr.write(f"creating the chroma equalization matrix ...\n")
  assert eval[ne-1] > 1.0e-6, "bug no chroma"
  dev_max = sqrt(eval[ne-1])
  H = numpy.zeros((nb,ne,))
  nh = 0; # Number of significant chroma dimensions.
  for ke in range(ne):
    ie = ne - 1 - ke
    sys.stderr.write(f"  eiegnval {ie:2d} = {eval[ie]:10.6f}")
    assert eval[ie] >= -0.00001, "bug sign chroma eigenvalue"
    dev_ie = sqrt(max(0,eval[ie]))
    sys.stderr.write(f" dev = {dev_ie:10.6f}")
    if dev_ie >= 0.01*dev_max:
      H[ie,ie] = 1/dev_ie
      nh += 1
      evec_ie_lum = numpy.sum(evec[:,ie])/nb
      sys.stderr.write(f" eigenvec lum = {evec_ie_lum:10.6f}\n")
      assert abs(evec_ie_lum) <= 1.0e-6, "bug eigenvector with luma"
    else:
      H[ie,ie] = 0.0
      sys.stderr.write(" ignored\n")
  sys.stderr.write(f"chroma subspace dimension = {nh}\n")
  assert nh > 0, "bug X/X != 1"
  assert nh <= ne-2, "bug residual background chroma signal"
  EHE = evec @ H @ numpy.transpose(evec)

  sys.stderr.write(f"applying equalization matrix ...\n")
  chromas = []
  for iy in range(ny):
    for ix in range(nx):
      dpix = D[iy,ix,:] @ EHE
      assert abs(numpy.sum(dpix)/nb) < 1.0e-6, "scaled chroma with luma"
      D[iy,ix,:] = dpix
      if M[iy,ix] != 0:
        chroma = numpy.linalg.norm(dpix)
        if chroma > 1.0e-6: chromas.append(chroma)
  nc = len(chromas)
  sys.stderr.write(f"collected {nc} chroma norms\n")
  assert nc >= 100, "bug too few non-gray pixels"
  
  frac_chroma = 0.002 # Acceptable fraction of outliers.
  mchr = numpy.quantile(numpy.array(chromas), 1-frac_chroma, axis=0)
  sys.stderr.write(f"chosen max chroma = {mchr:10.6}\n")

  sys.stderr.write(f"clipping chroma outliers ...\n")
  dchr_hi = 0.95*mchr # Max chroma norm to be preserved.
  dchr_max = -inf     # Max chroma norm befoe clipping.
  nclip = 0
  for iy in range(ny):
    for ix in range(nx):
      dpix = D[iy,ix,:]
      assert abs(numpy.sum(dpix)/nb) < 1.0e-6, "chroma with luma"
      dchr = numpy.linalg.norm(dpix)
      if dchr > dchr_hi:
        dchr_max = max(dchr, dchr_max)
        nclip += 1
        gchr = map_value_squashed(dchr, 0, dchr_hi, 0, dchr_hi, 0, mchr)
        D[iy,ix,:] = (gchr/dchr)*dpix
      
  sys.stderr.write(f"{nclip} chromas clipped/squashed to {mchr:10.6}")
  sys.stderr.write(f" original max = {dchr_max:10.6}\n")
 
  return mchr
  # ----------------------------------------------------------------------