#!/bin/bash

# gp_wrap.sh simplifies the Gnuplot interface and makes plotting
# data from CSV files easier.
# Copyright (C) 2021  Remington Furman

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.


# The command line arguments form a mini language that modify the
# current state of the plot parameters which are used when the -p
# (plot) argument is given to add a line to the plot.  By default, the
# plot parameters will plot the first column of the CSV file with no
# labels or title, and subsequent plots will plot the next column (-y
# will increment with each -p).

# Examples:
# ./gp_wrap.sh -f file.csv -p -p -p      # Plot columns 1, 2, 3 vs row number.
# ./gp_wrap.sh -f file.csv -h -p -p -p   # Same, but use headers from first row.
# ./gp_wrap.sh -h -p -p -p < file.csv    # Same, but read from stdin.
# ./gp_wrap.sh -f file.csv -x 2 -y 3 -p  # Plot col3 vs col2.


usage () {
    cat <<EOF
Usage:

    -f file    Plot data from File (default: stdin)
    -w file    Write output to file
    -F string  Use 'string' as Field separator (default: ",")
    -x n       Use column 'n' for X coordinates (default: use row number)
    -y n       Use column 'n' for Y coordinates (default: 1)
    -r range   Use Range for x and y axes (default: none)
    -a string  Label x axis (Abscissa) with 'string' (default: none)
    -o string  Label y axis (Ordinate) with 'string' (default: none)
    -t string  Label next plots with 'string' (default: filename)
    -T string  Set Title to 'string' (default: none)
    -s string  Style next plots with 'string' (default: lines)
    -k string  Place plot key at 'string' (default: top left)
    -i         Increment current column after each plot (default)
    -I         Don't Increment current column after each plot
    -m         Format Math in plot titles with TeX (default: off)
    -p         Make a Plot with the current parameters
    -l         List current parameters
    -h         Use first row (Header) as plot labels
    -c string  Add Custom Gnuplot commands to preamble (default: none)
    -u         Display Usage
    -d         Debug (print Gnuplot input)
EOF
}

# Set default parameters
STDIN=true                      # Read from stdin.
FILE=$(mktemp -u)
OUTPUT="./plot.png"
FIELD_SEP=","
X_COL=
Y_COL=1
RANGE=
X_LABEL=
Y_LABEL=
TITLE=
TITLE_MATH="set key noenhanced"
STYLE="with lines"
KEY_LOC="set key left top"
HEADER=
INCREMENT=true
CUSTOM=
PLOTS=
DEBUG=

# Read command line arguments
optstring="f:w:F:x:y:r:a:o:t:T:s:k:c:iImplhud"
while getopts ${optstring} arg; do
    case ${arg} in
        f)
            FILE="${OPTARG}"
            STDIN=false
            ;;
        w)
            OUTPUT="${OPTARG}"
            ;;
        F)
            FIELD_SEP="${OPTARG}"
            ;;
        x)
            # Note the trailing : which is necessary for specifying
            # the x and y coordinates.
            X_COL="${OPTARG}:"
            ;;
        y)
            Y_COL="${OPTARG}"
            ;;
        r)
            RANGE="${OPTARG}"
            ;;
        a)
            X_LABEL="${OPTARG}"
            ;;
        o)
            Y_LABEL="${OPTARG}"
            ;;
        t)
            PLOT_TITLE="title '${OPTARG}'"
            ;;
        T)
            TITLE="${OPTARG}"
            ;;
        s)
            STYLE="${OPTARG}"
            ;;
        k)
            KEY_LOC="set key ${OPTARG}"
            ;;
        i)
            INCREMENT=true
            ;;
        I)
            INCREMENT=false
            ;;
        m)
            TITLE_MATH=""
            ;;
        p)
            # Add a new plot command with current parameters.
            if [ "${STDIN}" = true -a \
                 -z "${HEADER}" -a    \
                 -z "${PLOT_TITLE}" ] ; then
                PLOT_TITLE="title 'stdin'"
            fi

            PLOTS="${PLOTS:+${PLOTS}, }" # Append comma if plots exist.
            PLOTS+="'${FILE}' using ${X_COL}${Y_COL} ${STYLE} ${PLOT_TITLE}"

            if "$INCREMENT" = true ; then
                # Move to next column.
                Y_COL=$((Y_COL+1))
            fi
            ;;
        l)
            echo FILE="$FILE"
            echo X_COL="$X_COL"
            echo Y_COL="$Y_COL"
            echo X_LABEL="$X_LABEL"
            echo Y_LABEL="$Y_LABEL"
            echo TITLE="$TITLE"
            echo HEADER="$HEADER"
            echo INCREMENT="$INCREMENT"
            echo CUSTOM="$CUSTOM"
            echo PLOTS="$PLOTS"
            ;;
        h)
            HEADER="set key autotitle columnheader"
            ;;
        c)
            CUSTOM+=$'\n'
            CUSTOM+="${OPTARG}"
            ;;
        u)
            usage
            exit
            ;;
        d)
            DEBUG=true
            ;;
    esac
done

if [ "${STDIN}" = true ] ; then
    # Save stdin to a file so that Gnuplot can read it for multiple
    # plots.
    cat /dev/stdin > "${FILE}"
    trap "rm ${FILE}" EXIT      # Make sure the temp file is deleted
                                # when this script exits.
    PLOT_TITLE="${PLOT_TITLE:-title 'stdin'}"
fi

if [ -z "${PLOTS}" ] ; then
    # Make at least one plot if none are requested.
    PLOTS="'${FILE}' using ${X_COL}${Y_COL} ${STYLE} ${PLOT_TITLE}"
fi

if [ "${DEBUG}" = true ] ; then
    COMMAND=cat
else
    COMMAND=gnuplot
fi

# Generate Gnuplot command and call Gnuplot
"${COMMAND}" <<-EOF
set terminal pngcairo
${CUSTOM}
set datafile separator ','
set output '${OUTPUT}'
set title '${TITLE}'
${TITLE_MATH}
${KEY_LOC}
${HEADER}
set xlabel '${X_LABEL}'
set ylabel '${Y_LABEL}'

plot ${RANGE} ${PLOTS}
EOF

if [ "${DEBUG}" = true ] ; then
    cat "${FILE}"
fi
