#! /usr/bin/python3
# Last edited on 2026-08-16 08:24:16 by stolfi

# Reads three narrow-band monochrome images for a given page, previosuly
# processed, extracting a specified rectangle. Combines them into a
# pseudocolor RGB image.
#
# COMMAND LINE ARGUMENTS
# 
# The command line arguments are
# 
#   {page} {walum0} {walum1} {walum2} {fgmask} {crop} {matrix} {bgmask} {border} {odir} {oname}
# 
# where
#
#   {page} is the normalized page's f-number like "008r1".
#
#   {walum0,walum1,walum2} are the wavelengths and light source
#   codes of the three band images to combine.
#
#   {fgmask} is the tag of the subset of the page (like "trace" for ink traces)
#   that was used for spectrum normalization.
#   
#   {crop} is an ImageMagik-like spec of a rectangle (relative 
#   to the "page.png" images).
#
#   {matrix} Transformation matrix to apply to the three images, or "YUV" or "RGB"
#
#   {bgmask} Mask for the subset of the page (like "parch" for blank
#   parchment) which should be considered forbrightness/contrast normalizaton.
#
#   {border} width of extra gray frame to add, in pixels.
# 
#   {odir} is the directory for output files.
#
#   {oname} is the prefix for names of the output files, sans extension.
#
# Each of {walum0,walum1,walum2} must be a pair "{wlen}-{illum}" where
# {wlen} is the dominant wavelength of the image to use (in nm, three
# digits,zero-padded) and {illum} is a light source code (0 = both
# sources, 1 = from south, 2 = from north).
#
# INPUT FILES
# 
# Assumes that there is a file "bands.txt" in the directory
# "MS/davis/{page}" that specifies the names of the bands avaliable for
# that page (e.g. "MB870IR_030_F") as well as their illumnation type
# {illum} ({0..3}) and main wavelength {wlen}.
#
# The input image for each wavelength-source pair "{wlen}-{illum}" is
# read from "{pdir???NO}/{page}/{wlen}-{illum}-{fgmask}.png". These
# images are assumed to have been suitably preprocessed, e.g. with
# spectrum normalized with respect to the {fgmask}.
#
# The program reads mask files "MS/davis/{page}/masks/page/{bgmask}.png"
# and "MS/davis/{page}/masks/page/{fgmask}.png". These files must
# contain blevel images specifying the pixels (black=no, white=yes) that
# are mostly blank parchment and ink traces, respectively. These masks
# are considered when normalizing images and color channels for
# brightness and constrast.
# 
# The program extracts from each file "page.png" and from the masks
# "page/{bgmask}.png" and "page/{fgmask}.png" a rectangle specified by
# the {crop} argument as float arrays with elements in {[0_1]}. It also
# adds a grey frame of width {border} all around
#
# OUTPUT FILES
#
# The output is the composite image is written as
# the RGB image file "{odir}/{oname}.png".
#
import sys, re, os
from sys import stderr as err
from math import sqrt, sin, cos, log, exp, floor, pi, inf
import numpy
from error_funcs import prog_error, arg_error, debug
from process_funcs import bash, run_command
import argparser

import vms_color_image_funcs as cfn
import vms_linear_gray_image_funcs as gfn
import bands_table_funcs as bfn

def main():

  page, walum0, walum1, walum2, fgmask, crop, matrix, bgmask, border, odir, oname = parse_options()

  # Gets the spectral bands tables:
  bfile = f"MS/davis/{page}/bands.txt"
  bands_tb = bfn.read_page_bands_tables(bfile)
  
  # Get the parchment mask image:
  M_bg = read_mask(page, bgmask, crop) # Background (e.g. blank parchment) mask
  M_fg = read_mask(page, fgmask, crop) # Foreground (e.g. ink traces) mask
  ny, nx = numpy.shape(M_bg)
  assert numpy.shape(M_fg) == (ny,nx,), f"mask shape mismatch {bgmask = } {fgmask = }"

  # Make and read the specified normalized narrow-band images:
  err.write(f"reading the specified bands ...\n")
  C = numpy.zeros((ny,nx,3))
  for ic in range(3):
    wlen, illum = ( walum0, walum1, walum2, )[ic]
    pdir = f"MS/davis/{page}/normalized"
    wfile = f"{pdir}/{wlen}-{illum}-{fgmask}.png"
    bash(f"cd MS/davis && make PAGE={page} WLEN={wlen} ILLUM={illum} FGMASK={fgmask} normalized-image")
    P = read_and_crop_page_image(wfile, crop)
    assert numpy.shape(P) == (ny,nx,), "image/mask size mismatch"
    C[:,:,ic] = P

  err.write(f"creating composite image ...\n")
  R = make_RGB_composite_image(C, M_bg, M_fg, matrix)
  
  err.write(f"statistics for RGB space:\n")
  show_image_statistics(R, M_bg, M_fg)
  
  err.write(f"statistics for YUV space:\n")
  show_image_statistics(cfn.YUV_image_from_RGB_image(R), M_bg, M_fg)
  
  if border >= 0:
    err.write(f"adding border to image ...\n")
    R = add_border_to_image(R, border)
  
  rfile = f"{odir}/{oname}.png"
  cfn.write_image_as_png(rfile, R)
  err.write(f"writing {rfile} ...\n")
  return None
  # ----------------------------------------------------------------------

