#! /usr/bin/python3
# Test program for module {palette}
# Last edited on 2021-02-14 19:52:01 by jstolfi

import palette
import color
import hacks
import rn
import pyx
import sys
from math import sqrt, hypot, floor, ceil, sin, cos, inf, nan, pi

def plot_sample(c, k, RGBk, ystart, np):
  #  Plot a sample rectangle with index {k} (in {0..np-1}) and color {RGBk}.
  # The sample has width 1 height 2; assumes that the row of samples starts at
  # height {ystart} with a margin 1 all around the sample row. 
  sys.stderr.write("k = %d RGBk = ( %5.3f %5.3f %5.3f )\n" % ((k,) + RGBk)) 
  xlo = 1 + k; xhi = xlo + 1;
  ylo = ystart + 1; yhi = ylo + 2
  pyck = pyx.color.rgb( RGBk[0], RGBk[1], RGBk[2] )
  hacks.plot_box(c, pyck, ((xlo,ylo), (xhi,yhi)))

def plot_tic(c, r, ystart, np):
  #  Plots a tic mark at position {r} between the centers
  # of samples {k=2} and {k=np-3}. 
  kf = r*(np-5) + 2 # Fractional palette index corresp. {r}
  xlo = 1 + kf; xhi = xlo + 1; xmd = (xlo + xhi)/2
  ylo = ystart + 1; yhi = ylo + 2; ymd = (ylo + yhi)/2
  yr0 = ymd - 0.2; yr1 = ymd + 0.2 
  hacks.plot_line(c, pyx.color.rgb.black, 0.05, (xmd, yr0), (xmd, yr1))

def test_smooth_single(c, np, ystart):
  # Plots samples of {smooth.single} on {c} starting at height {ystart}. Returns 
  # the {ystart} for the next stuff.
  
  sys.stderr.write("--- testing{smooth_single} --------------\n")

  x0 = 17; RGB0 = ( 0.000, 0.500, 0.000 )
  x1 = -3; RGB1 = ( 1.000, 0.000, 1.000 )
  blog = 2
  
  # Testing {smooth_single}:
  for ulog in False, True:
    for k in range(np):
      rk = (k-2)/(np-5)
      xk = (1-rk)*x0 + rk*x1
      RGBk = palette.smooth_single(xk, x0,RGB0, x1,RGB1, blog*int(ulog))
      plot_sample(c, k, RGBk, ystart, np)
    plot_tic(c, 0.0, ystart, np)
    plot_tic(c, 1.0, ystart, np)
    if ulog: plot_tic(c, abs(blog/(x1 - x0)), ystart, np)
    ystart += 3

  return ystart
  # ----------------------------------------------------------------------
  
def test_smooth_double(c, np, ystart):
  # Plots samples of {smooth.double} on {c} starting at height {ystart}. Returns 
  # the {ystart} for the next stuff.
  
  sys.stderr.write("--- testing{smooth_double} --------------\n")
  
  xmin =  -9; RGBmin = ( 0.000, 0.167, 1.000 )            # Y = 0.200
  xzer =  +4; RGBzer = ( 0.350, 0.350, 0.350 ); eps = 0.5 # Y = 0.350
  xmax = +10; RGBmax = ( 1.000, 0.500, 0.000 )            # Y = 0.600

  # Testing {smooth_double}:
  for ulog in False, True:
    for k in range(np):
      rk = (k-2)/(np-5)
      xk = (1-rk)*xmin + rk*xmax
      RGBk = palette.smooth_double(xk, xmin, RGBmin, xzer, eps, RGBzer, xmax, RGBmax, ulog)
      plot_sample(c, k, RGBk, ystart, np)
    plot_tic(c, 0.0, ystart, np)
    plot_tic(c, 1.0, ystart, np)
    plot_tic(c, (xzer-eps-xmin)/(xmax - xmin), ystart, np)
    plot_tic(c, (xzer-xmin)/(xmax - xmin), ystart, np)
    plot_tic(c, (xzer+eps-xmin)/(xmax - xmin), ystart, np)
    ystart += 3

  return ystart
  # ----------------------------------------------------------------------
  
def test_discrete(c, np, ystart):
  # Plots samples of {palette.discrete} on {c} starting at height {ystart}. Returns 
  # the {ystart} for the next stuff.
  
  sys.stderr.write("--- testing{discrete} --------------\n")
  
  def do_test_palette(Ymin, Ymax, Smin, Smax, ystart):
    sys.stderr.write("Y = [ %5.3f _ %5.3f ] S = [ %5.3f _ %5.3f ]\n" % (Ymin,Ymax,Smin,Smax))
    for k in range(np):
      RGBk = palette.discrete(k, Ymin, Ymax, Smin, Smax)
      plot_sample(c, k, RGBk, ystart, np)
      YUVk = color.YUV_from_RGB(RGBk)
      sys.stderr.write("YUVk = ( %5.3f %5.3f %5.3f )\n" % YUVk) 
      YHSk = color.YHS_from_YUV(YUVk)
      sys.stderr.write("YHSk = ( %5.3f %5.3f %5.3f )\n" % YHSk) 
      assert Ymin-1.0e-6 <= YUVk[0] and YUVk[0] <= Ymax+1.0e-6
    plot_tic(c, 0.0, ystart, np)
    return ystart+3

  tests = (
      (0.300, 0.600, 1.000, 1.000),
      (0.300, 0.600, 0.700, 1.000), 
      (0.300, 0.600, 0.500, 0.700), 
      (0.500, 0.700, 1.000, 1.000), 
      (0.500, 0.700, 0.700, 1.000), 
      (0.500, 0.700, 0.500, 0.700), 
      (0.100, 0.500, 0.700, 1.000),
      (0.100, 0.800, 0.700, 1.000),
    )
    
  # Testing {discrete}:
  for Ymin, Ymax, Smin, Smax in tests:
    ystart = do_test_palette(Ymin, Ymax, Smin, Smax, ystart)
    
  return ystart
  # ----------------------------------------------------------------------

c = pyx.canvas.canvas()
pyx.unit.set(uscale=1.00, wscale=1.00, vscale=1.00)

np = 25 # Number of color samples per row

ystart = 0
ystart = test_smooth_single(c, np, ystart)
ystart = test_smooth_double(c, np, ystart)
ystart = test_discrete(c, np, ystart)

dp = (0,0)
szx = 1 + np + 1
szy = ystart + 1
hacks.plot_frame(c, pyx.color.rgb.white, 0.10, dp, szx, szy, -0.06)

hacks.write_plot(c, "tests/out/palette_TST")