(**************************************************************************
 * You can add new test cases by adding new elements to the following lists
 * Format is:
 * TEST<X>ARG(<weight>, <function_name>, <arg1>, <arg2>, ..., <argX>)
 *
 * <X> is the number of argument that the function being tested takes.
 **************************************************************************)
open Mp5common
open Mp5eval

let evalAll a b = (eval_exp (LetExp (a,(AppExp ( AppExp (VarExp "list_all",FnExp ("x", BinOpAppExp (EqOp,(VarExp "x"), ConstExp (IntConst 1)))),(BinOpAppExp (ConsOp,ConstExp (IntConst 1),BinOpAppExp (ConsOp,ConstExp (IntConst 2),ConstExp NilConst)))))), [])) = (eval_exp (LetExp (b,(AppExp ( AppExp (VarExp "list_all",FnExp ("x", BinOpAppExp (EqOp,(VarExp "x"), ConstExp (IntConst 1)))),(BinOpAppExp (ConsOp,ConstExp (IntConst 1),BinOpAppExp (ConsOp,ConstExp (IntConst 2),ConstExp NilConst)))))), []))

let evalAllB a b = (eval_exp (LetExp (a,(AppExp ( AppExp (VarExp "list_all",FnExp ("x", BinOpAppExp (GreaterOp,(VarExp "x"), ConstExp (IntConst 1)))),(BinOpAppExp (ConsOp,ConstExp (IntConst 100),BinOpAppExp (ConsOp,ConstExp (IntConst 200),ConstExp NilConst)))))), [])) = (eval_exp (LetExp (b,(AppExp ( AppExp (VarExp "list_all",FnExp ("x", BinOpAppExp (GreaterOp,(VarExp "x"), ConstExp (IntConst 1)))),(BinOpAppExp (ConsOp,ConstExp (IntConst 100),BinOpAppExp (ConsOp,ConstExp (IntConst 200),ConstExp NilConst)))))), []))


let evalAllC a b = (eval_exp (LetExp (a,(AppExp ( AppExp (VarExp "list_all",FnExp ("x", BinOpAppExp (GreaterOp,(VarExp "x"), ConstExp (IntConst 1)))),ConstExp NilConst))), [])) = (eval_exp (LetExp (b,(AppExp ( AppExp (VarExp "list_all",FnExp ("x", BinOpAppExp (EqOp,(VarExp "x"), ConstExp (IntConst 1)))),ConstExp NilConst))), []))

let ackermann =
LetExp(Rec("a", "b",
 IfExp(BinOpAppExp(GreaterOp,VarExp "b", ConstExp(IntConst 1)),
   LetExp(Seq(Val("c", AppExp(VarExp "a",
                       BinOpAppExp(IntMinusOp, VarExp "b", ConstExp(IntConst 1)))),
              Val("d", AppExp(VarExp "a",
                       BinOpAppExp(IntMinusOp, VarExp "b", ConstExp(IntConst 2))))),
           BinOpAppExp(IntPlusOp, VarExp "c", VarExp "d")),
   ConstExp(IntConst 1))),
   AppExp(VarExp "a", ConstExp(IntConst 12)))

let e =
Seq(Rec("a","b",FnExp("i",IfExp(BinOpAppExp(EqOp,VarExp "i", ConstExp(IntConst 0)),
    ConstExp(BoolConst true),
    AppExp(VarExp "b", BinOpAppExp(IntMinusOp,VarExp "i", ConstExp(IntConst 1)))))),
    Rec("b","i", IfExp(BinOpAppExp(EqOp,VarExp "i", ConstExp(IntConst 0)),
    ConstExp(BoolConst false),
    AppExp(AppExp(VarExp "a", VarExp "b"),
          BinOpAppExp(IntMinusOp,VarExp "i", ConstExp(IntConst 1))))))
    

(* ======= Some relatively generic functions on lists and options ======= *)

let rec first p l =
    match l with [] -> None
    | (x::xs) -> if p x then Some x else first p xs