def read_mask(page, mask, crop):
  assert mask != None
  mfile = f"MS/davis/{page}/masks/page/{mask}.png"
  M = gfn.read_mask_image(mfile, crop)
  return M;
  # ----------------------------------------------------------------------

def show_image_statistics(R, M_bg, M_fg):
  ny,nx,nc = numpy.shape(R)
  ravg = None
  for M, tag in ((M_bg, "avg"), (M_fg, "dev")):
    err.write(f"  statistics for subset '{tag}':\n")
    for ic in range(nc):
      cavgs = [ \
          R[iy,ix,ic] \
          for ix in range(nx) \
            for iy in range(ny) \
              if M is None or M[iy,ix] != 0 \
        ]
      savg = numpy.mean(cavgs)
      if tag == "avg": ravg = savg
      cdevs = [ \
          R[iy,ix,ic] - ravg \
          for ix in range(nx) \
            for iy in range(ny) \
              if M is None or M[iy,ix] != 0 \
        ]
      sdev = numpy.std(cdevs)
      err.write(f"    channel {ic} samples = {len(cavgs)} avg = {savg:.4f} dev = {sdev:.4f}\n") 
  return
  # ----------------------------------------------------------------------

def read_and_crop_page_image(wfile, crop):
  # Reads the image {wfile} and extracts the rectange specified by {crop}. 
  # Returns it as a {numpy} array {P} of floats with shape
  # {(ny,nx,)} with float samples in {[0 _ 1]}.

  maxval = 65535
  B = gfn.read_gray_png_image(wfile, maxval, crop)
  assert len(numpy.shape(B)) == 2, "band image should be monochrome"
  return B
  # ----------------------------------------------------------------------
   
def make_RGB_composite_image(C, M_bg, M_fg, matrix):
  # Assumes that {C} is a {numpy} array with shape {(ny,nx,3)}
  # containing three grayscale images of {ny} rows and {nx} columns.
  #
  # The {matrix} argument may be a string containing a color space key
  # "YUV" or "RGB" optionally followed by nine numbers separated by
  # spaces. These are interpreted as the elements of a 3x3 matrix {F} used
  # to combine the three images {C[:,:,ki} for {ki} in {0..2}.
  #
  # If the color space key is "RGB", the first three numbers will be
  # used as the weights of the three grayscale images to compute the red
  # channel; the next three will be used for green, and the next three
  # for the blue channel.
  #
  # If the color space key is "YUV", the first three numbers will be the
  # weights to compute the Y channel (luminance), the next three will
  # define the U channel (blue-yellow chroma), and the next three will
  # define the V channel (green-red chroma). In this case, each channel
  # will then be normalzed so that the Y axis has mean 0.5 over the
  # subset selected by {M_bg}, and deviation 0.25 on the {M_fg}
  # subset, hard-clipped to {[0.2_0.8]}; while the U,V axes have mean 0
  # over {M-avg} and deviation 0.5 over {M_fg}. The U,V axes are then
  # multiplied by the Y value. The result is then converted to RGB.
  #
  # In any case, and each pixel is soft-clipped to the {[0 _ 1]}
  # cube.
  #
  # If the nine numbers are omitted, assumes the identity matrix.
  #
  ny,nx,nc = numpy.shape(C)
  assert nc == 3, "C should have exactly three channels"
  ckey, F = parse_matrix(matrix)
  if ckey == "YUV":
    RGB = make_RGB_image_via_YUV(C, M_bg, M_fg, F)
  elif ckey == "RGB":
    RGB = make_RGB_image_direct(C, M_bg, M_fg, F)
  else:
    assert False, f"invalid color key '{ckey}'"
  return RGB
  # ----------------------------------------------------------------------

