#! /usr/bin/python3 
# Last edited on 2025-08-23 12:03:49 by stolfi

# Reads a set of spectral components and combines a set of multispectral 
# images according to a specified subset of those components
# into an RGB image.

import sys, re, os
from math import sin, cos, log, exp, floor, pi, inf
import numpy
from PIL import Image
from error_funcs import prog_error, arg_error, debug
from process_funcs import bash, run_command
from bands_table_funcs import read_page_bands_tables
from image_funcs import read_mask_image, read_clip_image
from spectral_analysis_funcs import read_int_param
import argparser

def main():

  page, vset, clip, mask = parse_options()
  
  # Read the spectral bands tables:
  pfile = f"MS/davis/{page}/bands.txt"
  bands_tb = read_page_bands_tables(pfile)

  # Read the spectral components:
  bands, E = read_principal_components(page, vset)
  nb, ne = numpy.shape(E)
  debug("nb", nb)
  debug("ne", ne)

  # Read the normalization mask image:
  mfile = f"pages/{page}/clips/{clip}/masks/{mask}.png"
  M = read_mask_image(mfile)
  
  # Read the clip images:
  maxval = bands_tb[band]['maxval']
  B = read_band_clip_images(page, clip, bands, maxval)

  # Create Grayscale eigenvector images:
  odir = f"pages/{page}/clips/{clip}/comps/{vset}"

  for ie in range(ne):
    gimg = make_gray_composite_image(B, M, E, ie)
    gfile = f"{odir}/{ie}.png"
    write_numpy_image_as_png(gfile, gimg)

  # Create selected composite RGB images:
  for ie0, ie1, ie2 in ((2,6,7), (2,6,8), (2,7,8), (6,7,8), ):
    cimg = make_rgb_composite_image(B, M, E, ie0, ie1, ie2)
    cfile = f"{odir}/{ie0}_{ie1}_{ie2}.png"
    write_numpy_image_as_png(cfile, cimg)

  return 0
  # ----------------------------------------------------------------------
   
def read_band_clip_images(page, clip, bands, bands_tb):
  # Reads the images of the clip {clip} in the spectral bands listed in
  # {bands}. The result is a {numpy} array {B} of floats with shape
  # {(ny,nx,nb)} with float samples in {[0 _ 1]}.

  B = 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']
    cfile = f"MS/davis/{page}/{band}/clips/{clip}.png"
    # Read the clip image:
    C = read_clip_image(cfile, maxval)
    assert len(numpy.shape(C)) == 2, "clip image should be monochrome"
    ny, nx = numpy.shape(C)
    if B is None:
      # Initialize the array:
      B = numpy.zeros((ny, nx, nb))
    B[:,:,ib] = C
  return B
  # ----------------------------------------------------------------------

def make_gray_composite_image(B, M, E, ie):
  # Creates a grayscale image by combining the clip images in {B}
  # with weights taken from column {ie} of the spectral component matrix
  # {E}.
  #
  # The result is then shifted and scaled so that that the pixls
  # selecetd by the mask will best fit the range {[0 _ 1]}.
  # 
  # The result is a {numpy} array {cimg} of floats with shape
  # {(ny,nx,1)} with float samples in {[0 _ 1]}.
  #
  ny,nx,nb = numpy.shape(B)
  assert numpy.shape(M) == (ny,nx), "bug: mask size"
  assert nb == numpy.size(E,0)
  ne = numpy.size(E, 1)
  G = numpy.zeros((ny,nx,1))
  for ib in range(nb):
    G[:,:,0] += E[ib,ie]*B[:,:,ib]
  normalize_gray_image(G, M)
  return G
  # ----------------------------------------------------------------------
   