let rec lookup data_list key =
    match data_list with [] -> None
  | (k,d) :: rem_data -> if k = key then Some d else lookup rem_data key

let app renaming name =
    match lookup renaming name with Some new_name -> new_name
    | None -> name

let rec mergesort list =
let split l =
  let rec split_aux l left right = 
    match l,left,right with
    | ([] | [_]),_,_ -> (List.rev left),right
    | (_::_::t),_,h::right_t -> split_aux t (h::left) right_t
    | _ -> assert false
  in
  split_aux l [] l
  in
let rec merge l1 l2 =
  match l1,l2 with
  | [],l | l,[] -> l
  | h1::t1,h2::t2 ->
    if h1 < h2  then h1::(merge t1 l2)
    else if h2 < h1 then h2::(merge l1 t2)
    else merge t1 l2
  in match list with
  | ([] | [_]) as l -> l
  | l ->  let left,right = split l in 
          merge (mergesort left) (mergesort right)

let max_opt no mo =
    match no with None -> mo
    | Some n -> (match mo with None -> no | Some m -> Some (max n m))

(* ========= normalization stuff =============*)

let rec max_free_cont_var_in_exp bks expCPS =
 match expCPS with VarCPS (k, x) -> max_free_cont_var_in_cont bks k
    | ConstCPS (k, c) -> max_free_cont_var_in_cont bks k
    | MonOpAppCPS (k,m,s) -> (max_free_cont_var_in_cont bks k)
    | BinOpAppCPS (k,b,r,s) -> (max_free_cont_var_in_cont bks k)
    | IfCPS (r,e1,e2) -> 
      max_opt (max_free_cont_var_in_exp bks e1) (max_free_cont_var_in_exp bks e2)
    | AppCPS (k,x1,x2) -> (max_free_cont_var_in_cont bks k)
    | FnCPS (k,x,c,e) ->
      max_opt (max_free_cont_var_in_cont bks k)
              (max_free_cont_var_in_exp (c::bks) e)
    | FixCPS (k,f,x,kx,e) ->
      max_opt (max_free_cont_var_in_cont bks k)
              (max_free_cont_var_in_exp (kx::bks) e)
and max_free_cont_var_in_cont bks kappa =
 match kappa with External -> None
    | ContVarCPS n -> if List.mem n bks then None else Some n
    | ContCPS(x, e) -> max_free_cont_var_in_exp bks e

let rec norm_aux_exp_cps freeVars nextk renumber rename expCPS =
 match expCPS with
   VarCPS (kappa, x)
   -> VarCPS ((norm_aux_cps_cont freeVars nextk renumber rename kappa),
             (app rename x))
 | ConstCPS(kappa, c)
   -> ConstCPS ((norm_aux_cps_cont freeVars nextk renumber rename kappa), c)
 | MonOpAppCPS (kappa, mop, x)
   -> MonOpAppCPS ((norm_aux_cps_cont freeVars nextk renumber rename kappa), mop,
                  (app rename x))
 | BinOpAppCPS (kappa, bop, x, y)
   -> BinOpAppCPS ((norm_aux_cps_cont freeVars nextk renumber rename kappa), bop,
                  (app rename x), (app rename y))
 | IfCPS (b, e1, e2)
   -> IfCPS ((app rename b),
            (norm_aux_exp_cps freeVars nextk renumber rename e1),
            (norm_aux_exp_cps freeVars nextk renumber rename e2))
 | AppCPS (kappa, f, x)
   -> AppCPS ((norm_aux_cps_cont freeVars nextk renumber rename kappa),
             (app rename f), (app rename x))
 | FnCPS (kappa, bv, kx, e) ->
   let new_bv = freshFor freeVars in
   let new_nextk = nextk + 1 in
      FnCPS ((norm_aux_cps_cont freeVars nextk renumber rename kappa),
      new_bv, nextk,
      (norm_aux_exp_cps (new_bv :: freeVars) new_nextk ((kx,nextk)::renumber)
                         ((bv, new_bv)::rename) e))
 | FixCPS (kappa,f,x,kx,e) ->
   let newf = freshFor freeVars in
   let newx = freshFor (newf :: freeVars) in
   let newFreeVars = newx :: newf :: freeVars in
   let new_nextk = nextk + 1 in
   FixCPS ((norm_aux_cps_cont freeVars nextk renumber rename kappa),
           newf, newx, nextk,
           (norm_aux_exp_cps newFreeVars new_nextk 
            ((kx,nextk)::renumber)  ((x,newx)::(f,newf)::rename) e))
