module Register = struct
  type t = {index: string}

  let compare a b = compare a.index b.index
end


module CodeMap = Map.Make(String)
module RegisterMap = Map.Make(Register)
module MemoryMap = Map.Make(Int)

module RISCArchitecture = struct
  type t = {
    code: RISC.RISCAssembly.risci list CodeMap.t;
    registers: int RegisterMap.t;
    memory: int MemoryMap.t;
    outputreg: Register.t;
  }
end

let convert (prg: RISC.RISCAssembly.t) : RISC.RISCAssembly.risci list CodeMap.t =
  let rec helper
      (prg: RISC.RISCAssembly.risci list)
      (current: RISC.RISCAssembly.risci list)
      (current_label: string)
      (map: RISC.RISCAssembly.risci list CodeMap.t)
    : (RISC.RISCAssembly.risci list CodeMap.t) =
    match prg with
    | [] -> (CodeMap.union
               (fun _ _ _ -> failwith "Two labels are the same")
               (CodeMap.singleton current_label current)
               map)
    | Label l :: tl -> helper tl ([]) l
                         (CodeMap.union
                            (fun _ _ _ -> failwith "Two labels are the same")
                            (CodeMap.singleton current_label current)
                            map)
    | instr :: tl -> helper tl (current @ [instr]) current_label map
  in
  match prg.code with
  | Label "main" :: tl -> helper tl [] "main" CodeMap.empty
  | _ -> failwith "Program should begind with label main"

let label_order (prg: RISC.RISCAssembly.t) : string list =
  let rec helper
      (prg: RISC.RISCAssembly.risci list)
    : string list =
    match prg with
      [] -> []
    | Label l :: tl -> l :: (helper tl)
    | _ :: tl -> (helper tl)
  in
  helper (prg.code)

let reduce_instructions (prg: RISCArchitecture.t) (lo: string list) : int =
  let match_operator_r (brop: RISC.RISCAssembly.brop) =
    match brop with
    | Add -> (+)
    | Sub -> (-)
    | Mult -> ( * )
    | Div -> (/)
    | Mod -> (mod)
    | Pow -> (Utility.pow)
    | And -> (Utility.int_and)
    | Or -> (Utility.int_or)
    | Eq -> (Utility.int_eq)
    | Less -> (Utility.int_less)
    | LessEq -> (Utility.int_less_eq)
    | More -> (Utility.int_more)
    | MoreEq -> (Utility.int_more_eq)
  in
  let match_operator_i (biop: RISC.RISCAssembly.biop) =
    match biop with
    | AddI -> (+)
    | SubI -> (-)
    | MultI -> ( * )
    | DivI -> (/)
    | ModI -> (mod)
    | PowI -> (Utility.pow)
    | AndI -> (Utility.int_and)
    | OrI -> (Utility.int_or)
    | EqI -> (Utility.int_eq)
    | LessI -> (Utility.int_less)
    | LessEqI -> (Utility.int_less_eq)
    | MoreI -> (Utility.int_more)
    | MoreEqI -> (Utility.int_more_eq)
  in

  let rec helper
      (prg: RISCArchitecture.t)
      (current: RISC.RISCAssembly.risci list)
      (current_label: string)
    : RISCArchitecture.t =
    match current with
    | [] -> (
        (* falls to the next label *)
        match List.find_index ((=) current_label) lo with
          None -> prg (* should never happen *)
        | Some i ->
          if i + 1 < (List.length lo) then
            helper prg (CodeMap.find (List.nth lo (i+1)) prg.code) (List.nth lo (i+1))
          else
            prg
      )
    | Nop :: tl ->
      helper prg tl current_label
    | BRegOp (brop, r1, r2, r3) :: tl -> (
        let n = (match_operator_r brop)
            (RegisterMap.find {index = r1.index} prg.registers)
            (RegisterMap.find {index = r2.index} prg.registers)
        in
        helper {prg with registers = RegisterMap.add {index = r3.index} n prg.registers} tl current_label
      )
    | BImmOp (biop, r1, i, r3) :: tl -> (
        let n = (match_operator_i biop)
            (RegisterMap.find {index = r1.index} prg.registers)
            i
        in
        helper {prg with registers = RegisterMap.add {index = r3.index} n prg.registers} tl current_label
      )
    | URegOp (urop, r1, r3) :: tl -> (
        match urop with
        | Copy -> (
            let n = RegisterMap.find {index = r1.index} prg.registers in
            helper {prg with registers = RegisterMap.add {index = r3.index} n prg.registers} tl current_label
          )
        | Not -> (
            let n = Utility.int_not
                (RegisterMap.find {index = r1.index} prg.registers)
            in
            helper {prg with registers = RegisterMap.add {index = r3.index} n prg.registers} tl current_label
          )
        | Rand -> (
            let n = Random.int
                (RegisterMap.find {index = r1.index} prg.registers)
            in
            helper {prg with registers = RegisterMap.add {index = r3.index} n prg.registers} tl current_label
          )
      )
    | Load (r1, r3) :: tl -> (
        let n =
          MemoryMap.find
            (RegisterMap.find {index = r1.index} prg.registers)
            prg.memory
        in
        helper {prg with registers = RegisterMap.add {index = r3.index} n prg.registers} tl current_label
      )
    | LoadI (i, r3) :: tl -> (
        let n = i
        in
        helper {prg with registers = RegisterMap.add {index = r3.index} n prg.registers} tl current_label
      )
    | Store (r1, r3) :: tl -> (
        let n = RegisterMap.find {index = r1.index} prg.registers in
        let n1 = RegisterMap.find {index = r3.index} prg.registers in
        helper {prg with memory = MemoryMap.add n1 n prg.memory} tl current_label
      )
    | Jump l :: _ -> helper prg (CodeMap.find l prg.code) l
    | CJump (r, l1, l2) :: _ -> (
        let br = (RegisterMap.find {index = r.index} prg.registers) > 0 in
        if br
        then
          helper prg (CodeMap.find l1 prg.code) l1
        else
          helper prg (CodeMap.find l2 prg.code) l2
      )
    | Label _ :: tl -> helper prg tl current_label
  in
  RegisterMap.find
    prg.outputreg
    (helper prg (CodeMap.find "main" prg.code) "main").registers


let reduce (prg: RISC.RISCAssembly.t) : int =
  reduce_instructions
    {code = convert prg;
     registers = (
       match prg.inputoutputreg with
       | None ->
         RegisterMap.singleton
           {index = "in"}
           (Option.value prg.inputval ~default:0)
       | Some (i, _) ->
         RegisterMap.singleton
           {index = i.index}
           (Option.value prg.inputval ~default:0)
     );
     memory = MemoryMap.empty;
     outputreg = (
       match prg.inputoutputreg with
       | None -> {index = "out"}
       | Some (_, o) -> {index = o.index}
     )
    }
    (label_order prg)