#! /usr/bin/python3
# Last edited on 2026-02-16 17:12:31 by stolfi

from sys import stderr as err, stdout as out
from math import sqrt, sin, cos, log, exp, floor, pi, inf
import random
from del_shift_add_funcs import simulate_del_shift_add
from process_funcs import bash

def main():
  max_val = 3.0
  pert_dev = 0.1
  min_sep = 0.5*pert_dev
  random.seed(477000477)
  err.write("testing {simulate_del_shift_add}")
  err.write(f"with {max_val = } {pert_dev = } {min_sep = }\n")
  bash("mkdir -p tests/out")
  for nA in range(2,6):
    # Generate a list {LA} of {nA} random positions: 
    max_val_squish_A = max_val - (nA-1)*min_sep
    assert max_val_squish_A > 0
    LA = [ random.uniform(0,max_val_squish_A) for iA in range(nA) ]
    LA.sort()
    LA = [ LA[iA] + iA*min_sep for iA in range(nA) ]

    nK_min = 2
    nK_max = nA
    for nK in range(nK_min, nK_max+1):
      ofile = f"tests/out/dsa-A{nA:02d}-K{nK:02d}.dat"; wr = open(ofile, "w")
      yplot = 0
      nB_min = nK
      nB_max = max(nB_min, 5*nA//4)
      for nB in range(nB_min, nB_max+1):
        KA, LB, KB = simulate_del_shift_add(LA, nK, nB, max_val, min_sep, pert_dev)
        write_dots(wr, f"A.{nB}", LA, KA, max_val, yplot, 0xff2200); yplot += 1
        write_dots(wr, f"B.{nB}", LB, KB, max_val, yplot, 0x008822); yplot += 1
      wr.close();
      err.write(f"=== {ofile} ===\n")
      bash(f"head -n 20 {ofile} 1>&2")
      err.write(f"===============\n")
      pfile = f"tests/out/dsa-A{nA:02d}-K{nK:02d}-plot.png"
      title = f"test of simulate_del_shift_add {nA = } {nK = }"
      bash(f"plot_pos_file.sh {ofile} '{title}' > {pfile}")
  return
  # ----------------------------------------------------------------------

def write_dots(wr, loc, L, K, max_val, yplot, color):
  # Writes the positions {L[..]} and indices {K} for {plot_pos_file.sh}.
  n = len(L)
  K_set = set(K)
  for i in range(n):
    wtyp = 1 if i in K_set else 0
    wr.write(f"{loc:<12s} {max_val:8.4f} {L[i]:8.4f} {wtyp} 0x{color:06x}\n")
  wr.flush()
  return
  # ----------------------------------------------------------------------

main()