open Printf
open Ast
open Util

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, None, 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 (Array _ as vtype, name, Some (BoolConst _  as v), loc) :: t
                | VarDec (Array _ as vtype, name, Some (FloatConst _ as v), loc) :: t
                | VarDec (Array _ as vtype, name, Some (IntConst _   as v), loc) :: t ->
                    let init = Some (ArrayInit (ArrayScalar v, vtype)) in
                    trav inits (VarDec (vtype, name, init, loc) :: t)

                (* Wrap ArrayConst in ArrayInit to pass dimensions *)
                | VarDec (Array _ as vtype, name, Some (ArrayConst _ as v), loc) :: t ->
                    let init = Some (ArrayInit (v, vtype)) in
                    trav inits (VarDec (vtype, name, init, loc) :: t)

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

                (* 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
        FunDef (export, ret_type, name, params, Block (move_inits body), loc)

    | node -> transform_children var_init node

let rec replace_var var replacement node =
    let trav = (replace_var var replacement) in
    match node with
    | Var (name, loc) when name = var ->
        Var (replacement, loc)
    | For (counter, start, stop, step, body, loc) when counter = var ->
        For (replacement, trav start, trav stop, trav step, trav body, loc)
    | node ->
        transform_children trav node

let for_to_while node =
    let rec traverse new_vars = function
        | FunDef (export, ret_type, name, params, body, loc) ->
            let new_vars = ref [] in
            let body = traverse new_vars body in
            let create_vardec name = VarDec (Int, name, None, noloc) in
            let new_vardecs = List.map create_vardec !new_vars in
            let _body = new_vardecs @ (flatten_blocks (block_body body)) in
            FunDef (export, ret_type, name, params, Block _body, loc)

        (* 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, None, start, locof start);
                Assign (_stop, None, stop, locof stop);
                Assign (_step, None, step, locof step);
                While (cond, traverse new_vars (Block (
                    block_body (replace_var counter _i body) @
                    [Assign (_i, None, Binop (Add, vi, vstep, noloc), noloc)]
                )), loc);
            ]

        | node -> transform_children (traverse new_vars) node
    in
    traverse (ref []) node

let rec array_init = function
    (* Transform scalar assignment into nested for-loops *)
    | Assign (name, None, ArrayInit (ArrayScalar value, Array (_, dims)), loc) ->
        let rec add_loop indices = function
            | [] ->
                Assign (name, Some indices, value, loc)
            | dim :: rest ->
                let counter = fresh_var "i" in
                let body = Block [add_loop (indices @ [Var (counter, noloc)]) rest] in
                For (counter, IntConst (0, noloc), dim, IntConst (1, noloc), body, noloc)
        in
        add_loop [] dims

    (* Transform array constant inisialisation into separate assign statements
     * for all entries in the constant array *)
    (* TODO: only allow when array dimensions are constant? *)
    | Assign (name, None, ArrayInit (ArrayConst _ as value, Array (_, dims)), loc) ->
        let ndims = list_size dims in
        let rec make_assigns depth i indices = function
            | [] -> []
            | hd :: tl ->
                let assigns = traverse depth (i :: indices) hd in
                make_assigns depth (i + 1) indices tl @ assigns
        and traverse depth indices = function
            | ArrayConst (values, _) ->
                make_assigns (depth + 1) 0 indices values
            | value when depth = ndims ->
                let indices = List.map (fun i -> IntConst (i, noloc)) indices in
                [Assign (name, Some (List.rev indices), value, loc)]
            | node ->
                let msg = sprintf
                    "dimension mismatch: expected %d nesting levels, got %d"
                    ndims depth
                in
                raise (NodeError (node, msg))
        in
        Block (List.rev (traverse 0 [] value))

    | node -> transform_children array_init node

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