#!/bin/bash
# Preprocessing Script (SPM Based - Standard PiB Method)
# Based on Centiloid Project Guidelines

set -e

if [ -z "$1" ]; then
    echo "Usage: $0 <subject_id>"
    echo "Example: $0 02"
    exit 1
fi

subject_id=$1
subject="sub-${subject_id}"

echo "Preprocessing Subject: ${subject} (SPM Based)"

if ! ml fsl/6.0.7.8 2>/dev/null; then
    ml fsl
fi

if ! ml spm12/r7771 2>/dev/null; then
    echo "Specific SPM12 module not found, trying default..."
    ml spm12
else
    echo "Loaded SPM12/r7771 module."
fi

SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"
PROJECT_ROOT="$(dirname "$(dirname "$SCRIPT_DIR")")"

cd ~/Desktop/derivatives/data/${subject}/pet/
pet_dir=$(pwd)

FRAMING_CSV="${PROJECT_ROOT}/data/references/framing_info.csv"
frame_info=$(grep "AD${subject_id}," "$FRAMING_CSV")
range=$(echo "$frame_info" | cut -d',' -f3)
start_frame=$(echo "$range" | cut -d'-' -f1)
end_frame=$(echo "$range" | cut -d'-' -f2)

start_idx=$((start_frame - 1))
num_frames=$((end_frame - start_frame + 1))

echo "Extracting frames for realignment..."
mkdir -p "frames"
fslsplit "${subject}_pet.nii.gz" "frames/vol_"

frame_list=""
for ((i=start_frame; i<=end_frame; i++)); do
    padded_num=$(printf "%04d" $i)
    gunzip -f "frames/vol_${padded_num}.nii.gz"
    frame_list="$frame_list '$(pwd)/frames/vol_${padded_num}.nii',"
done

frame_list="${frame_list%,}"

echo "Generating SPM Realignment Job..."
job_file="${pet_dir}/realign_job.m"

files_matlab=""
for f in $(ls "${pet_dir}/frames"/vol_*.nii | sort); do
    files_matlab="${files_matlab}'${f},1';"$'\n'
done

cat <<EOF > "$job_file"
% SPM Realignment Job
spm('defaults', 'PET');
spm_jobman('initcfg');

matlabbatch{1}.spm.spatial.realign.estwrite.data = {
    {
$files_matlab
    }
}';
matlabbatch{1}.spm.spatial.realign.estwrite.eoptions.quality = 0.9;
matlabbatch{1}.spm.spatial.realign.estwrite.eoptions.sep = 4;
matlabbatch{1}.spm.spatial.realign.estwrite.eoptions.fwhm = 5;
matlabbatch{1}.spm.spatial.realign.estwrite.eoptions.rtm = 1;
matlabbatch{1}.spm.spatial.realign.estwrite.eoptions.interp = 2;
matlabbatch{1}.spm.spatial.realign.estwrite.eoptions.wrap = [0 0 0];
matlabbatch{1}.spm.spatial.realign.estwrite.eoptions.weight = '';
matlabbatch{1}.spm.spatial.realign.estwrite.roptions.which = [2 1];
matlabbatch{1}.spm.spatial.realign.estwrite.roptions.interp = 4;
matlabbatch{1}.spm.spatial.realign.estwrite.roptions.wrap = [0 0 0];
matlabbatch{1}.spm.spatial.realign.estwrite.roptions.mask = 1;
matlabbatch{1}.spm.spatial.realign.estwrite.roptions.prefix = 'r';

spm_jobman('run', matlabbatch);
exit;
EOF

echo "Running SPM Batch..."
if command -v spm12 &> /dev/null; then
    spm12 batch "$job_file"
elif command -v matlab &> /dev/null; then
    matlab -nodisplay -nosplash -nodesktop -r "run('$job_file');"
else
    echo "ERROR: Neither spm12 nor matlab command found."
    exit 1
fi

echo "Locating Mean Realigned Image..."
mean_realigned_image=$(find "${pet_dir}/frames" -name "meanvol_*.nii" | head -n 1)

if [ -z "$mean_realigned_image" ]; then
    echo "ERROR: Mean realigned image not found. Motion correction failed."
    exit 1
fi

echo "Motion Correction Complete."
echo "Mean Realigned Image: $mean_realigned_image"

mv "$mean_realigned_image" "${subject}_pet_avg.nii"
rm -rf "frames" "$runner_script"

