open Ast
open Util

let block_body = function
    | Block nodes -> nodes
    | _ -> raise InvalidNode

let rec replace_var var replacement = function
    | Var (name, loc) when name = var -> Var (replacement, loc)
    | node -> transform_children (replace_var var replacement) node

let for_to_while node =
    let new_vars = ref [] in
    let rec traverse = function
        (* Do not traverse into local functions (already done by var_init) *)
        | FunDef (_, _, _, _, _, _) as node -> node

        (* Transform for-loops to while-loops *)
        | For (counter, start, stop, step, body, loc) ->
            let _i = fresh_var counter in
            let _stop = fresh_var "stop" in
            let _step = fresh_var "step" in
            new_vars := !new_vars @ [_i; _stop; _step];

            let vi = Var (_i, noloc) in
            let vstop = Var (_stop, locof stop) in
            let vstep = Var (_step, locof step) in
            let cond = Cond (
                Binop (Gt, vstep, IntConst (0, noloc), noloc),
                Binop (Lt, vi, vstop, noloc),
                Binop (Gt, vi, vstop, noloc),
                noloc
            ) in
            Block [
                Assign (_i, start, locof start);
                Assign (_stop, stop, locof stop);
                Assign (_step, step, locof step);
                While (cond, traverse (Block (
                    (* TODO: check for illegal assigments of counter in body *)
                    block_body (replace_var counter _i body) @
                    [Assign (_i, Binop (Add, vi, vstep, noloc), noloc)]
                )), loc);
            ]

        | node -> transform_children traverse node
    in
    let node = traverse node in
    (node, new_vars)

let rec var_init = function
    (* Move global initialisations to __init function *)
    | Program (decls, loc) ->
        let decls = flatten_blocks (List.map var_init decls) in
        let rec trav assigns = function
            | [] -> (assigns, [])
            | (Assign (_, _, _) as h) :: t -> trav (assigns @ [h]) t
            | h :: t ->
                let (assigns, decls) = trav assigns t in
                (assigns, (h :: decls))
        in
        let (assigns, decls) = trav [] decls in (
            match assigns with
            | [] -> Program (decls, loc)
            | assigns ->
                let init_func = FunDef (true, Void, "__init", [], Block assigns, noloc) in
                Program (init_func :: decls, loc)
        )

    (* Move global variable initialisations to exported __init function *)
    | GlobalDef (export, ctype, name, Some init, loc) ->
        Block [GlobalDef (export, ctype, name, None, loc);
               Assign (name, init, locof init)]

    (* Split local variable initialisations in declaration and assignment *)
    | FunDef (export, ret_type, name, params, Block body, loc) ->
        let move_inits body =
            let rec trav inits node = match node with
                (* translate scalar array initialisation to ArrayScalar node,
                 * for easy replacement later on *)
                | VarDec (ArrayDef (_, _) as vtype, name,
                          Some ((BoolConst  (_, l)) as v), loc) :: t
                | VarDec (ArrayDef (_, _) as vtype, name,
                          Some ((FloatConst (_, l)) as v), loc) :: t
                | VarDec (ArrayDef (_, _) as vtype, name,
                          Some ((IntConst   (_, l)) as v), loc) :: t ->
                    trav inits (VarDec (vtype, name, Some (ArrayScalar (v, l)), loc) :: t)

                | VarDec (ctype, name, init, loc) :: t ->
                    (* array definition: create __allocate statement *)
                    let alloc = match ctype with
                        | ArrayDef (_, dims) -> [Allocate (name, dims, loc)]
                        | _ -> []
                    in
                    (* initialisation: create assign statement *)
                    let add = match init with
                        | Some value -> alloc @ [Assign (name, value, loc)]
                        | None -> alloc
                    in
                    VarDec (ctype, name, None, loc) :: (trav (inits @ add) t)

                (* initialisations need to be placed after local functions *)
                | (FunDef (_, _, _, _, _, _) as h) :: t ->
                    (var_init h) :: (trav inits t)

                (* rest of function body: recurse *)
                | rest -> inits @ (List.map var_init rest)
            in
            flatten_blocks (trav [] body)
        in
        let params = flatten_blocks (List.map var_init params) in
        let (body, new_vars) = for_to_while (Block (move_inits body)) in
        let create_vardec name = VarDec (Int, name, None, noloc) in
        let new_vardecs = List.map create_vardec !new_vars in
        let stats = new_vardecs @ (flatten_blocks (block_body body)) in
        FunDef (export, ret_type, name, params, Block stats, loc)

    | node -> transform_children var_init node

(*
let rec array_init = function
    (* transform scalar assignment into nested for loops *)
    | Assign (name, ArrayScalar (value)) ->
        let rec add_loop indices = function
            | [] ->
                Assign (Deref (name, indices), value)
            | dim :: rest ->
                let counter = fresh_var "counter" in
                let ind = (indices @ [Var counter]) in
                For (counter, IntConst 0, dim, IntConst 1, add_loop ind rest)
        in
        add_loop [] dims

    | Assign (name, ArrayConst (dims)) -> Block []

    | node -> transform array_init node
*)

let rec phase input =
    prerr_endline "- Desugaring";
    match input with
    | Ast (node, args) ->
        Ast (var_init node, args)
    | _ -> raise (InvalidInput "desugar")