open Cfg

module SimpleStatements = struct
  type t =
    | SimpleSkip
    | SimpleAssignment of Types.variable * simpleArithmetic
    | SimpleGuard of simpleBoolean
  and simpleBoolean =
    | SimpleBoolean of bool
    | SimpleBAnd of simpleBoolean * simpleBoolean
    | SimpleBOr of simpleBoolean * simpleBoolean
    | SimpleBNot of simpleBoolean
    | SimpleBCmp of simpleArithmetic * simpleArithmetic
    | SimpleBCmpLess of simpleArithmetic * simpleArithmetic
    | SimpleBCmpLessEq of simpleArithmetic * simpleArithmetic
    | SimpleBCmpGreater of simpleArithmetic * simpleArithmetic
    | SimpleBCmpGreaterEq of simpleArithmetic * simpleArithmetic
  and simpleArithmetic =
    | SimpleVariable of Types.variable
    | SimpleInteger of int
    | SimplePlus of simpleArithmetic * simpleArithmetic
    | SimpleMinus of simpleArithmetic * simpleArithmetic
    | SimpleTimes of simpleArithmetic * simpleArithmetic
    | SimpleDivision of simpleArithmetic * simpleArithmetic
    | SimpleModulo of simpleArithmetic * simpleArithmetic
    | SimplePower of simpleArithmetic * simpleArithmetic
    | SimplePowerMod of simpleArithmetic * simpleArithmetic * simpleArithmetic
    | SimpleRand of simpleArithmetic

  let pp (ppf: out_channel) (c: t) : unit =
    let rec helper_c (ppf) (c: t) : unit =
      match c with
      | SimpleSkip -> Printf.fprintf ppf "Skip"
      | SimpleAssignment (v, a) -> Printf.fprintf ppf "Assignment {%s, %a}" v helper_a a
      | SimpleGuard (b) -> Printf.fprintf ppf "Guard {%a}" helper_b b
    and helper_b (ppf) (c: simpleBoolean) : unit =
      match c with
      | SimpleBoolean b -> Printf.fprintf ppf "%b" b
      | SimpleBAnd (b1, b2) -> Printf.fprintf ppf "{%a && %a}" helper_b b1 helper_b b2
      | SimpleBOr (b1, b2) -> Printf.fprintf ppf "{%a || %a}" helper_b b1 helper_b b2
      | SimpleBNot b -> Printf.fprintf ppf "{not %a}" helper_b b
      | SimpleBCmp (a1, a2) -> Printf.fprintf ppf "{%a == %a}" helper_a a1 helper_a a2
      | SimpleBCmpLess (a1, a2) -> Printf.fprintf ppf "{%a < %a}" helper_a a1 helper_a a2
      | SimpleBCmpLessEq (a1, a2) -> Printf.fprintf ppf "{%a <= %a}" helper_a a1 helper_a a2
      | SimpleBCmpGreater (a1, a2) -> Printf.fprintf ppf "{%a > %a}" helper_a a1 helper_a a2
      | SimpleBCmpGreaterEq (a1, a2) -> Printf.fprintf ppf "{%a >= %a}" helper_a a1 helper_a a2
    and helper_a (ppf) (c: simpleArithmetic) : unit =
      match c with
      | SimpleVariable (v) -> Printf.fprintf ppf "%s" v
      | SimpleInteger (i) -> Printf.fprintf ppf "%d" i
      | SimplePlus (a1, a2) -> Printf.fprintf ppf "{%a + %a}" helper_a a1 helper_a a2
      | SimpleMinus (a1, a2) -> Printf.fprintf ppf "{%a - %a}" helper_a a1 helper_a a2
      | SimpleTimes (a1, a2) -> Printf.fprintf ppf "{%a * %a}" helper_a a1 helper_a a2
      | SimpleDivision (a1, a2) -> Printf.fprintf ppf "{%a / %a}" helper_a a1 helper_a a2
      | SimpleModulo (a1, a2) -> Printf.fprintf ppf "{%a %% %a}" helper_a a1 helper_a a2
      | SimplePower (a1, a2) -> Printf.fprintf ppf "{%a ^ %a}" helper_a a1 helper_a a2
      | SimplePowerMod (a1, a2, a3) -> Printf.fprintf ppf "{powmod %a %a %a}" helper_a a1 helper_a a2 helper_a a3
      | SimpleRand (a) -> Printf.fprintf ppf "{rand %a}" helper_a a
    in
    helper_c ppf c

  let pplist (ppf: out_channel) (c: t list) : unit =
    List.iter (fun x -> pp ppf x; Printf.printf "; ") c
end

module SSCfg  = Cfg.Make(SimpleStatements)

