#! /bin/sh

usage() {
    echo "EXAMPLE (whole tree)   : $0 | neato -T png > ps.png && open ps.png"
    echo "EXAMPLE (user clusters): $0 | sdp   -T png > ps.png && open ps.png"
}


compile() {
    awk -v kernel="$(uname -v)" -v whoami="$(whoami)" \
    '
        function num_scale(src_cur, src_max, dst_min, dst_max) {
            return dst_min + ((src_cur * (dst_max - dst_min)) / src_max)
        }

        function vert_print(v,    _color, _fontcolor) {
            _color =\
                num_scale(\
                    child2nice[v] + 20,
                    20 + 20,
                    COLORSCHEME_MIN,
                    COLORSCHEME_MAX\
                )
            _fontcolor = \
                _color == COLORSCHEME_MAX || _color == COLORSCHEME_MIN \
                ? sprintf("/%s/%d", COLORSCHEME, COLORSCHEME_MID) \
                : sprintf("/%s/%d", "greys9", 9)
            printf(\
                "\"%d\"\
                [ fontsize=8 \
                , style=filled \
                , fontname=Helvetica \
                , shape=ellipse \
                , label=\"%s\n%d\" \
                , color=\"/%s/%d\" \
                , fontcolor=\"%s\" \
                ];",
                v,
                child2cmd[v],
                v,
                COLORSCHEME,
                _color,
                _fontcolor\
            )
        }

        function edge_print(child,    _parent, _color, _colorscheme) {
            _parent = child2parent[child]
            _colorscheme = "greys9"
            _color = 3
            printf(\
                "\"%s\" -> \"%s\"\
                [ fontsize=8 \
                , fontname=Helvetica \
                , len=2.0 \
                , colorscheme=%s \
                , color=%d \
                ];\n",
                _parent,
                child,
                _colorscheme,
                _color\
            )
        }

        BEGIN {
            # Hot->Cold gradual colorschemes:
            # - rdbu11
            # - rdbu9
            # - rdbu8

            # Light->Dark gradual colorschemes:
            # - reds9
            # - blues9
            # - orrd9
            # - oranges9
            # - bupu9
            # - greys9

            COLORSCHEME_MIN = 1
            COLORSCHEME_MID = 5
            COLORSCHEME_MAX = 9
            COLORSCHEME = sprintf("rdbu%d", COLORSCHEME_MAX)

            child2cmd[0] = "swapper/sched"
        }

        NR > 1 {
            parent2child_count[$2]++
            max_children = \
                parent2child_count[$2] > max_children\
                ? parent2child_count[$2]\
                : max_children
            child2parent[$1]    = $2
            child2user_id[$1]   = $3
            child2user_name[$1] = $4
            child2nice[$1]      = $5
            child2cmd[$1]       = $6
            user_names[$4]      = 1
        }

        END {
            print "strict digraph G {";

            print "start=0;";
            print "fontsize=8;";
            print "fontname=Helvetica;";
            print "label=\"" kernel "\";";
            print "fontcolor=\"/greys9/9\";"

            ##### Vertices (clustered by user)
            for (user_name in user_names) {
                printf "subgraph \"cluster_%s\" {\n", user_name
                printf "label=\"%s\"\n", user_name
                for (c in child2parent)
                    if (child2user_name[c] == user_name)
                        vert_print(c)
                print "}"
            }

            ##### Vertices (without a user)
            for (c in child2cmd)
                if (!child2user_name[c])
                    vert_print(c)

            ##### Edges (across clusters)
            for (c in child2parent)
                edge_print(c)

            print "}";
        }
    '
}


procs() {
    if [ "$(uname)" = 'Linux' ]; then
        ps -eo 'pid,ppid,euid,euser,nice,comm'
    else
        ps -eco 'pid,ppid,euid,euser,nice,comm'
    fi
}


main() {
    case "$1" in
        '--help') usage
        ;;
        *) procs | grep "$1" | compile
        ;;
    esac
}


main "$1"