#!/bin/bash
# Removes possible Windows carriage returns at runtime
sed -i 's/\r$//' "$0"

set -euo pipefail
# ----------- ARGUMENT VALIDATION------------
sujeto="$1"
base_path="$2"
zip_t1="$3"
zip_dti="$4"

if [ $# -ne 4 ]; then
  echo "Use: $0 <subject_identifier> <base_path> <zip_T1> <zip_DTI>"
  exit 1
fi

for file in "$zip_t1" "$zip_dti"; do
  if [ ! -f "$file" ]; then
    echo "ERROR: File not found: $file"
    exit 1
  fi
done

# ----------- DEFINE DINAMIC PATHS------------
sujeto_path="${base_path}/${sujeto}"
datos_dir="${sujeto_path}/data"
t1_dir="${datos_dir}/t1"
dti_dir="${datos_dir}/dti"
nii_dir="${sujeto_path}/nii"
fs_dir="${sujeto_path}/freesurfer"
m2m_dir="${sujeto_path}/m2m_${sujeto}"
dwi_dir="${m2m_dir}/DWI"
prep_dir="${m2m_dir}/dMRI_prep"
resultados="${sujeto_path}/results"

mkdir -p "$t1_dir" "$dti_dir" "$nii_dir" "$resultados" "$fs_dir"

# ----------- 1. DCM2NIIX T1 ------------
echo "Unzipping  T1..."
unzip "$zip_t1" -d "$t1_dir"

echo "Converting T1 DICOM a NIfTI..."
dcm2niix -o "$nii_dir" -f MRI_T1 "$t1_dir"

# ----------- 2. FREESURFER ------------
echo "Running FreeSurfer..."
set +u
set +e
export FREESURFER_HOME=/usr/local/freesurfer/8.1.0
source $FREESURFER_HOME/SetUpFreeSurfer.sh
set -e
set -u

export SUBJECTS_DIR="$fs_dir"
recon-all -i "${nii_dir}/MRI_T1.nii" -s "$sujeto" -all

# ----------- 3. BINARIZATION------------
cd "${fs_dir}/${sujeto}" || exit 1
aparc="mri/aparc+aseg.mgz"

binarize() {
  mri_binarize --i "$aparc" --match "${@:3}" --o "$2"
}

echo "Segmenting brain regions..."
binarize "$aparc" whitematter.nii 2 41 77
binarize "$aparc" greymatter.nii 28 60 13 52 12 51 11 50 26 58 1001 1002 1003 1005 1007 1008 1009 1010 1011 1012 1013 1014 1015 1016 1017 1018 1019 1020 1021 1022 1023 1024 1025 1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 2001 2002 2003 2005 2007 2008 2009 2010 2011 2012 2013 2014 2015 2016 2017 2018 2019 2020 2021 2022 2023 2024 2025 2026 2027 2028 2029 2030 2031 2032 2033 2034 2035
binarize "$aparc" ventricles.nii 4 5 14 15 43 44
binarize "$aparc" entorhinalc.nii 1006 2006
binarize "$aparc" hippocampus.nii 17 53
binarize "$aparc" brainstem.nii 16
binarize "$aparc" cerebellumcortex.nii 8 47
binarize "$aparc" cerebellumwhitematter.nii 7 46
binarize "$aparc" thalamus.nii 49 10
binarize "$aparc" csf.nii 24
binarize "$aparc" amygdala.nii 18 54
binarize "$aparc" cc.nii 251 252 253 254 255

declare -A etiquetas=(
  [ventricles]=1 [whitematter]=2 [greymatter]=3 [entorhinalc]=4 [hippocampus]=5
  [brainstem]=6 [cerebellumcortex]=7 [cerebellumwhitematter]=8 [thalamus]=9
  [csf_dil]=10 [amygdala]=11 [cc]=12
)

for region in "${!etiquetas[@]}"; do
  if [ -f "${region}.nii" ]; then
    fslmaths "${region}.nii" -mul "${etiquetas[$region]}" "${region}_label.nii"
  fi
done

fslmaths ventricles_label.nii -add whitematter_label.nii -add greymatter_label.nii -add entorhinalc_label.nii \
  -add hippocampus_label.nii -add brainstem_label.nii -add cerebellumcortex_label.nii \
  -add cerebellumwhitematter_label.nii -add thalamus_label.nii -add amygdala_label.nii \
  -add cc_label.nii inside_csf.nii

fslmaths csf.nii -dilM -dilM -dilM csf_dil.nii.gz
fslmaths csf_dil.nii.gz -sub inside_csf.nii -thr 0 -bin csf_dil.nii.gz
fslmaths csf_dil.nii -mul 10 csf_dil_label.nii

fslmaths ventricles_label.nii -add whitematter_label.nii -add greymatter_label.nii -add entorhinalc_label.nii \
  -add hippocampus_label.nii -add brainstem_label.nii -add cerebellumcortex_label.nii \
  -add cerebellumwhitematter_label.nii -add thalamus_label.nii -add csf_dil_label.nii \
  -add amygdala_label.nii -add cc_label.nii brain_labels.nii.gz

# ----------- 4. SIMNIBS ------------
echo "Running SimNIBS..."
charm "$sujeto" "${nii_dir}/MRI_T1.nii" --fs-dir "${fs_dir}/${sujeto}" --
mv -f "${fs_dir}/${sujeto}/m2m_${sujeto}" "${sujeto_path}/m2m_${sujeto}"
meshmesh "${fs_dir}/${sujeto}/brain_labels.nii.gz" "${m2m_dir}/${sujeto}_mesh.msh" --voxsize_meshing 0.5
gmsh "${m2m_dir}/${sujeto}_mesh.msh" -format nas -o "${m2m_dir}/${sujeto}_mesh.bdf" -save

cp "${m2m_dir}/${sujeto}_mesh.bdf" "${resultados}/${sujeto}_mesh.bdf"



# ----------- 5. DCM2NIIX DTI ------------
echo "Unzipping DTI..."
unzip "$zip_dti" -d "$dti_dir"

mkdir -p "$dwi_dir"

echo "Converting DTI a NIfTI..."
dcm2niix -o "$dwi_dir" -f DTI "$dti_dir"
dcm2niix -o "$nii_dir" -f DTI "$dti_dir"

# ----------- 6. dwi2cond and vector field ------------
cd "$sujeto_path" || exit 1

export FSLDIR="${base_path}/fsl"

dwi2cond --all "$sujeto" "${dwi_dir}/DTI.nii" "${dwi_dir}/DTI.bval" "${dwi_dir}/DTI.bvec"

nii2msh -ev "${m2m_dir}/DTI_coregT1_tensor.nii.gz" "${m2m_dir}/${sujeto}.msh" "${prep_dir}/first_eig_for_check.msh"
gmsh "${prep_dir}/first_eig_for_check.msh" -format msh2 -save -o "${prep_dir}/first_eig_for_check_ascii.msh"

sed '/EndNodes/,$d' "${prep_dir}/first_eig_for_check_ascii.msh" | tail -n +6 | cut -d' ' -f2- | tr -s ' ' | tr ' ' ',' > "${prep_dir}/first_eig_for_check_csv.csv"

get_fields_at_coordinates -s "${prep_dir}/first_eig_for_check_csv.csv" -m "${prep_dir}/first_eig_for_check.msh" --method assign

LC_ALL=C awk -F, 'BEGIN{OFS=","} {norm=0; for(i=1;i<=NF;i++){gsub(",",".",$i); a[i]=$i; norm+=a[i]*a[i]} norm=sqrt(norm); for(i=1;i<=NF;i++) printf "%s%.6f",(i>1?OFS:""), a[i]/norm; printf "\n"}' "${prep_dir}/first_eig_for_check_csv_max_eig.csv" > "${prep_dir}/first_eig_for_check_csv_max_eig_unit.csv"

paste -d',' "${prep_dir}/first_eig_for_check_csv.csv" "${prep_dir}/first_eig_for_check_csv_max_eig_unit.csv" > "${resultados}/vector_field.csv"

# ----------- END ------------
echo "Processing completed for '$sujeto'."
echo "Results saved in: $resultados"