and norm_aux_cps_cont freeVars nextk renumber rename kappa =
 match kappa with External -> External
 | ContVarCPS kx -> ContVarCPS (app renumber kx)
 | ContCPS (x, e) ->
   let newx = freshFor freeVars in
   ContCPS (newx, (norm_aux_exp_cps (newx :: freeVars) nextk renumber rename e))

let normalize_exp_cps e =
    norm_aux_exp_cps (freeVarsInExpCPS e)
    (match max_free_cont_var_in_exp [] e with None -> 0 | Some n -> n + 1) [] [] e
let normalize_cps_cont k =
     norm_aux_cps_cont (freeVarsInContCPS k) 
    (match max_free_cont_var_in_cont [] k with None -> 0 | Some n -> n + 1) [] [] k

let k0 = ContVarCPS (~-2)
let k1 = ContCPS ("a", FnCPS(ContVarCPS 0, "b", 1,
                  AppCPS(ContVarCPS 1, "a", "b")))
let k2 = ContCPS ("b", FnCPS(ContVarCPS 0, "b", 1,
                  AppCPS(ContVarCPS 1, "a", "b")))

(* ========= Specific normalize versions of out code and the students code ======*)
let freeVarsInExp_stu exp = mergesort (Student.freeVarsInExp exp)
let freeVarsInExp_sol exp = mergesort (Solution.freeVarsInExp exp)

let freeAndBindingVarsInDec_stu dec = 
    match Student.freeAndBindingVarsInDec dec
     with (fv,bv) -> (mergesort fv, mergesort bv)

let freeAndBindingVarsInDec_sol dec =
    match Solution.freeAndBindingVarsInDec dec
     with (fv,bv) -> (mergesort fv, mergesort bv)

let cps_exp_stu e k n =  normalize_exp_cps (fst(Student.cps_exp e k n))
let cps_exp_sol e k n = normalize_exp_cps (fst(Solution.cps_exp e k n))

let cps_dec_stu d e n = normalize_exp_cps (fst(Student.cps_dec d e n))
let cps_dec_sol d e n = normalize_exp_cps (fst(Solution.cps_dec d e n))

