#!/bin/bash
# Visual QC Snapshot Generator

set -e

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

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

echo "Generating Visual QC for Subject: ${subject}"

BASE_DIR=~/Desktop/derivatives
SUB_DIR="${BASE_DIR}/data/${subject}"
ANAT_DIR="${SUB_DIR}/anat"
PET_DIR="${SUB_DIR}/pet"
QC_DIR="${BASE_DIR}/QC/${subject}"

mkdir -p "$QC_DIR"

if ! command -v fsleyes &> /dev/null; then
    if ! ml fsl/6.0.7.8 2>/dev/null; then
        ml fsl
    fi
fi

T1="${ANAT_DIR}/${subject}_T1w.reoriented.nii"
PET_MNI="${PET_DIR}/${subject}_pet_to_MNI_smoothed.nii.gz"
AVG_PET="${PET_DIR}/${subject}_pet_avg.nii"

MNI_TEMPLATE="${FSLDIR}/data/standard/MNI152_T1_2mm_brain.nii.gz"

SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"
PROJECT_ROOT="$(dirname "$(dirname "$SCRIPT_DIR")")"
MASK_DIR="${PROJECT_ROOT}/data/references/centiloid_masks"
MASK_CEREB="${MASK_DIR}/voi_WhlCbl_2mm.nii"
MASK_CTX="${MASK_DIR}/voi_ctx_2mm.nii"

if [ ! -f "$T1" ] || [ ! -f "$PET_MNI" ]; then
    echo "ERROR: Missing required input files for QC."
    echo "Expected: $T1 and $PET_MNI"
    exit 1
fi

echo "Snapshot 1/3: Coregistration..."
fsleyes render --outfile "${QC_DIR}/qc_coreg.png" \
    --size 1200 400 \
    --scene ortho \
    "$T1" --overlayType volume --name "T1" \
    "$AVG_PET" --overlayType volume --name "PET" --cmap hot --alpha 50 || echo "WARNING: Coregistration snapshot failed"

echo "Snapshot 2/3: Normalization..."
fsleyes render --outfile "${QC_DIR}/qc_norm.png" \
    --size 1200 400 \
    --scene ortho \
    "$MNI_TEMPLATE" --overlayType volume --name "MNI152" \
    "$PET_MNI" --overlayType volume --name "PET_MNI" --cmap hot --alpha 50 || echo "WARNING: Normalization snapshot failed"

echo "Snapshot 3/3: Mask Alignment..."
fsleyes render --outfile "${QC_DIR}/qc_masks.png" \
    --size 1200 400 \
    --scene ortho \
    "$PET_MNI" --overlayType volume --name "PET_MNI" --cmap gray \
    "$MASK_CEREB" --overlayType volume --name "Cerebellum" --cmap Blue --alpha 40 \
    "$MASK_CTX" --overlayType volume --name "Cortex" --cmap Red --alpha 40 || echo "WARNING: Mask alignment snapshot failed"

echo "Visual QC Complete for $subject. Images saved to $QC_DIR"
