#!/bin/tcsh

# help?

if( $#argv < 2 )then
  echo
  echo '-----------------------------------------------------------------'
  echo
  echo 'This script will make a JPEG image showing the edges of an'
  echo 'EPI dataset overlay-ed on an anatomical dataset.  The purpose is'
  echo 'to let the user (you) judge the quality of the 3D registration.'
  echo
  echo 'Three images from each of the coronal, axial, and sagittal'
  echo 'AFNI image viewers are used, laid out in a 3x3 grid.'
  echo
  echo '@snapshot_volreg works by running the AFNI GUI inside a "virtual"'
  echo 'X11 display server program named "Xvfb", and saving images from'
  echo 'that copy of AFNI.  The script also uses programs from the netpbm11'
  echo 'software library to put the saved images together into a pleasing'
  echo 'layout.  If the script cannot find the netpbm11 software, it will'
  echo 'not run :('
  echo
  echo '-----------------------------------------------------------------'
  echo 'Usage: @snapshot_volreg ANATdataset EPIdataset [jname] [xdisplay]'
  echo
  echo 'Sample (from an afni_proc.py results directory):'
  echo
  echo '  @snapshot_volreg anat_final.sub-10506+tlrc      \'
  echo '                   pb02.sub-10506.r01.volreg+tlrc sub-10506'
  echo
  echo 'The output file from this example is "sub-10506.jpg".'
  echo '-----------------------------------------------------------------'
  echo
  echo 'Do NOT put a sub-brick index (such as "[0]") on the EPIdataset'
  echo 'name -- the script will automatically only use the "[0]" volume.'
  echo
  echo '(( Although the original use was for visualizing how well EPI ))'
  echo '(( and anatomical datasets were aligned by align_epi_anat.py, ))'
  echo '(( it is also useful to see how well 3dQwarp aligned an       ))'
  echo '(( anatomical dataset to a template dataset.                  ))'
  echo
  echo 'The optional third argument is the name of the output JPEG'
  echo 'file -- if it does not end in ".jpg", that suffix will be added.'
  echo 'If you do NOT supply a 3rd argument, the script will invent a name:'
  echo 'it is probably better for you to supply a 3rd argument.'
  echo
  echo 'The fourth (and very optional) argument is the display number'
  echo 'of an ALREADY RUNNING copy of Xvfb, as in'
  echo '  Xvfb :88 -screen 0 1024x768x24 &'
  echo 'If you do NOT supply this number (88 in the example), then'
  echo 'the script will start its own Xvfb (on a display of its choosing),'
  echo 'use it once, and then stop it. If you are going to run this script'
  echo 'many times in a row, starting and stopping your own Xvfb'
  echo 'instance will speed things up a little. Normally, you do not'
  echo 'need to use this 4th argument.'
  echo
  echo '-----------------------------------------------------------------'
  echo
  echo 'The edges from a typical EPI dataset are usually broken up and'
  echo 'do not completely outline sulci, ventricles, etc.  In judging'
  echo 'the quality of alignment, I usually start by looking at the'
  echo 'outlines of the large lateral ventricles -- if those are very'
  echo 'wrong, the alignment is not good.  After that, I look at the'
  echo 'sulci in the superior part of the brain -- if the EPI edges'
  echo 'there seem to be mostly aligned with the sulci, then I am'
  echo 'usually happy.  The base of the brain, where lots of EPI'
  echo 'dropout happens, often does not not show good edge alignment'
  echo 'even when the rest of the brain alignment looks good.'
  echo
  echo '-----------------------------------------------------------------'
  echo
  echo 'If this script crashes, then it might leave behind files with'
  echo 'names that start with "zzerm".  Delete these files.'
  echo 'It is also possible that the Xvfb program will still be running'
  echo 'if this script crashes.  A command such as that below can'
  echo 'be used to see if you have any stray Xvfb programs running:'
  echo
  echo '  ps X | grep Xvfb | grep -v grep'
  echo
  echo 'If there are any such programs, the command below can be used'
  echo 'to kill all of them:'
  echo
  echo '  killall Xvfb'
  echo
  echo '-------------- Author: The Madd Allineator ----------------------'
  echo
  exit 0
endif

# control variables

setenv AFNI_DONT_LOGFILE YES
unset noclobber

# check if all the needed auxiliary programs exist

set nerr = 0
set errm = "** ERROR:"

set plist = ( Xvfb djpeg cjpeg pnmcat pbmtext pnmscale pbmtopgm )
foreach pppp ( $plist )
  set wwww = `which $pppp`
  if( $status != 0 )then
    @ nerr++
    set errm = "$errm $pppp"
  endif
end

# we can use either pamcomp or pnmcomp, so only need to
# find one of the twain

set pcprog = pamcomp
set wwww = `which $pcprog`
if( $status != 0 )then
  set pcprog = pnmcomp
  set wwww = `which $pcprog`
  if( $status != 0 )then
    @ nerr++
    set errm = "$errm (pnmcomp OR pamcomp)"
  endif
endif

if( $nerr > 0 )then
  echo "$errm -- not found in path -- @snapshot_volreg fails"
  echo "** WARNING: this script cannot run without installing package netpbm11"
  exit 1
endif

# set the prefix for the anat and epi datasets

set adset = $argv[1]
set abase = $adset:t
set anat  = `basename $abase .gz`
set anat  = `basename $anat  .nii`
set anat  = `basename $anat  .HEAD`
set anat  = `basename $anat  .BRIK`
set anat  = `basename $anat  +tlrc`
set anat  = `basename $anat  +orig`
set anat  = `basename $anat  +acpc`
set anat  = `basename $anat  +tlrc.`
set anat  = `basename $anat  +orig.`
set anat  = `basename $anat  +acpc.`
if( $abase == $anat.nii.gz )then
  set asuff = .nii.gz
else if( $abase == $anat.nii )then
  set asuff = .nii
else
  set asuff = ""
endif

set edset = $argv[2]
set epi   = `basename $edset .gz`
set epi   = `basename $epi   .nii`
set epi   = `basename $epi   .HEAD`
set epi   = `basename $epi   .BRIK`
set epi   = `basename $epi   +orig`
set epi   = `basename $epi   +acpc`
set epi   = `basename $epi   +tlrc`
set epi   = `basename $epi   +tlrc.`
set anat  = `basename $anat  +tlrc.`
set anat  = `basename $anat  +tlrc`
set anat  = `basename $anat  +orig.`
set anat  = `basename $anat  +acpc.`

# set output image prefix

if( $#argv > 2 )then
  set jnam = $argv[3]
  set jnam = `basename "$jnam" .jpg`
  set jnam = `basename "$jnam" .JPG`
  set jnam = `basename "$jnam" .jpeg`
  set jnam = `basename "$jnam" .JPEG`
else
  set jnam = $anat.$epi
endif

# Are we re-using Xvfb?

unset xdisplay
unset killX
if( $#argv > 3 )then
  set xdisplay = $argv[4]
  if( ! -e /tmp/.X${xdisplay}-lock )then
    echo "** ERROR: it doesn't look like Xvfb is running with display $xdisplay"
    exit 1
  endif
endif

set exad = `3dinfo -exists $adset`
set exep = `3dinfo -exists $edset`

if( $exad == 0 || $exep == 0 )then
  if( $exad == 0 ) echo "** ERROR: @snapshot_volreg can't find $adset"
  if( $exep == 0 ) echo "** ERROR: @snapshot_volreg can't find $edset"
  exit 1
endif

# set some AFNI GUI things

setenv AFNI_NOSPLASH          YES
setenv AFNI_SPLASH_MELT       NO
setenv AFNI_LEFT_IS_LEFT      YES
setenv AFNI_IMAGE_LABEL_MODE  5
setenv AFNI_IMAGE_LABEL_SIZE  2
setenv AFNI_ENVIRON_WARNINGS  NO
setenv AFNI_COMPRESSOR        NONE

# start the X virtual frame buffer on display #xdisplay

set ranval = `count -dig 1 1 999999 R1`

if( $?xdisplay == 0 )then
  set killX     = 1
  set ntry      = 1
  set Xnotfound = 1
  while( $Xnotfound )
    set xdisplay = `count -dig 1 3 999 R1`
    if( -e /tmp/.X${xdisplay}-lock ) continue
    echo " -- trying to start Xvfb :${xdisplay}"
    Xvfb :${xdisplay} -screen 0 1024x768x24 >& /dev/null &
    sleep 1
    jobs > zzerm.$ranval.txt
    grep -q Xvfb zzerm.$ranval.txt
    set Xnotfound = $status
    \rm -f zzerm.$ranval.txt
    if( $Xnotfound == 0 )break ;
    @ ntry++
    if( $ntry > 99 )then
      echo "** ERROR: can't start Xvfb -- exiting"
      exit 1
    endif
  end
endif

setenv DISPLAY :${xdisplay}

# quasi-random temp filename prefix

set zpref = zzerm.X${xdisplay}-${ranval}

# crop the input anat to a standard size

set cdset = ${zpref}.acrop.nii

3dAutobox -npad 17 -prefix $cdset $adset

# resample the EPI to the anat grid

3dAllineate -input  ${edset}'[0]'     \
            -master ${cdset}          \
            -prefix ${zpref}.epiR.nii \
            -1Dparam_apply IDENTITY   \
            -final cubic

# find edges in the EPI

3dedge3 -input ${zpref}.epiR.nii -prefix ${zpref}.epiE.nii

# get the EPI automask and apply it to the edgized EPI

3dAutomask -q -prefix ${zpref}.epiM.nii -clfrac 0.333 -dilate 5 ${zpref}.epiR.nii
3dcalc -a ${zpref}.epiE.nii -b ${zpref}.epiM.nii -expr 'a*b' -prefix ${zpref}.epiEM.nii

# get the lower and upper thresholds for the edgized EPI

set epp = `3dBrickStat -non-zero -percentile 20 60 80 ${zpref}.epiEM.nii`
set eth = $epp[2]
set emx = $epp[4]

# run AFNI to make the 3 overlay images

set anatCM = ( `3dCM $cdset` )

set anatNN = `3dinfo -nijk $cdset`

set astep = `ccalc -int "cbrt($anatNN)/6.111"`

if( $astep < 2 ) set astep = 2

afni -noplugins -no_detach                                       \
     -com "SWITCH_UNDERLAY ${cdset}"                             \
     -com "SWITCH_OVERLAY ${zpref}.epiEM.nii"                    \
     -com "SET_DICOM_XYZ $anatCM"                                \
     -com "SET_PBAR_ALL +99 1 Plasma"                            \
     -com "SET_FUNC_RANGE $emx"                                  \
     -com "SET_THRESHNEW $eth *"                                 \
     -com "SEE_OVERLAY +"                                        \
     -com "SET_XHAIRS OFF"                                       \
     -com "OPEN_WINDOW sagittalimage opacity=6 mont=3x1:$astep"  \
     -com "OPEN_WINDOW axialimage opacity=6 mont=3x1:$astep"     \
     -com "OPEN_WINDOW coronalimage opacity=6 mont=3x1:$astep"   \
     -com "SAVE_JPEG sagittalimage ${zpref}.sag.jpg blowup=2"    \
     -com "SAVE_JPEG coronalimage  ${zpref}.cor.jpg blowup=2"    \
     -com "SAVE_JPEG axialimage    ${zpref}.axi.jpg blowup=2"    \
     -com "QUITT"                                                \
     ${cdset} ${zpref}.epiEM.nii

# convert the output JPEGs to PNMs for manipulation

djpeg ${zpref}.sag.jpg > ${zpref}.sag.pnm
djpeg ${zpref}.cor.jpg > ${zpref}.cor.pnm
djpeg ${zpref}.axi.jpg > ${zpref}.axi.pnm

# cat them together, make a text label, overlay it, make output JPEG

pnmcat -tb -jcenter -black ${zpref}.sag.pnm ${zpref}.axi.pnm ${zpref}.cor.pnm           > ${zpref}.pnm
pbmtext -builtin fixed "$jnam"   | pnmcrop | pbmtopgm 1 1 | pnmscale 2                  > ${zpref}.t1.pgm
pbmtext -builtin fixed "<- Left"           | pbmtopgm 1 1                               > ${zpref}.t2.pgm
$pcprog -align=right -valign=bottom ${zpref}.t1.pgm ${zpref}.pnm                        > ${zpref}.t3.pnm
$pcprog -align=right -valign=top    ${zpref}.t2.pgm ${zpref}.t3.pnm | cjpeg -quality 95 > "$jnam.jpg"

# delete the trash data

\rm -f ${zpref}.*

echo "@snapshot_volreg output image = $jnam.jpg"

# stop Xvfb if we started it ourselves

if( $?killX ) kill %1

exit 0