# Implementation of module {seam}
# Last edited on 2021-03-21 11:56:46 by jstolfi

import seam
import contact
import move
import move_parms
import path
import block
import hacks
import rn
import pyx 
from math import nan, inf, sqrt
import sys

class Seam_IMP:
  # A block contact {sm} is between{Block} objects {bc0=sm.bc[0]} and {bc1=sm.bc[1]}, 
  # The field {sm.cts} is a tuple of {Contact} objects (trace-trace contacts)
  # where side {i} of each contact is a trace in some choice of block {sm.bc[i]}
  #
  # The field {sm.drc} field is a boolean that indicates whether the
  # chosen alternative for block {bc0} must be executed before the
  # chosen alternative of {bc1}, whatever those choices may be.
  # 

  def __init__(self, bc0, bc1, CTS, drc):
    # Read-only fields:
    self.bc = (bc0, bc1)
    self.CTS = tuple(CTS)
    self.drc = drc
    # Mutable fields for the {hotpath} heuristic:
    self.iscov = [False, False] # Status of block contact rel to current {Q}.

def make(bc0, bc1, CTS, drc):
  assert isinstance(bc0, block.Block)
  assert isinstance(bc1, block.Block)
  assert type(CTS) is tuple or type(CTS) is list
  assert type(drc) is bool
  assert bc0 != bc1
  return seam.Seam(bc0, bc1, CTS, drc)
  # ----------------------------------------------------------------------
  
def get_contacts(sm):
  assert isinstance(sm, seam.Seam)
  return sm.CTS
  # ----------------------------------------------------------------------

def side(sm, i):
  return sm.bc[i]
  # ----------------------------------------------------------------------

def which_side(bc, sm):
  assert isinstance(bc, block.Block)
  assert isinstance(sm, seam.Seam)
  for i in range(2):
    if sm.bc[i] == bc:
      return i
  return None 

def is_directed(sm):
  assert isinstance(sm, seam.Seam)
  return sm.drc
  # ----------------------------------------------------------------------

def bbox(SMS):
  B = None
  for sm in SMS:
    B = rn.box_join(B, contact.bbox(sm.CTS))
  return B
  # ----------------------------------------------------------------------


# def min_max_tcool(sm, mp_jump):
#   bc0 = sm.bc[0]
#   bc1 = sm.bc[1]
#   tc_max_min = +inf
#   for ich0 in range(block.nchoices(bc0)):
#     oph0 = block.choice(bc0,ich0)
#     omv0 = path.elem(oph0, path.nelems(oph0)-1) # Last move of {oph0}.
#     for ich1 in range(block.nchoices(bc1)):
#       oph1 = block.choice(bc1,ich1)
#       omv0 = path.elem(oph0, 0) # First move of {oph1}.
#       use_jump = False
#       use_link = False
#       tconn = connector_extime(omv0, omv1, use_jump, use_link, mp_jump)
#       # Compute max cooling time among contacts of {sm} that 
#       # are closed by {oph0} and {oph1}:
#       tc_max = pair_max_tcool(oph0, tconn, oph1, sm.CTS)
#       if tc_max < tc_max_min: tc_max_min = tc_max
#   return tc_max_min
#   # ----------------------------------------------------------------------

