type instruction =
  | NOOP (* no operation *)
  | PUSH of int (* push integer constant onto stack *)
  | ADD (* addition *)
  | SUB (* subtraction *)
  | EQ (* equal *)
  | LT (* less than *)
  | JMP of int (* relative jump *)
  | JMPZ of int (* relative jump if zero on the stack *)
  | EXIT (* stop execution, returning the top element of the stack *)
  | READ (* push read integer *)
  | DUP (* duplicate top element on the stack *)

type machine_state = {
  input : int list;
  rom : int -> instruction;
  stack : int list;
  instruction_pointer : int;
}

let pop state =
  match state.stack with
  | x :: stack -> (x, { state with stack })
  | [] -> failwith "empty stack"

let read state =
  match state.input with
  | x :: input -> (x, { state with input })
  | [] -> failwith "empty input"

let double_pop state =
  let x1, state = pop state in
  let x2, state = pop state in
  (x1, x2, state)

let push x state = { state with stack = x :: state.stack }

let increment_instruction_pointer i state =
  { state with instruction_pointer = state.instruction_pointer + i }

let string_of_instruction = function
  | NOOP -> "NOOP"
  | PUSH i -> "PUSH " ^ string_of_int i
  | ADD -> "ADD"
  | SUB -> "SUB"
  | EQ -> "EQ"
  | LT -> "LT"
  | JMP i -> "JMP " ^ string_of_int i
  | JMPZ i -> "JMPZ " ^ string_of_int i
  | EXIT -> "EXIT"
  | READ -> "READ"
  | DUP -> "DUP"

let print_debug_info { input; rom; stack; instruction_pointer } =
  Printf.printf "INFO: IP = %i (%s), stack = [%s], input = [%s]\n"
  instruction_pointer
  (string_of_instruction @@ rom instruction_pointer)
  (String.concat ", " @@ List.map string_of_int stack)
  (String.concat ", " @@ List.map string_of_int input)

let rec run (state : machine_state) =
  print_debug_info state;
  let incr = increment_instruction_pointer 1 in
  match state.rom state.instruction_pointer with
  | NOOP -> run state
  | _ -> failwith "not implemented"

(* sum *)
let romA i = [| PUSH 0; READ; DUP; JMPZ 3; ADD; JMP (-4); JMPZ 1; EXIT |].(i)
(* abs sum *)
let romB i = [| PUSH 0; READ; DUP; JMPZ 7; DUP; PUSH 0; LT; JMPZ 5; ADD; JMP (-8); JMPZ 1; EXIT; PUSH 0; SUB; ADD; JMP (-14) |].(i)

let computerA input =
  run { rom = romA; stack = []; input; instruction_pointer = 0 }

let computerB input =
  run { rom = romB; stack = []; input; instruction_pointer = 0 }

