using Nemerle.Collections;
using Nemerle.Utility;

[Record]
class Variable {
  public name : string;
  public max_value : int;
  public mutable current_value : int;
}

class UnfoldVars {
  label_under : Hashtable [string, list [Variable]] = Hashtable (100);
  variables : Hashtable [string, Variable] = Hashtable (100);
  mutable stmts : list [Stmt];

  error (msg : string) : void
  {
    throw UserErrorException (msg);
  }
  
  public this (stmts : list [Stmt])
  {
    this.stmts = stmts;
  }

  scan_labels () : void
  {
    def scan (vars, stmt) {
      match (stmt) {
        | Stmt.Label (name) =>
          if (label_under.Contains (name))
            error ("redef of label " + name)
          else
            label_under.Add (name, vars)
        | Stmt.If (_, l1, l2) =>
          scanl (vars, l1);
          scanl (vars, l2);
        | Stmt.Vars (newvars, body) =>
          mutable vrs = vars;
          List.Iter (newvars, fun (name, size) {
            def v = Variable (name, size - 1, 0);
            if (variables.Contains (name))
              error ("redef of variable " + name)
            else
              variables.Add (name, v);
            vrs = v :: vrs;
          });
          scanl (vrs, body)
        | Stmt.Goto
        | Stmt.Action => ()
      }
    } and scanl (vars, stmts) {
      List.Iter (stmts, fun (s) { scan (vars, s) }) 
    };
    scanl ([], stmts);
  }

  no_actions (expr : BooleanFormula) : void
  {
    | BooleanFormula.Cond => error ("condition ripped off")
    | BooleanFormula.Dummy_true
    | BooleanFormula.Const => ()
    | BooleanFormula.Not (e) => no_actions (e)
    | BooleanFormula.And (e1, e2)
    | BooleanFormula.Or (e1, e2) => no_actions (e1); no_actions (e2)
  }

  can_fold (expr : BooleanFormula) : bool
  {
    | BooleanFormula.Cond => false
    | BooleanFormula.Dummy_true
    | BooleanFormula.Const => true
    | BooleanFormula.Not (x) => can_fold (x)
    | BooleanFormula.And (x1, x2) 
    | BooleanFormula.Or (x1, x2) => can_fold (x1) && can_fold (x2)
  }

  fold_expr (expr : BooleanFormula) : BooleanFormula
  {
    | BooleanFormula.Dummy_true
    | BooleanFormula.Cond => expr
    | BooleanFormula.Const (c) => BooleanFormula.Const (ConstantExpr.Const (eval_const (c)))
    | BooleanFormula.Not (x) =>
      match (fold_expr (x)) {
        | BooleanFormula.Const (ConstantExpr.Const (0)) =>
          BooleanFormula.Const (ConstantExpr.Const (1))
        | BooleanFormula.Const (ConstantExpr.Const (1)) =>
          BooleanFormula.Const (ConstantExpr.Const (0))
        | BooleanFormula.Const =>
          error ("invalid argument to !"); null
        | x => BooleanFormula.Not (x)
      }
    | BooleanFormula.And (x1, x2) =>
      match (fold_expr (x1)) {
        | BooleanFormula.Const (ConstantExpr.Const (0)) => no_actions (x2); BooleanFormula.Const (ConstantExpr.Const (0))
        | BooleanFormula.Const (ConstantExpr.Const (1)) => fold_expr (x2)
        | BooleanFormula.Const =>
          error ("invalid left argument to &&"); null
        | x1 =>
          match (fold_expr (x2)) {
            | BooleanFormula.Const (ConstantExpr.Const (0)) => no_actions (x1); BooleanFormula.Const (ConstantExpr.Const (0))
            | BooleanFormula.Const (ConstantExpr.Const (1)) => x1
            | BooleanFormula.Const =>
              error ("invalid right argument to &&"); null
            | x2 => BooleanFormula.And (x1, x2)
          }
      }

    | BooleanFormula.Or (x1, x2) =>
      match (fold_expr (x1)) {
        | BooleanFormula.Const (ConstantExpr.Const (0)) => fold_expr (x2)
        | BooleanFormula.Const (ConstantExpr.Const (1)) => no_actions (x2); BooleanFormula.Const (ConstantExpr.Const (1))
        | BooleanFormula.Const =>
          error ("invalid left argument to ||"); null
        | x1 =>
          match (fold_expr (x2)) {
            | BooleanFormula.Const (ConstantExpr.Const (1)) => no_actions (x1); BooleanFormula.Const (ConstantExpr.Const (1))
            | BooleanFormula.Const (ConstantExpr.Const (0)) => x1
            | BooleanFormula.Const =>
              error ("invalid right argument to ||"); null
            | x2 => BooleanFormula.Or (x1, x2)
          }
      }
  }

