#! /usr/bin/python3
# Last edited on 2025-12-21 08:02:25 by stolfi

from math import hypot, pi, inf, nan, isnan
import numpy, sys, os
from sys import stderr as err, stdout as out

from process_funcs import bash
import vms_linear_gray_image_funcs as gfn

def fit_background(R, M):
  # Given a multichannel image {R} and a binary mask{M},
  # fits a "background" image {U} to those pixels 
  # where {M} is nonzero, and extrapolates it to those pixels
  # where {M} is zero.
  
  ny, nx, nc = numpy.shape(R)
  U = numpy.zeros((ny,nx,nc))
  for ic in range(nc):
    fit_gray_background(R[:,:,ic], M, U[:,:,ic])
  return U
  # ----------------------------------------------------------------------
  
def fill_holes_by_nearest(R, M, U):
  # Copies the grayscale image {R} into {U}
  # where {M} is nonzero, and fills the other pixels
  # by nearest neighbor propagation.
  ny, nx = numpy.shape(R)
  assert numpy.shape(M) == (ny,nx)
  assert numpy.shape(U) == (ny,nx)
  U = numpy.full((ny,nx), nan)
  # The queue {Q} holds the pixels that have been set 
  # whose neighbors need to be inspected.
  Q = [ None ]*(nx*ny)
  # Fill the queue with those pixels where {M} is nonzero:
  nq = 0
  for iy in range(ny):
    for ix in range(nx):
      if M[iy,ix] != 0:
        U[iy,ix] = R[iy,ix]
        Q[nq] = (ix,iy);
        nq += 1
  # Now propagate from those pixels to the neighbors:
  iq = 0
  while iq < nq:
    ix, iy = Q[iq]; iq += 1
    smp = U[iy,ix]
    assert not isnan(smp)
    for dy in (-1, 0, +1):
      for dx in (-1, 0, +1):
        jx = ix + dx; jy = iy + dy
        if jx >= 0 and jx < nx and jy >= 0 and jy < ny and isnan(U[jx,jy]):
          U[jy,jx] = smp
          Q[nq] = (jx,jy); nq += 1
  retun
  # ----------------------------------------------------------------------

def fit_gray_background(R, M, U):
  # Like {U = fit_background(R, M)}, but expects that {R} and {U}
  # are monochrome images.
  ny, nx = numpy.shape(R)
  fill_holes_by_nearest(R, M, U)
  for ip in range(3):
    relax_quadratic(R,M,U)
  return
  # ----------------------------------------------------------------------
  
def relax_quadratic(R, M, U):
  # Applies local relaxing to every pixel {p = [iy,ix]}.
  # Alternates 
  
  ny, nx = numpy.shape(R)
  for ky in 0, 2, 1, 3:
    for kx in 0, 2, 1, 3:
      for iy in range(ky,ny,4):
        for ix in range(kx,nx,4):
          relax_quadratic_locally(ix, iy, R, M, U)
  return
  # ----------------------------------------------------------------------
  
