#! /usr/bin/python3 
# Last edited on 2025-08-24 05:38:29 by stolfi

# Reads a set of {nb} narrow-band monochrome images {C[0:ny,0:nx,0:nb]}
# for a given {page} and {clip}. Reads a set of spectral analysis data
# for {nm} subsets of thay clip selected by {nm} mask images, including
# the brightest spectrum {smax[0:nb,0:nm]} and the darkest spectrum
# {smin[0:nb:,0:nm]}. 
#
# For each {mask} index {im} in {0..nm-1}, writes a grayscale image
# {ssim[0:ny,0:nx]} where the value {ssim[iy,ix]} of each pixel tells
# whether spectrum {spix[iy,ix,:]} of that pixel in the multispectral
# clip images is close to {smax[:,im}} (output pixel value 1.0) or to
# {smin[:,im]} (output pixel value 0.0). If {spix} is equidistant from
# both of these spectra, the output sample {ssim[iy,ix]} will be 0.5.
# 
# These images are written to disk as as file "{mdir}/simil.png" where
# {mdir} is "pages/{page}/clips/{clip}/masks/{mask}".
#
import sys, re, os
from math import sqrt, sin, cos, log, exp, floor, pi, inf, isfinite
import numpy
from error_funcs import prog_error, arg_error, debug
from bands_table_funcs import read_page_bands_tables
from process_funcs import bash, run_command
from image_funcs import read_clip_image, write_image_as_gray_png
import argparser

def main():

  page, clip, masks = parse_options()
  nm = len(masks)

  # masks the spectral bands tables:
  bfile = f"MS/davis/{page}/bands.txt"
  bands_tb = read_page_bands_tables(bfile)

  # Read the spectral analysis results for the given masks:
  bands, smin, smax = read_analysis_data(page, clip, masks)
  nb = len(bands)
  debug("nb", nb)
  assert numpy.shape(sref) == (nb,)
  assert numpy.shape(sink) == (nb,nm,)
   
  # Read the clip images:
  C = read_clip_images(page, clip, bands)
  ny, nx = numpy.shape(C[:,:,0])
  debug("ny", ny)
  debug("nx", nx)
  assert numpy.shape(C) == (ny,nx,nb,)
 
  # Create spectral similarity images:
  U, V = create_spectral_similarity_images(C, sref, sink)

  # Write the images:
  write_spectral_similarity_images(page, clip, masks, U, V)

  return 0
  # ----------------------------------------------------------------------

def ???create_spectral_similarity_images(C, sref, sink):
  # Given a multipsectral clip image set {C}, a "background" spectrum
  # {sref[0:nb]}, and {nm} "ink" spectra {sink[o:nb,0:nm]}, creates {nm+1} images
  # that show the mixing coefficients of those {nm+1} spectra that 
  # best account for the spectrum of each pixel. Here {nb} is the number of
  # spectral bands used in the analysis that determined the ink spectra.
  # 
  # Expects {C}, {sref}, and {sink} to have shapes {(ny,nx,nb)}, {(nb,)}, and {(nb,nm)},
  # respectively.
  # 
  # Returns an array {S} with shape {(ny,nx,nm+1)} such that {S[:,:,ii]}
  # is the mixing coefficient of ink {ii} in each pixel, and {S[:,:,nm]}
  # is the mixing coeff of the "background".  Beware that these numbers
  # may be less than 0.0 or greater than 1.0.
  #
  # Also returns an array {R} withshape {(ny,nx,nb)} such that
  # {R[ix,iy,0:nb]} is the part of the spectrum of each pixel that
  # cannot be obtained by mixing the given "ink" spectra. 
  # This spectrum residual is normally in {[-1 _ +1]}.
  # 
  ny, nx, nb = numpy.shape(C)
  debug("nb", nb)
  nm = numpy.size(sink, 1)
  debug("sref.shape", numpy.shape(sref))
  assert numpy.shape(sref) == (nb,), "bug bad sref shape"
  assert numpy.shape(sink) == (nb,nm), "bug bad sink shape"
  nd = nm + 1
  
  SU = numpy.zeros((nm,nb)) 
  for ii in range(nm):
    SU[ii,:] = sink[:,ii] - sref
    
  US = numpy.linalg.pinv(SU)
  assert numpy.shape(US) == (nb,nm,), "bug pinv"
  
  S = numpy.zeros((ny,nx,nd))
  R = numpy.zeros((ny,nx,nb))
  for iy in range(ny):
    for ix in range(nx):
      cpix = C[iy,ix,:]
      dpix = (cpix - sref) @ US
      assert numpy.shape(dpix) == (nm,), "bug pinv"
      S[iy,ix,0:nm] = dpix
      S[iy,ix,nm] = 1 - numpy.sum(dpix,0)
      rpix = cpix - (dpix @ SU + sref)
      assert numpy.shape(rpix) == (nb,), "bug pinv"
      R[iy,ix,0:nb] = rpix
  return S, R
  # ----------------------------------------------------------------------

