#!/bin/bash -ex

# install-npu-software.sh
#
# Copyright (c) 2025 Technologic Systems, Inc. (dba embeddedTS)

FIRMWARE_URL=https://github.com/nxp-imx/ethos-u-firmware.git
DRIVER_REPO_URL=https://github.com/nxp-imx/ethos-u-driver-stack-imx.git
DELEGATE_REPO_URL=https://github.com/nxp-imx/tflite-ethosu-delegate-imx
TENSORFLOW_REPO_URL=https://github.com/nxp-imx/tensorflow-imx
VELA_REPO_URL=https://github.com/nxp-imx/ethos-u-vela.git

SOURCE_DEST="$HOME/src/NPU-$(uname -r)"
ETHOS_KERNEL_BRANCH="lf-6.6.52_2.2.0"

# The "minimal example" referenced by NXP's ML guide is useful for
# development, but not needed to demonstrate functionality.
build_minimal_example=false

# ==================== Fetch ====================

function fetch-git {
    local REPO_URL=$1
    local BRANCH=$2
    local SOURCE_DEST_DIR=$3
    local SOURCE_DEST_PARENT_DIR="$(dirname ${SOURCE_DEST_DIR})"
    local SOURCE_DEST_BASENAME="$(basename ${SOURCE_DEST_DIR})"
    local rc
    
    if ! [ -e "${SOURCE_DEST_PARENT_DIR}" ] ; then
        echo "Creating: ${SOURCE_DEST_PARENT_DIR}"
        mkdir "${SOURCE_DEST_PARENT_DIR}"
    fi

    pushd ${SOURCE_DEST_PARENT_DIR}
    git clone ${REPO_URL} -b ${BRANCH} || return $?
    if ! [ -e "${SOURCE_DEST_BASENAME}" ] ; then
       echo "Not found: ${SOURCE_DEST_BASENAME}" >&2
       return 1
    fi
    return 0
}

fetch-git "${FIRMWARE_URL}" "${ETHOS_KERNEL_BRANCH}" "${SOURCE_DEST}/ethos-u-firmware"

fetch-git "${DRIVER_REPO_URL}" "${ETHOS_KERNEL_BRANCH}" "${SOURCE_DEST}/ethos-u-driver-stack-imx"

fetch-git "${DELEGATE_REPO_URL}" "${ETHOS_KERNEL_BRANCH}" "${SOURCE_DEST}/tflite-ethosu-delegate-imx"

fetch-git "${TENSORFLOW_REPO_URL}" "${ETHOS_KERNEL_BRANCH}" "${SOURCE_DEST}/tensorflow-imx"

# Vela is needed to convert Tensorflow Lite models for Ethos-U
# - Note: A specific version of Vela is tightly coupled with a
#   specific version of the Ethos-U driver. The compatibility between
#   different Vela versions is not guaranteed.
fetch-git "${VELA_REPO_URL}" "${ETHOS_KERNEL_BRANCH}" "${SOURCE_DEST}/ethos-u-vela"

# ==================== Build and Install ====================

if ! [ -e /swapfile ] ; then
    created_swap=true
    fallocate -l 2G /swapfile
    chown 600 /swapfile
    mkswap /swapfile
    swapon /swapfile
else
    created_swap=false
fi

apt-get update
# needed by most:
apt-get install cmake build-essential libflatbuffers-dev -y
# needed by delegate:
apt-get install python3-dev python3-venv -y
# needed to unpack a model:
apt-get install unzip

# firmware
echo "Building ethos-u-firmware"
pushd "${SOURCE_DEST}/ethos-u-firmware"
mkdir -p ${ROOT}/lib/firmware
cp -p ethosu_firmware ${ROOT}/lib/firmware/ethosu_firmware
popd

# ethos-uu driver/libraries (will be used to build tflist-ethosu-delegate-imx)
echo "Building ethos-u-driver-stack-imx"
pushd "${SOURCE_DEST}/ethos-u-driver-stack-imx"
mkdir build
cd build
cmake .. && make -j2 && make install
echo "Adding /usr/local/lib to /etc/ld.so.conf.d/ethosu.conf"
mkdir -p ${ROOT}/etc/ld.so.conf.d/
echo "/usr/local/lib" > ${ROOT}/etc/ld.so.conf.d/ethosu.conf && ldconfig
popd

# tflite-ethosu-delegate-imx
echo "Building tflite-ethosu-delegate-imx"
pushd "${SOURCE_DEST}/tflite-ethosu-delegate-imx"
mkdir build
cd build
cmake -DPython_EXECUTABLE=python3 .. && make -j2 && make install
popd

# tensorflow-imx
pushd "${SOURCE_DEST}/tensorflow-imx"
python3 -m venv ../env

. ../env/bin/activate
curl -o ../env/bin/bazelisk-linux-arm64 -LO https://github.com/bazelbuild/bazelisk/releases/latest/download/bazelisk-linux-arm64
chmod +x ../env/bin/bazelisk-linux-arm64
ln ../env/bin/bazelisk-linux-arm64 ../env/bin/bazelisk

pip install "numpy<2"
./configure <<EOF




n


EOF

popd

# ethos-u-vela
# continue sharing the venv from tensorflow-imx
pushd "${SOURCE_DEST}/ethos-u-vela"
pip install .
vela --version
popd

# build-minimal-example -- see tensorflow-imx and the NXP's ML User's Guide
if ${build_minimal_example} ; then
    ! [ -e ${SOURCE_DEST}/build-minimal-example ] && mkdir ${SOURCE_DEST}/build-minimal-example
    pushd "${SOURCE_DEST}/build-minimal-example"
    cmake -DTFLITE_ENABLE_XNNPACK=on \
          -DTFLITE_ENABLE_RUY=on \
          ../tensorflow-imx/tensorflow/lite/examples/minimal
    make -j2 && make install
fi

# ==================== Final cleanup ====================

if ${created_swap} ; then
    echo "Remember to remove the swap space added earlier"
fi