def add_border_to_image(R, border):
  # Adds a gray frame of width {border} all around.
  ny,nx,nc = numpy.shape(R)
  assert nc == 3, "bad nc"
  S = numpy.full((ny+2*border,nx+2*border,3), 0.50)
  S[border:border+ny,border:border+nx,:] = R
  return S
  # ----------------------------------------------------------------------

def make_RGB_image_via_YUV(C, M_bg, M_fg, F):
  ny, nx, nc = numpy.shape(C)
  assert nc == 3, "bad nc"
  assert numpy.shape(F) == (3,3)

  YUV = C @ F.T

  err.write("normalizing Y,U,V channels ...\n")
  Yavg = 0.60
  Ydev = 0.25
  UVdev = 0.50
  cfn.normalize_YUV_image(YUV, M_bg, M_fg, Yavg, Ydev, UVdev)

  err.write("clipping Y channel ...\n")
  Ymin = 0.2
  Ymax = 0.8
  cfn.clip_Y_in_YUV_image(YUV, Ymin, Ymax)
  
  err.write("computing best remix of U,V data ...\n")
  E = remix_UV_coords_of_image(YUV, M_fg)
  
  err.write("rescaling U,V of pixels according to {Y} ...\n")
  for iy in range(ny):
    for ix in range(nx):
      Y = YUV[iy,ix,0];
      assert Y >= 0.999*Ymin and Y <= 1.001*Ymax, f"invalid {Y = }"
      aUV = 4*Y*(1-Y)
      YUV[iy,ix,1] *= aUV
      YUV[iy,ix,2] *= aUV

  # Rescale the U,V coordinates to fit in the RGB cube:
  err.write("global rescaling of U,V to fit in RGB cube ...\n")
  sUV = cfn.choose_UV_scale(YUV, M_fg)
  err.write(f"  rescaling factor = {sUV}\n")
  YUV[:,:,1] *= sUV
  YUV[:,:,2] *= sUV

  err.write("converting YUV image to RGB ...\n")
  RGB = cfn.RGB_image_from_YUV_image(YUV)
  return RGB
  #----------------------------------------------------------------------
  
def remix_UV_coords_of_image(YUV, M):
  # Returns a 2x2 matrix that should post-multiply the U,V coordinatess for 
  # maximum color contrast.
  ny, nx, nc = numpy.shape(YUV)
  assert nc == 3, "image {YUV} must have three channels"
  if M is not None: assert numpy.shape(M) == (ny, nx)
  
  err.write(f"extracting UV coords of selected pixels ...\n")
  sel_UV_vecs = [ 
      YUV[iy,ix,1:3] \
      for iy in range(ny) \
        for ix in range(nx) \
          if M is None or M[iy,ix] != 0 \
    ]
  UV_sel = numpy.array(sel_UV_vecs)
  ns, ne = numpy.shape(UV_sel)
  assert ne == 2
  
  err.write(f"computing the UV averages of selected pixels ...\n")
  UV_avg = numpy.zeros((ne,))
  for ie in range(ne):
    UV_avg[ie] = numpy.average(UV_sel[:,ie])
  err.write(f"  averages: {UV_avg}\n")
  
  err.write(f"computing the UV covariance matrix ...\n")
  COV = numpy.zeros((ne,ne))
  for ks in range(ns):
    for ie in range(ne):
      di = UV_sel[ks,ie] - UV_avg[ie]
      for je in range(ne):
        dj = UV_sel[ks,je] - UV_avg[je]
        COV[ie,je] += di*dj
  COV /= ns
  err.write(f"  covariances:\n{COV}\n")
   
  # Eigens of covariance matrix:
  err.write(f"computing the eigenvalues and eigenvectors ...\n")
  E = numpy.linalg.eigh(COV)
  E_vec = E.eigenvectors
  err.write(f"  eigenvectors:\n{E_vec}\n")
  E_val = E.eigenvalues
  err.write(f"  eigenvalues:\n{E_val}\n");
  assert numpy.shape(E_val) == (ne,), "bug: eigenvalues shape"
  assert numpy.shape(E_vec) == (ne,ne,), "bug: eigenvectors shape"
  
  D = numpy.array([[1.0/sqrt(E_val[0]), 0], [0, 1.0/sqrt(E_val[1])]])

  err.write(f"remapping UV coordinates ...\n")
  for iy in range(ny):
    for ix in range(nx):
      YUV[iy,ix,1:3] =  E_vec @ (YUV[iy,ix,1:3] - UV_avg) @ D

  return
  # ----------------------------------------------------------------------

