type simpleStatements =
  | 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 printSingleStatement (ppf) (c: simpleStatements) : unit =
  let rec helper_c (ppf) (c: simpleStatements) : 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 printSimpleStatements (ppf) (c: simpleStatements list) : unit =
  List.iter (fun x -> printSingleStatement ppf x; Printf.printf "; ") c



let globalIdNode = ref 0;

module Node = struct
  type t = {
    id: int
  }
  let compare a b = compare a.id b.id

  let newNode () =
    globalIdNode := !globalIdNode + 1;
    {id = !globalIdNode}
end
;;

module NodeMap = Map.Make(Node)
module NodeSet = Set.Make(Node)

module Cfg = struct
  type t = {
    empty: bool;
    nodes: NodeSet.t;
    edges: (Node.t * (Node.t option)) NodeMap.t;
    initial: Node.t option;
    terminal: Node.t option;
    code: (simpleStatements list) NodeMap.t
  }

  let newCfg () =
    { empty = true;
      nodes = NodeSet.empty;
      edges = NodeMap.empty;
      initial = None;
      terminal = None;
      code = NodeMap.empty }

  let mergeCfg (cfg1: t) (cfg2: t) (entryNode: Node.t) (exitNode: Node.t) : t =
    match (cfg1.empty, cfg2.empty) with
      true, _ -> cfg2
    | _, true -> cfg1
    | false, false ->
      let cfg1initial = Option.get cfg1.initial in
      let cfg2initial = Option.get cfg2.initial in
      let cfg1terminal = Option.get cfg1.terminal in
      let cfg2terminal = Option.get cfg2.terminal in
      { empty = false;
        nodes = NodeSet.union cfg1.nodes cfg2.nodes |>
                NodeSet.add entryNode |>
                NodeSet.add exitNode;
        edges = NodeMap.union (fun _ -> failwith "Failed merging edges of cfg.")
            cfg1.edges cfg2.edges |>
                NodeMap.add entryNode (cfg1initial, Some cfg2initial) |>
                NodeMap.add cfg1terminal (exitNode, None) |>
                NodeMap.add cfg2terminal (exitNode, None);
        initial = Some entryNode;
        terminal = Some exitNode;
        code = NodeMap.union (fun _ -> failwith "Failed merging code of cfg.")
            cfg1.code cfg2.code
      }

  let concatCfg (cfg1: t) (cfg2: t) : t =
    match (cfg1.empty, cfg2.empty) with
      true, _ -> cfg2
    | _, true -> cfg1
    | false, false ->
      let cfg1initial = Option.get cfg1.initial in
      let cfg2initial = Option.get cfg2.initial in
      let cfg1terminal = Option.get cfg1.terminal in
      let cfg2terminal = Option.get cfg2.terminal in
      { empty = false;
        nodes = NodeSet.union cfg1.nodes cfg2.nodes;
        edges = NodeMap.union (fun _ -> failwith "Failed merging edges of cfg.")
                              cfg1.edges cfg2.edges |>
                NodeMap.add cfg1terminal (cfg2initial, None);
        initial = Some cfg1initial;
        terminal = Some cfg2terminal;
        code = NodeMap.union (fun _ -> failwith "Failed merging code of cfg.")
            cfg1.code cfg2.code
      }

  let addToLastNode (newcode: simpleStatements) (cfg: t) : t =
    match cfg.empty with
    | true -> let newnode = Node.newNode () in
      { empty = false;
        nodes = NodeSet.singleton newnode;
        edges = NodeMap.empty;
        initial = Some newnode;
        terminal = Some newnode;
        code = NodeMap.singleton newnode [newcode]
      }
    | false ->
      let prevcfgterminal = Option.get cfg.terminal in
      { cfg with
        code = (NodeMap.add_to_list
                  prevcfgterminal
                  newcode
                  cfg.code) }

  let pp (ppf) (c: t) : unit =
    Printf.fprintf ppf "Nodes' ids: ";
    List.iter (fun (x : Node.t) -> Printf.fprintf ppf "%d " x.id) (NodeSet.to_list c.nodes);
    Printf.fprintf ppf "\n";

    Printf.fprintf ppf "Nodes' edges:\n";
    List.iter (fun ((n, (a, b)) : (Node.t * (Node.t * Node.t option))) : unit ->
        match b with None -> Printf.fprintf ppf "\t%d -> %d\n" n.id a.id
                   | Some b -> Printf.fprintf ppf "\t%d -> %d, %d\n" n.id a.id b.id
      ) (NodeMap.to_list c.edges);
    Printf.fprintf ppf "\n";

    Printf.fprintf ppf "Initial node's id: ";
    Printf.fprintf ppf "%d" ((Option.get c.initial).id);
    Printf.fprintf ppf "\n";

    Printf.fprintf ppf "Terminal node's id: ";
    Printf.fprintf ppf "%d" ((Option.get c.terminal).id);
    Printf.fprintf ppf "\n";

    Printf.fprintf ppf "Code:\n";
    List.iter (fun ((n, stms) : Node.t * simpleStatements list) : unit ->
        Printf.fprintf ppf "\tid %d --> %a\n%!" n.id printSimpleStatements (List.rev stms)
      ) (NodeMap.to_list c.code);
    Printf.fprintf ppf "\n";
end
;;


let rec convert_c (prevcfg: Cfg.t) (prg: Types.c_exp) : Cfg.t =
  match prg with
  | Skip -> prevcfg |> Cfg.addToLastNode SimpleSkip
  | Assignment (x, a) -> prevcfg |> Cfg.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 (Cfg.newCfg ()) c1 in
    let cfg2 = convert_c (Cfg.newCfg ()) c2 in
    let entrynode = Node.newNode () in
    let exitnode = Node.newNode () in
    let newcfg = Cfg.mergeCfg cfg1 cfg2 entrynode exitnode in
    let mergedcfg = Cfg.concatCfg prevcfg newcfg in
    { mergedcfg with
      code = mergedcfg.code |>
             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 (Cfg.newCfg ()) c in
    let cfginitial = Option.get cfg.initial in
    let cfgterminal = Option.get cfg.terminal in
    let entrynode = Node.newNode () in
    let guardnode = Node.newNode () in
    let exitnode = Node.newNode () 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);
      initial = Some entrynode;
      terminal = Some exitnode;
      code = NodeMap.add_to_list guardnode (SimpleGuard (convertedb)) cfg.code |>
             NodeMap.add_to_list exitnode (SimpleSkip)
    } |> Cfg.concatCfg prevcfg
  | For (assignment, guard, increment, body) ->
    let cfgassignment = convert_c (Cfg.newCfg ()) assignment in
    let convertedguard = convert_b guard in
    let cfgincrement = convert_c (Cfg.newCfg ()) increment in
    let cfgbody = convert_c (Cfg.newCfg ()) body in

    let prevassignment = Cfg.concatCfg prevcfg cfgassignment in
    let bodyincrement = Cfg.concatCfg cfgbody cfgincrement in

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

    let guardnode = Node.newNode () in
    let exitnode = Node.newNode () 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);
      initial = Some guardnode;
      terminal = Some exitnode;
      code = NodeMap.add_to_list guardnode (SimpleGuard (convertedguard)) bodyincrement.code |>
             NodeMap.add_to_list exitnode (SimpleSkip)
    } |> Cfg.concatCfg prevassignment

and convert_b (prg: Types.b_exp) : 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) : 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) : Cfg.t =
  match prg with
  | Main (_, _, exp) ->
    convert_c (Cfg.newCfg ()) exp