#! /usr/bin/python3
# Test program for module {best_elem_path}
# Last edited on 2021-12-28 15:05:37 by stolfi

import best_elem_path
import paper_example_B
import raster
import raster_example
import path
import path_example
import move 
import move_parms
import contact
import hacks
import palette
import txt_read
import job_parms

import rn
import rmxn

import pyx
import sys
import time
from math import sqrt, sin, cos, floor, ceil, inf, nan, pi

# Move parameters matching those in the paper:
acc = 3000
sp_trace = 40
sp_jump = 130
udp = 0.05

wd_cont = 0.30; mp_cont = move_parms.make(wd_cont, acc, sp_trace, 0.0)
wd_fill = 0.40; mp_fill = move_parms.make(wd_fill, acc, sp_trace, 0.0)
wd_link = 0.30; mp_link = move_parms.make(wd_link, acc, sp_trace, 0.0)
mp_jump = move_parms.make(0.0, acc, sp_jump, udp)

wd_axes = 0.15*min(wd_fill,wd_cont)

def plot(OPHS, CTS, tag):
  # Plots the paths in {OPHS} and the contacts in {CTS}.
  # 
  # Plot files are "tests/out/best_elem_path_TST_{tag}.{ext}"
  # where {ext} is "png", "jpg", "eps".
  
  nph = len(OPHS)

  LKS = set()
  for oph in OPHS:
    for olk in list(path.get_links(oph) + path.get_links(path.rev(oph))):
      lk, dr = path.unpack(olk)
      LKS.add(lk)
  LKS = list(LKS)
  nlk = len(LKS)

  # Compute the bounding box of it all:
  B = None
  B = rn.box_join(B, path.bbox(OPHS))
  B = rn.box_join(B, path.bbox(LKS))
  B = hacks.round_box(B,0.5)
  
  dp = (0,0)
  
  frame = False
  grid = True
  c,szx,szy = hacks.make_canvas(B, dp, None, frame, grid, 1, 1)

  # PLOT THE FULLPATH
  
  # Paint the material traces:

  rwd_ph = 0.80
  axes_ph = False
  dots_ph = False
  arrows_ph = False
  matter_ph = True
  path.plot_standard(c, OPHS, dp, 0, None, rwd_ph, wd_axes, axes_ph, dots_ph, arrows_ph, matter_ph)

  # Paint the nominal traces of the links:

  if nlk > 0:
    CLRS_lk = hacks.link_colors(nlk, None)
    rwd_lk = 0.60
    axes_lk = True
    dots_lk = True
    arrows_lk = False
    matter_lk = False
    path.plot_standard(c, LKS, dp, None, CLRS_lk, rwd_lk, wd_axes, axes_lk, dots_lk, arrows_lk, matter_lk)

  # Paint the nominal traces of the path:

  if len(OPHS) == 1:
    clr_ph = pyx.color.rgb( 0.000, 0.800, 1.000 )
    CLRS_ph = [clr_ph,]
  else:
    CLRS_ph = hacks.trace_colors(nph, None)
  rwd_ph = 0.80
  axes_ph = False
  dots_ph = False
  arrows_ph = True
  matter_ph = False
  path.plot_standard(c, OPHS, dp, None, CLRS_ph, rwd_ph, wd_axes, axes_ph, dots_ph, arrows_ph, matter_ph)
    
  if len(CTS) > 0:
    clr_ct = pyx.color.rgb( 1.000, 0.000, 0.700 )
    wd_ct = 1.5*wd_axes
    dashpat_ct = None
    sz_tics_ct = 0
    arrows_ct = False
    for ct in CTS:
      ep = contact.endpoints(ct)
      dep = rn.dist(ep[0], ep[1])
      len_ct = max(wd_ct, dep - 4*wd_ct)
      ext_ct = (len_ct - dep)/2 
      contact.plot_single(c, ct, dp, clr_ct, dashpat_ct, ext=ext_ct, wd=wd_ct, sz_tic=sz_tics_ct, arrow=arrows_ct)

  fname = "tests/out/best_elem_path_TST" "_" + tag
  hacks.write_plot(c, fname)
  return
  # ----------------------------------------------------------------------