(* This list is for regular problems *)
let rubric =
[
  TEST1ARG(1, import_list, [("a",1);("b",2);("c",3)]);
  TEST1ARG(1, import_list, []);
  TEST1ARG(1, import_list, [("",10000000);("()",-2);("()",-10000000);("()",10000000);("()",2);("()",101000)]);
  TEST1ARG(1, import_list, [("",1);("(((((((())))))))",2);("$$$$$$$$$$$$$$$$$$$$$$$$",3);("avreheh",1);("568h",2);("gg",3);("",1);("(((((((())))))))",2);("$$$$$$$$$$$$$$$$$$$$$$$$",3);("avre`vheh",1);("568h",2);("gg",3)]);
 
  TEST0ARG(1, list_all);
  TEST0ARG(1, (evalAll (Student.list_all) (Solution.list_all)));
  TEST0ARG(1, (evalAllB (Student.list_all) (Solution.list_all)));
  TEST0ARG(1, (evalAllC (Student.list_all) (Solution.list_all)));
  

  TEST1ARG(1, cal_max_exp_height, (BinOpAppExp (ConsOp,
BinOpAppExp (CommaOp, ConstExp (StringConst "a"), ConstExp (IntConst 1)),
BinOpAppExp (ConsOp,
BinOpAppExp (CommaOp, ConstExp (StringConst "b"), ConstExp (IntConst 2)),
BinOpAppExp (ConsOp,
BinOpAppExp (CommaOp, ConstExp (StringConst "c"), ConstExp (IntConst 3)),
ConstExp NilConst)))));
  TEST1ARG(1, cal_max_exp_height, (VarExp "a"));
  TEST1ARG(1, cal_max_exp_height, (ConstExp (IntConst 5)));
  TEST1ARG(1, cal_max_exp_height, (IfExp (BinOpAppExp(EqOp, VarExp "a", ConstExp (StringConst "b")), MonOpAppExp(HdOp, AppExp(VarExp "ww", ConstExp NilConst)), FnExp("x", VarExp "a"))));
  TEST1ARG(1, cal_max_exp_height, (LetExp(Seq(Rec("g","w",VarExp "q"), Val("a", ConstExp UnitConst)),
AppExp(AppExp(VarExp "g", VarExp "a"), VarExp "d"))));

  TEST1ARG(1, freeVarsInExp, (VarExp "x"));
  TEST1ARG(1, freeVarsInExp, (ConstExp NilConst));

  TEST1ARG(1, freeVarsInExp, (MonOpAppExp (TlOp, VarExp "v")));
  TEST1ARG_TWOFUN(1, freeVarsInExp_sol, freeVarsInExp_stu, (BinOpAppExp (CommaOp, VarExp "v", VarExp "dv")));
  TEST1ARG_TWOFUN(1, freeVarsInExp_sol, freeVarsInExp_stu, (MonOpAppExp (IntNegOp, BinOpAppExp (ConcatOp, VarExp "fg", VarExp "yz"))));
  TEST1ARG_TWOFUN(1, freeVarsInExp_sol, freeVarsInExp_stu, (BinOpAppExp (ConsOp, MonOpAppExp (TlOp, VarExp "v"), MonOpAppExp (TlOp, VarExp "w"))));
  TEST1ARG_TWOFUN(1, freeVarsInExp_sol, freeVarsInExp_stu, (IfExp(VarExp "a",VarExp "b",VarExp "c")));
  TEST1ARG_TWOFUN(1, freeVarsInExp_sol, freeVarsInExp_stu, (AppExp(VarExp "d",VarExp "e")));
  TEST1ARG_TWOFUN(1, freeVarsInExp_sol, freeVarsInExp_stu, (IfExp(AppExp(VarExp "a",VarExp "b"),VarExp "c",VarExp "f")));
  TEST1ARG_TWOFUN(1, freeVarsInExp_sol, freeVarsInExp_stu, (AppExp(VarExp "d",IfExp(BinOpAppExp(GreaterOp, VarExp "silly", ConstExp (BoolConst true)), VarExp "x", VarExp "y"))));

  TEST1ARG(1, freeVarsInExp, (FnExp("x", VarExp "x")));
  TEST1ARG(1, freeVarsInExp, (FnExp("x", VarExp "y")));
  TEST1ARG_TWOFUN(1, freeVarsInExp_sol, freeVarsInExp_stu, (AppExp(FnExp("x", FnExp("z",AppExp(VarExp "y", VarExp "z"))), VarExp "x")));

  TEST1ARG(1, freeAndBindingVarsInDec, (Val ("",ConstExp NilConst)));
  TEST1ARG(1, freeAndBindingVarsInDec, (Val ("q",VarExp "r")));
  TEST1ARG(1, freeAndBindingVarsInDec, (Val ("q",VarExp "q")));

  TEST1ARG(1, freeAndBindingVarsInDec,
              (Rec ("f", "x", BinOpAppExp(IntPlusOp, VarExp "x",
                                         AppExp(VarExp  "f", VarExp "y")))));

  TEST1ARG_TWOFUN(1, freeAndBindingVarsInDec_stu, freeAndBindingVarsInDec_sol,
    (Seq(Val("15", VarExp "12"), Val("0", VarExp "4"))));
  TEST1ARG_TWOFUN(1, freeAndBindingVarsInDec_stu, freeAndBindingVarsInDec_sol,
    (Seq(Val("15", VarExp "12"), Val("0", VarExp "15"))));
  TEST1ARG_TWOFUN(1, freeAndBindingVarsInDec_stu, freeAndBindingVarsInDec_sol,
    (Seq(Val("0", VarExp "12"), Val("12", VarExp "15"))));
 
  TEST1ARG(1, freeVarsInExp, (LetExp(Val("x", VarExp "y"), VarExp "x")));
  TEST1ARG_TWOFUN(1, freeVarsInExp_sol, freeVarsInExp_stu,
    (LetExp(Val("x", AppExp(VarExp "x", VarExp "y")), VarExp "z")));
  TEST1ARG_TWOFUN(1, freeVarsInExp_sol, freeVarsInExp_stu,
    (LetExp(Val("x", VarExp "y"), AppExp(VarExp "x", VarExp "y"))));

  TEST1ARG_TWOFUN(1, freeAndBindingVarsInDec_sol, freeAndBindingVarsInDec_stu,
    (Seq(Rec("g", "h", VarExp "q"), Val("x", AppExp(AppExp(VarExp "h", VarExp "g"),AppExp(VarExp "x", VarExp "q"))))));
  TEST1ARG_TWOFUN(1, freeAndBindingVarsInDec_sol, freeAndBindingVarsInDec_stu,
    (Seq (Val ("aa", LetExp (Seq(Val ("a", VarExp "b"), Rec("f", "x", AppExp(AppExp(VarExp "f", VarExp "a"), VarExp "x"))), AppExp(AppExp(VarExp "f", VarExp "a"), VarExp "aa"))), Val("bb", LetExp(Rec ("g", "y", AppExp(AppExp(VarExp "aa", VarExp "a"), AppExp(VarExp "g", VarExp "y"))), AppExp(AppExp(VarExp "g", VarExp "y"), AppExp(VarExp "f",VarExp "x")))))));

  TEST3ARG(1, cps_exp, (VarExp "x"), (ContVarCPS 0), 1);
  TEST3ARG(1, cps_exp, (VarExp "y"), External, 10);
  TEST3ARG(1, cps_exp, (ConstExp (StringConst "x")), (ContVarCPS 0), 1);
  TEST3ARG(1, cps_exp, (ConstExp (BoolConst true)), k1, 2);

  TEST3ARG(1, cps_exp, (IfExp (VarExp "b", ConstExp (IntConst 2),
                                         ConstExp (IntConst 5))),
                                 (ContVarCPS 0), 1);

  TEST3ARG_TWOFUN(1, cps_exp_sol, cps_exp_stu, 
                  (IfExp (ConstExp(BoolConst false), VarExp "a", VarExp "b")),
                                 k1, 2);
  TEST3ARG_TWOFUN(1, cps_exp_sol, cps_exp_stu,
                  (IfExp (IfExp(VarExp "b", ConstExp (IntConst 2),
                          ConstExp (IntConst 5)),
                             IfExp (ConstExp(BoolConst false), VarExp "a",
                                         VarExp "aa"),
                               IfExp(VarExp "e", ConstExp (RealConst 5.5),
                                         VarExp "c"))),
                                 k0, (~-1));

  TEST3ARG(1, cps_exp, (AppExp (VarExp "f", VarExp "x")),
                         (ContVarCPS 0), 1);
  TEST3ARG_TWOFUN(1, cps_exp_sol, cps_exp_stu,
              (AppExp(AppExp(VarExp "a", ConstExp UnitConst), VarExp "c")), k2, 2);
  TEST3ARG_TWOFUN(1, cps_exp_sol, cps_exp_stu,
              (AppExp(AppExp(VarExp "b", ConstExp UnitConst), VarExp "c")), k1, 2);

  TEST3ARG(1, cps_exp, (BinOpAppExp (IntPlusOp, ConstExp(IntConst 5),
                                              ConstExp(IntConst 1))),
                          (ContVarCPS 3), 4);
  TEST3ARG_TWOFUN(1, cps_exp_sol, cps_exp_stu,
  (BinOpAppExp(ConcatOp, BinOpAppExp(IntMinusOp, VarExp "a", VarExp "b"),
              VarExp "c")), k2, 2);
  TEST3ARG_TWOFUN(1, cps_exp_sol, cps_exp_stu,
  (BinOpAppExp(ConsOp,BinOpAppExp(CommaOp,VarExp "b", ConstExp UnitConst),
               VarExp "a")), k1, 2);

  TEST3ARG(1, cps_exp, (MonOpAppExp (HdOp, ConstExp NilConst)), (ContVarCPS 0), 1);
  TEST3ARG_TWOFUN(1, cps_exp_sol, cps_exp_stu,
                  (MonOpAppExp (HdOp, VarExp "b")), k2, 2);
  TEST3ARG_TWOFUN(1, cps_exp_sol, cps_exp_stu,
                  (MonOpAppExp (HdOp, VarExp "a")), k1, 2);

  TEST3ARG(1, cps_exp, (FnExp ("x", VarExp "x")), (ContVarCPS 0), 1);
  TEST3ARG_TWOFUN(1, cps_exp_sol, cps_exp_stu,
   (FnExp ("a", ConstExp (IntConst 2))), k0, 1);
  TEST3ARG_TWOFUN(1, cps_exp_sol, cps_exp_stu, (FnExp ("a", VarExp "a")), k2, 1);

  TEST3ARG(1, cps_dec, (Val ("x", ConstExp(IntConst 2))),
                                 (VarCPS (ContVarCPS 0, "x")), 1);
  TEST3ARG_TWOFUN(1, cps_dec_sol, cps_dec_stu,
                  (Val ("a", AppExp(VarExp "a", VarExp "b"))),
                   (AppCPS(k2,"a","b")), 2);
  TEST3ARG_TWOFUN(1, cps_dec_sol, cps_dec_stu,
                  (Val ("b", AppExp(VarExp "a", VarExp "b"))),
                   (AppCPS(k1,"a","b")), 2);

  TEST3ARG(1, cps_dec, (Seq (Val ("x", ConstExp(IntConst 2)),
                                       Val ("y", VarExp "x"))),
                                 (VarCPS (ContVarCPS 0, "x")), 1);
  TEST3ARG_TWOFUN(1, cps_dec_sol, cps_dec_stu,
                     (Seq (Val ("a", ConstExp(IntConst 5)),
                            Val ("b", VarExp "c"))),
                      (AppCPS(k1,"a","b")), 2);
  TEST3ARG_TWOFUN(1, cps_dec_sol, cps_dec_stu,
                     (Seq (Val ("a", ConstExp(IntConst 5)),
                            Val ("b", VarExp "a"))),
                      (AppCPS(k2,"a","b")), 2);

  TEST3ARG(1, cps_exp, (LetExp (Val ("x", ConstExp(IntConst 2)), VarExp "x")),
                      (ContVarCPS 0), 1);
  TEST3ARG_TWOFUN(1, cps_exp_sol, cps_exp_stu,
   (LetExp (Val ("a", VarExp "a"), LetExp(Val ("b", AppExp(VarExp "a", VarExp "b")),
            BinOpAppExp(ConsOp, VarExp "b", VarExp "a")))), k0, (~-1));
  TEST3ARG_TWOFUN(1, cps_exp_sol, cps_exp_stu,
   (LetExp(Seq(Val ("a", VarExp "b"), Val("c", VarExp "d")),
           AppExp (VarExp "a", VarExp "b"))), k1, 2)
]


let extra_rubric = [ 
  TEST3ARG(1, cps_exp, (LetExp (Rec("f", "x",VarExp "x"),
                                  ConstExp (IntConst 4))),
                                 (ContVarCPS 1), 2);
  TEST3ARG_TWOFUN(1, cps_exp_sol, cps_exp_stu, ackermann, External, 0);
  TEST3ARG_TWOFUN(1, cps_dec_sol, cps_dec_stu, e, (AppCPS(k1,"a","b")), 2)
]