  eval_const (expr : ConstantExpr) : int
  {
    match (expr) {
      | ConstantExpr.Const (x) => if (x < 0) { error ("negative literal"); 0 } else x
      | ConstantExpr.Ref (name) =>
        match (variables.Get (name)) {
          | Some (v) =>
            v.current_value
          | None => error ("undef variable " + name); 0
        }
      | ConstantExpr.Binary (op, e1, e2) =>
        def e1 = eval_const (e1);
        def e2 = eval_const (e2);
        match (op) {
          | BinaryOperator.Plus => e1 + e2
          | BinaryOperator.Minus =>
            if (e1 - e2 < 0) { error ("negative `-' result"); 0 }
            else e1 - e2
          | BinaryOperator.Equal =>
            if (e1 == e2) 1 else 0
          | BinaryOperator.Less_than => 
            if (e1 < e2) 1 else 0
          | BinaryOperator.More_than => 
            if (e1 > e2) 1 else 0
        }

      | ConstantExpr.Not (e) =>
        def e = eval_const (e);
        if (e > 1) { error ("non boolean arg to !"); 0 }
        else if (e == 0) 1 else 0
    }
  }

  mutable unique_id : int;
  
  unfolds_and (s : list [Stmt]) : list [Stmt]
  {
    List.Concat (List.Map (s, unfold_and))
  }

  unfold_and (s : Stmt) : list [Stmt]
  {
    match (s) {
      | Stmt.Vars (newvars, body) =>
        def label = "__varsend_" + unique_id.ToString ();
        ++unique_id;
        def body = unfolds_and (body + [Stmt.Goto ([], label)]);
        [Stmt.Vars (newvars, body), Stmt.Label (label)]
        
      | Stmt.If (expr, l1, l2) =>
        match (expr) {
          | c when can_fold (c) =>
            [Stmt.If (c, unfolds_and (l1), unfolds_and (l2))]
          | BooleanFormula.And (e1, e2) =>
            def label = "__and_" + unique_id.ToString ();
            ++unique_id;
            unfold_and (Stmt.If (e1, [Stmt.If (e2, l1, [Stmt.Goto ([], label)])], Stmt.Label (label) :: l2))
          | BooleanFormula.Or (e1, e2) =>
            def label = "__or_" + unique_id.ToString ();
            ++unique_id;
            unfold_and (Stmt.If (e1, Stmt.Label (label) :: l1, [Stmt.If (e2, [Stmt.Goto ([], label)], l2)]))
          | BooleanFormula.Not (e) =>
            unfold_and (Stmt.If (e, l2, l1))
          | c =>
            [Stmt.If (c, unfolds_and (l1), unfolds_and (l2))]
        }
        
      | Stmt.Goto
      | Stmt.Label
      | Stmt.Action => [s]
    }
  }
 
  unfolds (s : list [Stmt]) : list [Stmt]
  {
    List.Concat (List.Map (s, unfold))
  }

