#! /usr/bin/python3
# Last edited on 2025-12-17 18:23:16 by stolfi

# Reads a set of {nb} narrow-band monochrome images for a given page and
# clip.  Removes variations due to lighting and writes them out,
# normalized with respect to a given mask.
#
# To remove lighting effects, it scales the {nb} samples of each pixel
# in the {nb} bands have average 0.4. This preliminary "light removal"
# step is requested by the {unlight} flag below.
#
# COMMAND LINE ARGUMENTS
# 
# The command line arguments are
# 
#   {page} {unlight} {mask} {odir}
# 
# where
#
#   {page} is the normalized page's f-number like "008r1".
#   
#   {unlight} is a flag ("True" or "False", or 0 or 1) that requests
#   light removal.
#
#   {mask} is the tag of a subset of the page (like "wr" for writing)
#   used for brightness/contrast normalization.
#
#   {odir} is the directory for output files.
#
# 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 images are read from "MS/davis/{page}/{band}/page.png", for
# each {band} listed in the "bands.txt" file above that was imaged in
# reflected light, with two-lamp frontal narrow-band illumination
# ({illum == 0} and name starts with "MB"). These images are assumed to
# have been converted from the TIFF files without any sample scaling or
# gamma mapping.
#
# If a file "MS/davis/{page}/masks/page/{mask}.png" exists, the program
# expects that it will contain a blevel image specifying the pixels 
# (black=no, white=yes) to consider when normalizing images and color
# channels for brightness and constrast.  If "None" or "NONE", all
# pixels are considered for that purpose.
#
# OUTPUT FILES
#
# The outputs are the {mb} "page.png" narrowband images, with lighting removed and
# normalized within the given {mask} (if it exists) written as {nb}
# grayscale image files "{odir}/page-{wlen}-{mask}.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

import cProfile

def main():

  page, unlight, mask, odir, oname = parse_options()

  # Gets the spectral bands tables:
  bfile = f"MS/davis/{page}/bands.txt"
  bands_tb = bfn.read_page_bands_tables(bfile)
  
  # Select the bands to use (only reflection two-lamp illum):
  bands = []      # Bands included in the analysis.
  for band in bands_tb.keys():
    illum = bands_tb[band]['illum']
    if illum == 0:
      assert band[0:2] == "MB", "bug: band image name"
      bands.append(band)
  nb = len(bands)
  debug("bands used", bands)
  assert nb > 0, "bug nb"

  # Read the page images:
  C = read_page_images(page, bands, bands_tb)
  assert numpy.size(C,2) == nb
  ny, nx = numpy.shape(C[:,:,0])
  
  # Get the mask image:
  if mask == None:
    M = None
  else:
    mfile = f"MS/davis/{page}/masks/page/{mask}.png"
    M = gfn.read_mask_image(mfile, None)
    Msh = numpy.shape(M)
    assert Msh == (ny,nx,), "shape mismatch {ny = } {nx = } {Msh = }"
  
  if unlight:
    # "Light removal": normalizes the samples of each pixel to 
    # get average albedo 0.4 across all bands:
    for iy in range(ny):
      for ix in range(nx):
        p = C[iy,ix,:]
        assert numpy.shape(p) == (nb,)
        pavg = numpy.average(p)
        p = 0.4*p/pavg
        for ib in range(nb):
          if p[ib] > 1.0:
            err.write(f"!! light removal: C[{iy},{ix},{ib}] from {C[iy,ix,ib]} to {p[ib]}\n")
        C[iy,ix,:] = p
  
  # Write the individual clipped and un-lighted images:
  for ib in range(nb):
    band = bands[ib]
    illum = 0
    wlen = bands_tb[band]['wlen']
    assert wlen >= 200 and wlen <= 999, "bad wlen"
    gimg = make_gray_band_image(C[:,:,ib], band, M)
    gfile = f"{odir}/page-{wlen}-{mask}.png"
    gfn.write_image_as_gray_png(gfile, gimg)

  return None
  # ----------------------------------------------------------------------

def read_page_images(page, bands, bands_tb):
  # Reads the image "{band}/page.png" for {band} in the lits {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]
    maxval = bands_tb[band]['maxval']
    pfile = f"MS/davis/{page}/{band}/page.png"
    # Read the page image, scale linearly {0..maxval} to {[0_1]} :
    B = gfn.read_gray_png_image(pfile, maxval, None)
    assert len(numpy.shape(B)) == 2, "band image should be monochrome"
    ny, nx = numpy.shape(B)
    if C is None:
      # Initialize the array:
      C = numpy.zeros((ny, nx, nb))
    C[:,:,ib] = B
  return C
  # ----------------------------------------------------------------------

def make_gray_band_image(P, band, M):
  # Returne the float image {P} with the samples
  # rescaled so as to best fit the range {[0 _ 1]}.
  #
  ny,nx = numpy.shape(P)
  G = gfn.map_image_to_unit_range(P, True, M, band) 
  return G
  # ----------------------------------------------------------------------
 
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()
  unlight = bool(get_arg())

  mask = get_arg(); 
  if mask == "None" or mask == "NONE": mask = None

  odir = get_arg()
  oname = get_arg()
  return page, unlight, mask, odir, oname
  # ----------------------------------------------------------------------

# cProfile.run('B = gfn.read_gray_png_image(f"MS/davis/008r1/MB365UV_007_F/page.png", 3000, None)')
main()