fslreorient2std "${subject}_pet_avg.nii" "${subject}_pet_avg_reoriented.nii.gz"
mv "${subject}_pet_avg_reoriented.nii.gz" "${subject}_pet_avg.nii.gz"
gunzip -f "${subject}_pet_avg.nii.gz"

fslreorient2std ${subject}_pet_avg.nii.gz ${subject}_pet_avg_reoriented.nii.gz
mv ${subject}_pet_avg_reoriented.nii.gz ${subject}_pet_avg.nii.gz
gunzip -f ${subject}_pet_avg.nii.gz

cd ~/Desktop/derivatives/data/${subject}/anat/
cp ${subject}_T1w.nii.gz ${subject}_T1w_preproc.nii.gz

fslreorient2std ${subject}_T1w_preproc.nii.gz ${subject}_T1w_preproc.reorient.nii.gz

mv ${subject}_T1w_preproc.reorient.nii.gz ${subject}_T1w.reoriented.nii.gz

gunzip -f ${subject}_T1w.reoriented.nii.gz

cd ~/Desktop/derivatives/data/${subject}/

mri_file="$(pwd)/anat/${subject}_T1w.reoriented.nii"
pet_file="$(pwd)/pet/${subject}_pet_avg.nii"
spm_script="spm_preprocessing_job.m"

cat > "$spm_script" <<EOF
% SPM Batch Script for Centiloid Preprocessing

spm('defaults', 'FMRI');
spm_jobman('initcfg');