def ???write_spectral_similarity_images(page, clip, masks, U, V): 
  ???
  nm = len(masks)
  ny, nx, ns = numpy.shape(S) # ns = number of mixing coeffs ("masks" plus background).
  assert ns == nm + 1, "bug ns"

  sdir = f"pages/{page}/clips/{clip}/separations/{sepa}"
  bash(f"mkdir -p {sdir}")
  for ks in range(ns): 
    ink = "BG" if ks == nm else masks[ks]
    idir = f"masks/{ink}"
    bash(f"mkdir -p {idir}")
    clo = 0.00; plo = 0.25
    chi = 1.00; phi = 0.75
    sfile = f"{sdir}/{ink}.png"
    write_image(sfile, S[:,:,ks], clo, chi, plo, phi)
  return
  # ----------------------------------------------------------------------



def parse_options():
  na = len(sys.argv)
  assert na >= 4, "insuff args"
  page = sys.argv[1]
  clip = sys.argv[2]    # name of clip to take pixels from.       
  masks = [ sys.argv[k] for k in range(4,na) ]
  return page, clip, masks
  # ----------------------------------------------------------------------

def band_error(band, msg):
  sys.stderr.write(f"** {band}: {msg}\n")
  assert False
  # ----------------------------------------------------------------------

main()

@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
  
def write_residual_spectra_images(page, clip, sepa, bands, R): 
  
  nb = len(bands)
  ny, nx, nr = numpy.shape(R)
  assert nr == nb # Number of bands in residual including the total one.

  rdir = f"pages/{page}/clips/{clip}/separations/{sepa}/residual"
  bash(f"mkdir -p {rdir}")
  clo = 0.00; plo = 0.25
  chi = 1.00; phi = 0.75
  for ir in range(nr):
    band = bands[ir]
    rfile = f"{rdir}/{band}.png"
    write_image(rfile, R[:,:,ir], clo, chi, plo, phi)
  return
  # ----------------------------------------------------------------------

???
 Namely, gets
  # {sref} as the {smax} from the data of mask {mref}, and 
  # {sink[0..nm-1]} as the {smin} from the data of masks {masks[0..nm-1]}.
  
def parse_options():
  na = len(sys.argv)
  assert na >= 5, "insuff args"
  page = sys.argv[1]
  clip = sys.argv[2]    # name of clip to take pixels from.       
  sepa = sys.argv[3]    # Name of separation run.
  mref = sys.argv[4]    # Name of background mask.
  masks = [ sys.argv[k] for k in range(5,na) ]
  return page, clip, sepa, mref, masks
  # ----------------------------------------------------------------------
  
def debug(name, arg):
  sys.stderr.write(f"{name}");
  if isinstance(arg, list) or isinstance(arg, tuple) or isinstance(arg, dict):
    sys.stderr.write(f" (length = {len(arg)}):\n")
    sys.stderr.write(r"  ")
  else:
    sys.stderr.write(r" = ")
  sys.stderr.write(f"{arg}\n")
  return
  # ----------------------------------------------------------------------

main()