#!/bin/bash
# Copyright (C) 2018, Raffaello Bonghi <raffaello@rnext.it>
# All rights reserved
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# 1. Redistributions of source code must retain the above copyright 
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
# 3. Neither the name of the copyright holder nor the names of its 
#    contributors may be used to endorse or promote products derived 
#    from this software without specific prior written permission.
# 
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND 
# CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, 
# BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS 
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, 
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; 
# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, 
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE 
# OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, 
# EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

# NVIDIA Identify version 
# reference: 
# https://devtalk.nvidia.com/default/topic/1014424/jetson-tx2/identifying-tx1-and-tx2-at-runtime/
# https://devtalk.nvidia.com/default/topic/996988/jetson-tk1/chip-uid/post/5100481/#5100481

if [ -f /sys/module/tegra_fuse/parameters/tegra_chip_id ]; then
    case $(cat /sys/module/tegra_fuse/parameters/tegra_chip_id) in
        64)
            JETSON_BOARD="TK1" ;;
        33)
            JETSON_BOARD="TX1" ;;
        24)
            JETSON_BOARD="TX2" ;;
        25)
            JETSON_BOARD="Xavier" ;;
        *)
            JETSON_BOARD="UNKNOWN" ;;
    esac
    if [ "$JETSON_BOARD" = "Xavier" ] ; then
        JETSON_TYPE="AGX Xavier"
    elif [ "$JETSON_BOARD" = "TX1" ] ; then
        JETSON_TYPE="NANO/TX1"
    else
        JETSON_TYPE="$JETSON_BOARD"
    fi
    # Make Jetson Description
    JETSON_DESCRIPTION="NVIDIA Jetson $JETSON_TYPE"
    # Export variables
    export JETSON_BOARD
    export JETSON_TYPE
    export JETSON_DESCRIPTION
fi

# NVIDIA Jetson version
# reference https://devtalk.nvidia.com/default/topic/860092/jetson-tk1/how-do-i-know-what-version-of-l4t-my-jetson-tk1-is-running-/
# https://stackoverflow.com/questions/16817646/extract-version-number-from-a-string
if [ -f /etc/nv_tegra_release ]; then
    # L4T string
    # First line on /etc/nv_tegra_release 
    # - "# R28 (release), REVISION: 2.1, GCID: 11272647, BOARD: t186ref, EABI: aarch64, DATE: Thu May 17 07:29:06 UTC 2018"
    JETSON_L4T_STRING=$(head -n 1 /etc/nv_tegra_release)

    # Load release and revision
    JETSON_L4T_RELEASE=$(echo $JETSON_L4T_STRING | cut -f 2 -d ' ' | grep -Po '(?<=R)[^;]+')
    JETSON_L4T_REVISION=$(echo $JETSON_L4T_STRING | cut -f 2 -d ',' | grep -Po '(?<=REVISION: )[^;]+')
    # unset variable
    unset JETSON_L4T_STRING
    
    # Write Jetson description
    JETSON_L4T="$JETSON_L4T_RELEASE.$JETSON_L4T_REVISION"

    # Write version of jetpack installed
    # https://developer.nvidia.com/embedded/jetpack-archive
    if [ "$JETSON_BOARD" = "Xavier" ] ; then
        case $JETSON_L4T in
            "31.1.0") 
                    JETSON_JETPACK="4.1.1" ;;
            "31.0.2") 
                    JETSON_JETPACK="4.1" ;;
            "31.0.1") 
                    JETSON_JETPACK="4.0" ;;
            *)
               JETSON_JETPACK="UNKNOWN" ;;
        esac
    elif [ "$JETSON_BOARD" = "TX2i" ] ; then 
        case $JETSON_L4T in
            "28.2.1") 
               JETSON_JETPACK="3.3 | 3.2.1" ;;
            "28.2" | "28.2.0" ) 
               JETSON_JETPACK="3.2" ;;
            *)
               JETSON_JETPACK="UNKNOWN" ;;
        esac        
    elif [ "$JETSON_BOARD" = "TX2" ] ; then
        case $JETSON_L4T in
            "28.2.1") 
                    JETSON_JETPACK="3.3 | 3.2.1" ;;
            "28.2" | "28.2.0" ) 
                    JETSON_JETPACK="3.2" ;;
            "28.1") 
                    JETSON_JETPACK="3.1" ;;
            "27.1") 
                    JETSON_JETPACK="3.0" ;;
            *)
               JETSON_JETPACK="UNKNOWN" ;;
        esac
    elif [ "$JETSON_BOARD" = "TX1" ] ; then
        case $JETSON_L4T in
            "32.1.0" ) 
                    JETSON_JETPACK="4.2" ;;
            "28.2" | "28.2.0" ) 
                    JETSON_JETPACK="3.2" ;;
            "28.1") 
                    JETSON_JETPACK="3.1" ;;
            "24.2.1") 
                    JETSON_JETPACK="3.0 | 2.3.1" ;;
            "24.2") 
                    JETSON_JETPACK="2.3" ;;
            "24.1") 
                    JETSON_JETPACK="2.2.1 | 2.2" ;;
            "23.2") 
                    JETSON_JETPACK="2.1" ;;
            "23.1") 
                    JETSON_JETPACK="2.0" ;;
            *)
               JETSON_JETPACK="UNKNOWN" ;;
        esac
    elif [ "$JETSON_BOARD" = "TK1" ] ; then
        case $JETSON_L4T in
            "21.5") 
                    JETSON_JETPACK="2.3.1 | 2.3" ;;
            "21.4") 
                    JETSON_JETPACK="2.2 | 2.1 | 2.0 | DP 1.2" ;;
            "21.3") 
                    JETSON_JETPACK="DP 1.1" ;;
            "21.2") 
                    JETSON_JETPACK="DP 1.0" ;;
            *)
               JETSON_JETPACK="UNKNOWN" ;;
        esac
    else
        # Unknown board
        JETSON_JETPACK="UNKNOWN"
    fi
    # Export Release L4T
    export JETSON_L4T_RELEASE
    export JETSON_L4T_REVISION
    export JETSON_L4T
    # Export Jetson Jetpack installed
    export JETSON_JETPACK
