#! /usr/bin/python3 
# Last edited on 2025-10-15 11:51:01 by stolfi

# Reads a set of multispectral images and spectral analysis data and
# splits the former into an affine combination of specified inks. Then
# writes images with those ink components.
# 
# Command line arguments are the {page} and {clip}, a tag {run_tag} for 
# this separation run, and two or more mask tags ("bg", "ch", etc).
# The first mask will be assumed to select pixels in the background 
# (blank parch). The {smax} spectrum of that mask will be taken to be the
# color of the fully illuminated substrate.  The {smin} spectra
# of that and subsequent masks will be taken to be the full-coverage
# colors of particle suspension inks applied over that substrate.
# The first mask may appear among the other masks.
#
# One set of  input files are the analysis files 
# "pages/{page}/clips/{clip}/masks/{mask}/analysis.txt" with the
# spectral analysis of the pixels in that clip sepected by each given{mask}.
#
# Another set of input files are the multispectral image clips 
# "pages/{page}/clips/{clip}/npy/{band}.npy" for each multispectral band
# used in that analysis file.
#
# For each given {mask} name, writes an image file
# "pages/{pages}/clips/{clip}/sepas/{run_tag}/inks/{mask}.png" with
# images of the inferred coverage fraction of the ink of that {mask} as
# a grayscale file.
#
# Also, for each spectral {band}, writes an image file
# "pages/{page}/clips/{clip}/sepas/{run_tag}/residual/{band}.png"
# with the part of the clip's 
# 
import sys, re, os
from math import sqrt, sin, cos, log, exp, floor, pi, inf, isfinite
import numpy
from PIL import Image
from error_funcs import prog_error, arg_error, debug
from process_funcs import bash, run_command
import argparser

import vms_linear_gray_image_funcs as gfn
import vms_multispectral_image_funcs as mfn
import vms_spectral_analysis_funcs as afn

def main():

  page, clip, run_tag, mref, inks = parse_options()
  ni = len(inks)
  debug("ni", ni)
  
  # Read the spectral bands tables:
  pfile = f"MS/davis/{page}/bands.txt"
  bands_tb = mfn.read_page_bands_tables(pfile)

  # Read the spectral analysis results for the given ref mask:
  bands, sref = read_mref_analysis_data(page, clip, mref)
  nb = len(bands)
  debug("nb", nb)
  assert numpy.shape(sref) == (nb,), "bug sref shape"
  
  # Read the spectral analysis results for the given inks:
  bands2, sink = read_inks_analysis_data(page, clip, inks)
  assert bands2 == bands, "bug inconsistent bands"
  assert numpy.shape(sink) == (nb,ni,), "bug sink shape"
  
  # Read the clip images:
  C = mfn.read_clip_npy_images(page, clip, bands)
  ny, nx = numpy.shape(C[:,:,0])
  debug("ny", ny)
  debug("nx", nx)
  assert numpy.shape(C) == (ny,nx,nb,)

  # Create ink separation images:
  S, R = create_ink_separation_images(C, sref, sink)
  
  # Write the images:
  write_ink_separation_images(page, clip, run_tag, inks, S)
  write_residual_spectra_images(page, clip, run_tag, bands, R)

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

def read_mref_analysis_data(page, clip, mref):
  # Reads the bands list {bands} and the analysis data for the given
  # and mask {mref}.  {page}, {clip}, Returns the list of {bands} and
  # and mask {mref}.  the {smax} of that mask.
  # 
  bands, illums, wlens, savg, smin, smax, rdev, tdev, P = \
    afn.read_mask_analysis_data(page, clip, mref)
  nb, ne = numpy.shape(P)
  assert numpy.shape(smax) == ((nb,)), "bug smax shape"
  assert len(bands) == nb, "bug inconsistent bands"
  sref = smax
  return bands, sref
  # ----------------------------------------------------------------------

def read_inks_analysis_data(page, clip, inks):
  # Reads the bands list {bands} and the analysis data for the given {page}, {clip},
  # and all the masks whose names are {inks[0:ni]}.  Namely, gets
  # {sink[:,ii]} as the {smin[:]} from the analysis file of mask {inks[ii]}.
  # 
  ni = len(inks)
  bands, illums, wlens, savg, smin, smax, rdev, tdev, P = \
    afn.read_mask_list_analysis_data(page, clip, inks)
  nb = len(bands)
  assert numpy.shape(smin) == (nb,ni,), "bug smin shape"
  sink = smin
  return bands, sink
  # ----------------------------------------------------------------------