def relax_quadratic_locally(ix, iy, R, M, U):
  # fits a quadratic {F(dx,dy)} to the values of {U[iy+dy,ix+dx]} in a
  # neighborhood of pixel {p = (ix,iy)}, and to the values of {R} in
  # that neighborhood where the mask is nonzero. Then replaces
  # {U[iy,ix]} by the value of {F(0,0)}.
  
  ny, nx = numpy.shape(R)
  assert numpy.shape(M) == (ny,nx)
  assert numpy.shape(U) == (ny,nx)

  # The data samples are {D[ks]} with corresponding weights {W[ks]}
  # for {ks} in {0..ns-1} where {ns} is the number
  # of pixels in the neighborhood used in the fit.
  #
  # The basis functions are {f[0..5][ks]} are the quadratic monomials basis
  #   { 1, dx, dy, dx^2, dx*dy, dy^2 }
  # where {(dx,dy)} are the displacements of the coords {(jx,jy)}
  # of sample pixel {ks} relative to pixel {(ix,iy)}

  # The neighborhood is the {5x5} square centered at pixel {(ix,iy)},
  # namely pixels {(ix+dx,iy+dy)} for {dx,dy} in {-2..+2}.
  # Thus {ns = 25} and the sample index {ks} is {5*(dy+2) + (dx+2)}.
  nb = 6
  ns = 25
  
  def basis(ks):
    # Computes the {nb} basis functions {f[0..nb-1][ks]} at 
    # sample index {ks}. Assumes {ks} is {5*dy + dx}.
    dx = (ks % 5) - 2
    dy = (ks // 5) - 2
    f = (1, dx, dy, dx*dx, dx*dy, dy*dy)
    return f
    # ....................................................................
  
  def sample(ks):
    # Relative weights for the quadratic error functional:
    WR = 3.0
    WU = 1.0

    # Delives the data value {D[ks]} for sample index {ks}
    # and the corresponding weight {W[ks]}.
    # Assumes {ks} is {5*dy + dx}.
    
    # The weight will be zero if {(dy,dx) = (±2,±2)}
    # or {(dy,dx) = (0,0)}.
    
    nonlocal nx, ny, ix, iy, R, M, U
    d2max = 6 # To exclude only {(dy,dx) = (±2,±2)}.
    D = 0; W = 0
    dy = (ks // 5) - 2
    jy = iy + dy
    if jy >= 0 and jy < ny:    
      dx = (ks % 5) - 2
      jx = ix + dx
      if ((dx != 0) or (dy != 0)) and jx >= 0 and jx < nx:  
        d2 = dx*dx + dy*dy
        if d2 <= d2max:
          err.write(f"    {ks = :3d} {dx = :3d} {dy = :3d} {jx = :3d} {jy = :3d}")
          err.write(f" R = {R[jy,jx]:12.8f} M = {M[jy,jx]} U = {U[jy,jx]:12.8f} ")
          if M[jy,jx] != 0:
            W = WR
            D = R[jy,jx]
          else:
            W = WU
            D = U[jy,jx]
          wd2 = 1 - d2/d2max; 
          wd4 = wd2*wd2
          W = W * wd4
          err.write(f" {D = :12.8f} {W = :12.8f}\n")
    return D, W
    # ......................................................................
  
  A, B = least_squares_system(nb, basis, ns, sample)
  c = numpy.linalg.solve(A, B)
  err.write(f"  {c = }\n")
  # Assumes {basis(kd)} is {(1,0,...0)} for the central sample.
  U[iy,ix] = c[0]
  return
  # ----------------------------------------------------------------------
  
def least_squares_system(nb, basis, ns, sample):
  # The least-squares system {A*c = B} finds the {nb}-coefficient vector
  # {c} that provides the best fit of a function {F(c)[ks]} to 
  # data {D[ks]} with weights {W[ks]}, for {ks} in {0..ns-1}.
  #
  # The error functional {Q(c)} is
  #   { Q(c) = SUM(W[ks]*(F(c)[ks] - D[ks])^2 : ks) }
  # {F(c)[ks]} is assumed to be a lineat combination of {nb} 
  # basis functions {f[jb][ks]} with coefficients {c}
  #   { F(c)[ks] = SUM(c[jb]*f[jb][ks] : jb) }
  # where {c[0..nb-1]} are the coeffs to be fitted
  #
  # Minimizing {Q(c)} means making {dQ(c)/dc[ib] = 0} for {ib} in {0..nb-1}
  # The derivative is 
  #   { dQ(c)/dc[ib] = SUM(2*W[ks]*(F(c)[ks] - D[ks])*dF(c)[ks]/dc[ib] : ks)
  # and 
  #   { dF(c)[ks]/dc[ib] = f[ib][ks] }
  # Therefore 
  #   { dQ(c)/dc[ib] = SUM(A[ib,jb]*c[jb] : jb) - B[ib])
  # where
  #   { A[ib,jb] = 2*SUM(W[ks]*f[jb][ks]*f[ib][ks] : ks) }
  #   { B[ib] = 2*SUM(W[ks]*f[ib]*D[ks] : ks) }
  # Then we determine the coefficients {C} by solving {A*c = B}.

  A = numpy.zeros((nb,nb,))
  B = numpy.zeros((nb,))
  for ks in range(ns):
    D, W = sample(ks)
    if W == 0: err.write(f"    {ks = :3d} {D = :12.8f} {W = :12.8f}\n")
    f = basis(ks)
    for ib in range(nb):
      B[ib] += W*D*f[ib]
      for jb in range(ib+1):
        A[ib,jb] += W*f[ib]*f[jb]
  # Complete A by symmetry:
  for ib in range(nb):
    for jb in range(ib):
      A[jb,ib] = A[ib,jb]
      
  err.write("  A,B:\n")
  err.write(f"{A}\n")
  err.write(f"{B}\n")
  return A, B
  # ----------------------------------------------------------------------

def make_test_M_full(nx, ny):
  # A mask that is 1 everywhere.
  M = numpy.full((ny,nx,), 1)
  return M
  # ----------------------------------------------------------------------

def make_test_S(nx, ny, which):
  if which == "constant":
    # An image that is a quadratic everywhere:
    S = numpy.full((ny,nx,), pi)
  else:
    S = numpy.zeros((ny,nx,))
    for iy in range(ny):
      for ix in range(nx):
        ux = ix - 5.7
        uy = iy - 6.2
        if which == "linear":
          S[iy,ix] = 0.5*ux + 0.3*uy
        elif which == "quadratic":
          S[iy,ix] = 0.5*ux*ux - 0.2*ux*uy + 0.4*uy*uy
        else:
          assert False
  return S
  # ----------------------------------------------------------------------
  
def test_relax_quadratic_locally(nx, ny):
  assert nx >= 10 and ny >= 10
  ix0 = 1; iy0 = ny//2
  
  # A quadratic image:
  S = make_test_S(nx, ny, "quadratic")
  test_show(S, "S")
  
  # The mask:
  M = make_test_M_full(nx,ny)
  for jy in range(ny):
    for jx in range(nx):
      d = (jx + jy) - (ix0 + iy0)
      if abs(d) <= 1: M[jy,jx] = 0
  
  # The data image {R}:
  R = S.copy()
  for jy in range(ny):
    for jx in range(nx):
      if M[jy,jx] == 0: R[jy,jx] = 0.5
  
  test_show(R, "R")
  
  # A computed background {U} with garbage at {(ix0,iy0)}:
  U = S.copy(); U[iy0,ix0] = 100000
  
  relax_quadratic_locally(ix0, iy0, R, M, U)
  error = abs(U[iy0,ix0] - S[iy0,ix0])
  err.write(f"  {error = :12.8f}\n")
  assert error <= 1.0e-8
  return
  # ----------------------------------------------------------------------
  
def test_show(S, tag):
  Swr = gfn.map_image_to_unit_range(S, True, None, "S")
  fname = f".test-{tag}.png"
  gfn.write_image_as_gray_png(fname, Swr)
  bash(f"display -filter Box -resize '1600%' {fname}")
  return
  # ----------------------------------------------------------------------
  
def test_it():
  ny = 30; nx = 40;
  test_relax_quadratic_locally(nx,ny)
  return
  # ----------------------------------------------------------------------

# test_it()