#!/bin/bash
# ------------------------------------------------------
# Script to set up a MATLAB kernel for Jupyterhub @ DKRZ 
# Author: Bendoukha Sofiane
# Date: 2025-07-04
# ------------------------------------------------------

set -e 

# ---------------------------
# 1. Load MATLAB & Python modules
# ---------------------------
MATLAB_MODULE=""
# Manual argument parser
while [[ $# -gt 0 ]]; do
    case "$1" in
        --matlab-module)
            MATLAB_MODULE="$2"
            shift 2
            ;;
        --help)
            echo "Usage: $0 --matlab-module <matlab/version>"
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            exit 1
            ;;
    esac
done

if [[ -z "$MATLAB_MODULE" ]]; then
    echo "ERROR: --matlab-module is required"
    exit 1
fi
module load $MATLAB_MODULE
module load python3

# ---------------------------
# 2. Create and activate Python virtual environment
# ---------------------------
ENV_DIR="$HOME/matlab-kernel-env"
python3 -m venv "$ENV_DIR"
source "$ENV_DIR/bin/activate"

# ---------------------------
# 3. Install matlab_kernel and register the Jupyter kernel
# ---------------------------
pip install --upgrade pip
pip install matlab_kernel

python -m matlab_kernel install --user

# ---------------------------
# 4. Update the kernel.json to use the venv's Python
# ---------------------------
KERNEL_NAME="matlab"
KERNEL_DIR="$HOME/.local/share/jupyter/kernels/$KERNEL_NAME"
KERNEL_JSON="$KERNEL_DIR/kernel.json"

if [[ -f "$KERNEL_JSON" ]]; then
    echo "Updating kernel.json to use Python from the venv: $ENV_DIR"
    # Replace the first entry in "argv" with full Python path from venv
    VENV_PYTHON="$ENV_DIR/bin/python"
    TMP_JSON="$KERNEL_JSON.tmp"

    jq --arg py "$VENV_PYTHON" '.argv[0] = $py' "$KERNEL_JSON" > "$TMP_JSON" && mv "$TMP_JSON" "$KERNEL_JSON"
else
    echo "ERROR: kernel.json not found at $KERNEL_JSON"
    exit 1
fi

# ---------------------------
# 5. Locate MATLAB and prepare to build MATLAB Engine API for Python
# ---------------------------
MATLAB_ROOT=$(which matlab | sed 's|/bin/matlab||')
ENGINE_DIR="$MATLAB_ROOT/extern/engines/python"

if [[ ! -d "$ENGINE_DIR" ]]; then
    echo "ERROR: Could not find MATLAB Engine API directory at $ENGINE_DIR"
    exit 1
fi

cd "$ENGINE_DIR"

# ---------------------------
# 6. Build the MATLAB Engine for Python into $HOME
# ---------------------------
python setup.py build -b "$HOME"

# ----------------------
# 7. Move/copy the directory to python venv
BUILD_OUTPUT_DIR="$HOME/lib/matlab"
SITE_PACKAGES_DIR=$(python -c "import site; print(site.getsitepackages()[0])")

cd ~

if [[ -d "$BUILD_OUTPUT_DIR" ]]; then
    echo "Moving $BUILD_OUTPUT_DIR to $SITE_PACKAGES_DIR/matlab"
    # Remove any old version first (optional, but avoids conflicts)
    rm -rf "$SITE_PACKAGES_DIR/matlab"
    mv "$BUILD_OUTPUT_DIR" "$SITE_PACKAGES_DIR/matlab"
else
    echo "ERROR: $BUILD_OUTPUT_DIR does not exist."
    exit 1
fi

# 8. Cleanup
if [[ -d "$SITE_PACKAGES_DIR/matlab" ]]; then
    echo "Cleaning up build dir at $BUILD_OUTPUT_DIR"
    rm -rf "$BUILD_OUTPUT_DIR"
fi


echo "MATLAB kernel enabled, it will be shortly available in jupyterlab launcher!"