  unfold (s : Stmt) : list [Stmt]
  {
    match (s) {
      | Stmt.Label (name) =>
        def var_name (v : Variable) {
          "_" + v.name + "=" + v.current_value.ToString ()
        };
        def vars = Option.UnSome (label_under.Get (name));
        [Stmt.Label (name + NString.Concat ("", List.Map (vars, var_name)))]
      
      | Stmt.Goto (assigns, name) =>
        def used = Hashtable ();
        def value (v : Variable) {
          mutable expr = null;
          if (List.Exists (assigns, fun (x, e) { 
                if (x == v.name) { expr = e; true } else false 
              })) {
            used.Set (v.name, null);
            def e = 
              match (fold_expr (expr)) {
                | BooleanFormula.Const (ConstantExpr.Const (x)) => x
                | _ => error ("evil constant folding result"); 0
              };
            e % (v.max_value + 1);
          } else
            v.current_value
        };
        def var_name (v : Variable) {
          "_" + v.name + "=" + value (v).ToString ()
        };
        match (label_under.Get (name)) {
          | Some (vars) =>
            def name = name + NString.Concat ("", List.Map (vars, var_name));
            List.Iter (assigns, fun (x, _) {
              unless (used.Contains (x))
                error ("unused goto assign to `" + x + "'")
            });
            [Stmt.Goto ([], name)]
          | None => error ("undef label " + name); null
        }
        
      | Stmt.Vars (newvars, body) =>
        def incr () {
          List.Exists (newvars, fun (name, _) {
            def v = Option.UnSome (variables.Get (name));
            if (v.current_value < v.max_value) {
              ++v.current_value;
              true
            } else {
              v.current_value = 0;
              false
            }
          })
        };
        def loop (acc) {
          def acc = unfolds (body) :: acc;
          if (incr ())
            loop (acc)
          else 
            List.Concat (List.Rev (acc))
        };
        loop ([])
        
      | Stmt.If (expr, l1, l2) =>
        match (fold_expr (expr)) {
          | BooleanFormula.Dummy_true
          | BooleanFormula.Const (ConstantExpr.Const (1)) =>
            [Stmt.If (BooleanFormula.Dummy_true (), unfolds (l1), unfolds (l2))]
            //unfolds (l1)
          | BooleanFormula.Const (ConstantExpr.Const (0)) =>
            [Stmt.If (BooleanFormula.Dummy_true (), unfolds (l2), unfolds (l1))]
            //unfolds (l2)
          | BooleanFormula.Const => error ("evil constant folding result to if"); null
          | BooleanFormula.And
          | BooleanFormula.Or
          | BooleanFormula.Not => error ("&&, || or ! survived!"); null
          | (BooleanFormula.Cond) as c =>
            [Stmt.If (c, unfolds (l1), unfolds (l2))]
        }
      
      | Stmt.Action => [s]
    }
  }
 
  public static Dump (stmts : list [Stmt]) : void
  {
    def w (i, s : string) {
      for (mutable j = 0; j < i; ++j)
        System.Console.Error.Write ("  ");
      System.Console.Error.WriteLine (s);
    };
    
    def dump (i, x) {
      match (x) {
        | Stmt.Label (name) => w (i - 1, name + ":");
        | Stmt.Goto (_, name) => w (i, "goto " + name)
        
        | Stmt.Vars => assert (false)
        
        | Stmt.If (BooleanFormula.Cond (NodeCondition.Pickup), l1, l2) =>
          w (i, "pickup {");
          dumps (i + 1, l1);
          w (i, "} nofood {");
          dumps (i + 1, l2);
          w (i, "}");

        | Stmt.If (BooleanFormula.Cond (NodeCondition.Move), l1, l2) =>
          w (i, "move {");
          dumps (i + 1, l1);
          w (i, "} blocked {");
          dumps (i + 1, l2);
          w (i, "}");

        | Stmt.If (BooleanFormula.Cond (NodeCondition.Flip (max)), l1, l2) =>
          w (i, "ifrand " + max.ToString () + " {");
          dumps (i + 1, l1);
          w (i, "} else {");
          dumps (i + 1, l2);
          w (i, "}");
          
        | Stmt.If (BooleanFormula.Cond (NodeCondition.Sense (c1, c2)), l1, l2) =>
          w (i, "if " + c1.ToString () + " " + c2.ToString () + " {");
          dumps (i + 1, l1);
          w (i, "} else {");
          dumps (i + 1, l2);
          w (i, "}");
          
        | Stmt.If (BooleanFormula.Dummy_true, l1, l2) =>
          w (i, "alwys {");
          dumps (i + 1, l1);
          w (i, "} junk {");
          dumps (i + 1, l2);
          w (i, "}");
          
        | Stmt.If (c, _, _) => 
          assert (false, c.ToString ())
        
        | Stmt.Action (NodeAction.Mark (c)) => w (i, "mark " + c.ToString ());
        | Stmt.Action (NodeAction.Unmark (c)) => w (i, "unmark " + c.ToString ());
        | Stmt.Action (NodeAction.Turn (true)) => w (i, "turn left");
        | Stmt.Action (NodeAction.Turn (false)) => w (i, "turn right");
        | Stmt.Action (NodeAction.Drop) => w (i, "drop");
      }
    } and dumps (i, s) { List.Iter (s, fun (x) { dump (i, x) }) };

    dumps (1, stmts);
  }

  public Execute () : list [Stmt]
  {
    stmts = unfolds_and (stmts);
    scan_labels ();
    stmts = unfolds (stmts);
    when (System.Array.IndexOf (System.Environment.GetCommandLineArgs (), "-dump") != -1)
      Dump (stmts);
    stmts
  }
}