#! /usr/bin/python3
# Last edited on 2021-03-25 16:47:16 by jstolfi

import sys
import random
from math import sqrt, sin, cos, log, exp, floor, ceil, inf, nan, pi

def main():
  random.seed(4615)
  ntmin = 100000
  for ns in (3,):
    for nt in ntmin, 4*ntmin:
      qmax = ns*ns
      # qmax = min(255, ns*ns)
      ends = True
      cook = True
      analyze_aa_err(ns, cook, nt, qmax, ends)
  return
  # ----------------------------------------------------------------------
  
def analyze_aa_err(ns, cook, nt, qmax, ends):
  # Analyzes the sampling and quantization errors of antialiasing 
  # a halfplane with {ns} by {ns} samples per pixel, as a function of
  # the reative coverage area.
  #
  # Writes to stadout a file with lines of the form
  #
  #   {qcov} {prob} {savg} {srms} {qavg} {qrms}
  #
  # where {qcov} is the quantized fraction of the pixel that is covered by the
  # halfplane, {prob} is the relative probability of that coverage fraction,
  # {srms} is the rms value of the sampling error {serr} before quatization,
  # and {qrms} is the same for the error {qerr} after quantization.
  #
  # The data is obtained by generating {nt} random halfplanes whose
  # boundary crosses the interior of the pixel, and sampling each of
  # them in a set of {ns*ns} sampling points inside the pixel. If {cook}
  # is false, these points form a regular grid. If {cook} is true, they
  # are perturbed in an attempt to improve the accuracy.
  #
  # Let {cov} be the actual fraction of the pixel that is inside a halfplane {H},
  # and {m} be the number of those sampling points that fall inside {H}.
  # The sampled coverage {scov} is {m/ns^2}, and the 
  # pre-quantization sampling error {serr = scov - cov}. 
  # 
  # The quantized coverage {qcov} is the value of {scov} converted to an integer in {0..qmax},
  # namely {enc(scov, ends)} for some suitable quantzation function {enc}.
  # The post-quantization sampling error is {qerr = dec(qcov,ends) - cov}, where  
  # {dec} is the appropriate decoding function.
  
  # Choose the sampling points in the unit square:
  smp = sampling_points(ns, cook)

  wr = open("tests/out/data_s%02d_c%d_t%07d_q%03d.txt" % (ns,int(cook),nt,qmax), "w")
  
  # Sums indexed by {qcov} in {0..qmax}:
  num_q = [0]*(qmax+1)  # Count of cases.
  sum_s = [0]*(qmax+1)  # Sum of {serr}
  sum_q = [0]*(qmax+1)  # Sum of {qerr}
  sum_s2 = [0]*(qmax+1) # Sum of {serr^2}
  sum_q2 = [0]*(qmax+1) # Sum of {qerr^2}
  
  # Accumulate stats:
  for it in range(nt):
    u, d, cov = random_halfplane()
    scov = antialias_halfplane(smp, u, d)
    serr = scov - cov
    qcov = enc(scov, qmax, ends)
    qerr = dec(qcov, qmax, ends) - cov
    
    assert 0 <= qcov and qcov <= qmax
    num_q[qcov] += 1
    sum_s[qcov] += serr; sum_s2[qcov] += serr*serr
    sum_q[qcov] += qerr; sum_q2[qcov] += qerr*qerr
    
  # Write stats:
  for qcov in range(qmax+1):
    prob = num_q[qcov]/nt
    savg = sum_s[qcov]/nt
    srms = sqrt(sum_s2[qcov]/nt)
    qavg = sum_q[qcov]/nt
    qrms = sqrt(sum_q2[qcov]/nt)
    wr.write("%3d %10.7f" % (qcov, prob))
    wr.write("  %10.7f %10.7f" % (savg, srms))
    wr.write("  %10.7f %10.7f\n" % (qavg, qrms))
  
  wr.close()
  return 
  # ----------------------------------------------------------------------
  
def random_halfplane():
  # Generates a random halfplane whose boundary 
  # intersects the interior of the pixel, assuemed to be the unit square {[0_1]^2}.
  # The underlying distribution is uniform in normal direction and position,
  # rejecting those that do not cross the pixel.
  #
  # The halfplane is described by the outwards normal {u} (a unit
  # vector perpendicular to the edge, pointing out) and the signed
  # distance {d} from in the direction {u} from the pixel center to the
  # edge.
  
  r = sqrt(0.5) # Circumradius of pixel.
  while True:
    # Generate a uniform random halfplane that cuts the circunscribed circle
    a = 2*pi*random.random()
    u = (cos(a), sin(a))
    d = (2*random.random() - 1)*r
    # If it cuts the pixel, stop:
    dmax = 0.5*(abs(u[0])+abs(u[1]))
    if abs(d) < dmax: break
    
  cov = halfplane_coverage(u, d)
  return u, d, cov
  # ----------------------------------------------------------------------
  
