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

# Reads two monochromatic images {G1,G2} taken with light coming
# from different directions. Tries to determine the local roughness of
# the vellum as the difference {(G1-G2)/(G1+G2)}, sample by sample. To
# account for variations in light intensity, applies a local
# normalization to each image first.
# 
# Command line arguments are the {page}, {clip}, the tags {band1} and {band2}, 
# of two spectral bands, and a tag {run_tag} for this run of the program.
#
# The original clip images {G1} and {G2} are read from the files
# "{cdir}/{band1}.npy" "{cdir}/{band2}.npy", where {cdir} is
# "pages/{page}/clips/{clip}/npy" The must have illumination types 1 and
# 2, respectively, as declared on the bands table file for the given
# {page}.
#
# The roughness map image images are written to files
# "{odir}/{band1}-{band2}-rough.png" where {odir} is 
# "pages/{page}/clips/{clip}/roughs/". 
#
# Also writes to "{odir}/{band1}-avg.txt" and "{odir}/{band2}-avg.txt" 
# images with the local averages of {G1} and {G2}, and 
# "{odir}/{band1}-dev.txt" and "{odir}/{band2}-dev.txt
# with the local deviations from local average.
# 
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 bands_table_funcs as bfn

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

import argparser

def main():

  page, clip, run_tag, bands = parse_options()
  
  # Read the spectral bands tables:
  pfile = f"MS/davis/{page}/bands.txt"
  sys.stderr.write(f"reading bands table {pfile} ...\n")
  bands_tb = bfn.read_page_bands_tables(pfile)

  # Read the clip images:
  sys.stderr.write(f"reading the two clip images ...\n")
  ny = None; nx = None;
  G = None # Shape {(ny,nx,2)}.
  for illum in 1, 2:
    ig = illum - 1
    band = bands[ig]
    assert bands_tb[band]['illum'] == illum, "images have wrong illum"
    gfile = f"pages/{page}/clips/{clip}/npy/{band}.npy"
    Gi = numpy.load(gfile)
    assert len(numpy.shape(Gi)) == 2, "clip image should be monochrome"
    if G is None: 
      ny, nx = numpy.shape(Gi)
      G = numpy.zeros((ny,nx,2))
    else:
      assert numpy.shape(Gi) == (ny, nx), "bug inconsistent clip shape"
    G[:,:,ig] = Gi
  debug("ny", ny)
  debug("nx", nx)
  
  odir = make_output_directory(page, clip, run_tag)
    
  sys.stderr.write(f"computing locally normalized images ...\n")
  G = normalize_locally(odir, G, bands)
    
  sys.stderr.write(f"writing locally normalized images ...\n")
  for illum in 1, 2:
    ig = illum - 1
    band = bands[ig]
    oname_nrm = f"{band}-nrm"
    write_monochrome_png_image(odir, oname_nrm, G[:,:,ig], True)
 
  sys.stderr.write(f"computing the roughness map ...\n")
  G_ruf = compute_roughness_map(G)
  
  sys.stderr.write(f"writing the roughness image ...\n")
  oname_ruf = f"{band[0]}-{band[1]}-rough.png"
  write_monochrome_png_image(odir, oname_ruf, G_ruf, True)

  return 0
  # ----------------------------------------------------------------------
  
def normalize_locally(odir, G, bands):
  # Given a set {G} of {ng} greyscale images of the same
  # clip with different illuminations, adjusts them so that 
  # locally the mean illumination is the same.
  #
  # The returned images should have values spanning {[0_1]}.
  
  ny, nx, ng = numpy.shape(G)

  sys.stderr.write(f"  obtaining the local average and variance images ...\n")
  G_avg, G_dev = compute_local_avg_dev(G)
  assert numpy.shape(G_avg) == (ny,nx,ng,), "incompatible shapes G, G_avg"
  assert numpy.shape(G_dev) == (ny,nx,ng,), "incompatible shapes G, G_dev"

  sys.stderr.write(f"  writing average images ...\n")
  for ig in range(ng):
    band = bands[ig]
    # The local avg should span {[0_1]} already so no need to normalize:
    oname_avg = f"{band}-avg"
    write_monochrome_png_image(odir, oname_avg, G_avg[:,:,ig], True)

  # The local dev images are usually too dark so we scale up for display:
  sys.stderr.write(f"  scaling deviation images for display ...\n")
  F = mfn.map_images_to_unit_range(G_dev, False, None, bands)

  sys.stderr.write(f"  writing (scaled) deviation images ...\n")
  for ig in range(ng):
    band = bands[ig]
    oname_dev = f"{band}-dev"
    write_monochrome_png_image(odir, oname_dev, F[:,:,ig], True)
  
  sys.stderr.write(f"  applying local normalization to the given images ...\n")
  N = numpy.zeros((ny,nx,ng,))
  eps = 0.001
  for ig in range(ng):
    for iy in range(ny):
      for ix in range(nx):
        val = G[iy,ix,ig]
        assert val >= 0 and val <= 1, "bad image value"
        avg = G_avg[iy,ix,ig]
        assert avg >= 0 and avg <= 1, "bad local avg"
        nrm = log((val + eps)/(avg + eps))
        N[iy,ix,ig] = nrm
        
  # Normalize both locally normalized images by same scale & shift to span {[0_1]}:
  N = mfn.map_images_to_unit_range(N, True, None, "nrm")
  return N
  # ----------------------------------------------------------------------