def print_contact_table(wr, BCS, CTS, MVS):
  nbc = len(BCS)
  nct = len(CTS)
  nmv = (0 if MVS == None else len(MVS))
  
  def fmt_mv(i, n):
    fmt = "M%0" + ("2" if n < 100 else "3")+ "d"
    return fmt % i
    
  def fmt_bc(i, j, n):
    fmt = "B%0" + ("2" if n < 100 else "3")+ "d:%d"
    return fmt % (i,j)

  def fmt_ct(i, j, n):
    fmt = "C%0" + ("2" if n < 100 else "3")+ "d:%d"
    return fmt % (i,j)

  # Table header:
  mvsp = " "*len(fmt_mv(0, nmv))
  bcds = "-"*len(fmt_bc(0, 0, nbc))
  ctsp = " "*len(fmt_ct(0, 0, nct))
  wr.write("%s %s " % (ctsp,mvsp))
  for ib in range(nbc):
    bci = BCS[ib]
    nchi = block.nchoices(bci)
    for jpi in range(nchi):
      wr.write(" " + (fmt_bc(ib,jpi, nbc)))
  wr.write("\n")
  
  # Contact entries:
  for i in range(nct):
    cti = CTS[i]
    for j in range(2):
      ctID = fmt_ct(i, j, nct)
      mvij = contact.side(cti,j)
      if MVS == None:
        mvID = mvsp
      else:
        r = MVS.index(mvij)
        mvID = fmt_mv(r, nmv)
      wr.write("%s %s " % (ctID,mvID))
      for k in range(nbc):
        bck = BCS[k]
        nchk = block.nchoices(bck)
        for l in range(nchk):
          ophkl = block.choice(bck, l)
          t = path.find(ophkl, mvij)
          if t == None:
            wr.write(" " + bcds)
          else:
            wr.write(" %*d" % (len(bcds), t))
      wr.write("\n")
    wr.write("\n")
  return
  # ----------------------------------------------------------------------

def show(wr, sm):
  wr.write("\nseam:\n")
  for i in range(2):
    pti = sm.pt[i]
    wr.write("  pt[%d] = ( %6.3f, %.3f ) drc = %d" % (i, pti[0], pti[1], int(sm.drc)))
  wr.write("\n")
  for i in range(2):
    bci = sm.bc[i]
    wr.write("  bc[%d] =" % i)
    wr.write(" ( %6.3f, %6.3f )" % move.pini(bci))
    wr.write(" --> ( %6.3f, %6.3f )" % move.pfin(bci))
    wr.write(" wd = %6.3f tex = %8.3f\n" % (move.width(bci), move.extime(bci)))
  return
  # ----------------------------------------------------------------------

# PLOTTING

def plot_to_files(fname, SMS, clr, BCS, CLRS, wd_axes, matter):

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

  dp = None
  axes = True
  dots = True
  arrows = True
  grid = True
  wd_line = 2*wd_axes
  plot_standard \
    ( c, SMS, dp, clr, wd_line, BCS, CLRS, wd_axes, \
      axes=axes, dots=dots, arrows=arrows, matter=matter, grid=grid
    )
  
  hacks.write_plot(c, fname)
  return
  # ----------------------------------------------------------------------

