#! /usr/bin/python3 
# Last edited on 2025-08-23 16:19:40 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 {sepa} 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}/separations/{sepa}/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}/separations/{sepa}/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
from bands_table_funcs import read_page_bands_tables
from image_funcs import read_mask_image, read_clip_image, write_image_as_gray_png
from spectral_analysis_funcs import read_int_param
import argparser

def main():

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

  # Read the spectral analysis results for the given mref, inks:
  bands, sref, sink = read_analysis_data(page, clip, mref, inks)
  nb = len(bands)
  debug("nb", nb)
  assert numpy.shape(sref) == (nb,)
  assert numpy.shape(sink) == (nb,ni,)
  
  # 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 ink separation images:
  S, R = create_ink_separation_images(C, sref, sink)
  
  # Write the images:
  write_ink_separation_images(page, clip, sepa, inks, S)
  write_residual_spectra_images(page, clip, sepa, bands, R)

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

def read_analysis_data(page, clip, mref, inks):
  # Reads the bands list {bands} and the analysis data for the given {page}, {clip},
  # and masks {mref} and {inks[0..ni-1]}.  Namely, gets
  # {sref} as the {smax} from the data of mask {mref}, and 
  # {sink[0..ni-1]} as the {smin} from the data of masks {inks[0..ni-1]}.
  # 
  ni = len(inks)
  nb = None # For now, reset later.
  ne = None # For now, reset later.
  sref = None; sink = None

  bands = []
  for ii in range (ni+1):
    mask = mref if ii == ni else inks[ii]
    afile = f"pages/{page}/clips/{clip}/masks/{mask}/analysis.txt"
    with open(afile, "r") as rd:
      if nb == None:
        nb = read_int_param(rd, "nb")
        ne = read_int_param(rd, "ne")
        sref = numpy.zeros((nb,))
        sink = numpy.zeros((nb,ni,))
      else:
        assert read_int_param(rd, "nb") == nb, "inconsistent nb"
        assert read_int_param(rd, "ne") == ne, "inconsistent ne"
      for ib in range(nb):
        line = rd.readline().strip()
        line = re.sub(r'[ ][ ]+', " ", line)
        flds = line.split(' ')
        assert len(flds) == 3 + 5 + ne, f"bad format: {len(flds)} fields"
        band = flds[0].strip(); 
        if ii == 0:
          bands.append(band)
        else:
          assert band == bands[ib], "bug: inconsistent bands"
        ifld = 1
        ilum = int(flds[ifld].strip()); ifld += 1
        wlen = float(flds[ifld].strip()); ifld += 1
        savg = float(flds[ifld].strip()); ifld += 1
        smin = float(flds[ifld].strip()); ifld += 1
        smax = float(flds[ifld].strip()); ifld += 1
        rdev = float(flds[ifld].strip()); ifld += 1
        tdev = float(flds[ifld].strip()); ifld += 1
        # for ie in range(ne):
        #   P[ib,ie,ii] = float(flds[ifld].strip()); ifld += 1
        if ii == ni:
          sref[ib] = smax
        else:
          sink[ib,ii] = smin
      rd.close()
  return bands, sref, sink
  # ----------------------------------------------------------------------
   
def read_clip_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 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, sepa, 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}/separations/{sepa}"
  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, 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
  # ----------------------------------------------------------------------
  
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")
  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.       
  sepa = 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, sepa, mref, inks
  # ----------------------------------------------------------------------
  
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()

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

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 = \
      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
  # ----------------------------------------------------------------------