def make_rgb_composite_image(B, M, E, ie0,ie1,ie2):
  # Creates a composite 3-channel image by combining the clip images in
  # {B} with weights taken from the spectral component matrix {E}.
  #
  # Specifically, channels {0,1,2} are obtained by combining the clips
  # with weights taken form columns {ie0,ie1,ie2} of the matrix {E},
  # respectively. 
  # 
  # The result is interepreted as a YUV image, shifted and scaled so
  # that Y axis has mean 0.5 and deviation 0.25, while the U,V axes have
  # mean 0 and deviation 0.5. The result is then converted to RGB and
  # clipped to the {[0 _ 1]} cube, while trying to preserve (1) the Y
  # value, and (2) the hue (direction in the U,V plane).
  #
  # The YUV and RGB normalizations consider only pixels that are nonzero
  # in the monochrome mask image {M},
  #
  ny,nx,nb = numpy.shape(B)
  assert numpy.shape(M) == (ny,nx), "bug: mask size"
  assert nb == numpy.size(E,0)
  ne = numpy.size(E, 1)
  YUV = numpy.zeros((ny,nx,3))
  for ib in range(nb):
    YUV[:,:,0] += E[ib,ie0]*B[:,:,ib]
    YUV[:,:,1] += E[ib,ie1]*B[:,:,ib]
    YUV[:,:,2] += E[ib,ie2]*B[:,:,ib]
  normalize_YUV_image(YUV, M)
  RGB = RGB_image_from_YUV_image(YUV)
  return RGB
  # ----------------------------------------------------------------------

def write_numpy_image_as_png(rfile, R):
  # Assumes {R} is a numpy array with shape {(ny,nx,nc)}
  # where {nc} is either 1 or 3.
  # 
  # Writes it out to file "{rfile}".
  #
  # The image should be a {numpy} array with shape {(ny,nx,nc)}
  # 
  ny, nx, nc = numpy.shape(R)
  assert nc == 1 or nc == 3, "bug {nc}"
  ext = "pgm" if nc == 1 else "ppm" # File name extension.
  magic = "P2" if nc == 1 else "P3" # Magic code for PGM or PPM file.
  inSpace = "linearGray" if nc == 1 else "linearRGB"
  outSpace = "Gray" if nc == 1 else "sRGB"
  tfile = f"/tmp/{os.getpid()}-R.{ext}"
  # Write an ascii PGM or PPM:
  maxval = 1023
  with open(tfile, "w") as wr:
    wr.write(f"{magic}\n")
    wr.write(f"{nx} {ny}\n")
    wr.write(f"{maxval}\n")
    for iy in range(ny):
      for ix in range(nx):
        if ix % 20 == 0: wr.write("\n")
        for ic in range(nc):
          fsmp = R[iy,ix,ic]
          ival = int(floor(maxval*fsmp + 0.5))
          if ival < 0 or ival > maxval: prog_error("int sample out of bounds")
          wr.write(f" {ival}")
      wr.write("\n")
    wr.close()
  
  # Convert to PNG:
  bash\
    ( f"convert {tfile}" + \
      f"    -set colorspace {inSpace}" + \
      f"    -depth 8" + \
      f"    -colorspace {outSpace}" +
      f"  {rfile}"
    )
  return
  # ----------------------------------------------------------------------

def read_principal_components(page, vset):
  # Reads the bands list {bands} and the corresponding principal
  # component matrix {E} for the component set {vset} of the given
  # {page}.
  # 
  efile = f"pages/{page}/vsets/{vset}/princ.txt"
  with open(efile, "r") as rd:
    nb = read_int_param(rd, "nb")
    ne = read_int_param(rd, "ne")
    bands = []
    E = numpy.zeros((nb,ne))
    for ib in range(nb):
      line = rd.readline().strip()
      flds = line.split(' ')
      assert len(flds) == ne+1, "bad format (cols)"
      band = flds[0].strip(); bands.append(band)
      for ie in range(ne):
        fval = float(flds[ie + 1].strip())
        E[ib,ie] = fval
    rd.close()
  return bands, E
  # ----------------------------------------------------------------------

def parse_options():
  page = sys.argv[1]      
  vset = sys.argv[2]  # Set of spectral components.        
  clip = sys.argv[3]  # Clip to use for composed images.              
  mask = sys.argv[4]  # Mask to use when normalizing the composed images.           
  return page, vset, clip, mask
  # ----------------------------------------------------------------------
  
def debug(name, arg):
  sys.stderr.write(f"{name}");
  if isinstance(arg, list) or isinstance(arg, tuple) or isinstance(arg, dict):
    sys.stderr.write(f" (length = {len(arg)}):\n")
    sys.stderr.write(r"  ")
  else:
    sys.stderr.write(r" = ")
  sys.stderr.write(f"{arg}\n")
  return
  # ----------------------------------------------------------------------

main()