def make_RGB_image_direct(C, M_bg, M_fg, F):
  RGB = C @ F.T
  
  err.write("normalizing RGB image ...\n")
  Ymin = 0.2
  Ymax = 0.8
  cfn.normalize_RGB_image(RGB, M_bg, M_fg, Ymin, Ymax)
  
  err.write("clipping composite image to RGB cube ...\n")
  cfn.clip_RGB_image_to_unit_cube(RGB)
  return RGB
  # ----------------------------------------------------------------------

def parse_options():
  na = len(sys.argv)
  ia = 1
  
  def get_arg(): 
    nonlocal na, ia
    assert ia < na, "insuff args"
    arg = sys.argv[ia]; ia += 1
    return arg
    # ....................................................................
  
  page = get_arg()
  walum0 = get_arg()
  walum1 = get_arg()
  walum2 = get_arg()
  fgmask = get_arg() 
  assert len(fgmask) == 5, "invalid mask code"
  crop = get_arg()
  matrix = get_arg()
  bgmask = get_arg()
  border = int(get_arg())
  odir = get_arg()
  oname = get_arg()
  return page, walum0, walum1, walum2, fgmask, crop, matrix, bgmask, border, odir, oname
  # ----------------------------------------------------------------------
  
def parse_matrix(matrix):
  ckey = None; F = None
  matrix = re.sub(r"[^-+.0-9a-zA-Z]", " ", matrix)
  if m := re.fullmatch(r"[ ]*([A-Z]+)([-+.0-9 ]*)", matrix):
    ckey = m.group(1)
    melems = m.group(2)
    melems = melems.strip()
    if melems == "":
      F = numpy.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])
    else:
      melems = melems.split()
      assert len(melems) == 9, f"invalid matrix size = {len(melems)}"
      F = numpy.zeros((3,3))
      for ic in range(3):
        for jc in range(3):
          F[ic,jc] = melems[3*ic + jc]
  else:
    assert False, f"invalid compositing matrix = '{matrix}'"
  return ckey, F
  # ----------------------------------------------------------------------
  
def test_remix_UV_coords_of_image():
  err.write(f" testing remix_UV_coords_of_image ...\n")
  YUV = numpy.zeros((9,1,3))
  YUV[:,0,:] = numpy.array( \
      (
        ( 0.1, 00.00 + -0.40 + -0.04, 1.00 + -0.80 + +0.02, ),
        ( 0.2, 00.00 + -0.40 + +0.04, 1.00 + -0.80 + -0.02, ),
        ( 0.3, 00.00 + -0.20 + -0.04, 1.00 + -0.40 + +0.02, ),
        ( 0.4, 00.00 + -0.20 + +0.04, 1.00 + -0.40 + -0.02, ),
        ( 0.5, 00.00 + 00.00 + 00.00, 1.00 + 00.00 + 00.00, ),
        ( 0.6, 00.00 + +0.20 + +0.04, 1.00 + +0.40 + -0.02, ),
        ( 0.7, 00.00 + +0.20 + -0.04, 1.00 + +0.40 + +0.02, ),
        ( 0.8, 00.00 + +0.40 + +0.04, 1.00 + +0.80 + -0.02, ),
        ( 0.9, 00.00 + +0.40 + -0.04, 1.00 + +0.80 + +0.02, ),
      )
    )
  err.write(f" before:\n {YUV = }\n")
  remix_UV_coords_of_image(YUV, None)
  err.write(f" after:\n {YUV = }\n")
  return
  # ----------------------------------------------------------------------
  
if sys.argv[1] == "test_remix_UV_coords_of_image":
  test_remix_UV_coords_of_image()
else:
  main()