def plot_standard \
  ( c, SMS, dp, clr, wd_line, BCS, CLRS, wd_axes, 
    axes, dots, arrows, matter, grid
  ):

  assert type(SMS) is list or type(SMS) is tuple
  nsm = len(SMS)
  assert type(BCS) is list or type(BCS) is tuple
  nbc = len(BCS)
  if CLRS == None: 
    CLRS = [ pyx.color.rgb(0.050, 0.800, 0.000), ] # Default trace color.
  else:
    assert type(CLRS) is list or type(CLRS) is tuple
  nclr = len(CLRS)
  assert nclr == 1 or nclr == nbc
  if dp == None: dp = (0,0)

  # Get a plotting bounding box {pbox} of all blocks and contacts:
  B = block.bbox(BCS)
  B = rn.box_join(B, bbox(SMS))
  pbox = hacks.round_box(B, 1)
  szx, szy = rn.box_size(pbox)

  # Choose the verical displacement {ystep} between sub-plots:
  ystep = 1 + szy

  # Make a list {S} of pairs {(ibc0, ibc1)} such that {BCS[S[ism][i]]} is
  # the block that is side {i} of the seam {SMS[ism]}:
  S = []
  for ism in range(len(SMS)):
    sm = SMS[ism]
    ibc = [None, None]
    for i in range(2):
      bc = side(sm, i)
      if not bc in BCS:
        sys.stderr.write("side %d of seam {SMS[%d]} not in list {BCS}" % (i, ism))
        assert False
      ibc[i] = BCS.index(bc)
    S.append(tuple(ibc))
  
  # Make a list {C} of pairs {(ism, ict)} where {ism} is the index of a 
  # seeam in {SMS} and {ict} is the index of a contact in {SMS[ism]}:
  C = []
  for ism in range(len(SMS)):
    sm = SMS[ism]
    CTSi = get_contacts(sm)
    for ict in range(len(CTSi)):
      C.append((ism,ict,))
  nct = len(C) # Number of contacts to plot.
  
  # Plot as many sub-plots as needed:
  nct_plotted = 0
  while len(C) > 0: 
    # Need at least one more sub-plot.
    
    if grid:
      wd_grid = 0.002*szx
      hacks.plot_grid(c, None, wd_grid, dp, pbox, -0.25, 1, 1)

      wd_frame = 0.003*szx
      hacks.plot_frame(c, None, wd_frame, dp, pbox, 0.25)
    
    if matter:
      # Plot the matter shadows of all blocks:
      for ibc in range(nbc):
        bc = BCS[ibc]
        block.plot_matter_shadow(c, bc, dp)

    # Decide which contacts and which choices of which blocks to plot:
    K, CS, CR = plot_standard_select_items(SMS, BCS, S, C)
    assert len(CS) > 0
    assert len(CS) + len(CR) == len(C)
    assert len(K) == nbc

    # Plot the block choices:
    for ibc in range(nbc):
      bc = BCS[ibc]
      # Plot the 
      k = K[ibc]
      if k != None:
        # Must plot choice {k} of {bc}
        oph = block.choice(bc, k)
        clr_bc = CLRS[ibc] if nclr == nbc else CLRS[0]
        path.plot_standard \
          ( c, [oph,], dp, None, [clr_bc,], wd_axes, 
            axes=axes, dots=dots, arrows=arrows, matter=False
          )
    
    # Now plot the contacts in {CS}:
    for ism, ict in CS:
      sm = SMS[ism]
      ct = get_contacts(sm)[ict]
      arrow = is_directed(sm)
      sz_tic = 0 if arrow else wd_line
      contact.plot_single(c, ct, dp, clr, wd_line, sz_tic, arrow)
      nct_plotted += 1

    # Prepare for next sub-plot:
    C = CR
    dp = ( dp[0], dp[1] + ystep )
  return
  # ----------------------------------------------------------------------

def plot_standard_select_items(SMS, BCS, S, C):
  # Selects which block choices and contacts to show in the next sub-plot of
  # {plot_standard}.  Receives
  #
  #   {SMS} a list of {Seam}  objects.
  #
  #   {BCS} a list of {Block} objects.
  #
  #   {S}   a list of pairs such that {S[ism]} is {(ibc0,ibc1)} if sides 0 and 1 of seam
  #         {SMS[ism]} are the blocks {BCS[ibc0]} and {BCS[ibc1]}, respecctively.
  #
  #   {C}   a list of contacts that have not been plotted yet, each representd by a pair {ism,ict}
  #         meaning the contact with index {ict} in seam {SMS[ism]}.
  # 
  # Returns three lists:
  # 
  #   {K}   a list with {nbc = len(BCS)} elements.  Elem {K[ibc]} will be 
  #         {k} if choice {k} of block {BCS[ibc]} is to be plotted in this sub-plot,
  #         or {None} if no choice of that block needs to be plotted.
  #
  #   {CS}  the pairs in {C} of the contacts that should be shown.
  #
  #   {CR}  the pairs {C} that were not included in {CS}.
  #
  nbc = len(BCS) # Number of blocks.
  K = [None]*nbc
  CS = [] # Elements of {C} that will be shown in this sub-plot.
  CR = [] # Elements of {C} that will not be shown in this sub-plot.
  for ism, ict in C:
    sm = SMS[ism]
    ibc = S[ism] # Indices in {BCS} of blocks that are sides of {sm}.
    ct = get_contacts(sm)[ict]
    BCP = (side(sm,0), side(sm,1)) # The blocks that are the sides {sm}.
    assert BCP[0] == BCS[ibc[0]] and BCP[1] == BCS[ibc[1]]
    KPchosen = (K[ibc[0]], K[ibc[1]]) # Which choice of {BCP[i]} has been selected for this sub-plot.
    KPneed = plot_standard_contact_is_plottable(ct, BCP, KPchosen)
    if KPneed[0] == None:
      # Contact {ct} cannot be shown in this sub-plot.  Leave it for later:
      assert KPneed[1] == None
      CR.append((ism,ict))
    else:
      # Contact {ct} can be shown in this sub-plot:
      assert KPneed[1] != None
      # Mark those choices of the two blocks as the ones to be plotted:
      for i in range(2):
        if K[ibc[i]] != None:
          assert KPneed[i] == K[ibc[i]]
        else:
          K[ibc[i]] = KPneed[i]
      # Add contact to list to be shown:
      CS.append((ism,ict))
  return K, CS, CR
  # ----------------------------------------------------------------------

