open Analysis

module Variable = struct
  type t = string

  let _pp (ppf: out_channel) (v: t) : unit =
    Printf.fprintf ppf "%s" v

  let _pplist (ppf: out_channel) (vv: t list) : unit =
    List.iter (Printf.fprintf ppf "%s, ") vv

  let compare a b =
    String.compare a b
end


module RISCCfg = CfgRISC.RISCCfg
module VariableMap = Map.Make(Variable)

let variables_frequency (instr : RISCCfg.elt) : (Variable.t * int) list =
  let add_one = (fun x -> match x with None -> Some 1 | Some x -> Some (x + 1)) in

  let helper (acc: int VariableMap.t) (instr: RISCCfg.elt) : int VariableMap.t =
    match instr with
    | Nop ->
      acc
    | BRegOp (_, r1, r2, r3) ->
      VariableMap.update r1.index add_one acc |>
      VariableMap.update r2.index add_one |>
      VariableMap.update r3.index add_one
    | BImmOp (_, r1, _, r3)
    | URegOp (_, r1, r3)
    | Load (r1, r3)
    | Store (r1, r3) ->
      VariableMap.update r1.index add_one acc |>
      VariableMap.update r3.index add_one
    | LoadI (_, r3) ->
      VariableMap.update r3.index add_one acc
  in

  helper VariableMap.empty instr |> VariableMap.to_list

let variables_all_frequency (instructions : RISCCfg.elt list) : (Variable.t * int) list =
  List.fold_left
    ( fun (acc: int VariableMap.t) (instr: RISCCfg.elt) ->
        VariableMap.union
          (fun _v x y -> Some (x + y))
          acc (variables_frequency instr |> VariableMap.of_list) )
    VariableMap.empty instructions |> VariableMap.to_list