% Auto-Reorient: Reset Origin to Center of Image
files_to_reorient = {'$mri_file', '$pet_file'};
for i = 1:numel(files_to_reorient)
    fname = files_to_reorient{i};
    v = spm_vol(fname);
    M = v.mat;
    com = (v.dim(1:3)' + 1) / 2;
    R = M(1:3, 1:3);
    T_new = -R * com;
    M_new = [R T_new; 0 0 0 1];
    spm_get_space(fname, M_new);
end

matlabbatch = {};

% Coregister: Estimate (Reference: MRI, Source: PET)
matlabbatch{1}.spm.spatial.coreg.estimate.ref = {'$mri_file,1'};
matlabbatch{1}.spm.spatial.coreg.estimate.source = {'$pet_file,1'};
matlabbatch{1}.spm.spatial.coreg.estimate.other = {''};
matlabbatch{1}.spm.spatial.coreg.estimate.eoptions.cost_fun = 'nmi';
matlabbatch{1}.spm.spatial.coreg.estimate.eoptions.sep = [4 2];
matlabbatch{1}.spm.spatial.coreg.estimate.eoptions.tol = [0.02 0.02 0.02 0.001 0.001 0.001 0.01 0.01 0.01 0.001 0.001 0.001];
matlabbatch{1}.spm.spatial.coreg.estimate.eoptions.fwhm = [7 7];

% Segment (Unified Segmentation)
matlabbatch{2}.spm.spatial.preproc.channel.vols = {'$mri_file,1'};
matlabbatch{2}.spm.spatial.preproc.channel.biasreg = 0.001;
matlabbatch{2}.spm.spatial.preproc.channel.biasfwhm = 60;
matlabbatch{2}.spm.spatial.preproc.channel.write = [0 1];
matlabbatch{2}.spm.spatial.preproc.tissue(1).tpm = {fullfile(spm('Dir'),'tpm','TPM.nii,1')};
matlabbatch{2}.spm.spatial.preproc.tissue(1).ngaus = 1;
matlabbatch{2}.spm.spatial.preproc.tissue(1).native = [1 0];
matlabbatch{2}.spm.spatial.preproc.tissue(1).warped = [0 0];
matlabbatch{2}.spm.spatial.preproc.tissue(2).tpm = {fullfile(spm('Dir'),'tpm','TPM.nii,2')};
matlabbatch{2}.spm.spatial.preproc.tissue(2).ngaus = 1;
matlabbatch{2}.spm.spatial.preproc.tissue(2).native = [1 0];
matlabbatch{2}.spm.spatial.preproc.tissue(2).warped = [0 0];
matlabbatch{2}.spm.spatial.preproc.tissue(3).tpm = {fullfile(spm('Dir'),'tpm','TPM.nii,3')};
matlabbatch{2}.spm.spatial.preproc.tissue(3).ngaus = 2;
matlabbatch{2}.spm.spatial.preproc.tissue(3).native = [1 0];
matlabbatch{2}.spm.spatial.preproc.tissue(3).warped = [0 0];
matlabbatch{2}.spm.spatial.preproc.tissue(4).tpm = {fullfile(spm('Dir'),'tpm','TPM.nii,4')};
matlabbatch{2}.spm.spatial.preproc.tissue(4).ngaus = 3;
matlabbatch{2}.spm.spatial.preproc.tissue(4).native = [1 0];
matlabbatch{2}.spm.spatial.preproc.tissue(4).warped = [0 0];
matlabbatch{2}.spm.spatial.preproc.tissue(5).tpm = {fullfile(spm('Dir'),'tpm','TPM.nii,5')};
matlabbatch{2}.spm.spatial.preproc.tissue(5).ngaus = 4;
matlabbatch{2}.spm.spatial.preproc.tissue(5).native = [1 0];
matlabbatch{2}.spm.spatial.preproc.tissue(5).warped = [0 0];
matlabbatch{2}.spm.spatial.preproc.tissue(6).tpm = {fullfile(spm('Dir'),'tpm','TPM.nii,6')};
matlabbatch{2}.spm.spatial.preproc.tissue(6).ngaus = 2;
matlabbatch{2}.spm.spatial.preproc.tissue(6).native = [0 0];
matlabbatch{2}.spm.spatial.preproc.tissue(6).warped = [0 0];
matlabbatch{2}.spm.spatial.preproc.warp.mrf = 1;
matlabbatch{2}.spm.spatial.preproc.warp.cleanup = 1;
matlabbatch{2}.spm.spatial.preproc.warp.reg = [0 0.001 0.5 0.05 0.2];
matlabbatch{2}.spm.spatial.preproc.warp.affreg = 'mni';
matlabbatch{2}.spm.spatial.preproc.warp.fwhm = 0;
matlabbatch{2}.spm.spatial.preproc.warp.samp = 3;
matlabbatch{2}.spm.spatial.preproc.warp.write = [1 1];

% Normalise: Write (Apply to PET)
matlabbatch{3}.spm.spatial.normalise.write.subj.def(1) = cfg_dep('Segment: Forward Deformations', substruct('.','val', '{}',{2}, '.','val', '{}',{1}, '.','val', '{}',{1}), substruct('.','fordef', '()',{':'}));
matlabbatch{3}.spm.spatial.normalise.write.subj.resample = {'$pet_file,1'};
matlabbatch{3}.spm.spatial.normalise.write.woptions.bb = [-90 -126 -72; 91 91 109];
matlabbatch{3}.spm.spatial.normalise.write.woptions.vox = [2 2 2];
matlabbatch{3}.spm.spatial.normalise.write.woptions.interp = 4;
matlabbatch{3}.spm.spatial.normalise.write.woptions.prefix = 'w';

% Smooth (8mm FWHM)
matlabbatch{4}.spm.spatial.smooth.data(1) = cfg_dep('Normalise: Write: Resampled Images (Subj 1)', substruct('.','val', '{}',{3}, '.','val', '{}',{1}, '.','val', '{}',{1}, '.','val', '{}',{1}), substruct('()',{1}, '.','files'));
matlabbatch{4}.spm.spatial.smooth.fwhm = [8 8 8];
matlabbatch{4}.spm.spatial.smooth.dtype = 0;
matlabbatch{4}.spm.spatial.smooth.im = 0;
matlabbatch{4}.spm.spatial.smooth.prefix = 's';


EOF

echo "Running SPM Batch..."
if command -v matlab &> /dev/null; then
    matlab -nodisplay -nosplash -nodesktop -r "run('$spm_script'); exit"
elif command -v spm12 &> /dev/null; then
    spm12 batch "$spm_script"
else
    echo "WARNING: Neither MATLAB nor SPM12 command found. The batch script '$spm_script' has been generated but not run."
    echo "Please run it manually in your SPM environment."
fi

if [ -f "pet/sw${subject}_pet_avg.nii" ]; then
    gzip -c "pet/sw${subject}_pet_avg.nii" > "pet/${subject}_pet_to_MNI_smoothed.nii.gz"
    echo "SPM Processing Complete. Output: pet/${subject}_pet_to_MNI_smoothed.nii.gz"
    
    echo "Cleaning up intermediate files..."
    rm -f "pet/sw${subject}_pet_avg.nii" "pet/w${subject}_pet_avg.nii" "anat/y_${subject}_T1w.reoriented.nii"
    rm -f "pet/${subject}_pet_crop.nii.gz" "pet/${subject}_pet_avg.nii.gz" "anat/${subject}_T1w_preproc.nii.gz"
    echo "Cleanup complete."
fi

echo "SPM Based Preprocessing Script Finished."