def plot_standard_contact_is_plottable(ct, BCP, KPchosen):
  # Decides whether a contact {ct} can be shown in a sub-plot of {plot_standard}.
  # The parameter {BCP} must be a pair of {Block} objects such that
  # {BCP[i]} contains side {i} of {ct}.  The parameter {KPchosen} must be a pair
  # such that {KPchosen[i]} is the choice of block {BCP[i]} that has been selected 
  # to be drawn in this sub-plot; or {None} if no choice of that block has been
  # selected yet.
  #
  # The procedure returns a pair {KPneed} such that {KPneed[i]} is the choice
  # of block {BCP[i]} that should be painted in order to show {ct}.
  # It is either {KPchosen[i]}, or some integer index if {KPchosen[i]} is {None}.
  # If the contact {ct} cannot be shown in this sub-plot, return {(None,None)}.

  showct = True      # Shall we show contact {ct} on this sub-plot?
  KPneed = [None, None] # To show {ct} we need to plot these choices of blocks {ibc[0..1]}
  for i in range(2):
    bc = BCP[i] # The block that is side {i} of {sm}.
    mv = contact.side(ct, i) # The trace that is side {i} of {ct}
    chosenk = KPchosen[i]
    if chosenk != None:
      # We already decided to show choice {chosenk} of block {BCP[i]}. 
      # See if that choice contains side {i} of {ct}:
      assert chosenk < block.nchoices(bc)
      oph = block.choice(bc, chosenk)
      imv = path.find(oph, mv)
      if imv == None:
        # Nope -- can't show this contact now:
        showct = False
      else:
        # So far,contact {ct} could be shown now:
        KPneed[i] = chosenk
    else:
      # We have not selected a choice for this block yet.
      # Find some choice that contains side {i} of {ct}:
      chosenk = block.find_choice_with_move(bc, mv)
      if chosenk == None:
        sys.stderr.write \
          ( "block {BCS[%d]} has no choice with side %d of contact %d of seam {SMS[%d]}"
             % (ibc[i],i,ict,ism)
          )
        assert False
      # To show {ct} we need to plot choice {chosenk} of block {BCS[ibc[i]]}:
      KPneed[i] = chosenk
  if showct:
    assert KPneed[0] != None and KPneed[1] != None
  else:
    KPneed = (None,None)
  return tuple(KPneed)
  # ----------------------------------------------------------------------

def plot_single(c, sm, dp, clr, oph0, oph1, wd_line):
  
  nch = block.nchoices(bc)
  
  arrow = is_directed(sm) 
  sz_tic = 0 if arrow else wd_line
      
  for ct in get_contacts(sm):
    ok = True
    for i in range(2):
      ophi = oph0 if i == 0 else oph1
      mvi = contact.side(ct, i)
      k = path.find_move(ophi, mvi)
      if k == None: ok = False
    if ok:
      contact.plot_single(c, ct, dp, clr, wd_line, sz_tic, arrow)
  return
  # ----------------------------------------------------------------------