#! /usr/bin/python3
# Last edited on 2025-12-19 04:20:52 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
# 
#   {pdir} {page} {crop} {mask} {wlen0} {wlen1} {wlen2} {border} {odir} {oname}
# 
# where
#
#   {pdir} is the directory were to find the preprocessed page images.
#
#   {page} is the normalized page's f-number like "008r1".
#   
#   {crop} is an ImageMagik-like spec of a rectangle (relative 
#   to the "page.png" images).
#
#   {mask} is the tag of a subset of the page (like "trace" for ink traces)
#   that was used for brightness/contrast normalization.
#
#   {wlen0,wlen1,wlen2} are the wavelengths of the three band 
#   images to combine.
#
#   {yuv} should the three wavelengths should be used as YUV not RGB?
#
#   {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.
#
# 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 {wlen} is read from 
# "{pdir}/{page}-{wlen}-{mask}.png". These images are assumed to
# have been suitably preprocessed, e.g. with lighting effects removed
# and normalized with respect to the {mask}.
#
# The program reads mask files "MS/davis/{page}/masks/page/parch.png"
# and "MS/davis/{page}/masks/page/{mask}.png". These fiels 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 {mask} may be "None" or "NONE", 
# implying that all pixels are considered for that purpose.
#
# The program extracts from each file "page.png" and from the masks
# "page/parch.png" and "page/{mask}.png" a rectangle specified by the
# {crop} argument and converts the samples affinely from {0..maxval} to
# floats in the range {[0 _ 1]} where the {maxval} is specified in the
# "bands.txt" for the band in question. 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():

  pdir, page, crop, mask, wlen0, wlen1, wlen2, yuv, 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:
  mfile_avg = f"MS/davis/{page}/masks/page/parch.png"
  M_avg = gfn.read_mask_image(mfile_avg, crop)
  ny, nx = numpy.shape(M_avg)
  
  # Get the mask with the pixels of interest:
  if mask == None:
    M_dev = None
  else:
    mfile_dev = f"MS/davis/{page}/masks/page/{mask}.png"
    M_dev = gfn.read_mask_image(mfile_dev, crop)
    assert numpy.shape(M_dev) == (ny,nx,), "mask shape mismatch {mfile = }"

  # Select the bands to use (only reflection two-lamp illum):
  err.write(f"reading the specified bands ...\n")
  C = numpy.zeros((ny,nx,3))
  for ic in range(3):
    wlen = ( wlen0, wlen1, wlen2, )[ic]
    wfile = f"{pdir}/page-{wlen}-{mask}.png"
    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_avg, M_dev, yuv)
  
  err.write(f"statistics for RGB space:\n")
  show_image_statistics(R, M_avg, M_dev)
  
  err.write(f"statistics for YUV space:\n")
  show_image_statistics(cfn.YUV_image_from_RGB_image(R), M_avg, M_dev)
  
  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 show_image_statistics(R, M_avg, M_dev):
  ny,nx,nc = numpy.shape(R)
  ravg = None
  for M, tag in ((M_avg, "avg"), (M_dev, "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_avg, M_dev, yuv):
  # Assumes that {C} is a {numpy} array with shape {(ny,nx,3)}
  #
  # If {yuv} is true, interprets it as as a YUV image. Normalizes each
  # channel so that the Y axis has mean 0.5 over the subset selected by {M_avg}, and deviation 0.25
  # on the {M_dev} 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_dev}. The U,V
  # axes are then multiplied by the Y value. The result is then
  # converted to RGB and each pixel is soft-clipped to the {[0 _ 1]}
  # cube.
  #
  # If {yuv} is false, merely soft-clips each pixel to the unit cube.
  #
  ny,nx,nc = numpy.shape(C)
  assert nc == 3, "C should have exactly three channels"
  if yuv:
    RGB = make_RGB_image_via_YUV(C, M_avg, M_dev)
  else:
    RGB = make_RGB_image_direct(C, M_avg, M_dev)
  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_avg, M_dev):
  ny, nx, nc = numpy.shape(C)
  assert nc ==3, "bad nc"

  YUV = numpy.copy(C)

  err.write("normalizing Y,U,V channels ...\n")
  Yavg = 0.60
  Ydev = 0.25
  UVdev = 0.50
  cfn.normalize_YUV_image(YUV, M_avg, M_dev, 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_dev)
  
  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_dev)
  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_avg, M_dev):
  RGB = numpy.copy(C)
  
  err.write("normalizing RGB image ...\n")
  Ymin = 0.2
  Ymax = 0.8
  cfn.normalize_RGB_image(RGB, M_avg, M_dev, 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
    # ....................................................................
  
  pdir = get_arg()
  page = get_arg()
  crop = get_arg()
  mask = get_arg(); 
  assert len(mask) == 5, "invalid mask code"
  wlen0 = int(get_arg())
  wlen1 = int(get_arg())
  wlen2 = int(get_arg())
  yuv_str = get_arg(); yuv = (yuv_str == "True" or yuv_str == "1" or yuv_str == "T")
  border = int(get_arg())
  odir = get_arg()
  oname = get_arg()
  return pdir, page, crop, mask, wlen0, wlen1, wlen2, yuv, border, odir, oname
  # ----------------------------------------------------------------------
  
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
  # ----------------------------------------------------------------------
  
# test_remix_UV_coords_of_image()

main()