#! /usr/bin/python3
# Last edited on 2025-08-30 18:19:33 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 monochome mask image
# {M[0>ny,0:nx]} for a given {mask}. Reads the spectral analysis file
# for that {page}, {clip}, and {mask}, specifically the brightest
# spectrum {smax[0:nb]} and the darkest spectrum {smin[0:nb]}.
# 
# Then extracts the spectra {C[iy,ix,:]} where {M[iy,ix]} is nonzero.
# Maps those spectra, as well as {smin,smax}, from {\RR^{nb}} to
# {\RR^3}. Displays those points with an interactive 3D interface.

import sys, re, os
from math import sqrt, hypot, log, exp, floor, ceil, comb, 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
from image_funcs import \
  mfn.read_clip_npy_images, gfn.read_mask_image, \
  gfn.quantize_image, gfn.write_uint_array_as_png, mfn.spectrum_to_color_matrix
from spectral_analysis_funcs import \
  afn.read_mask_analysis_data, afn.extract_masked_pixel_spectra
import argparser

def main():

  page, clip, mask = parse_options()

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

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

  # Reads the multispectral clip images:
  sys.stderr.write(f"reading images of clip {clip} ...\n")
  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,)

  # Extract the pixel spectra:
  sys.stderr.write(f"extracting the masked pixel spectra ...\n")
  S = afn.extract_masked_pixel_spectra(C, M)
  ns = numpy.size(S,0)
  assert numpy.shape(S) == (ns,nb,)
  
  # Compute blending weights:
  sys.stderr.write(f"computing the spectrum to color matrix ...\n")
  RfromS = mfn.spectrum_to_color_matrix(bands, bands_tb)

  sys.stderr.write(f"converting spectra to colors ...\n")
  # Project to {\RR^3}:
  R = S @ RfromS
  
  Rmin = smin @ RfromS 
  Rmax = smax @ RfromS 
  
  sys.stderr.write(f"computing color histogram ...\n")
  hist_maxval = 127 if ns > 10000 else 255
  RH = compute_rgb_histogram(R, hist_maxval)
  
  write_points_for_view_point_cloud(page, clip, mask, RH, Rmin, Rmax)

  write_points_as_rgb_png(page, clip, mask, R, Rmin, Rmax)
  
  return 0
  # ----------------------------------------------------------------------
  
def read_analysis_data(page, clip, mask):
  mfile = f"pages/{page}/clips/{clip}/masks/{mask}/mask.png"
  sys.stderr.write(f"reading {mfile} ...\n")
  M = gfn.read_mask_image(mfile)
  sys.stderr.write(f"reading analysis data ...\n")
  bands, illums, wlens, savg, smin, smax, rdev, tdev, P = \
    afn.read_mask_analysis_data(page, clip, mask)
  return M, bands, smin, smax
  # ----------------------------------------------------------------------

def compute_rgb_histogram(R, maxval):
  # Expects {R} to be an array with shape {(ns,3)} containing {ns} RGB colors.
  # Reduces the colors to the range {0..maxval} and computes their 
  # histogram.  Returns an array {RH} with shape {(nh,4,)}
  # where the first three columns are a reduced RGB color and the fourth 
  # column in a count for the same.
  
  ns, nc = numpy.shape(R)
  assert nc == 3, "bad R channels"
  assert maxval < 256, "bad maxval"
  mult = maxval + 1
  
  L = []
  for ks in range(ns):
    Rval = 0
    for ic in range(nc):
      Rval = mult*Rval + int(floor(maxval*R[ks,ic] + 0.5))
    L.append(Rval)
  L.sort()
  LH = []
  oval = -1; ct = 0;
  for val in L:
    if oval != val and ct != 0:
      LH.append((oval, ct,))
      oval = val; ct = 0
    ct += 1
  if ct != 0: LH.append((oval, ct,))
  
  nh = len(LH)
  sys.stderr.write(f"found {nh} distinct colors among {ns} pixels\n")

  RH = numpy.zeros((nh,4,))
  for ih in range(nh):
    val, ct = LH[ih]
    for ic in range(nc):
      RH[ih,nc-1-ic] = (val % mult) / maxval
      val = val // mult
    RH[ih,nc] = float(ct)
  
  return RH
  # ----------------------------------------------------------------------

def write_points_for_view_point_cloud(page, clip, mask, RH, Rmin, Rmax):
  nh = numpy.size(RH,0)
  assert numpy.shape(RH) == (nh, 4,), "bug RH shape"
  
  gdir = f"pages/{page}/clips/{clip}/masks/{mask}"
  bash(f"mkdir -p {gdir}")
  gfile = f"{gdir}/ptcloud.png"
  sys.stderr.write(f"writing {gfile} ...\n")
  with open(gfile, "w") as wr:
  
    def wrpt(tag, cx, cy, cz, rad):
      nonlocal wr
      wr.write("%s %7.4f %7.4f %7.4f  %7.4f" % (tag, cx,cy,cz, rad))
      R = min(1.0, max(0.0, cx))
      G = min(1.0, max(0.0, cy))
      B = min(1.0, max(0.0, cz))
      wr.write("  %5.3f %5.3f %5.3f\n" % (R, G, B))
      return 
      # ....................................................................

    rbox = 0.01
    wr.write("-box 0.300 0.300 0.300\n")
    for bx in range(2):
      for by in range(2):
        for bz in range(2):
          wrpt("", bx, by, bz, 0.02)
          
    for ih in range(nh):
      cx, cy, cz, ct = RH[ih,:]
      rad = 0.005*sqrt(sqrt(sqrt(ct)))
      wrpt("-point", cx, cy, cz, rad)
    
    wr.write("-segment 0.700 0.700 0.700\n")
    for Rext in Rmin, Rmax:
      cx, cy, cz = Rext[:]
      rad = 0.03
      wrpt("", cx, cy, cz, rad)
      
    wr.close()

  bash(f"view_point_cloud.sh {gfile}")
  return
  # ----------------------------------------------------------------------

def write_points_as_rgb_png(page, clip, mask, R, Rmin, Rmax):
  
  ns, nc = numpy.shape(R)
  assert nc == 3, "bad R channels"
  
  np = ns + 2; # Num of pixels to generate.
  ny = int(ceil(sqrt(np) + 1.0e-6))
  nx = (np + ny - 1)//ny
  
  P = numpy.zeros((ny,nx,3,))
  ks = 0
  for iy in range(ny):
    for ix in range(nx):
      if ks < ns:
        P[iy,ix,:] = R[ks,:]
        ks += 1
  P[ny-1,nx-2,:] = Rmin
  P[ny-1,nx-1,:] = Rmax
  P = gfn.quantize_image(P, 255)
  pdir = f"pages/{page}/clips/{clip}/masks/{mask}"
  bash(f"mkdir -p {pdir}")
  pfile = f"{pdir}/colors.png"
  sys.stderr.write(f"writing {pfile} ...\n")
  gfn.write_uint_array_as_png(pfile, P)
  bash(f"display -title '%d' {pfile}")
  return 
  # ----------------------------------------------------------------------

def parse_options():
  na = len(sys.argv)
  assert na >= 3, "insuff args"
  page = sys.argv[1]
  clip = sys.argv[2]    # name of clip to take pixels from.       
  mask = sys.argv[3]    # name of mask that selects the desied pixels.    
  return page, clip, mask
  # ----------------------------------------------------------------------

main()