#!/usr/bin/env bash

# ================================

SCRIPT=$(readlink -m $(type -p $0))

SCRIPTDIR=$(dirname $SCRIPT)

testDataDir=$SCRIPTDIR/data

prismaDir=$SCRIPTDIR/data/Prisma

testDir=$SCRIPTDIR/

zipFile=dtitest_Siemens_SC.zip

# ================================

echo Downloading Prisma data

if [ ! -f $testDataDir/$zipFile ]

then

    wget http://people.cas.sc.edu/rorden/SW/dcm2niix/$zipFile -P $testDataDir

else

    echo $testDataDir/$zipFile exists, unzipping $zipFile

fi

rm -rf $prismaDir

unzip -o -d $testDataDir $testDataDir/$zipFile

pushd .

cd $prismaDir

dcm2niix -o . -f %z -z y .

IMAGELIST=$testDataDir/imagelist.txt

MASKLIST=$testDataDir/masklist.txt

B0LIST=$testDataDir/b0list.txt

WarpedB0LIST=$testDataDir/warped_b0list.txt

WarpedMASKLIST=$testDataDir/warped_masklist.txt

CASELIST=$testDataDir/caselist.txt

if [ -f $IMAGELIST ]

then

    rm $IMAGELIST $MASKLIST $B0LIST $WarpedB0LIST $WarpedMASKLIST $CASELIST

fi

echo 'Generating b0 image and corresponding brain mask for test data ...'

for i in `ls *.nii.gz`

do

    IFS=., read -r prefix _ _ <<< $i

    if [ -f $prefix.bval ] && [ -f $prefix.bvec ]

    then

        echo $i

        b0=${prefix}_b0.nii.gz

        mask=${prefix}

        fslroi $i $b0 0 1

        bet $b0 $mask -m -n       

        echo `pwd`/$i >> $IMAGELIST

        echo `pwd`/${prefix}_mask.nii.gz >> $MASKLIST

        echo `pwd`/${prefix}_b0.nii.gz >> $B0LIST

        echo `pwd`/${prefix}_b0-Warped.nii.gz >> $WarpedB0LIST

        echo `pwd`/${prefix}_b0-Warped-mask.nii.gz >> $WarpedMASKLIST

        echo $prefix >> $CASELIST

    fi

done

popd

# ================================

# prediction based on provided model

../pipeline/dwi_masking.py -i $IMAGELIST -nproc 5 -f ../model_folder/

../pipeline/dwi_masking.py -i $IMAGELIST -nproc 5 -f ../model_folder/ -filter scipy

../pipeline/dwi_masking.py -i $B0LIST -nproc 5 -f ../model_folder/

# ================================

# training

../src/registration.py -b0 $B0LIST -mask $MASKLIST -ref ../model_folder/IITmean_b0_256.nii.gz

../src/preprocess_b0.py -i $WarpedB0LIST

../src/preprocess_mask.py -i $WarpedMASKLIST

echo -e "

[COMMON]

save_model_dir = \"$testDataDir/model_folder_test\"

log_dir = \"$testDataDir/log_dir\"

[DATA]

data_dir = \"$prismaDir\"

train_data_file = \"sagittal-traindata-dwi.npy\"

train_label_file = \"sagittal-traindata-mask.npy\"

[TRAINING]

principal_axis = \"sagittal\"

learning_rate = 1e-3

train_batch_size = 4

validation_split = 0.2

num_epochs = 1

shuffle_data = \"True\"

" > $testDataDir/settings.ini

export COMPNET_CONFIG=$testDataDir/settings.ini

../src/train.py

sed -i "s/sagittal/coronal/g" $testDataDir/settings.ini

../src/train.py

sed -i "s/coronal/axial/g" $testDataDir/settings.ini

../src/train.py

# ================================

# prediction based on trained model

cp ../model_folder/IITmean_b0_256.nii.gz $testDataDir/model_folder_test/

../pipeline/dwi_masking.py -i $IMAGELIST -nproc 5 -f $testDataDir/model_folder_test/