let rec convert_c (prevcfg: SSCfg.t) (prg: Types.c_exp) : SSCfg.t =
  let open SimpleStatements in
  match prg with
  | Skip -> prevcfg |> SSCfg.addToLastNode SimpleSkip
  | Assignment (x, a) -> prevcfg |> SSCfg.addToLastNode (SimpleAssignment (x, convert_a a))
  | Sequence (c1, c2) ->
    let cfg1 = convert_c prevcfg c1 in
    let cfg2 = convert_c cfg1 c2 in
    cfg2
  | If (b, c1, c2) ->
    let convertedb = convert_b b in
    let cfg1 = convert_c SSCfg.empty c1 in
    let cfg2 = convert_c SSCfg.empty c2 in
    let entrynode = Node.create () in
    let exitnode = Node.create () in
    let newcfg = SSCfg.merge cfg1 cfg2 entrynode exitnode in
    let mergedcfg = SSCfg.concat prevcfg newcfg in
    { mergedcfg with
      content = mergedcfg.content |>
                NodeMap.add_to_list entrynode (SimpleGuard convertedb) |>
                NodeMap.add_to_list exitnode (SimpleSkip) }
  | While (b, c) ->
    let convertedb = convert_b b in
    let cfg = convert_c SSCfg.empty c in
    let cfginitial = Option.get cfg.initial in
    let cfgterminal = Option.get cfg.terminal in
    let entrynode = Node.create () in
    let guardnode = Node.create () in
    let exitnode = Node.create () in
    { empty = false;
      nodes = cfg.nodes |>
              NodeSet.add entrynode |>
              NodeSet.add guardnode |>
              NodeSet.add exitnode;
      edges = cfg.edges |>
              NodeMap.add entrynode (guardnode, None) |>
              NodeMap.add guardnode (cfginitial, Some exitnode) |>
              NodeMap.add cfgterminal (guardnode, None);
      reverseEdges = cfg.reverseEdges |>
                     NodeMap.add_to_list guardnode entrynode |>
                     NodeMap.add_to_list cfginitial guardnode |>
                     NodeMap.add_to_list exitnode guardnode |>
                     NodeMap.add_to_list guardnode cfgterminal;
      inputVal = prevcfg.inputVal;
      inputOutputVar = prevcfg.inputOutputVar;
      initial = Some entrynode;
      terminal = Some exitnode;
      content = NodeMap.add_to_list guardnode (SimpleGuard (convertedb)) cfg.content |>
                NodeMap.add_to_list exitnode (SimpleSkip)
    } |> SSCfg.concat prevcfg
  | For (assignment, guard, increment, body) ->
    let cfgassignment = convert_c SSCfg.empty assignment in
    let convertedguard = convert_b guard in
    let cfgincrement = convert_c SSCfg.empty increment in
    let cfgbody = convert_c SSCfg.empty body in

    let prevassignment = SSCfg.concat prevcfg cfgassignment in
    let bodyincrement = SSCfg.concat cfgbody cfgincrement in

    let cfginitial = Option.get bodyincrement.initial in
    let cfgterminal = Option.get bodyincrement.terminal in

    let guardnode = Node.create () in
    let exitnode = Node.create () in
    { empty = false;
      nodes = bodyincrement.nodes |>
              NodeSet.add guardnode |>
              NodeSet.add exitnode;
      edges = bodyincrement.edges |>
              NodeMap.add guardnode (cfginitial, Some exitnode) |>
              NodeMap.add cfgterminal (guardnode, None);
      reverseEdges = bodyincrement.reverseEdges |>
                     NodeMap.add_to_list cfginitial guardnode |>
                     NodeMap.add_to_list exitnode guardnode |>
                     NodeMap.add_to_list guardnode cfgterminal;
      inputVal = prevcfg.inputVal;
      inputOutputVar = prevcfg.inputOutputVar;
      initial = Some guardnode;
      terminal = Some exitnode;
      content = NodeMap.add_to_list guardnode (SimpleGuard (convertedguard)) bodyincrement.content |>
                NodeMap.add_to_list exitnode (SimpleSkip)
    } |> SSCfg.concat prevassignment

and convert_b (prg: Types.b_exp) : SimpleStatements.simpleBoolean =
  match prg with
  | Boolean (b) -> SimpleBoolean b
  | BAnd (b1, b2) -> SimpleBAnd (convert_b b1, convert_b b2)
  | BOr (b1, b2) -> SimpleBOr (convert_b b1, convert_b b2)
  | BNot (b) -> SimpleBNot (convert_b b)
  | BCmp (a1, a2) -> SimpleBCmp (convert_a a1, convert_a a2)
  | BCmpLess (a1, a2) -> SimpleBCmpLess (convert_a a1, convert_a a2)
  | BCmpLessEq (a1, a2) -> SimpleBCmpLessEq (convert_a a1, convert_a a2)
  | BCmpGreater (a1, a2) -> SimpleBCmpGreater (convert_a a1, convert_a a2)
  | BCmpGreaterEq (a1, a2) -> SimpleBCmpGreaterEq (convert_a a1, convert_a a2)

and convert_a (prg: Types.a_exp) : SimpleStatements.simpleArithmetic =
  match prg with
  | Variable x -> SimpleVariable x
  | Integer n -> SimpleInteger n
  | Plus (a1, a2) -> SimplePlus (convert_a a1, convert_a a2)
  | Minus (a1, a2) -> SimpleMinus (convert_a a1, convert_a a2)
  | Times (a1, a2) -> SimpleTimes (convert_a a1, convert_a a2)
  | Division (a1, a2) -> SimpleDivision (convert_a a1, convert_a a2)
  | Modulo (a1, a2) -> SimpleModulo (convert_a a1, convert_a a2)
  | Power (a1, a2) -> SimplePower (convert_a a1, convert_a a2)
  | PowerMod (a1, a2, a3) -> SimplePowerMod (convert_a a1, convert_a a2, convert_a a3)
  | Rand (a) -> SimpleRand (convert_a a)

let convert (prg: Types.p_exp) : SSCfg.t =
  match prg with
  | Main (i, o, exp) ->
    {(convert_c SSCfg.empty exp) with inputOutputVar = Some (i, o)}

let convert_io (prg: Types.p_exp) (i: int) : SSCfg.t =
  {(convert prg) with inputVal = Some i}