{

local cur prev words cword

_init_completion -s || return

local prototxts='@(prototxt)'

local caffemodels='@(caffemodel,binaryproto)'

local solverstates='@(solverstate)'

local caffefiles='@(prototxt|caffemodel|solverstate)'

local flags='-gpu -iterations -model -snapshot -solver -weights -sighup_effect -sigint_effect -level -stage -phase'

if [[ $cword -eq 1 ]]; then

COMPREPLY=( $( compgen -W 'train test time device_query' -- "$cur" ) )

return 0

fi

if [[ $cword -eq 2 ]]; then

case ${words[1]} in

train|test|device_query|time)

COMPREPLY=( $( compgen -W "$flags" -- "$cur") )

return 0

;;

*)

return 0

;;

esac

fi

case $prev in

-gpu|-iterations|-version|-level|-stage)

return 0

;;

-solver|-model)

_filedir $prototxts

return 0

;;

-weights)

_filedir $caffemodels

return 0

;;

-snapshot)

_filedir $solverstates

return 0

;;

-sighup_effect|-sigint_effect)

COMPREPLY=( $( compgen -W 'snapshot stop none' -- "$cur") )

return 0

;;

-phase)

COMPREPLY=( $( compgen -W 'TRAIN TEST' -- "$cur") )

return 0

;;

*)

COMPREPLY=( $( compgen -W "$flags" -- "$cur") )

return 0

;;

esac

# file completion on relevant files

_filedir "$caffefiles"

return 0