#!/bin/bash -e

wd=$(pwd)
jobid=$(squeue --me | head -2 | tail -n1 | awk '{print $1}')

#
# Example assume allocation was created, e.g.:
# N=1 ; salloc -p standard-g  --threads-per-core 1 --exclusive -N $N --gpus $((N*8)) -t 4:00:00 --mem 0
#

set -x

SIF=/pfs/lustref1/appl/local/laifs/containers/lumi-multitorch-u24r70f21m50t210-20260415_130625/lumi-multitorch-full-u24r70f21m50t210-20260415_130625.sif 

mkdir -p $wd/deps
mkdir -p $wd/deps/opt/venv-extra

# Install profiler dependencies from the distro
singularity exec \
  -B $wd:/workdir \
  $SIF bash -c -eux '
    
    cd /workdir/deps ;
    
    for i in libdw1t64 rocprofiler-systems rocprofiler-compute ; do
      apt-get download $i ;
      dpkg -x $i*.deb . ;
      rm -rf $i*.deb ;
    done
  '

# Install venv with rocprofiler-compute requirements
singularity exec \
  -B $wd:/workdir \
  -B $wd/deps/opt/venv-extra:/opt/venv-extra \
  $SIF bash -c -eux '
    \
    rm -rf /opt/venv-extra/* ;
    python -m venv --system-site-packages /opt/venv-extra ;
    source /opt/venv-extra/bin/activate ;

    pip install -r /workdir/deps/opt/rocm-7.0.2/libexec/rocprofiler-compute/requirements.txt ;

    #pip install "numpy<1.24" ;
    #pip install "pandas<2" ;
  '

# Adjust amd-smi to be compatible with LUMI
singularity exec \
  -B $wd:/workdir \
  $SIF bash -c -eux '

    mkdir -p /workdir/deps/opt/rocm-7.0.2/libexec/amdsmi_cli ;

    cat /opt/rocm-7.0.2/libexec/amdsmi_cli/amdsmi_helpers.py | \
      sed "s/required_groups =.*/required_groups = set()/g" > \
      /workdir/deps/opt/rocm-7.0.2/libexec/amdsmi_cli/amdsmi_helpers.py ; 
  '

# Clone rocm-systems
singularity exec \
  -B $wd:/workdir \
  $SIF bash -c -eux '
    
    cd /workdir/deps/ ;

    rm -rf rocprofiler-systems rocm-systems ;    
    git clone -b rocm-7.2.3 --recursive https://github.com/ROCm/rocm-systems ;
    ln -s rocm-systems/projects/rocprofiler-systems . ;
  '

# Build rocprof-systems with Python support
singularity exec \
  -B $wd:/workdir \
  $SIF bash -c -eux '
    \
    mkdir -p /workdir/deps/rocprofiler-systems/external/timemory/I_MPI_SUBSTITUTE_INSTALLDIR/include ;
    rm -rf /workdir/deps/rocprof-sys-build ;
    cmake \
      -B /workdir/deps/rocprof-sys-build \
      -D CMAKE_INSTALL_PREFIX=/workdir/deps/opt/rocm-7.2.3-extra \
      -D ROCPROFSYS_USE_ROCM=ON \
      -D ROCPROFSYS_USE_PYTHON=ON \
      -D ROCPROFSYS_USE_OMPT=OFF \
      -D ROCPROFSYS_USE_MPI_HEADERS=ON \
      -D ROCPROFSYS_USE_MPI=OFF \
      -D ROCPROFSYS_BUILD_PAPI=ON \
      -D ROCPROFSYS_BUILD_LIBUNWIND=ON \
      -D ROCPROFSYS_BUILD_DYNINST=ON \
      -D ROCPROFSYS_BUILD_TBB=ON \
      -D ROCPROFSYS_BUILD_BOOST=ON \
      -D ROCPROFSYS_BUILD_ELFUTILS=ON \
      -D ROCPROFSYS_BUILD_LIBIBERTY=ON \
      -D CMAKE_CXX_COMPILER=g++ \
      -D CMAKE_C_COMPILER=gcc \
      -S /workdir/deps/rocprofiler-systems
  '

singularity exec \
  -B $wd:/workdir \
  $SIF bash -c -eux '
    \
    cmake --build /workdir/deps/rocprof-sys-build --target all --parallel 64 ;
    cmake --build /workdir/deps/rocprof-sys-build --target install ;

    cd /workdir/deps/ ;
    rm -rf rocprofiler-systems rocm-systems rocprof-sys-build; 
  '

# Create squashfs image
rm -rf $wd/deps.sqsh
mksquashfs $wd/deps $wd/deps.sqsh -xattrs-exclude lustre.lov 