def halfplane_coverage(u, d):
  # Compute the fraction of the unit square covered by the 
  # halfplane {H}  with parameters {u,d}.
  
  ux = abs(u[0])
  uy = abs(u[1])
  dmax = 0.5*(ux + uy)
  dmed = 0.5*abs(ux - uy)
  tmin = min(ux, uy)/max(ux,uy)
  r = (dmax - abs(d))/(dmax - dmed)
  if d <= -dmax:
    cov = 0
  elif d <= -dmed:
    cov = 0.5*tmin*r*r
  elif d <= +dmed:
    s = 0.5*(d/dmed + 1)
    cov = 0.5*tmin + s*(1-tmin)
  elif d <= +dmax:
    cov = 1 - 0.5*tmin*r*r
  else:
    cov = 1
  return cov
  # ----------------------------------------------------------------------

def sampling_points(ns, cook):
  # Returns a list of {ns*ns} sampling points to use for 
  # computing the antialiasing.  Normally they are 
  # a regular grid, but if {cook} is true they are cooked
  # in an attempt to get more even error.
  smp = []
  for ix in range(ns):
    x = (ix + 0.5)/ns; vx = x - 0.5
    for iy in range(ns):
      y = (iy + 0.5)/ns; vy = y - 0.5;
      if cook:
        assert ns == 3, "cooking not implemented for this {ns}"
        bar = -0.20
        if (x == 0 or x == 2) and (y == 0 or y == 2):
          vx = (1+bar)*vx; x = 0.5 + vx
          vy = (1+bar)*vy; y = 0.5 + vy
        elif x != 0 and y != 0:
          vx = (1+bar/2)*vx; x = 0.5 + vx
          vy = (1+bar/2)*vy; y = 0.5 + vy
      smp.append((x,y))
  return smp
  # ----------------------------------------------------------------------

def antialias_halfplane(smp, u, d):
  # Returns the coverage fraction of the unit pixel by the halfplane
  # {H} with parameters {u,d}, estimated by sampling at the points
  # {smp[i]}.
  
  np = len(smp)
  m = 0; # Samples inside {H}.
  
  for i in range(np):
    x, y = smp[i]
    vx = x - 0.5
    vy = y - 0.5
    c = vx*u[0] + vy*u[1]
    if c <= d:
      # Inside {H}:
      m += 1
  return m/np
  # ----------------------------------------------------------------------
  
def enc(scov, qmax, ends):
  # Encodes the float coverage {scov} as an integer in {0..qmax}.
  #
  # If {ends} is false, each output value represents
  # an equal interval of {scov} valies, with width {1/(qmax+1)}.
  #
  # If {ends} is true, each output value represents an equal
  # interval of {scov} valies, with width {1/qmax}; except
  # that values 0 and {qmax} represent half-intervals adjacent 
  # to the ends of {[0_1]}.
  #
  # Halfway cases are rounded arbitrarily.  Hopefully they 
  # are too rare to matter.
  
  if ends:
    qcov = int(floor(scov*qmax + 0.5))
  else:
    qcov = int(floor(scov*(qmax+1)))
  if qcov > qmax: qcov = qmax
  if qcov < 0: qcov = 0
  return qcov
  # ----------------------------------------------------------------------
  
def dec(qcov, qmax, ends):
  # Returns the nominal float coverage value that corresponds to the 
  # the quantized coverage {qcov}.  See {enc}.
  
  if ends:
    scov = qcov/qmax
  else:
    scov = (qcov + 0.5)/(qmax+1)
  return scov
  # ----------------------------------------------------------------------

def test_halfplane_coverage():
  wr = open("tests/out/test.txt", "w")
  a = 0.418
  u = (cos(a), sin(a))
  nd = 50
  for id in range(nd+1):
    d = (2*id/nd - 1)*sqrt(0.5)
    cov = halfplane_coverage(u, d)
    wr.write("%+10.7f %10.8f\n" % (d, cov))
  wr.close()
  return
  # ----------------------------------------------------------------------

main()
# test_halfplane_coverage()