fi

# Read CUDA version
if [ -f /usr/local/cuda/version.txt ]; then
    JETSON_CUDA=$(cat /usr/local/cuda/version.txt | sed 's/\CUDA Version //g')
else
    JETSON_CUDA="NOT_INSTALLED"
fi
# Jetson CUDA version
export JETSON_CUDA

# Read from OpenCV if is installed CUDA
opencv_read_cuda()
{
    # Red if use CUDA or not
    OPENCV_VERSION_VERBOSE=$(opencv_version --verbose | grep "Use Cuda" )

    if [ ! -z "$OPENCV_VERSION_VERBOSE" ]; then
        # Read status of CUDA
        OPENCV_CUDA_FLAG=$(echo $OPENCV_VERSION_VERBOSE | cut -f2 -d ':' | cut -f2 -d ' ' )
        
        if [ $OPENCV_CUDA_FLAG=="NO" ] ; then
            # Set NO if cuda is not installed
	    echo "NO"
        else
            # Set YES if cuda is installed
	    echo "YES"
        fi
        # Remove variable
        unset OPENCV_CUDA_FLAG
    else
        # read NVIDIA CUDA version
        OPENCV_VERSION_VERBOSE=$(opencv_version --verbose | grep "NVIDIA CUDA" )
        # get information
        OPENCV_CUDA_FLAG=$(echo $OPENCV_VERSION_VERBOSE | cut -f2 -d ':')
        OPENCV_CUDA_FLAG=${OPENCV_CUDA_FLAG//[[:blank:]]/}
	# Remove variable
        unset OPENCV_CUDA_FLAG
        # Set YES if cuda is installed
	echo "YES"
    fi
    # Remove variable
    unset OPENCV_VERSION_VERBOSE
    return
}

if hash opencv_version 2>/dev/null; then
    JETSON_OPENCV="$(opencv_version)"
    # Read information about cuda status
    JETSON_OPENCV_CUDA=$(opencv_read_cuda)
else
    JETSON_OPENCV="NOT_INSTALLED"
    JETSON_OPENCV_CUDA="NO"
fi
# Opencv variables
export JETSON_OPENCV
export JETSON_OPENCV_CUDA

# Write CUDA architecture
# https://developer.nvidia.com/cuda-gpus
# https://devtalk.nvidia.com/default/topic/988317/jetson-tx1/what-should-be-the-value-of-cuda_arch_bin/
case $JETSON_BOARD in
    "Xavier") 
        JETSON_CUDA_ARCH_BIN="7.2" ;;
    "TX2" | "TX2i") 
        JETSON_CUDA_ARCH_BIN="6.2" ;;
    "TX1")
        JETSON_CUDA_ARCH_BIN="5.3" ;;
    "TK1")
        JETSON_CUDA_ARCH_BIN="3.2" ;;
    * )
        JETSON_CUDA_ARCH_BIN="NONE" ;;
esac
# Export Jetson CUDA ARCHITECTURE
export JETSON_CUDA_ARCH_BIN

# Extract cuDNN version
JETSON_CUDNN=$(dpkg -l 2>/dev/null | grep -m1 "libcudnn")
if [ ! -z "$JETSON_CUDNN" ] ; then
    JETSON_CUDNN=$(echo $JETSON_CUDNN | sed 's/.*libcudnn[0-9] \([^ ]*\).*/\1/' )
else
    JETSON_CUDNN="NOT_INSTALLED"
fi
# Export NVIDIA CuDNN Library
export JETSON_CUDNN

# Extract TensorRT version
JETSON_TENSORRT=$(dpkg -l 2>/dev/null | grep -m1 " tensorrt ")
if [ ! -z "$JETSON_TENSORRT" ] ; then
    JETSON_TENSORRT=$(echo $JETSON_TENSORRT | sed 's/.*tensorrt \([^ ]*\).*/\1/' )
else
    JETSON_TENSORRT="NOT_INSTALLED"
fi
# Export NVIDIA CuDNN TensorRT
export JETSON_TENSORRT

# Extract Visionworks version
JETSON_VISIONWORKS=$(dpkg -l 2>/dev/null | grep -m1 "libvisionworks")
if [ ! -z "$JETSON_VISIONWORKS" ] ; then
    JETSON_VISIONWORKS=$(echo $JETSON_VISIONWORKS | sed 's/.*libvisionworks \([^ ]*\).*/\1/' )
else
    JETSON_VISIONWORKS="NOT_INSTALLED"
fi
# Export NVIDIA CuDNN VisionWorks
export JETSON_VISIONWORKS

# TODO Add enviroments variables:
# - UID -> https://devtalk.nvidia.com/default/topic/996988/jetson-tk1/chip-uid/post/5100481/#5100481
# - GCID, BOARD, EABI

