module PASPProc

include("../asplang/ASPLang.jl")
using ASPLang

export tree_string,
    filter_tree,
    filter_children,
    has_child,
    node_findreplace,
    has_annotation_P,
    repl_annotation_disjunction,
    derived

tree_string(node::AnonymousVariable; level=0, indent="\t") = "$(indent^level)Anonymous: $(node.args |> string)"
tree_string(node::Functor; level=0, indent="\t") = "$(indent^level)Functor: $(node.args |> string)"
tree_string(node::Annotation; level=0, indent="\t") = "$(indent^level)Annotation: $(node.args |> string)"
tree_string(node::Variable; level=0, indent="\t") = "$(indent^level)Variable: $(node.args |> string)"
tree_string(node::StringConstant; level=0, indent="\t") = "$(indent^level)String: $(node.args |> string)"
tree_string(node::SymbolicConstant; level=0, indent="\t") = "$(indent^level)Symbol: $(node.args |> string)"
tree_string(node::NumberConstant; level=0, indent="\t") = "$(indent^level)Number: $(node.args |> string)"
#tree_print(node::Terminal; level=0, indent="\t") = node |> typeof |> string
function tree_string(node::ASPLang.NonTerminal; level=0, indent="\t")
    node_head = node |> typeof |> string
    sub_results = [ "$(indent^level)$node_head" ] 
    if isa(node.args, AbstractArray) && length(node.args) > 0
        for child in node.args 
            child_tree = if isa(child, ASPLang.Node)
                tree_string(child, level = level + 1, indent=indent)
            else
                "$(indent^(level+1))$(child |> string)"
            end
            push!(sub_results, child_tree)
        end
    else
        push!(sub_results, tree_string(node.args, level=level+1, indent=indent))
    end
    join(sub_results, "\n")
end
tree_string(node; level=0, indent="\t") = "$(indent^level)$(node |> typeof |> string):$(node |> string)"


function filter_tree(node::NonTerminal, prop)
    results = if prop(node)
        [node]
    else
        []
    end
    if isa(node.args, AbstractArray)
        for d in node.args
            append!(results, filter_tree(d, prop))
        end
    else
        if prop(node.args)
            push!(results, node.args)
        end
    end
    results
end

function filter_children(node::NonTerminal, prop)
    if isa(node.args, AbstractArray)
        filter(prop, node.args)
    else
        if prop(node.args)
            node.args
        else
            nothing
        end
    end
end

function has_child(prop)
    node::NonTerminal -> any(prop, node.args)
end

function node_findreplace(node::NonTerminal, prop, action)
    if prop(node)
        action(node)
    elseif isa(node.args, AbstractArray)
        new_args = []
        for arg in node.args
            push!(new_args, node_findreplace(arg, prop, action))
        end
        typeof(node)(new_args)
    else
        node
    end
end

has_annotation_P = n -> has_child(isa(n, ASPLang.Annotation))
function repl_annotation_disjunction(n::NonTerminal)
    a = [xi for xi in n.args if !isa(xi, Annotation)][1]
    Head(
        Disjunction(
            [a, Negated(a)]
        ))
end

function derived(p::ASPLang.Program)
    node_findreplace(p,
        has_annotation_P,
        repl_annotation_disjunction)
end

end # PASPproc