def compute_roughness_map(G):
  # Given a set {G} of {ng=2} greyscale images of the same
  # clip with different illuminations, computes a grayscale
  # image {R} that hopefully shows the effect of illumination 
  # canceling out the albedo.
  
  ny, nx, ng = numpy.shape(G)
  assert ng == 2, "should be two images"
  
  sys.stderr.write(f"  computing raw roughness ...\n")
  R = numpy.zeros((ny,nx,))
  for iy in range(ny):
    for ix in range (nx):
      for ig in range(ng):
        assert G[iy,ix,ig] >= 0 and G[iy,ix,ig] <= 1, "bug G values"
      dif = abs(G[iy,ix,0] - G[iy,ix,1])
      sum = G[iy,ix,0] + G[iy,ix,1]
      ruf = 0 if sum < 1.0e-4 else dif/sum
      assert ruf >= 0 and ruf <= 1, f"bug ruf {ruf}"
      R[iy,ix] = ruf

  return R
  # ----------------------------------------------------------------------

def compute_local_avg_dev(G):
  # Given a set {G} of {ng} grayscale images, with shape {(ny,nx,ng)},
  # returns sets {G_avg}, {G_dev} with same shape,
  # such that {G_avg[iy,ix,:]} is a windowed average of {G[jy,jx,:]}
  # for {jy,jx} in the neighborhood of pixel {iy,ix};
  # and {G_dev[iy,ix,:]} is a windowed RMS average
  # of {G[jy,jx,:] - G_avg[jy,jx,:]} in such a neighborhood.
  
  ny, nx, ng = numpy.shape(G)
  assert ng == 2, "should be two images"
  
  sys.stderr.write(f"  generating the local average images ...\n")
  G_avg = numpy.zeros((ny,nx,ng,))
  for ig in range(ng):
    G_avg[:,:,ig] = gfn.blur_image(G[:,:,ig])
  
  sys.stderr.write(f"  generating the local deviation images ...\n")
  G_dev = numpy.zeros((ny,nx,ng,))
  eps = 0.0005
  for ig in range(ng):
    G_var = numpy.power(G[:,:,ig] - G_avg[:,:,ig], 2)
    assert numpy.shape(G_var) == (ny,nx,)
    B_var = gfn.blur_image(G_var)
    for iy in range(ny):
      for ix in range (nx):
        dev = sqrt(B_var[iy,ix] + eps*eps)
        G_dev[iy,ix,ig] = dev
  
  return G_avg, G_dev
  # ----------------------------------------------------------------------
  
def make_output_directory(page, clip, run_tag):
  # Constructs the name {bdir} of the output directory, namely
  # "pages/{page}/clips/{clip}/roughs/{run_tag}".
  # Also creates that directory.
  
  odir = f"pages/{page}/clips/{clip}/roughs/{run_tag}"
  sys.stderr.write(f"creating directory {odir}/ ...\n")
  bash(f"mkdir -p {odir}")
  return odir
  # ----------------------------------------------------------------------
 
def write_monochrome_png_image(bdir, bname, B, display):
  # Given a monochrome image {B} with shape {(ny,nx,)}
  # and values in {[0_1]}, writes it to file
  # "{bdir}/{bname}.png"
  #
  
  ny, nx = numpy.shape(B)
  
  pfile = f"{bdir}/{bname}.png"
  sys.stderr.write(f"writing {pfile} ...\n")
  maxval = 65535
  P = gfn.quantize_image(B, maxval)
  gfn.write_uint_array_as_png(pfile, P)
  if display:
    dtag = re.sub(r"^.*[/]", "", bdir)
    bash(f"display -title '{dtag}/{bname}.png' {pfile}")
  return
  # ----------------------------------------------------------------------
    
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()       # Name of page in LFD multiscale set.  
  clip = get_arg()       # Name of clip to take pixels from.       
  run_tag = get_arg()    # Tag identifying the run of this program.
  bands = [ get_arg(), get_arg() ]
  assert ia == na, "spurious args"
  return page, clip, run_tag, bands
  # ----------------------------------------------------------------------

main()