def do_test_solve(OPHS, tag, Delta):

  sys.stderr.write("--- testing {solve} tag = %s ---\n" % tag)
  
  debug = False
   
  # Set the contact time limits:
  CTS = set()
  for oph in OPHS:
    for isd in 0,1:
      for ct in path.get_contacts(oph,isd):
        CTS.add(ct)
        contact.set_tcool_limit(ct, Delta)
  CTS = list(CTS)
  
  move_parms.show(sys.stderr, "mp_jump = ", mp_jump, "\n")

  org = None
  best_elem_path.describe_and_check_input(sys.stderr, OPHS, org, mp_jump)
  tbeg = time.clock_gettime(time.CLOCK_THREAD_CPUTIME_ID)
  fph, ncalls = best_elem_path.solve(OPHS, org, mp_jump)
  tend = time.clock_gettime(time.CLOCK_THREAD_CPUTIME_ID)
  Tcpu = tend - tbeg
  
  sys.stderr.write("=== results ===\n")
  if fph == None:
    sys.stderr.write("!! {best_elem_path.solve} returned {None}\n")
  else:
    sys.stderr.write("tag = %s\n" % tag)

    # Compute input element fabtime and jumptime:
    udp = move_parms.ud_penalty(mp_jump)
    Tfab_in = 0
    Tjump_in = 0
    for oph in OPHS:
      Tfab_in += path.fabtime(oph)
      nmv = path.nelems(oph)
      for imv in range(nmv):
        omv = path.elem(oph, imv)
        if move.is_jump(omv):
          assert imv != 0 and imv != nmv-1
          Tjump_in += move.fabtime(omv) + 2*udp
    Ttrace_in = Tfab_in - Tjump_in
    sys.stderr.write("Tfab (input)   = %.6f\n" % Tfab_in)
    if Tjump_in > 0:
      sys.stderr.write("Ttrace (input) = %.6f\n" % Tfab_in)
      sys.stderr.write("Tjump (input)  = %.6f\n" % Tfab_in)
    sys.stderr.write("\n")
   
    Tfab_ot = path.fabtime(fph)
    sys.stderr.write("Tfab (output)  = %.6f\n" % Tfab_ot)
    assert Tfab_ot >= Tfab_in

    # Compute jump and link time in output (excluding input jumps if any):
    Tjump_ot = 0
    nmv = path.nelems(fph)
    for imv in range(nmv):
      omv = path.elem(fph, imv)
      if move.is_jump(omv):
        Tjump_ot += move.fabtime(omv)
        if imv != 0: Tjump_ot += udp
        if imv != nmv-1: Tjump_ot += udp
    Tjump_ot -= Tjump_in
    Tlink_ot = Tfab_ot - Tfab_in - Tjump_ot
    assert Tjump_ot >= 0.0
    assert Tlink_ot >= -1.0e-8
    Tlink_ot = max(Tlink_ot, 0)
    sys.stderr.write("Tlink (output) = %.6f\n" % Tlink_ot)
    sys.stderr.write("Tjump (output) = %.6f\n" % Tjump_ot)
    sys.stderr.write("\n")
    
    # Find and print coldest contacts:
    ncold = 3 # Number of coldest contacts to show.
    CTScold = contact.coldest(fph, CTS, ncold)
    contact.show_times(sys.stderr, [fph,], CTScold)

    # Computation tine of {best_elem_path.solve}:
    sys.stderr.write("\n")

  sys.stderr.write("Tcpu = %.2f s\n" % Tcpu)
  sys.stderr.write("ncalls = %d s\n" % ncalls)
  sys.stderr.write("===============\n")

  if fph != None:
    plot([fph,], CTScold, tag)
  return
  # ----------------------------------------------------------------------
  
def test_raster_example_comb(nm, np, ns, dr, do, Delta):
  # Tests {best_elem_path.solve} with the {raster_example.make_comb} dataset.

  tag_num = "nm%02d_np%02d_ns%02d" % (nm, np, ns)
  tag_dim = "dr%05d_do%05d" % (int(dr*1000), int(do*1000))
  tag = "raster_example_comb_%s_%s_D%06.0f" % (tag_num, tag_dim, 1000*Delta)

  OCRS, OPHS, OLKS, CTS = raster_example.bicomb(nm, np, ns, dr, do, mp_cont,mp_fill,mp_link) 

  do_test_solve(OPHS, tag, Delta)
  
  return 
  # ----------------------------------------------------------------------
  
def test_trihex(Delta):
  # Tests {best_elem_path.solve} a small sample, forcing to enumerate all paths.

  tag = "raster_example_trihex_D%05d" % int(Delta*1000)
  
  Rp = wd_fill
  Rt = 4.0
  contacts = True
  OPHS, CTS = path_example.trihex(Rp, Rt, contacts, mp_fill, mp_jump)
  do_test_solve(OPHS, tag, Delta)
  # plot(OPHS, CTS, tag)
  
  return 
  # ----------------------------------------------------------------------

### test_trihex(0.010)
### test_trihex(9.000)
test_trihex(2.000)

test_raster_example_comb(3, 1, 0, 6*wd_fill, 2*wd_fill, 0.345)
test_raster_example_comb(3, 1, 1, 6*wd_fill, 2*wd_fill, 0.210)
test_raster_example_comb(3, 1, 1, 6*wd_fill, 2*wd_fill, 0.240)
test_raster_example_comb(3, 1, 1, 6*wd_fill, 2*wd_fill, 0.280)
test_raster_example_comb(3, 1, 1, 6*wd_fill, 2*wd_fill, 0.400)
test_raster_example_comb(3, 1, 1, 6*wd_fill, 2*wd_fill, 0.500)