using Nemerle
using Nemerle.Utility
using Nemerle.Logging
using Nemerle.Collections
using Nemerle.Imperative
using Nemerle.Profiling
using Nemerle.Text

using System
using System.Console
using System.Text.RegularExpressions

set namespace Fx7

class Hash [K, V] : Hashtable[K, V] where V : new ()
  private check(k : K) : void
    when (!Contains(k)) base[k] = V()
  
  public new Item [k : K] : V
    get
      check(k)
      base[k]
    set
      check(k)
      base[k] = value

/*
  STATUS

  Fairly stable.

  NOTES

  - I rely on the grinder to flatten \/ and /\ and to remove
    duplicate children of those operators.
  

  TODO

  - Check the effect of transforming an assert into an assume
    on the query.

  - Improve speed.

 */

public class Pruner
  core : Core
  pool : Term.Pool
  star : Term
  False : Term
  True : Term

  Old : Hashtable[Term, Term] = Hashtable() // maps new to old terms
  New : Hashtable[Term, Term] = Hashtable() // maps old to new terms

  oldParents : Hash[Term, Hash[list[string], int]] = Hash()
  newParents : Hash[Term, Hash[list[string], int]] = Hash()
  oldIds : Hashtable[Term, int] = Hashtable()
  newIds : Hashtable[Term, int] = Hashtable()

  sortedTerm : Hashtable[Term, Term] = Hashtable()
  expliesCache : Hashtable[Term, Hashtable[Term, bool]] = Hashtable()

  static quantVarPrefix : string = "qVar"
  static skolemFunctionPrefix : string = "$@fn"
  
  public this (c : Core)
    core = c
    pool = c.TermPool
    star = pool.Get("*", [])
    False = pool.Get("false", [])
    True = pool.Get("true", [])

  regex : Regex = Regex (@".*[:_]\d+\.\d.*")
  locationCache : Hashtable [Term, bool] = Hashtable ()

  Substitute(t : Term, s : Hashtable[Term, Term]) : Term
    if (s.Contains(t))
      s[t]
    else
      t.MapChildren(Substitute(_, s))

  // Renames the quantified variables in |t| with "qv<dl>",
  // where <dl> is the depth level and "(forall x y" is
  // interpreted as "(forall x (forall y". (roughly)
  [Profile]\
  RenameQuantifiedVar(t : Term) : Term
    def cache = Hashtable ()
    def RecRQV(d, m, t)
      if (d == 0 && cache.Contains(t)) cache[t]
      else
        def res =
          match (t.Name, t.Children)
            | _ when t.Arity == 0 && m.Contains(t.Id) => m.Get(t.Id)
            | ("forall", vars::_) \
            | ("exists", vars::_) =>
                def addRename(v, (m, d))
                  (m.Replace(v.Id, pool.Get($"$quantVarPrefix$d", [])), d + 1)
                def (m, d) = vars.Children.Fold((m, d), addRename)
                t.MapChildren(RecRQV(d, m, _))
            | (name, _ :: _) when name.StartsWith(skolemFunctionPrefix) =>
              pool.Get ("%skol", pool.Get (name, []) :: t.Children)
            | _ =>
                t.MapChildren(RecRQV(d, m, _))
        when (d == 0) cache[t] = res
        res
          
    RecRQV(0, IntTree.Empty(), t)

  [Profile]\
  static IsVariable(a : Term) : bool
    // TODO: perhaps add here classes from the java API?
    | V ("true") | V ("false") | V ("null") => false
    | V (name) =>
        !(LinearTheory.IsNumber(name)) &&\
        !name.StartsWith(quantVarPrefix)
    | _ => false


  // Check if |a| is implied by |b|
  [Profile]\
  Explies(a : Term, b : Term) : bool
    ++mems
    if (a === b)
      true
    else
      ++mems
      match (a.Name)
        | "and" =>
            mems += a.Arity
            a.Children.ForAll(Explies(_, b))
        | "or" =>
            mems += a.Arity
            a.Children.Exists(Explies(_, b))
        | _ => false


  MkAndOr(a : list [Term], one : Term, _zero : Term, name : string) : Term
    if (a.Exists(_ === one))
      one
    else
      pool.Get(name, a)
      /* for debug, makes it easier to compare a.smt to b.smt
      match (a.Filter(_ !== zero))
        | [] => zero
        | [t] => t
        | lst => pool.Get(name, lst)
       */
        
  MkAnd(a : list [Term]) : Term
    MkAndOr(a, False, True, "and")
    
  MkOr(a : list [Term]) : Term
    MkAndOr(a, True, False, "or")

  [Profile]\
  RemoveDuplicates(tl : list [Term]) : list [Term]
    def h = Hashtable()
    def f(t, r)
      if (h.Contains(t))
        log(PRUNE, "remove duplicate")
        r
      else
        h[t] = t
        t :: r
    tl.Fold([], f).Rev()

  mutable mems : int = 0

  [Profile] \
  // a is disjunction of conjuncts (DNF)
  Flatten (a : list[list[Term]]) : list[list[Term]]
    def flattenOr(x : list [Term])
      | [E("or", c)] => c.Map(x => [x])
      | _ => [x]
    def flattenAnd(x : Term)
      | E("and", c) => c
      | _ => [x]
    def a = a.Fold([], (x, y) => flattenOr(x) + y)
    a.Map(x => x.Fold([], (y, z) => flattenAnd(y) + z))
    
  // Returns a term equisatisfiable to |b| assuming that |a| is unsat
  [Profile]\
  Prune(a : Term, b : Term) : Term
    mems = 0
    def RecPrune(a, b)
      def smems = mems
      _ = smems
      log(PRUNE, "{")
      log(PRUNE, $"old=$a")
      log(PRUNE, $"new=$b")
      def r = RecPruneNoLog(a, b)
      log(PRUNE, $"prn=$r}")
      log(PRUNE, $"mems=$(mems-smems)")
      log(PRUNE, $"L1xL2 = $(a.Length * a.Fold(0,(t,s)=>t.Length+s))")
      r
    and RecPruneNoLog(a : list [list [Term]], b : Term)
      ++mems
      if (a is [])
        b
      else
        def a = Flatten(a)
        match (b.Name)
          | "and" =>
              // commonAB = terms that appear in both A and B
              // prunedA = A - commonA
  
              //log (TEMP, $ "len: $(b.Arity) $(a.Length) $(a.Fold(0,(l,a) => a+l.Length))")
              //def time = Core.read_tsc()
  
              // TODO: handle DISTINCT specially
              // This part must be linear time, otherwise the speed sucks
              def allB = Hashtable()
              mems += b.Arity
              b.IterChildren(bt => allB[bt] = bt)
              def commonAB = Hashtable()
              def nestIter(lst, f)
                lst.Iter(l => l.Iter(f))
              def recCommon(t)
                ++mems
                when (allB.Contains(t))
                  ++mems
                  commonAB[t] = t
              nestIter(a, recCommon)
  
              def prune(l, s)
                l.Filter(x => !s.Contains(x))
              mems += a.Length + a.Fold(0, (t, s) => t.Length + s)
              def prunedA = a.Filter(x => x.Exists(y => commonAB.Contains(y)))
              mems += prunedA.Length + prunedA.Fold(0, (t, s) => t.Length + s)
              def prunedA = prunedA.Map(prune(_, commonAB))
              
              
              //def time = (Core.read_tsc() - time) / 1000
              //log (TEMP, $ "done: -> $(prunedA.Fold(0,(l,a) => a+l.Length)) $time ")

              ++mems
              if (prunedA.Contains([]))
                log(PRUNE, $"throw (True) $b")
                False
              else
                def filter(bc)
                  ++mems
                  if (commonAB.Contains(bc))
                    bc
                  else
                    RecPrune(prunedA, bc)
                mems += b.Arity + b.Arity
                MkAnd(RemoveDuplicates(b.Children.Map(filter)))
            
          | "or" =>
              mems += b.Arity
              MkOr(b.Children.Map(RecPrune(a, _)))
  
          | "distinct" =>
              log(PRUNE, $"b=$b a=$a")
              def allB = Hashtable()
              b.IterChildren(bc => allB[bc] = bc)
              def expliesDistinct(at : list[Term])
                | [E("distinct", ac)] => ac.ForAll(allB.Contains(_))
                | _ => false
              if (a.Exists(expliesDistinct))
                log(PRUNE, $"throw (Distinct) $b")
                False
              else
                b
              
          | _ =>
              mems += a.Length + a.Fold(0, (t, s) => t.Length + s)
              if (a.Exists(x => x.ForAll(Explies(_, b))))
                log(PRUNE, $"throw (Explies) $b")
                False
              else
                b
  
    log(PRUNE, "start the work")
    def a = RenameQuantifiedVar(a)
    def b = RenameQuantifiedVar(b)
    
    CollectVariables(a, b)
    Unify()
    def a = Substitute(a, New)

    log(PRUNE, $"same = $(a === b)")
    def a = SortTerm(a)
    def b = SortTerm(b)

    //System.IO.File.WriteAllText("q1-ren.smt", a.ToPrettyString())
    //System.IO.File.WriteAllText("q2-ren.smt", b.ToPrettyString())
    
    mutable newDist = []
    def collectDistinct(t)
      match (t.Name)
        | "and" | "or" | "not" | "forall" | "exists" => t.IterChildren(collectDistinct)
        | "distinct" => newDist ::= t
        | _ => {}
    collectDistinct(b)
    def strengthenDist(t)
      match (t.Name)
        | "and" | "or" | "not" | "forall" | "exists" =>
          t.MapChildren(strengthenDist)
        | "distinct" =>
          log (PRUNE, $"process: $t")
          mutable best = null
          mutable best_add = 0
          foreach (t' in newDist)
            def nd = Hashtable()
            t'.IterChildren(t => nd[t] = t)
            when (t.Children.ForAll(t => nd.Contains(t)) &&
                  (best == null || nd.Count - t.Arity < best_add))
              best = t'
              best_add = nd.Count - t.Arity
          if (best != null && best !== t)
            log (PRUNE, $ "strengthen dist: $t -> $best")
            best
          else t
        | _ => t
    def a = strengthenDist(a)
    
    System.IO.File.WriteAllText("a.smt", b.ToPrettyString())
    def b = RecPrune([[a]], b)
    System.IO.File.WriteAllText("b.smt", b.ToPrettyString())
    log(PRUNE, $"total_mems=$mems")
    b

  static sizeCache : Hashtable[Term, int] = Hashtable()
  static TermSize(t : Term) : int
    use_cache(sizeCache[t])
      t.Children.Fold(1, (t, s) => TermSize(t) + s)

  // Uses the element comparator |elCmp| to make a list comparator
  public static Compare['a,'b](x : list ['a], y : list ['b], elCmp : ('a * 'b) -> int) : int
    match (x, y)
      | ([], []) => 0
      | ([], _::_) => -1
      | (_::_, []) => +1
      | (xh::xt, yh::yt) =>
          def h = elCmp(xh, yh)
          if (h == 0)
            Compare(xt, yt, elCmp)
          else
            h

  // Prepares for `similarity' heuristic
  // TODO: perhaps memoize this to make it faster.
  [Profile]\
  CollectVariables(a : Term, b : Term) : void
    mutable termId = 0;
    def RecCV(t, tp, parents, ids)
      mutable pos =
        match (t.Name)
          | "and" | "or" | "distinct" | "=" | "+" | "*" =>
            -1
          | _ => 0
      foreach (c in t.Children)
        def name = 
          if (pos == -1) t.Name 
          else 
            ++pos
            $ "t.Name/$pos"
        RecCV (c, name :: tp, parents, ids)
      when (IsVariable(t))
        parents[t][tp]++
        when (!ids.Contains(t))
          ids[t] = termId
          termId++
    termId = 0
    RecCV(a, [], oldParents, oldIds)
    termId = 0
    RecCV(b, [], newParents, newIds)

  public static Filter (this s : string, f : char -> bool) : string
    def res = System.Text.StringBuilder ()
    foreach (c in s)
      when (f (c)) _ = res.Append (c)
    res.ToString ()

  [Profile]\
  LCS(a : string, b : string) : int
    def a = a.Filter (x => ! char.IsDigit (x))
    def b = b.Filter (x => ! char.IsDigit (x))
    mutable tab1 = array (b.Length + 1)
    mutable tab2 = array (b.Length + 1)
    for (mutable i = 1; i <= a.Length; ++i)
      for (mutable j = 1; j <= b.Length; ++j)
        def x = if (a[i-1] == b[j-1]) 1 + tab2[j-1] else 0
        tab1[j] = Math.Max(x, Math.Max(tab2[j], tab1[j-1]))
      (tab1, tab2) = (tab2, tab1)
    tab2[b.Length]

  #if false
  SimilarityStats (a : Term, old : bool) : void
    def ap = (if (old) oldParents else newParents)[a]
    log (PRUNE, $"similarity: $a")
    foreach (x in ap)
      when (x.Value != 0)
        log (PRUNE, $"$(x.Key) -> $(x.Value)")
  #endif

  [Profile]\
  Similarity(a : Term, b : Term) : int
    def ap = oldParents[a]
    def bp = newParents[b]
    mutable cp = 0
    foreach (x in ap) 
      def v1 = x.Value
      mutable v2
      _ = bp.TryGetValue (x.Key, out v2)
      cp += 2*Math.Min(v1, v2) - Math.Abs(v1 - v2)
    cp + LCS(a.Name, b.Name)

  // Finds a minimum weighted bipartite matching. See `selfish cities'
  // at http://topos.ucd.ie/rgrig/programs/index.html
  [Profile]\
  HungarianAlgorithm(w : array[2, int]) : (array[int] * array[int])
    def inf = System.Int32.MaxValue
    def m = w.GetLength(0)
    def n = w.GetLength(1)
    assert (m <= n)
    def row_mate = array(n)
    def col_mate = array(m)
    def parent_row = array(n)
    def unchosen_rows = array(m)
    def row_dec = array(m)
    def col_inc = array(n)
    def slack = array(n)
    mutable t = 0  // count of unchosen rows
    mutable q = 0
    mutable delta = 0
    mutable r = 0  // current row of interest
    mutable c = 0  // current column of interest
    mutable cc = 0 // second column of interest
    mutable unmatched_rows = m
    def W(r, c)
      w[r, c] - row_dec[r] + col_inc[c]
    def mate(r, c)
      row_mate[c] = r
      col_mate[r] = c
    for (r = 0; r < m; ++r) col_mate[r] = -1
    for (c = 0; c < n; ++c) row_mate[c] = -1
    while (unmatched_rows > 0)
      // Prepare
      t = 0
      for (r = 0; r < m; ++r) when (col_mate[r] == -1)
        unchosen_rows[t] = r
        ++t
      for (c = 0; c < n; ++c)
        parent_row[c] = -1
        slack[c] = inf

      // Do BFS
      doneBFS:
        for (q = 0; q < t; ++q)
          r = unchosen_rows[q]
          for (c = 0; c < n; ++c) when (W(r, c) < slack[c])
            slack[c] = W(r, c)
            when (slack[c] == 0)
              parent_row[c] = r
              when (row_mate[c] == -1) doneBFS()
              unchosen_rows[t] = row_mate[c]
              ++t

      if (q == t) // No improving path found
        // Adjust |row_dec| and |col_inc|
        delta = inf
        for (c = 0; c < n; ++c) when (parent_row[c] == -1)
          delta = Math.Min(delta, slack[c])
        for (c = 0; c < n; ++c) when (parent_row[c] != -1)
          col_inc[c] += delta
        for (q = 0; q < t; ++q)
          row_dec[unchosen_rows[q]] += delta
      else
        // Improve matching
        r = parent_row[c]
        while (col_mate[r] != -1)
          cc = col_mate[r]
          mate(r, c)
          c = cc
          r = parent_row[c]
        mate(r, c)
        --unmatched_rows

    // The result
    (col_mate, row_mate)
        

  // Populates |New| and |Old|.
  [Profile]\
  Unify() : void
    // TODO: try with doubles (instead of integers)
    def oi = if (oldIds.Count <= newIds.Count) 0 else 1 // old index
    def ni = 1 - oi // new index
    def m = Math.Min(oldIds.Count, newIds.Count)
    def n = Math.Max(oldIds.Count, newIds.Count)
    def idxToTerm = array(2, n)
    def w = array (m, n)  // max_similarity - similarity
    
    mutable max = 0
    foreach (x in oldIds) idxToTerm[oi, x.Value] = x.Key
    foreach (x in newIds) idxToTerm[ni, x.Value] = x.Key
    for (mutable i = 0; i < m; ++i)
      for (mutable j = 0; j < n; ++j)
        def (ii, jj) = if (oi == 0) (i, j) else (j, i)
        w[i, j] = Similarity(idxToTerm[oi, ii], idxToTerm[ni, jj])
        max = Math.Max(max, w[i, j])
    for (mutable i = 0; i < m; ++i)
      for (mutable j = 0; j < n; ++j)
        w[i, j] = max - w[i, j]
    def ha = HungarianAlgorithm(w)
    def mate = if (oi == 0) ha[0] else ha[1]
    assert (mate.Length == oldIds.Count)
    for (mutable i = 0; i < mate.Length; ++i)
      if (mate[i] != -1)
        New[idxToTerm[oi, i]] = idxToTerm[ni, mate[i]]
        Old[idxToTerm[ni, mate[i]]] = idxToTerm[oi, i]
      else
        log(PRUNE, $"not mapping $(idxToTerm[oi, i].Name)")
    
            
  // Sorts terms lexicographically
  [Profile]\
  SortTerm(t : Term) : Term
    use_cache(sortedTerm[t])
      def comp(a, b)
        if (a === b) 0
        else
          def nc = a.Name.CompareTo(b.Name)
          if (nc != 0) nc
          else Compare(a.Children, b.Children, comp)
      match (t.Name)
        // TODO: add other commutative functions
        | "and" | "or" | "distinct" | "=" => 
            pool.Get(t.Name, t.Children.Map(SortTerm).Sort(comp))
        | "exists" | "forall" =>
            match (t.Children)
              | [qv, pats, form] =>
                pool.Get(t.Name, [qv, pats, SortTerm(form)])
              | [qv, form] =>
                pool.Get(t.Name, [qv, SortTerm(form)])
              | _ => assert (false)
        | _ =>
            t.MapChildren(SortTerm)

  DNFSize (t : Term) : double
    match (t.Name)
      | "and" =>
        t.Children.FoldLeft (1.0, (e, a) => a * DNFSize (e))
      | "or" =>
        t.Children.FoldLeft (0.0, (e, a) => a + DNFSize (e))
      | _ => 1.0
  
  [Profile] \
  public Run (prev : Term, cur : Term) : void
    def grinder = Grinder(core)
    def prev = grinder.Skolemize(prev)
    def cur = grinder.Skolemize(cur)

    log(PRUNE, $"old_sz=$(TermSize(prev))")
    log(PRUNE, $"new_sz=$(TermSize(cur))")

    //System.IO.File.WriteAllText("q1.smt", SortTerm(prev).ToPrettyString())
    //System.IO.File.WriteAllText("q2.smt", SortTerm(cur).ToPrettyString())

    def newQuery = Prune(prev, cur) //MkAnd([pool.Get("not", [prev]), cur])

    New.Iter((k, v) => when (k !== v) log(PRUNE, $"New[$k]=$v"))

    //log(TEMP, $"old: $(TermSize(prev))/$(DNFSize(prev)), "
    //           "new: $(TermSize(cur))/$(DNFSize(cur)), "
    //           "pruned: $(TermSize(newQuery))/$(DNFSize(newQuery))")

    System.IO.File.WriteAllText("pruner-prev.smt", prev.ToPrettyString())
    System.IO.File.WriteAllText("pruner-cur.smt", cur.ToPrettyString())
    System.IO.File.WriteAllText("pruner-pruned.smt", newQuery.ToPrettyString())