#! /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, _shape, _state, _size, _height, _label, _label_base, _label_ext) {
            _state = child2state[v]
            _style = "filled,solid"

            # -----------------------------------------------------------------
            # Sleeping/idling
            # -----------------------------------------------------------------
            # D uninterruptible sleep (usually IO)
            if (_state == "D") {
                _shape = "circle"
            # I Idle kernel thread
            } else if (_state == "I") {
                _shape = "circle"
            # S interruptible sleep (waiting for an event to complete)
            } else if (_state == "S") {
                _shape = "circle"
            # -----------------------------------------------------------------
            # Running
            # -----------------------------------------------------------------
            # R running or runnable (on run queue)
            } else if (_state == "R") {
                _shape = "rarrow"
            # -----------------------------------------------------------------
            # Stopped
            # -----------------------------------------------------------------
            # T stopped by job control signal
            } else if (_state == "T") {
                _shape = "square"
            # t stopped by debugger during the tracing
            } else if (_state == "t") {
                _shape = "square"
            # -----------------------------------------------------------------
            # Dead
            # -----------------------------------------------------------------
            # Z defunct ("zombie") process, terminated but not reaped by its parent
            } else if (_state == "Z") {
                _shape = "Msquare"
                _style = "solid"
            # -----------------------------------------------------------------
            # UNKNOWN STATE
            # -----------------------------------------------------------------
            } else {
                _shape = "doublecircle"
            }

            _color =\
                num_scale(\
                    child2cpu[v],
                    max_cpu,
                    VERT_COLORSCHEME_MAX,
                    VERT_COLORSCHEME_MIN\
                )
            _size =\
                num_scale(\
                    child2mem[v],
                    max_mem,
                    1,
                    4\
                ) / 4
            _height = _size
            _width  = _size
            _fontcolor = \
                _color == VERT_COLORSCHEME_MAX || _color == VERT_COLORSCHEME_MIN \
                ? sprintf("/%s/%d", VERT_COLORSCHEME, VERT_COLORSCHEME_MID) \
                : sprintf("/%s/%d", "greys9", 9)
            _fontcolor = \
                _size < 0.5 \
                ? sprintf("/%s/%d", "greys9", 9) \
                : _fontcolor
            _label_base = \
                sprintf("%s\n%d", child2comm[v], v)
            _label_ext = \
                _size >= 0.5 \
                ? sprintf("\ncpu: %.1f%%\nmem: %.1f%%", child2cpu[v], child2mem[v]) \
                : ""
            _label = _label_base _label_ext
            printf(\
                "\"%d\"\
                [ fontsize=8 \
                , fixedsize=true \
                , height=%f \
                , width=%f \
                , border=1 \
                , style=\"%s\" \
                , fontname=Helvetica \
                , label=\"%s\" \
                , shape=\"%s\" \
                , fillcolor=\"/%s/%d\" \
                , fontcolor=\"%s\" \
                ];",
                v,
                _height,
                _width,
                _style,
                _label,
                _shape,
                VERT_COLORSCHEME,
                _color,
                _fontcolor\
            )
        }

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

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

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

            VERT_COLORSCHEME_MIN = 1
            VERT_COLORSCHEME_MID = 4
            VERT_COLORSCHEME_MAX = 8
            VERT_COLORSCHEME = "rdylgn10"

            EDGE_COLOR = "/ylorbr9/3"

            child2comm[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
            child2state[$1]     = $6
            child2cpu[$1]       = $7
            child2mem[$1]       = $8
            child2comm[$1]      = $9
            user_names[$4]      = 1
            max_cpu = $7 > max_cpu ? $7 : max_cpu
            max_mem = $8 > max_mem ? $8 : max_mem
        }

        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 child2comm)
                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,s,%cpu,%mem,comm'
    else
        ps -eco 'pid,ppid,euid,euser,nice,s,%cpu,%mem,comm'
    fi
}


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


main "$1"