#!/bin/bash

# This script uses imagemagick to create grids of images and automatically adds captions to them.
# There are various grid styles which can be made outlined below

# Style 1 - n by m grid with caption for every row
# for a 3 by 2 grid it will look like:
#       ------       ------
#      |      |     |      |
#      |      |     |      |
#       ------       ------
#       ------       ------
#      |      |     |      |
#      |      |     |      |
#       ------       ------
#       ------       ------
#      |      |     |      |
#      |      |     |      |
#       ------       ------
#       

####
#Help
####
Help(){
    echo "fig_grid combines images into a grid, useful for creating complex figures that can be moved between pandoc/latex and other sources such as word or powerpoint."
    echo
    echo "Syntax: fig_grid [-h|-o filename|-c captionfile | -s size | -b | -p padding | -n] style inputs"
    echo "options:"
    echo "-h: Prints this help"
    echo "-o: Specify output file as filename"
    echo "-c: Path to caption file. Caption file must have specific syntax, please check documentation"
    echo "-s: Font size for captions in pixels"
    echo "-b: Add thin black border around images"
    echo "-p: Add padding around images before combining them "
    echo "-n: Prevents resizing the images"
    echo "style: Represents the style of the output figure. Currently only accepts:"
    echo "          - nxm (example: 3x2): Arranges a grid into n rows and m columns populating the rows first"
    echo "inputs: Input files to be processed"
}

####
#Error
####
function fail {
    printf 'Error: %s\n' "$1" >&2
    exit "${2-1}"
}


####
#Add caption to image
# First argument should be image file path
# Second argument should be the message
# Third argument is the padding to add to the image 
# Fourth argument is the size of the font
####
function add_caption {
    file=$1
    msg=$2
    pad=$3
    size=$4
    width=$(identify -format '%w' $file)
    height=$(identify -format '%h' $file)
    (( height = height+pad ))
    convert -extent x"$height" $file $file
    captmp="$RANDOM$RANDOM$RANDOM$RANDOM.png"
    convert -background " #FFF " -fill black -font Helvetica -size x$size label:"$msg" $captmp
    convert -append -gravity center $file $captmp $file
    rm $captmp
}

args=()
argnum=0
csize=0
cpad=0
border=0
pad=0
resize=1
output="output.png"
while [ $# -gt 0 ]; do
    unset OPTIND
    unset OPTARG
    #Parse Flags
    while getopts ':ho:c:s:bp:n' options; do 
        case $options in
            h)
                Help
                break;;
            o)
                output=$OPTARG;;
            c)
                captionfile=$OPTARG;;
            s)
                csize=$OPTARG;;
            b)
                border=2;;
            p)
                pad=$OPTARG;;
            n)
                resize=0;;
        esac
    done
    shift $((OPTIND-1))
    args+=("$1")
     (( argnum++ ))
     shift
 done

if [ $argnum -eq 0 ]; then
    Help
    exit
fi

#Specify tmp file prefix
tmp=$(echo $RANDOM$RANDOM$RANDOM)
#First calculate the number of inputs to be expected, then check to make sure enough images are provided
style=${args[0]}
calc_style="${style/x/\*}"
expfiles=$(echo $calc_style | bc)

if [ $expfiles -ne $(( argnum -1 )) ]; then
    fail "expected $expfiles files but only got $(( argnum - 1 )) files"
fi

#Now read the caption file into an array if it exists
if [ "$captionfile" != "" ]; then
    mapfile -t captions < $captionfile
fi

#Finally generate the correct images. By default this will size all images to the size of the largest image to create regular spacings. It does this by scaling and the extent command

#First find max image size
max_x=0
max_y=0

for i in $(seq 1 $((argnum-1)) ); do
    if [ ! -f ${args[$i]} ]; then
        fail "File ${args[$i]} doesn't exist"
    fi
    x=$(identify -format '%w' ${args[$i]})
    y=$(identify -format '%h' ${args[$i]})
    if [ "$x" -gt "$max_x" ]; then
        max_x=$x
    fi
    if [ "$y" -gt "$max_y" ]; then
        max_y=$y
    fi
done


#Get caption size 

if [ $csize -eq 0 ]; then
    csize=$(echo "scale=0; $max_y*10/100" | bc)
fi

#Get padding if necessary
if [ $pad -eq 0 ]; then 
    pad=$(echo "scale=0; $max_y*5/100" | bc)
    cpad=$(echo "scale=0; $csize/2" | bc)
fi
#Now resize all of the images
if [ $resize -eq 1 ]; then 
    echo "Resizing all images to $max_x $max_y"
    for i in $(seq 1 $((argnum-1)) ); do
        info=$(identify ${args[$i]} ) || fail "Failed to identify image ${args[$i]}"
        x=$(identify -format '%w' ${args[$i]})
        y=$(identify -format '%h' ${args[$i]})

        diffx=$(echo " $max_x - $x" | bc )
        diffy=$(echo "$max_y - $y "  | bc )

        if [ $diffx -eq 0 ]; then 
            cp ${args[$i]} $tmp$i.png 
        elif [ $diffy -eq 0 ]; then 
            cp ${args[$i]} $tmp$i.png 
        elif [ $diffy -gt $diffx ]; then
            convert -scale "$max_x"x ${args[$i]} $tmp$i.png || fail "Failed to scale image"
        elif [ $diffx -gt $diffy ]; then
            convert -scale x"$max_y" ${args[$i]} $tmp$i.png || fail "Failed to scale image"
        fi
        convert -extent "$max_x"x"$max_y" -gravity center -border "$border"x$border -bordercolor black $tmp$i.png $tmp$i.png || fail "Failed to add padding to image"
    done
else
    for i in $(seq 1 $((argnum-1)) ); do
        cp ${args[$i]} $tmp$i.png
    done
fi

cap_num=1

if [ "${captions[0]}" = "each" ]; then
    for i in $(seq 1 $((argnum-1)) ); do
        add_caption $tmp$i.png "${captions[$cap_num]}" $cpad $csize || fail "Couldn't add caption"
        (( cap_num++ ))
    done
fi

k=$(( i+1 ))
#Now build the actually images based on style
if [[ "$style" == *"x"* ]]; then
    vertappend=""
    rows=${style%x*}
    columns=${style#*x}
    echo "Combining images into a $rows by $columns grid"
    l=1
    for i in $(seq 1 $rows); do 
        horizappend=""
        for j in $(seq 1 $columns); do 
            horizappend="$horizappend ( -border $pad -bordercolor white $tmp$l.png )"
            (( l++ ))
        done
        convert $horizappend +append $tmp$k.png
        if [ "${captions[0]}" = "rows" ]; then 
            add_caption $tmp$k.png "${captions[$cap_num]}" $cpad $csize || fail "Couldn't add caption"
            (( cap_num++ ))
        fi 
        vertappend="$vertappend $tmp$k.png"
        (( k++ ))
    done
    convert -gravity center $vertappend -append $output
fi

rm $tmp*.png