open Ast
open Bytecode

(* Stack layout just after "Ent":

              <-- SP
   Local n
   ...
   Local 0
   Saved FP   <-- FP
   Saved PC
   Arg 0
   ...
   Arg n *)

let execute_prog prog =
  let stack = Array.make 1024 "0"
  and globals = Array.make prog.size_globals "0" in

  let rec exec fp sp pc = match prog.text.(pc) with
    Lit i  -> stack.(sp) <- string_of_int i ; exec fp (sp+1) (pc+1)
  | Stg s  -> stack.(sp) <- s ; exec fp (sp+1) (pc+1)
  | Max    -> stack.(sp) <- "Max"; exec fp (sp+1) (pc+1)
  | Acc    -> if stack.(sp-3) = "Max" then
               let rows = (int_of_string stack.(sp-4)) and 
                   cols = (int_of_string stack.(sp-5)) and
                   r = (int_of_string stack.(sp-1)) and
                   c = (int_of_string stack.(sp-2)) in (
               stack.(sp-5-(rows*cols)) <- stack.(sp-6-(rows*cols)+((r-1)*cols)+c);
               exec fp (sp-4-(rows*cols)) (pc+1))
  | Drp    -> exec fp (sp-1) (pc+1)
  | Bin op -> if stack.(sp-1) = "Max" || stack.(sp-2) = "Max" then ((
              match op with
              Add | Sub | Equal | Neq -> (
                let rows = (int_of_string stack.(sp-2)) and 
                    cols = (int_of_string stack.(sp-3)) in (
                for i = 4 to (3+rows*cols) do
                  let op1 = (int_of_string stack.(sp-3-i-rows*cols)) and
                      op2 = (int_of_string stack.(sp-i)) in 
                  stack.(sp-i) <- (let boolean i = if i then "1" else "0" in
                  match op with
                    Add   -> string_of_int (op1 + op2)
                  | Sub   -> string_of_int (op1 - op2)
                  | Equal -> boolean (op1 = op2)
                  | Neq   -> boolean (op1 != op2))
                done))
            | Mult  -> (
                if stack.(sp-1) <> "Max" then (
                  let rows = (int_of_string stack.(sp-3)) and
                      cols = (int_of_string stack.(sp-4)) and
                      const = (int_of_string stack.(sp-1)) in (
                  stack.(sp-1) <- "Max";
                  stack.(sp-2) <- string_of_int rows;
                  stack.(sp-3) <- string_of_int cols;
                  for i = 5 to (4+rows*cols) do
                    stack.(sp-i+1) <- string_of_int (const*(int_of_string stack.(sp-i)))
                  done))
                else (
                  let rows = (int_of_string stack.(sp-2)) and
                      cols = (int_of_string stack.(sp-3)) in (
                  if stack.(sp-rows*cols-4) <> "Max" then (
                    for i = 4 to (3+rows*cols) do
                      stack.(sp-i) <- string_of_int (int_of_string
                      stack.(sp-rows*cols-4)*(int_of_string stack.(sp-i)))
                    done)
                  else (
                    let rows2 = (int_of_string stack.(sp-5-rows*cols)) and 
                        cols2 = (int_of_string stack.(sp-6-rows*cols)) in (
                    if cols2 != rows then (raise (Failure(
                      "Operators for * do not satisfy matrix multiplication criteria")))
                    else (
                      let ops1 = ref [] and ops2 = ref [] and
                          res = ref [] and sum = ref 0 in (
                      let count = ref 0 and
                          m = (Array.sub stack (sp-rows*cols-6-rows2*cols2) (rows2*cols2)) in
                      for i = 0 to (Array.length m - 1) do
                        count := (!count + 1);
                        ops1 := (!ops1 @ [Array.get m i]);
                        if !count = cols2 then (
                          count := 0;
                          let len = List.length !ops1 in
                          for j = cols2 downto 1 do
                            ops1 := (!ops1 @ [List.nth !ops1 (len - j)])
                        done)
                      done;
                      let count = ref 0 and 
                          m = (Array.sub stack (sp-rows*cols-3) (rows*cols)) in (
                      for r = 1 to rows2 do
                        for c = 1 to cols do
                          for i = 0 to (Array.length m - 1) do
                            count := (!count + 1);
                            if !count = c then (ops2 := (!ops2 @ [Array.get m i]));
                            if !count = cols then (count := 0)
                          done
                        done
                      done;
                      count := 0;
                      stack.(sp-2) <- string_of_int rows2;
                      for i = 0 to (List.length !ops1 - 1) do
                        count := (!count + 1);
                        sum := (!sum + ((int_of_string (List.nth !ops1 i))*(
                                int_of_string (List.nth !ops2 i))));
                        if !count = cols2 then (
                          res := (!res @ [!sum]); count := 0; sum := 0;)
                        done;
                        for i = 0 to (List.length !res - 1) do
                          stack.(sp-3-(rows2*cols)+i) <- (string_of_int (List.nth !res i))
                        done)
                    )))))
                )));
                (match op with
                | Add | Sub | Mult -> ()
                | Equal -> let rows = (int_of_string stack.(sp-2)) and 
                               cols = (int_of_string stack.(sp-3)) in
                           stack.(sp-1) <- (let cmp = (
                             Array.fold_left (fun s e -> s + (int_of_string e)) 0 (
                               Array.sub stack (sp-3-rows*cols) (rows*cols))) in
                             if cmp = (rows*cols) then "1" else "0")
                | Neq -> let rows = (int_of_string stack.(sp-2)) and 
                             cols = (int_of_string stack.(sp-3)) in
                         stack.(sp-1) <- (let cmp = (
                           Array.fold_left (fun s e -> s + (int_of_string e)) 0 (
                             Array.sub stack (sp-3-rows*cols) (rows*cols))) in
                           if cmp > 0 then "1" else "0"));
                exec fp sp (pc+1))
              else (
                let op1 = (int_of_string stack.(sp-2)) and 
                    op2 = (int_of_string stack.(sp-1)) in (
                stack.(sp-2) <- (let boolean i = if i then "1" else "0" in
                match op with
                  Add     -> string_of_int (op1 + op2)
                | Sub     -> string_of_int (op1 - op2)
                | Mult    -> string_of_int (op1 * op2)
                | Div     -> string_of_int (op1 / op2)
                | Equal   -> boolean (op1 =  op2)
                | Neq     -> boolean (op1 != op2)
                | Less    -> boolean (op1 <  op2)
                | Leq     -> boolean (op1 <= op2)
                | Greater -> boolean (op1 >  op2)
                | Geq     -> boolean (op1 >= op2));
                exec fp (sp-1) (pc+1)))
  | Lod i   -> if globals.(i-1) = "Max" then
                 let rows = (int_of_string globals.(i-2)) and 
                     cols = (int_of_string globals.(i-3)) in (
                 for j = 1 to (3+rows*cols) do
                    stack.(sp+(3+rows*cols)-j) <- globals.(i-j)
                 done;
                 exec fp (sp+(rows*cols+3)) (pc+1))
               else (
                 stack.(sp) <- globals.(i-1);
                 exec fp (sp+1) (pc+1))
  | Str i   -> if stack.(sp-1) = "Max" then
                 let rows = (int_of_string stack.(sp-2)) and 
                     cols = (int_of_string stack.(sp-3)) in (
                 for j = 1 to (3+rows*cols) do
                    globals.(i-j) <- stack.(sp-j)
                 done)
               else (globals.(i-1) <- stack.(sp-1));
               exec fp sp (pc+1)
  | Lfp i   -> if i < 0 then let rec f1 = (fun x offset -> 
                 if offset = (-2) then x
                 else if stack.(fp+x) = "Max" then 
                   f1 (x-3-(int_of_string stack.(fp+x-1))*(
                     int_of_string stack.(fp+x-2))) (offset+1)
                   else f1 (x-1) (offset+1)) in (
                 if stack.(fp+(f1 (-2) i)) = "Max" then
                   let rows = (int_of_string stack.(fp+(f1 (-2) i)-1)) and
                       cols = (int_of_string stack.(fp+(f1 (-2) i)-2)) in (
                   for j = 1 to (rows*cols+3) do
                     stack.(sp+(rows*cols+3)-j) <- stack.(fp+(f1 (-2) i)-j+1)
                   done;
                   exec fp (sp+rows*cols+3) (pc+1))
                 else (
                   stack.(sp) <- stack.(fp+i);
                   exec fp (sp+1) (pc+1)))
               else if stack.(fp+i-1) = "Max" then (
                 let rows = (int_of_string stack.(fp+i-2)) and 
                     cols = (int_of_string stack.(fp+i-3)) in (
                 for j = 1 to (rows*cols+3) do
                   stack.(sp+(rows*cols+3)-j) <- stack.(fp+i-j)
                 done;
                 exec fp (sp+(if stack.(fp+i-1) = "Max" then (rows*cols+3) else 1)) (pc+1)))
               else (
                 stack.(sp) <- stack.(fp+i);
                 exec fp (sp+1) (pc+1))
  | Sfp i   -> if stack.(sp-1) = "Max" then (
                 let rows = (int_of_string stack.(sp-2)) and 
                     cols = (int_of_string stack.(sp-3)) in (
                 for j = 1 to (rows*cols+3) do
                   stack.(fp+i-j) <- stack.(sp-j)
                 done))
               else (stack.(fp+i) <- stack.(sp-1));
               exec fp sp (pc+1)
  | Jsr(-1) -> if stack.(sp-1) = "Max" then (
                 let rows = (int_of_string stack.(sp-2)) and 
                     cols = (int_of_string stack.(sp-3)) in (
                 for i = rows downto 1 do
                   Array.iter (fun e -> Printf.printf "%s " e) 
                     (Array.sub stack (sp-3-i*cols) cols);
                   Printf.printf "\n"
                 done))
               else (print_endline stack.(sp-1));
               exec fp sp (pc+1)
  | Jsr(-2) -> if stack.(sp-1) = "Max" then
                 let rows = (int_of_string stack.(sp-2)) and 
                     cols = (int_of_string stack.(sp-3)) in (
                 stack.(sp-3-(cols*rows)) <- string_of_int rows;
               exec fp (sp-2-(cols*rows)) (pc+1))
  | Jsr(-3) -> if stack.(sp-1) = "Max" then
                 let rows = (int_of_string stack.(sp-2)) and 
                     cols = (int_of_string stack.(sp-3)) in (
                 stack.(sp-3-(cols*rows)) <- string_of_int cols;
               exec fp (sp-2-(cols*rows)) (pc+1))
  | Jsr(-4) -> let oc = open_out stack.(sp-1) in (
                 if stack.(sp-2) = "Max" then
                   let rows = (int_of_string stack.(sp-3)) and 
                       cols = (int_of_string stack.(sp-4)) in (
                   for i = rows downto 1 do
                     Array.iter (fun e -> Printf.fprintf oc "%s " e) (
                     Array.sub stack (sp-4-i*cols) cols);
                     Printf.fprintf oc "\n"
               done)
               else Printf.fprintf oc "%s" stack.(sp-2));
               exec fp sp (pc+1)
  | Jsr i   -> stack.(sp) <- string_of_int (pc + 1);
               exec fp (sp+1) i
  | Ent i   -> stack.(sp) <- string_of_int fp;
               exec sp (sp+i+1) (pc+1)
  | Rts i   -> let j = (if i > 0 then (let rec f1 = (
                 fun x offset -> 
                   if offset = 0 then x
                   else if stack.(fp+x) = "Max" then 
                     f1 (x-3-(int_of_string stack.(fp+x-1))*(
                       int_of_string stack.(fp+x-2))) (offset+1)
                    else f1 (x-1) (offset+1)) in (f1 (-2) (-i)))
                 else i) in (
               let new_fp = int_of_string stack.(fp) and 
                   new_pc = int_of_string stack.(fp-1) in
               if stack.(sp-1) = "Max" then (
                 let rows = (int_of_string stack.(sp-2)) and 
                     cols = (int_of_string stack.(sp-3)) in 
                 for x = 1 to (rows*cols+3) do
                   stack.(fp-j-x) <- stack.(sp-x)
                 done)
               else stack.(fp-j-1) <- stack.(sp-1);
               exec new_fp (fp-j) new_pc)
  | Beq i   -> exec fp (sp-1) (pc + if (int_of_string stack.(sp-1)) =  0 then i else 1)
  | Bne i   -> exec fp (sp-1) (pc + if (int_of_string stack.(sp-1)) != 0 then i else 1)
  | Bra i   -> exec fp sp (pc+i)
  | Hlt     -> ()

  in exec 0 0 0