(*
# run { rom = (fun i -> [| PUSH 42; EXIT |].(i)); stack = []; input = [1; 2]; instruction_pointer = 0 } ;;
INFO: IP = 0 (PUSH 42), stack = [], input = [1; 2]
INFO: IP = 1 (EXIT), stack = [42], input = [1; 2]
- : int = 42

# run { rom = (fun i -> [| READ; EXIT |].(i)); stack = []; input = [1; 2]; instruction_pointer = 0 } ;;
INFO: IP = 0 (READ), stack = [], input = [1; 2]
INFO: IP = 1 (EXIT), stack = [1], input = [2]
- : int = 1

# run { rom = (fun i -> [| ADD; EXIT |].(i)); stack = [-1; 1; 2]; input = []; instruction_pointer = 0 } ;;
INFO: IP = 0 (ADD), stack = [-1; 1; 2], input = []
INFO: IP = 1 (EXIT), stack = [0; 2], input = []
- : int = 0

# run { rom = (fun i -> [| SUB; EXIT |].(i)); stack = [-1; 1; 2]; input = []; instruction_pointer = 0 } ;;
INFO: IP = 0 (SUB), stack = [-1; 1; 2], input = []
INFO: IP = 1 (EXIT), stack = [-2; 2], input = []
- : int = -2

# run { rom = (fun i -> [| EQ; EXIT |].(i)); stack = [-1; 1; 2]; input = []; instruction_pointer = 0 } ;;
INFO: IP = 0 (EQ), stack = [-1; 1; 2], input = []
INFO: IP = 1 (EXIT), stack = [0; 2], input = []
- : int = 0

# run { rom = (fun i -> [| EQ; EXIT |].(i)); stack = [1; 1; 2]; input = []; instruction_pointer = 0 } ;;
INFO: IP = 0 (EQ), stack = [1; 1; 2], input = []
INFO: IP = 1 (EXIT), stack = [1; 2], input = []
- : int = 1

# run { rom = (fun i -> [| EQ; EXIT |].(i)); stack = [1; 1; 2]; input = []; instruction_pointer = 0 } ;;
INFO: IP = 0 (EQ), stack = [1; 1; 2], input = []
INFO: IP = 1 (EXIT), stack = [1; 2], input = []
- : int = 1

# run { rom = (fun i -> [| LT; EXIT |].(i)); stack = [-1; 1; 2]; input = []; instruction_pointer = 0 } ;;
INFO: IP = 0 (LT), stack = [-1; 1; 2], input = []
INFO: IP = 1 (EXIT), stack = [1; 2], input = []
- : int = 1

# run { rom = (fun i -> [| JMP 2; READ; EXIT |].(i)); stack = []; input = []; instruction_pointer = 0 } ;;
INFO: IP = 0 (JMP 2), stack = [], input = []
INFO: IP = 2 (EXIT), stack = [], input = []
Exception: Failure "empty stack".

# run { rom = (fun i -> [| JMP 2; READ; EXIT |].(i)); stack = [100]; input = []; instruction_pointer = 0 } ;;
INFO: IP = 0 (JMP 2), stack = [100], input = []
INFO: IP = 2 (EXIT), stack = [100], input = []
- : int = 100

# run { rom = (fun i -> [| JMPZ 2; READ; EXIT |].(i)); stack = [0; 100]; input = [-1]; instruction_pointer = 0 } ;;
INFO: IP = 0 (JMPZ 2), stack = [0; 100], input = [-1]
INFO: IP = 2 (EXIT), stack = [100], input = [-1]
- : int = 100

# run { rom = (fun i -> [| JMPZ 2; READ; EXIT |].(i)); stack = [3; 100]; input = [-1]; instruction_pointer = 0 } ;;
INFO: IP = 0 (JMPZ 2), stack = [3; 100], input = [-1]
INFO: IP = 1 (READ), stack = [100], input = [-1]
INFO: IP = 2 (EXIT), stack = [-1; 100], input = []
- : int = -1

# run { rom = (fun i -> [| DUP; EXIT |].(i)); stack = [3; 100]; input = []; instruction_pointer = 0 } ;;
INFO: IP = 0 (DUP), stack = [3; 100], input = []
INFO: IP = 1 (EXIT), stack = [3; 3; 100], input = []
- : int = 3

# computerA [42; -2; 60; 0; 1337];;
INFO: IP = 0 (PUSH 0), stack = [], input = [42; -2; 60; 0; 1337]
INFO: IP = 1 (READ), stack = [0], input = [42; -2; 60; 0; 1337]
INFO: IP = 2 (DUP), stack = [42; 0], input = [-2; 60; 0; 1337]
INFO: IP = 3 (JMPZ 3), stack = [42; 42; 0], input = [-2; 60; 0; 1337]
INFO: IP = 4 (ADD), stack = [42; 0], input = [-2; 60; 0; 1337]
INFO: IP = 5 (JMP -4), stack = [42], input = [-2; 60; 0; 1337]
INFO: IP = 1 (READ), stack = [42], input = [-2; 60; 0; 1337]
INFO: IP = 2 (DUP), stack = [-2; 42], input = [60; 0; 1337]
INFO: IP = 3 (JMPZ 3), stack = [-2; -2; 42], input = [60; 0; 1337]
INFO: IP = 4 (ADD), stack = [-2; 42], input = [60; 0; 1337]
INFO: IP = 5 (JMP -4), stack = [40], input = [60; 0; 1337]
INFO: IP = 1 (READ), stack = [40], input = [60; 0; 1337]
INFO: IP = 2 (DUP), stack = [60; 40], input = [0; 1337]
INFO: IP = 3 (JMPZ 3), stack = [60; 60; 40], input = [0; 1337]
INFO: IP = 4 (ADD), stack = [60; 40], input = [0; 1337]
INFO: IP = 5 (JMP -4), stack = [100], input = [0; 1337]
INFO: IP = 1 (READ), stack = [100], input = [0; 1337]
INFO: IP = 2 (DUP), stack = [0; 100], input = [1337]
INFO: IP = 3 (JMPZ 3), stack = [0; 0; 100], input = [1337]
INFO: IP = 6 (JMPZ 1), stack = [0; 100], input = [1337]
INFO: IP = 7 (EXIT), stack = [100], input = [1337]
- : int = 100

# computerB [42; -2; 60; 0; 1337];;
INFO: IP = 0 (PUSH 0), stack = [], input = [42; -2; 60; 0; 1337]
INFO: IP = 1 (READ), stack = [0], input = [42; -2; 60; 0; 1337]
INFO: IP = 2 (DUP), stack = [42; 0], input = [-2; 60; 0; 1337]
INFO: IP = 3 (JMPZ 7), stack = [42; 42; 0], input = [-2; 60; 0; 1337]
INFO: IP = 4 (DUP), stack = [42; 0], input = [-2; 60; 0; 1337]
INFO: IP = 5 (PUSH 0), stack = [42; 42; 0], input = [-2; 60; 0; 1337]
INFO: IP = 6 (LT), stack = [0; 42; 42; 0], input = [-2; 60; 0; 1337]
INFO: IP = 7 (JMPZ 5), stack = [1; 42; 0], input = [-2; 60; 0; 1337]
INFO: IP = 8 (ADD), stack = [42; 0], input = [-2; 60; 0; 1337]
INFO: IP = 9 (JMP -8), stack = [42], input = [-2; 60; 0; 1337]
INFO: IP = 1 (READ), stack = [42], input = [-2; 60; 0; 1337]
INFO: IP = 2 (DUP), stack = [-2; 42], input = [60; 0; 1337]
INFO: IP = 3 (JMPZ 7), stack = [-2; -2; 42], input = [60; 0; 1337]
INFO: IP = 4 (DUP), stack = [-2; 42], input = [60; 0; 1337]
INFO: IP = 5 (PUSH 0), stack = [-2; -2; 42], input = [60; 0; 1337]
INFO: IP = 6 (LT), stack = [0; -2; -2; 42], input = [60; 0; 1337]
INFO: IP = 7 (JMPZ 5), stack = [0; -2; 42], input = [60; 0; 1337]
INFO: IP = 12 (PUSH 0), stack = [-2; 42], input = [60; 0; 1337]
INFO: IP = 13 (SUB), stack = [0; -2; 42], input = [60; 0; 1337]
INFO: IP = 14 (ADD), stack = [2; 42], input = [60; 0; 1337]
INFO: IP = 15 (JMP -14), stack = [44], input = [60; 0; 1337]
INFO: IP = 1 (READ), stack = [44], input = [60; 0; 1337]
INFO: IP = 2 (DUP), stack = [60; 44], input = [0; 1337]
INFO: IP = 3 (JMPZ 7), stack = [60; 60; 44], input = [0; 1337]
INFO: IP = 4 (DUP), stack = [60; 44], input = [0; 1337]
INFO: IP = 5 (PUSH 0), stack = [60; 60; 44], input = [0; 1337]
INFO: IP = 6 (LT), stack = [0; 60; 60; 44], input = [0; 1337]
INFO: IP = 7 (JMPZ 5), stack = [1; 60; 44], input = [0; 1337]
INFO: IP = 8 (ADD), stack = [60; 44], input = [0; 1337]
INFO: IP = 9 (JMP -8), stack = [104], input = [0; 1337]
INFO: IP = 1 (READ), stack = [104], input = [0; 1337]
INFO: IP = 2 (DUP), stack = [0; 104], input = [1337]
INFO: IP = 3 (JMPZ 7), stack = [0; 0; 104], input = [1337]
INFO: IP = 10 (JMPZ 1), stack = [0; 104], input = [1337]
INFO: IP = 11 (EXIT), stack = [104], input = [1337]
- : int = 104
*)