def create_ink_separation_images(C, sref, sink):
  # Given a multipsectral clip image set {C}, a "background" spectrum
  # {sref[0:nb]}, and {ni} "ink" spectra {sink[o:nb,0:ni]}, creates {ni+1} images
  # that show the mixing coefficients of those {ni+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,ni)},
  # respectively.
  # 
  # Returns an array {S} with shape {(ny,nx,ni+1)} such that {S[:,:,ii]}
  # is the mixing coefficient of ink {ii} in each pixel, and {S[:,:,ni]}
  # 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)
  ni = numpy.size(sink, 1)
  debug("sref.shape", numpy.shape(sref))
  assert numpy.shape(sref) == (nb,), "bug bad sref shape"
  assert numpy.shape(sink) == (nb,ni), "bug bad sink shape"
  nd = ni + 1
  
  SU = numpy.zeros((ni,nb)) 
  for ii in range(ni):
    SU[ii,:] = sink[:,ii] - sref
    
  US = numpy.linalg.pinv(SU)
  assert numpy.shape(US) == (nb,ni,), "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) == (ni,), "bug pinv"
      S[iy,ix,0:ni] = dpix
      S[iy,ix,ni] = 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_ink_separation_images(page, clip, run_tag, inks, S): 
  
  ni = len(inks)
  ny, nx, ns = numpy.shape(S) # ns = number of mixing coeffs ("inks" plus background).
  assert ns == ni + 1, "bug ns"

  sdir = f"pages/{page}/clips/{clip}/sepas/{run_tag}"
  bash(f"mkdir -p {sdir}")
  for ks in range(ns): 
    ink = "BG" if ks == ni else inks[ks]
    idir = f"inks/{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 write_residual_spectra_images(page, clip, run_tag, 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}/sepas/{run_tag}/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
  # ----------------------------------------------------------------------
  
def write_image(pfile, C, clo, chi, plo, phi):
  # Writes {C} as an grayscale linear PNG file,
  # mapping {[clo _ chi]} to {[plo _ phi]} and clipping
  # the result to {[0 _ 1]}.
  sys.stderr.write(f"writing {pfile} ...\n")
  ny, nx = numpy.shape(C)
  P = numpy.zeros((ny,nx))
  nbig = 0; cval_max = -inf; pval_max = -inf
  nsma = 0; cval_min = +inf; pval_min = +inf
  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)
        pval = plo + (cval - clo)*(phi - plo)/(chi - clo);
        pval_min = min(pval_min, pval)
        pval_max = max(pval_max, pval)
        if pval < 0.0:
          nsma += 1; pval = 0.0
        elif pval > 1.0:
          nbig += 1; pval = 1.0
      else:
        pval = 0.5
        nbad += 1
      P[iy,ix] = pval
  sys.stderr.write("input range [%+8.5f _ %+8.5f]" % (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} mapped samples were below 0.00\n")
    sys.stderr.write(f"{nbig:6d} mapped samples were above 1.00\n")
    sys.stderr.write(f"{nbad:6d} input samples were INF or NAN\n")
  gfn.write_image_as_gray_png(pfile, P)
  return 
  # ----------------------------------------------------------------------
    
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.       
  run_tag = sys.argv[3]    # Name of separation run.
  mref = sys.argv[4]    # Name of background mask.
  inks = [ sys.argv[k] for k in range(5,na) ]
  return page, clip, run_tag, mref, inks
  # ----------------------------------------------------------------------

main()

@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@

def ???read_analysis_data(page, clip, mref, masks):
  # Reads the bands list {bands} and the analysis data for the given {page}, {clip},
  # and masks {mref} and {masks[0..nm-1]}. 
  # 
  nm = len(masks)
  nb = None # For now, reset later.
  ne = None # For now, reset later.
  bands = None; illum = None; wlen = None; 
  sref = None; sink = None;  # For now, reset later.
  for im in range (nm + 1):
    mask = mref if im == nm else masks[im]
    bds, savg, smin, smax, rdev, tdev, P = \
      afn.read_mask_analysis_data(page, clip, mask)
    if bands == None:
      ne = numpy.size(P,1)
      bands = bds
      nb = len(bands)
      sref = numpy.zeros((nb,))
      sink = numpy.zeros((nb,nm,))
    else:
      assert bds == bands, "bug inconsistent {bands}"
      assert numpy.size(P,1) == ne, "bug inconsistent {ne}"
    assert numpy.shape(savg) == (nb,)
    assert numpy.shape(smin) == (nb,)
    assert numpy.shape(smax) == (nb,)
    assert numpy.shape(rdev) == (nb,)
    assert numpy.shape(tdev) == (nb,)
    assert numpy.shape(P) == (nb,ne,)
    if im == nm:
      sref = smax
    else
      sink[:,im] = smin
  return sref, sink
  # ----------------------------------------------------------------------