let reduceregisters (n: int) (cfg: RISCCfg.t) : RISCCfg.t =
  (if n < 4 then (failwith "ReduceRegisters: number of registers too small") else ());

  (* we get all the variables with associated frequency (only syntactic use) *)
  let all_variables = List.fold_left
      (fun acc (_, code) ->
         Utility.unique_union_assoc (fun _n x y -> x + y) acc (variables_all_frequency code))
      []
      (Cfg.NodeMap.to_list cfg.content)
  in

  let all_variables =
    match cfg.inputOutputVar with
    | None -> all_variables
    | Some (i, _o) -> (
      match List.assoc_opt i all_variables with
      | None -> (i, 1) :: all_variables
      | Some f -> (i, f+1) :: (List.remove_assoc i all_variables)
    )
  in

  let all_variables =
    match cfg.inputOutputVar with
    | None -> all_variables
    | Some (_i, o) -> (
      match List.assoc_opt o all_variables with
      | None -> (o, 1) :: all_variables
      | Some f -> (o, f+1) :: (List.remove_assoc o all_variables)
    )
  in

  let replaceregisters
      (remappedregisters: Variable.t VariableMap.t)
      (memorymap: int VariableMap.t)
      (temporaryregisters: Variable.t list)
      (code: RISCCfg.elt list)
    : RISCCfg.elt list =

    let tmpreg1: CfgRISC.RISCSimpleStatements.register =
      {index = List.nth temporaryregisters 0} in
    let tmpreg2: CfgRISC.RISCSimpleStatements.register =
      {index = List.nth temporaryregisters 1} in

    let aux (instruction: RISCCfg.elt) : RISCCfg.elt list =
      match instruction with
      | Nop -> [Nop]
      | BRegOp (brop, r1, r2, r3) -> (
          match ( VariableMap.find_opt r1.index remappedregisters,
                  VariableMap.find_opt r2.index remappedregisters,
                  VariableMap.find_opt r3.index remappedregisters,
                  VariableMap.find_opt r1.index memorymap,
                  VariableMap.find_opt r2.index memorymap,
                  VariableMap.find_opt r3.index memorymap )
          with
          | Some r1, Some r2, Some r3, _, _, _ ->
            [BRegOp (brop, {index = r1}, {index = r2}, {index = r3})]
          | Some r1, None, Some r3, _, Some m2, _ ->
            [LoadI (m2, tmpreg2);
             Load (tmpreg2, tmpreg2);
             BRegOp (brop, {index = r1}, tmpreg2, {index = r3})]
          | None, Some r2, Some r3, Some m1, _, _ ->
            [LoadI (m1, tmpreg1);
             Load (tmpreg1, tmpreg1);
             BRegOp (brop, tmpreg1, {index = r2}, {index = r3})]
          | None, None, Some r3, Some m1, Some m2, _ ->
            [LoadI (m1, tmpreg1);
             Load (tmpreg1, tmpreg1);
             LoadI (m2, tmpreg2);
             Load (tmpreg2, tmpreg2);
             BRegOp (brop, tmpreg1, tmpreg2, {index = r3})]
          | Some r1, Some r2, None, _, _, Some m3 ->
            [BRegOp (brop, {index = r1}, {index = r2}, tmpreg2);
             LoadI (m3, tmpreg1);
             Store (tmpreg2, tmpreg1)]
          | Some r1, None, None, _, Some m2, Some m3 ->
            [LoadI (m2, tmpreg2);
             Load (tmpreg2, tmpreg2);
             BRegOp (brop, {index = r1}, tmpreg2, tmpreg2);
             LoadI (m3, tmpreg1);
             Store (tmpreg2, tmpreg1)]
          | None, Some r2, None, Some m1, _, Some m3 ->
            [LoadI (m1, tmpreg1);
             Load (tmpreg1, tmpreg1);
             BRegOp (brop, tmpreg1, {index = r2}, tmpreg2);
             LoadI (m3, tmpreg1);
             Store (tmpreg2, tmpreg1)]
          | None, None, None, Some m1, Some m2, Some m3 ->
            [LoadI (m1, tmpreg1);
             Load (tmpreg1, tmpreg1);
             LoadI (m2, tmpreg2);
             Load (tmpreg2, tmpreg2);
             BRegOp (brop, tmpreg1, tmpreg2, tmpreg2);
             LoadI (m3, tmpreg1);
             Store (tmpreg2, tmpreg1)]
          | _ -> [BRegOp (brop, {index = r1.index}, {index = r2.index}, {index = r3.index})]
        )
      | BImmOp (biop, r1, i, r3) -> (
          match ( VariableMap.find_opt r1.index remappedregisters,
                  VariableMap.find_opt r3.index remappedregisters,
                  VariableMap.find_opt r1.index memorymap,
                  VariableMap.find_opt r3.index memorymap )
          with
          | Some r1, Some r3, _, _ ->
            [BImmOp (biop, {index = r1}, i, {index = r3})]
          | Some r1, None, _, Some m3 ->
            [BImmOp (biop, {index = r1}, i, tmpreg2);
             LoadI (m3, tmpreg1);
             Store (tmpreg2, tmpreg1)]
          | None, Some r3, Some m1, _ ->
            [LoadI (m1, tmpreg1);
             Load (tmpreg1, tmpreg1);
             BImmOp (biop, tmpreg1, i, {index = r3})]
          | None, None, Some m1, Some m3 ->
            [LoadI (m1, tmpreg1);
             Load (tmpreg1, tmpreg1);
             BImmOp (biop, tmpreg1, i, tmpreg2);
             LoadI (m3, tmpreg1);
             Store (tmpreg2, tmpreg1)]
          | _ -> failwith ("ReduceRegisters: fail to partition a set, some" ^
                           " registers have no binding.")
      )
      | URegOp (urop, r1, r3) ->(
          match ( VariableMap.find_opt r1.index remappedregisters,
                  VariableMap.find_opt r3.index remappedregisters,
                  VariableMap.find_opt r1.index memorymap,
                  VariableMap.find_opt r3.index memorymap )
          with
          | Some r1, Some r3, _, _ ->
            [URegOp (urop, {index = r1}, {index = r3})]
          | Some r1, None, _, Some m3 ->
            [URegOp (urop, {index = r1}, tmpreg2);
             LoadI (m3, tmpreg1);
             Store (tmpreg2, tmpreg1)]
          | None, Some r3, Some m1, _ ->
            [LoadI (m1, tmpreg1);
             Load (tmpreg1, tmpreg1);
             URegOp (urop, tmpreg1, {index = r3})]
          | None, None, Some m1, Some m3 ->
            [LoadI (m1, tmpreg1);
             Load (tmpreg1, tmpreg1);
             URegOp (urop, tmpreg1, tmpreg2);
             LoadI (m3, tmpreg1);
             Store (tmpreg2, tmpreg1)]
          | _ -> failwith ("ReduceRegisters: fail to partition a set, some" ^
                           " registers have no binding.")
        )
      | Load (r1, r3) -> (
          match ( VariableMap.find_opt r1.index remappedregisters,
                  VariableMap.find_opt r3.index remappedregisters,
                  VariableMap.find_opt r1.index memorymap,
                  VariableMap.find_opt r3.index memorymap )
          with
          | Some r1, Some r3, _, _ ->
            [Load ({index = r1}, {index = r3})]
          | Some r1, None, _, Some m3 ->
            [Load ({index = r1}, tmpreg2);
             LoadI (m3, tmpreg1);
             Store (tmpreg2, tmpreg1)]
          | None, Some r3, Some m1, _ ->
            [LoadI (m1, tmpreg1);
             Load (tmpreg1, tmpreg1);
             Load (tmpreg1, {index = r3})]
          | None, None, Some m1, Some m3 ->
            [LoadI (m1, tmpreg1);
             Load (tmpreg1, tmpreg1);
             Load (tmpreg1, tmpreg2);
             LoadI (m3, tmpreg1);
             Store (tmpreg2, tmpreg1)]
          | _ -> failwith ("ReduceRegisters: fail to partition a set, some" ^
                           " registers have no binding.")
        )
      | LoadI (i, r3) -> (
          (* we want to store an integer in memory immediately (strange, but
             unless better heuristic to choose the variables to replace we are
             stuck) *)
          match ( VariableMap.find_opt r3.index remappedregisters,
                  VariableMap.find_opt r3.index memorymap )
          with
          | Some r3, _ ->
            [LoadI (i, {index = r3})]
          | None, Some m3 ->
            [LoadI (i, tmpreg2);
             LoadI (m3, tmpreg1);
             Store (tmpreg2, tmpreg1)]
          | _ -> failwith ("ReduceRegisters: fail to partition a set, some" ^
                           " registers have no binding.")
        )
      | Store (r1, r3) -> (
          (* we want to maybe store an address in memory (very confusing, don't
             think can happen) *)
          match ( VariableMap.find_opt r1.index remappedregisters,
                  VariableMap.find_opt r3.index remappedregisters,
                  VariableMap.find_opt r1.index memorymap,
                  VariableMap.find_opt r3.index memorymap )
          with
          | Some r1, Some r3, _, _ ->
            [Store ({index = r1}, {index = r3})]
          | Some r1, None, _, Some m3 ->
            [Store ({index = r1}, tmpreg2);
             LoadI (m3, tmpreg1);
             Store (tmpreg2, tmpreg1)]
          | None, Some r3, Some m1, _ ->
            [LoadI (m1, tmpreg1);
             Load (tmpreg1, tmpreg1);
             Store (tmpreg1, {index = r3})]
          | None, None, Some m1, Some m3 ->
            [LoadI (m1, tmpreg1);
             Load (tmpreg1, tmpreg1);
             Store (tmpreg1, tmpreg2);
             LoadI (m3, tmpreg1);
             Store (tmpreg2, tmpreg1)]
          | _ -> failwith ("ReduceRegisters: fail to partition a set, some" ^
                           " registers have no binding.")
      )
    in

    List.map aux code |> List.concat
  in


  let aux (cfg: RISCCfg.t) (all_variables: (string * int) list) =
    (* we keep the first two variables free for immediate use *)
    let most_frequent, least_frequent =
      List.sort (fun (_a, fa) (_b, fb) -> Int.compare fb fa) all_variables
      |> Utility.take (n-2)
    in
    let most_frequent = fst (List.split most_frequent) in
    let least_frequent = fst (List.split least_frequent) in

    (* we map the most frequent to new registers, so that the first two are
       always free *)
    let most_frequent_mapping = (* +3 because starts at 0, but we want to start
                                   at 1*)
      List.mapi (fun n v -> (v, (string_of_int (n+3): Variable.t))) most_frequent
      |> VariableMap.of_list
    in
    (* we map the least to memory *)
    let least_frequent_mapping =
      List.mapi (fun n v -> (v, (n: int))) least_frequent
      |> VariableMap.of_list
    in

    (* we need to replace both at the same time, because we might have mapped
       some registers to already used registers, so a double pass might not
       differentiate the two *)
    (* special care must be taken for the in and out registers *)
    let newcfg = {
      cfg with
      content = Cfg.NodeMap.map
          (fun x -> replaceregisters most_frequent_mapping least_frequent_mapping ["1"; "2"] x)
          cfg.content}
    in

    match newcfg.inputOutputVar with
    | None -> newcfg (* if no input or output variables we ignore *)
    | Some (i, o) -> (
        match (VariableMap.find_opt i most_frequent_mapping,
               VariableMap.find_opt o most_frequent_mapping,
               VariableMap.find_opt i least_frequent_mapping,
               VariableMap.find_opt o least_frequent_mapping )
        with (*we check if in and out are simply remapped or are put in memory*)
        | Some i, Some o, _, _ ->
          { newcfg with inputOutputVar = Some (i, o) }
        | Some i, None, _, Some mo -> ( (* since the output simbol is in memory
                                           we need to first retrive it and then
                                           put the result in a temporary
                                           register *)
            match newcfg.terminal with (* we check for the terminal node, if not
                                          present we are very confused and dont
                                          modify the out variable *)
            | None -> { newcfg with inputOutputVar = Some (i, o)}
            | Some n -> (
                let terminalcontent = (
                  match Cfg.NodeMap.find_opt n newcfg.content with
                  | None -> []
                  | Some x -> x
                ) @ [LoadI (mo, {index = "2"});
                     Load ({index = "2"}, {index = "2"})]
                in
                let content = Cfg.NodeMap.add n terminalcontent newcfg.content in
                { newcfg with
                  inputOutputVar = Some (i, "2");
                  content = content
                }
              )
          )
        | None, Some o, Some mi, _ -> ( (* the input simbol should be stored in
                                           memory *)
            match newcfg.initial with
            | None -> { newcfg with inputOutputVar = Some (i, o) }
            | Some n -> (
                let initialcontent =
                  [(LoadI (mi, {index = "2"}) : RISCCfg.elt);
                   Store ({index = "1"}, {index = "2"})] @ (
                    match Cfg.NodeMap.find_opt n newcfg.content with
                    | None -> []
                    | Some x -> x
                  )
                in
                let content = Cfg.NodeMap.add n initialcontent newcfg.content in
                { newcfg with
                  inputOutputVar = Some ("1", o);
                  content = content
                }
              )
          )
        | None, None, Some mi, Some mo -> ( (* both simbols should be in
                                               memory *)
            match newcfg.initial, newcfg.terminal with
            | None, None -> { newcfg with inputOutputVar = Some (i, o) }
            | None, Some n -> (
                let terminalcontent = (
                  match Cfg.NodeMap.find_opt n newcfg.content with
                  | None -> []
                  | Some x -> x
                ) @ [LoadI (mo, {index = "2"});
                     Load ({index = "2"}, {index = "2"})]
                in
                let content = Cfg.NodeMap.add n terminalcontent newcfg.content in
                { newcfg with
                  inputOutputVar = Some (i, "2");
                  content = content
                }
              )
            | Some n, None -> (
                let initialcontent =
                  [(LoadI (mi, {index = "2"}) : RISCCfg.elt);
                   Store ({index = "1"}, {index = "2"})] @ (
                    match Cfg.NodeMap.find_opt n newcfg.content with
                    | None -> []
                    | Some x -> x
                  )
                in
                let content = Cfg.NodeMap.add n initialcontent newcfg.content in
                { newcfg with
                  inputOutputVar = Some ("1", o);
                  content = content
                }
              )
            | Some ni, Some no -> (
                let initialcontent =
                  [(LoadI (mi, {index = "2"}) : RISCCfg.elt);
                   Store ({index = "1"}, {index = "2"})] @ (
                    match Cfg.NodeMap.find_opt ni newcfg.content with
                    | None -> []
                    | Some x -> x
                  )
                in
                let terminalcontent = (
                  match Cfg.NodeMap.find_opt no newcfg.content with
                  | None -> []
                  | Some x -> x
                ) @ [LoadI (mo, {index = "2"});
                     Load ({index = "2"}, {index = "2"})]
                in
                let content = Cfg.NodeMap.add ni initialcontent newcfg.content in
                let content = Cfg.NodeMap.add no terminalcontent content in
                { newcfg with
                  inputOutputVar = Some ("1", "2");
                  content = content
                }
              )
          )
        | _ -> failwith ("ReduceRegisters: fail to partition a set, some" ^
                           " registers have no binding.")
      )
  in


  ( if List.length all_variables <= n
    then cfg
    else aux cfg all_variables )