#!/bin/bash
# Data Organization Script - Converts DICOM to NIfTI and structures data
set -e

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

SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"
subject_id=$1
subject="sub-${subject_id}"

echo "Processing Subject: ${subject}"

mkdir -p ~/Desktop/derivatives
mkdir -p ~/Desktop/derivatives/data/${subject}/anat
mkdir -p ~/Desktop/derivatives/data/${subject}/pet

MRI_ZIP="$HOME/Downloads/AD-100_MR.zip"
PET_ZIP="$HOME/Downloads/AD_PET_01-25.zip"

extract_zip() {
    local zip_file=$1
    local pattern=$2
    local output_dir=$3
    
    echo "Extracting $zip_file..."
    if unzip -n "$zip_file" "$pattern" -d "$output_dir"; then
        echo "Extraction successful."
        return 0
    else
        echo "ERROR: Extraction failed for $zip_file"
        return 1
    fi
}

echo "Extracting data for ${subject_id}..."

if [ -f "$MRI_ZIP" ]; then
    extract_zip "$MRI_ZIP" "*AD${subject_id}*" ~/Desktop/derivatives/
else
    echo "WARNING: MRI zip not found at $MRI_ZIP"
fi

if [ -f "$PET_ZIP" ]; then
    extract_zip "$PET_ZIP" "*AD${subject_id}*" ~/Desktop/derivatives/
else
    echo "WARNING: PET zip not found at $PET_ZIP"
fi

cd ~/Desktop/derivatives/

if ! ml dcm2niix/v1.0.20240202 2>/dev/null; then
    ml dcm2niix
fi

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

pet_dirs=$(find . -maxdepth 5 -type d -name "AD${subject_id}*_PET_DICOM")
if [ -z "$pet_dirs" ]; then
    pet_dirs=$(find . -maxdepth 5 -type d -name "AD${subject_id}*" ! -name "*MR_DICOM*" | head -n 1)
fi
anat_dir=$(find . -maxdepth 5 -type d -name "AD${subject_id}_MR_DICOM*" | head -n 1)

if [ -d "$anat_dir" ]; then
    echo "Converting Anatomical: $anat_dir"
    dcm2niix -o "$HOME/Desktop/derivatives/data/${subject}/anat" -f "${subject}_T1w" -z y -ba y -v y "$anat_dir"
    
    if [ -f "$HOME/Desktop/derivatives/data/${subject}/anat/${subject}_T1wa.nii.gz" ] && [ ! -f "$HOME/Desktop/derivatives/data/${subject}/anat/${subject}_T1w.nii.gz" ]; then
        mv "$HOME/Desktop/derivatives/data/${subject}/anat/${subject}_T1wa.nii.gz" "$HOME/Desktop/derivatives/data/${subject}/anat/${subject}_T1w.nii.gz"
        mv "$HOME/Desktop/derivatives/data/${subject}/anat/${subject}_T1wa.json" "$HOME/Desktop/derivatives/data/${subject}/anat/${subject}_T1w.json"
    fi
else
    echo "ERROR: Anatomical DICOM directory not found for ${subject_id}"
    exit 1
fi

if [ -n "$pet_dirs" ]; then
    echo "Found PET directories: $pet_dirs"
    
    SAVEIFS=$IFS
    IFS=$'\n'
    i=1
    for p_dir in $pet_dirs; do
        echo "Converting PET from: $p_dir (Part $i)"
        dcm2niix -o "$HOME/Desktop/derivatives/data/${subject}/pet" -f "${subject}_pet_part${i}" -z y -ba y -v y "$p_dir"
        i=$((i+1))
    done
    IFS=$SAVEIFS
    
    cd "$HOME/Desktop/derivatives/data/${subject}/pet"
    
    rm -f "${subject}_pet_merged.nii.gz" 2>/dev/null || true
    
    files=($(ls ${subject}_pet*.nii.gz 2>/dev/null | grep -v "${subject}_pet.nii.gz" | grep -v "merged" | sort))
    num_files=${#files[@]}
    
    if [ "$num_files" -gt 0 ]; then
        echo "Detected split PET series ($num_files files). Merging..."
        fslmerge -t "${subject}_pet_merged.nii.gz" "${files[@]}"
        
        if [ -f "${subject}_pet_merged.nii.gz" ]; then
            echo "Merge successful."
            rm "${files[@]}"
            mv "${subject}_pet_merged.nii.gz" "${subject}_pet.nii.gz"
        else
            echo "ERROR: fslmerge failed."
            exit 1
        fi
    else
        echo "Single PET file detected. No merging needed."
    fi
    
    cd ~/Desktop/derivatives/
else
    echo "ERROR: PET DICOM directory not found for ${subject_id}"
    exit 1
fi

cd data/
if [ ! -f dataset_description.json ]; then
    echo "{ \"Name\": \"PET_NeuroProject\", \"BIDSVersion\": \"1.9.0\" }" > dataset_description.json
fi

echo "Data organization complete for ${subject}"
