using System.Text

using SCG = System.Collections.Generic

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

set namespace Fx7

public partial class Term
  // -------------------------------------------------------------------------
  // Pool
  [CreateMemento] \
  public partial class Pool
    [Record] \
    struct NameArity : System.IComparable [NameArity]
      public name : string
      public arity : int

      public CompareTo (o : NameArity) : int
        if (o.arity == arity)
          string.CompareOrdinal (o.name, name)
        else o.arity - arity

    [Copy] mutable current_pool : ULongTree [Term] = ULongTree.Empty ()
    mutable master_pool : ULongTree [Term] = ULongTree.Empty ()

    mutable fv_count : int
    internal core : Core

    mutable current_name_id : int
    mutable current_apply_id : int
    mutable id_by_name : Map [NameArity, int] = Map ()
    internal mutable name_by_id : array [string] = array (1000)
    internal mutable arity_by_id : array [int] = array (1000)

    internal mutable current_level : int = 0
    [Copy] mutable rollback_queue : list [Term] = []
    [Copy] mutable distinct_stack_ptr = 0
    [Copy] internal mutable MatchHeight : int
    internal distinct_stack : array [Proof * list [Term]] = array (64)
    

    public PushState () : void
      current_level++
      SaveMemento ()
      rollback_queue = []

    public PopState () : void
      current_level--
      foreach (r in rollback_queue)
        r.Rollback ()
      RestoreMemento ()
      //log (TEMP, $ "pop state --> $distinct_stack_ptr")

    public QueueRollback (t : Term) : void
      rollback_queue ::= t

    public DumpStats () : void
      master_pool.IterValues (fun (t) {
        when (t.Stats != null)
          t.Stats.Dump (t)
      })

    static sig (left : int, right : int) : ulong
      unchecked
        (((right :> uint) :> ulong) << 32) | (left :> uint)

    public AssertDistinct (terms : list [Term], proof : Proof) : void
      assert (distinct_stack_ptr < distinct_stack.Length)
      def mask = 1UL << distinct_stack_ptr
      distinct_stack [distinct_stack_ptr] = (proof, terms)
      distinct_stack_ptr++
      def ht = Hashtable ()
      foreach (t in terms)
        if (ht.Contains (t.Root))
          core.Refute (Proof.Rule3 ("direct-distinct",
                                    Proof.EqProof (t, t.Root),
                                    Proof.EqProof (ht [t.Root], t.Root),
                                    proof))
          return
        else
          ht [t.Root] = t

      //log (TEMP, $ "assert distinct: $(distinct_stack_ptr - 1)")
      foreach (t in terms)
        //log (TEMP, $ "mask=$mask $(t.Root.distinct_mask) $t $(t.Root)")
        assert (t.Root.distinct_mask & mask == 0)
        t.Root.WillWrite ()
        t.Root.distinct_mask |= mask

    internal this (c : Core)
      core = c
      EmptySubst = MakeSubstLite ([])
      True = Get ("$@true", [])
      True.active = true
      PredTrue = Get ("true", [])
      PredFalse = Get ("false", [])
      Star = Get ("$@*", [])
      Forall2 = GetConst ("forall", 2)
      Forall3 = GetConst ("forall", 3)
      core.OnPush += PushState
      core.OnPop += PopState

    public True : Term
    public Star : Term
    public PredTrue : Term
    public PredFalse : Term

    internal Forall2 : Term
    internal Forall3 : Term

    public FreshVar () : Term
      ++fv_count
      while (id_by_name.Contains (NameArity ($"$$@$fv_count", 0)))
        ++fv_count
      Get ($"$$@$fv_count", [])

    public FreshFn (arity : int) : string
      ++fv_count
      while (id_by_name.Contains (NameArity ($"$$@fn$fv_count", arity)))
        ++fv_count
      $"$$@fn$fv_count"

    internal GetApply (left : Term, right : Term, skip_parents : bool) : Term \
      ensures value != null

      def rewritten = TryRewrites (left.Id, right)
      if (rewritten != null)
        rewritten
      else
        def k = sig (left.Id, right.Id)
        mutable res = null
        if (master_pool.TryGetValue (k, out res))
          res
        else
          --current_apply_id
          def t = Term (this, current_apply_id, left, right, skip_parents)
          t.WillWrite ()
          log (VTERM, $"master_pool[$k] = $t")
          master_pool = master_pool.Add (k, t)
          t

    GetConst (name : string, arity : int) : Term \
      ensures value != null

      match (id_by_name.Find (NameArity (name, arity)))
        | Some (res) =>
          master_pool.Get (sig (res, 0))
        | None =>
          ++current_name_id

          def id = current_name_id

          when (id >= arity_by_id.Length)
            System.Array.Resize (ref arity_by_id, id * 4)
            System.Array.Resize (ref name_by_id, id * 4)
          arity_by_id [id] = arity
          name_by_id [id] = name
          
          def t = Term (this, id, null, null, skip_parents = false)
          t.WillWrite ()
          
          id_by_name = id_by_name.Add (NameArity (name, arity), id)
          master_pool = master_pool.Add (sig (id, 0), t)
          log (VTERM, $"master_pool[$(sig(id,0))] = $t")

          t

    internal IsParentSkipping (name : string) : bool
      match (name)
        | "=" | "PATS" | "MPAT" | "NOPATS" | "ite" | "and" | "or" | "forall" | "not" => true
        | _ => false
    
    public Get (name : string, children : list [Term]) : Term
      log (VTERM, $"get: $name $children")

      // skip marking logical connectives as parents, as they are not going
      // to be merged with anything anyway
      def skip_parents = IsParentSkipping (name)

      def loop (_)
        | [x] => x
        | x :: xs =>
          GetApply (x, loop (xs), skip_parents)
        | [] => assert (false)

      loop (GetConst (name, children.Length) :: children)

    internal CheckSigHit (u :  Term) : option [Term * Proof]
      def k =
        if (u.left == null)
          sig (u.Id, 0)
        else
          sig (u.left.Root.Id, u.right.Root.Id)
      mutable term = null
      log (VTERM, $"look for sig($k)")
      if ((current_pool.TryGetValue (k, out term) || 
           master_pool.TryGetValue (k, out term)) && term.active)
        log (VTERM, $"found sig hit ($k), $u -> $term")
        if (term.SameAs (u))
          None ()
        else
          log (TERM, $"will merge ($k), $u -> $term ") //$(System.Environment.StackTrace)")
          def proof = 
             Proof.Rule2 ("fun-congr",
                          Proof.EqProof (u.left, term.left), 
                          Proof.EqProof (u.right, term.right))
          Some (term, proof)
      else
        log (VTERM, $"not found $u, won't merge")
        current_pool = current_pool.Add (k, u)
        None ()

   
    [Profile] \
    SetLastMerge (a : Term) : void
      def set_last_merge (a)
        log (TERM, $"set_last_merge: $(a.active) $(a.last_merge)/$MatchHeight $a $(a.ToExactString()) ")
        def a = a.Root
        when (a.active && a.last_merge != MatchHeight)
          a.WillWrite ()
          a.last_merge = MatchHeight
          when (Driver.FnLastMerge)
            a.UpdateFnLastMerge ()
          a.IterEqc (p => p.parents.Iter (set_last_merge))

      a.IterEqc (set_last_merge)

    [Profile] \
    public Merge (orig_a : Term, orig_b : Term, proof : Proof) : void
      assert (orig_a.active && orig_b.active)

      log (TERM, $"merge $orig_a = $orig_b by $proof at $MatchHeight")
      def a = orig_a.Root
      def b = orig_b.Root
      log (TERM, $"merge2 $(a.ToExactString()) = $(b.ToExactString()) fa=$(a.forbidden_merges) fb=$(b.forbidden_merges)")
      
      if (a : object == b)
        log (TERM, "already done")
      #if DETERMINISTIC
      else if (orig_a.LexiCmp (orig_b) < 0)
      #else
      // else if (a.Id > b.Id)
      else if (a.depth > b.depth || (a.depth == b.depth && a.Id > b.Id))
      #endif
        Merge (orig_b, orig_a, proof)
      else
        orig_a.WillWrite ()
        orig_b.WillWrite ()
        a.WillWrite ()
        b.WillWrite ()

        //log (TEMP, $"merge $a $b $(System.Environment.StackTrace)")
        def proof' = orig_a.CannotMergeProof (orig_b)
        if (proof' != null)
          core.Refute (Proof.Rule2 ("impossible-merge", proof', proof))
        else
          orig_b.ReverseProofPath ()
          orig_b.proof_parent = orig_a
          orig_b.proof = proof

          // now check which theories we need to inform
          when (a.IsReal && b.IsReal)
            core.TellEquality (a, b)

          a.forbidden_merges += b.forbidden_merges
          a.distinct_mask |= b.distinct_mask

          def delay = Queue ()
          
          log (VTERM, $"b.eqc_par: $b -> $(b.parents)")

          SetLastMerge (a)
          SetLastMerge (b)

          b.IterEqc (fun (p) {
            p.WillWrite ();
            p.Root = a
          })

          def check_parents (p)
            foreach (u in p.parents)
              log (VTERM, $"consider: $u ($(u.active))")
              when (u.active)
                match (CheckSigHit (u))
                  | Some ((term, proof)) =>
                    delay.Push (u, term, proof)
                  | None => {}

          b.IterEqc (check_parents)
          b.eqc_next <-> a.eqc_next

          while (! delay.IsEmpty)
            Merge (delay.Take ())


    #region rewrite rule handling
    // rewrite_rules[c] is a list of pairs (X, something), where c(X) => something
    // is a rewrite rule
    mutable rewrite_rules : IntTree [list [Term * Term]] = IntTree.Empty ()

    public AddRewriteRule (from : Term, to : Term) : void
      assert (from.left != null && from.left.left == null)
      //log (TEMP, $ "Add rule $from -> $to $(from.left) $(from.right)")
      mutable current = []
      _ = rewrite_rules.TryGetValue (from.left.Id, out current)
      rewrite_rules = rewrite_rules.Add (from.left.Id, (from.right, to) :: current)

    FindMatch (pat : Term, t : Term, sub : IntTree [Term]) : IntTree [Term]
      if (sub == null || pat : object == t) sub
      else
        if (pat.left == null && pat.Name [0] == '$')
          mutable now = null
          if (sub.TryGetValue (pat.Id, out now))
            if (now : object == t) sub
            else null
          else
            sub.Add (pat.Id, t)
        else if (pat.left != null && t.left != null)
          FindMatch (pat.right, t.right, FindMatch (pat.left, t.left, sub))
        else null

    TryRewrites (id : int, t : Term) : Term
      mutable rules = null
      if (rewrite_rules.TryGetValue (id, out rules))
        foreach ((from, to) in rules)
          //log (TEMP, $"check for match $from at $t")
          def sub = FindMatch (from, t, IntTree.Empty ())
          when (sub != null)
            def r = to.Apply (sub)
            //log (TEMP, $"rewritten: $t -> $r")
            return (r)
        null
      else null
    #endregion