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

# Reads a set of {nb} narrow-band monochrome images {C[0:ny,0:nx,0:nb]}
# for a given {page} and {clip}. Reads a set of spectral analysis data
# for {nm} subsets of thay clip selected by {nm} mask images, including
# the brightest spectrum {smax[0:nb,0:nm]} and the darkest spectrum
# {smin[0:nb:,0:nm]}. 
#
# For each {mask} index {im} in {0..nm-1}, writes three grayscale images:
# {slon[0:ny,0:nx]}, {slat[0:ny,0:nx]}, and {spol[0:ny,0:nx]}.  
#
# The value {slon[iy,ix]} tells the position of the pixel spectrum
# {spix[iy,ix,:]} projected on the line {L} from {smin[:,im]} to
# {smax[:,im}. Namely it tends to 1 as the projection tends to infinity
# beyond {smax}, and to 0 as it tends to infinity beyond {smin}. If
# {spix} is equidistant {smin} and {smax}, then {slon[iy,ix]} will be
# 0.5.
#
# The value of {slat[iy,ix]} is 1.0 if the pixel spectrum
# {spix[iy,ix,:]} lies on the line {L}, and tends to zero as {spix}
# moves away from that line.
#
# The value of {spol[iy,ix]} is {0.5 + slat[iy,ix]*(slon[iy,ix] - 0.5)}.
# It tends to 0.5 away from the line {L}, and on the line {L} it 
# is the same as {slon}.
# 
# These images are written to disk as as file "{mdir}/slon.png",
# "{mdir}/slat.png", and "{mdir}/spol.png" where {mdir} is
# "pages/{page}/clips/{clip}/masks/{mask}".
#

import sys, re, os
from math import sqrt, hypot, sin, cos, log, exp, floor, pi, inf, isfinite
import numpy
from error_funcs import prog_error, arg_error, debug
from bands_table_funcs import mfn.read_page_bands_tables
from process_funcs import bash, run_command
import argparser

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

def main():

  page, clip, masks = parse_options()
  nm = len(masks)

  # Reads the spectral bands tables:
  bfile = f"MS/davis/{page}/bands.txt"
  bands_tb = mfn.read_page_bands_tables(bfile)

  # Reads the spectral analysis results for the given masks:
  bands, smin, smax = read_analysis_data(page, clip, masks)
  nb = len(bands)
  debug("nb", nb)
  assert numpy.shape(smin) == (nb,nm,)
  assert numpy.shape(smax) == (nb,nm,)

  # Reads the multispectral clip images:
  C = mfn.read_clip_npy_images(page, clip, bands)
  ny, nx = numpy.shape(C[:,:,0])
  debug("ny", ny)
  debug("nx", nx)
  assert numpy.shape(C) == (ny,nx,nb,)
  
  for im in range(nm):
    mask = masks[im]

    # Create spectral similarity images:
    slon, slat = create_spectral_similarity_images(C, smin[:,im], smax[:,im])

    # Write the images:
    write_spectral_similarity_images(page, clip, mask, slon, slat)

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

def read_analysis_data(page, clip, masks):
  # Reads analysis data for the given {masks}.
  nm = len(masks)
  bands, illums, wlens, savg, smin, smax, rdev, tdev, P = \
    afn.read_mask_list_analysis_data(page, clip, masks)
  return bands, smin, smax
  # ----------------------------------------------------------------------

def create_spectral_similarity_images(C, smin, smax):
  # Given a multipsectral clip image set {C} with shape {(ny,nx,nb)}
  # and two spectra {smin[0:nb],smax[0:nb]};
  # creates two images {slon[0:ny,0:nx]} and {slat[0:ny,0:nx]}
  # that show the position of each pixel spectrum relative to 
  # {smin} and {smax}.  Returns the arrays {slon,slat}.
  # 
  ny, nx, nb = numpy.shape(C)
  debug("nb", nb)

  debug("smin.shape", numpy.shape(smin))
  assert numpy.shape(smin) == (nb,), "bug bad smin shape"
  assert numpy.shape(smax) == (nb,), "bug bad smax shape"
  
  slon = numpy.zeros((ny,nx,))
  slat = numpy.zeros((ny,nx,))
  
  # Compute unit vector {u} para to subspace:
  u = smax - smin;
  u = u / numpy.linalg.norm(u)
  # Compute distance {d} between two centers:
  d = numpy.dot(smax - smin, u) # D
  bgrd = 0.15  
  bsim = 0.05
  for iy in range(ny):
    for ix in range(nx):
      spix = C[iy,ix,:]
      # Compute coord {t} of {spix-smin} in subspace:
      t = numpy.dot(spix - smin, u)
      # Compute distance {p} of {spix} to subspace:
      p = numpy.linalg.norm(spix - (smin + t * u))
      # Define the {slon} image:
      r = hypot(bgrd, (t/d)/2)
      s = hypot(bgrd, (1-t/d)/2)
      slon[iy,ix] = 0.5 + r - s
      # Define the {slat} image: 
      slat[iy,ix] = bsim/hypot(bsim, p/d)
  return slon, slat  
  # ----------------------------------------------------------------------

def write_spectral_similarity_images(page, clip, mask, slon, slat): 
  ny, nx = numpy.shape(slon) 
  assert numpy.shape(slat) == (ny,nx,), "bug inconsistent shape slon, slat"

  sdir = f"pages/{page}/clips/{clip}/masks/{mask}"
  bash(f"mkdir -p {sdir}")
  gfn.write_image_as_gray_png(f"{sdir}/slon.png", slon)
  gfn.write_image_as_gray_png(f"{sdir}/slat.png", slat)
  spol = numpy.zeros((ny,nx,))
  for iy in range(ny):
    for ix in range(nx):
      spol[iy,ix] = 0.5 + slat[iy,ix] * (slon[iy,ix] - 0.5)
  gfn.write_image_as_gray_png(f"{sdir}/spol.png", spol)
  return
  # ----------------------------------------------------------------------
  
def parse_options():
  na = len(sys.argv); ia = 1;
  
  def get_opt():
    nonlocal ia
    if ia >= na: arg_error(f"missing arg {ia}")
    arg = sys.argv[ia]; ia += 1
    return arg 
    # ......................................................................
  
  page = get_opt();    # Name of page.
  clip = get_opt();    # Name of clip to take pixels from.       
  masks = [ get_opt() for k in range(ia,na) ]
  return page, clip, masks
  # ----------------------------------------------------------------------

main()