commonrc

# Common rc
# Things I want in both bash and zsh, because I don't necessarily use zsh all the time.

# Base16 colors
# Credit: https://github.com/chriskempson/base16-shell
if [[ -f $HOME/.base16colors ]]; then
    source $HOME/.base16colors
fi

# If cuda isn't already loaded
if [[ "$OSTYPE" == "linux-gnu"* ]] && [[ -f "$(which nvidia-smi)" ]] && [[ ! -f "$(which nvcc)" ]]; then
    CUDA_PATHS=( 
        "/usr/local/cuda-12.3"
        "/usr/local/cuda-12.2"
        "/usr/local/cuda"
    )
    for cuda_path in ${CUDA_PATHS[*]}; do
        if [[ -d "${cuda_path}" ]] && [[ -f "${cuda_path}/bin/nvcc" ]] && [[ -d "${cuda_path}/lib64" ]]; then
            export PATH=${cuda_path}/bin${PATH:+:${PATH}}
            export LD_LIBRARY_PATH=${cuda_path}/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
            break
        fi
    done
fi

# General Exports
export EDITOR="vim"

# Load local binaries
export PATH=$HOME/.local/bin:$PATH
export MANPATH=$HOME/.local/man:$MANPATH
if [[ -d "$HOME/.local/extras/coreutils" ]]; then
  export PATH=$HOME/.local/extras/coreutils/bin:$PATH
  export MANPATH=$HOME/.local/extras/coreutils/man:$MANPATH
fi

# TODO: the perl build is buggy
if [[ -d "$HOME/.local/lib/perl5/" ]]; then
  export PERL5LIB=$HOME/.local/lib/perl5/
fi

export LD_LIBRARY_PATH=$HOME/.ncurses/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
export LD_LIBRARY_PATH=$HOME/.local/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}

# Bat config
export BAT_CONFIG_PATH="$HOME/.config/bat/bat.conf"

# ls color scheme
export LS_COLORS="$(cat ~/.lscolors)"

##########
# Aliases
##########

# Replace ls with lsd
if [[ -f "$(which lsd)" ]]; then
    alias ls="lsd --group-directories-first --color=auto"
    alias lc="lsd --group-directories-first -lh --color=auto"
    alias la='lsd --group-directories-first -lah --color=auto'
    alias lm='lsd --group-directories-first -lah --sort modified --color=auto'
    alias ll='lsd --group-directories-first -lah --sort size --color=auto'
elif [[ -f "$(which exa)" ]]; then
    alias ls="exa --color=auto"
    alias lc="exa -lh --color=auto"
    alias la='exa -lah -s name --color=auto'
    alias lm='exa -lah -s modified --color=auto'
    alias ll='exa -lah -s size --color=auto'
fi

# Add color and make human readable
alias df='df -h'
# Grep to have color, give line number, don't tell me it can't access restricted files (sudo), and don't process binary files (garbage output ):
alias grep='grep --color=always --line-number --no-messages --binary-files=without-match'


alias prs='ps -o uid,user,pid,%cpu,%mem,tt,state,start,time,command'

# Tmux
# One command to attach to an existing session, or create one if none exist.
mux () {
    tmux attach $@ || tmux
}

# Show nvidia gpus, their archtag, and memory
if [[ -f "$(which nvidia-smi)" ]]; then
  nv-arch () {
      nvidia-smi --query-gpu=name,memory.total,compute_cap --format=csv
  }

  nv-name () {
    full_name=$(nvidia-smi -i 0 --query-gpu=gpu_name --format=csv,noheader)
    name_tmp=${full_name//NVIDIA/}
    name_tmp=${name_tmp//Tesla/}
    name_tmp=${name_tmp//Ampere/}
    name_tmp=${name_tmp//GeForce/}
    name_tmp=${name_tmp//RTX/}
    name_tmp=${name_tmp// /}
    name_final=${name_tmp:0:8}
    echo $name_final
  }

  export HAS_NV_SMI=1
fi

# OSX-specific aliases for my mac
if [[ "$OSTYPE" == "darwin"* ]]; then
    if [[ -d $HOME/.brew ]]; then
        export PATH="$HOME/.brew/bin:$HOME/.brew/sbin:$PATH"
    fi

    # Zathura in background
    zat () {
        zathura --page=1 "$@" > /dev/null 2>&1 &
    }

    ## Replace the mac sed with GNU sed:
    #alias sed="gsed"
    
    # GUI applications
    alias zoom="open -a zoom.us"
    alias spotify="open -a Spotify"
    alias slack="open -a Slack"
    alias settings="open -a System\ Settings"
    alias amph="open -a Amphetamine -g"
    alias finder="open ~/Documents"

    # Vimwiki markdown to original markdown converter
    w2md () {
        if [[ -f ~/.vimwiki2md ]]; then
            srcf=$1
            srcfmd=$(echo $srcf|sed -e "s/\.wiki$/.md/")
            destf=${2:-srcfmd}
            sed -f ~/.vimwiki2md $srcf > $destf
        else
            echo "Conversion from vimwiki markdown to md failed! ~/.vimwiki2md not found."
        fi
    }
fi

# SSH-Agent (from scripts/)
alias sagent='source sagent'

# Login vs compute node rules
# Login nodes should automatically try and set up an ssh agent; compute nodes shouldn't.
# Compute nodes should try and find conda and other such tools, not login nodes.
if [[ -z $SLURM_JOB_ID ]]; then
    # Start/attach to existing ssh agent
    source sagent
fi
# Attempt to load anaconda on compute nodes only
if [[ $HAS_NV_SMI = 1 ]]; then
  source cnda
fi

# Modules
if [[ -f /etc/profile.d/modules.sh ]]; then
    source /etc/profile.d/modules.sh
    # Slurm
    if [[ "$(module avail)" == *"slurm"* ]]; then
        module load slurm
    fi

    # Bright cluster manager shell
    if [[ "$(module avail)" == *"cmsh"* ]]; then
        module load cmsh
    fi
fi

# Custom wandb tmp directory
MY_WANDB_DIR="$HOME/.tmp/wandb/"
[ ! -d "$MY_WANDB_DIR" ] && mkdir -p $MY_WANDB_DIR
export WANDB_DIR="$MY_WANDB_DIR"

# Custom ENROOT config path
MY_ENROOT="$HOME/.enroot"
export ENROOT_CONFIG_PATH="$MY_ENROOT"

# My functions

# Who has processes?
alias whodat='ps haeo user | sort -u'

# IP address lookup
alias lip="hostname -I | cut -d' ' -f1"
alias gip="curl -s http://ipinfo.io/ip | cut -d'%' -f1"

# Clear out those latex compile files for goodness sake
cleartex () {
  if [[ $# -eq 1 ]]; then
      rm -f $1.log $1.aux $1.bbl $1.blg $1.out $1.dvi $1.ps $1.ps.gz \
          $1.synctex.gz $1.fls $1.fdb_latexmk $1.synctex\(busy\) \
          $1.bcf $1.run.xml
      rm -f texput.log
  else
      echo "No argument passed! cleartex expects exactly 1 argument: file name."
      echo "For example: cleartex main will remove latex cache files associated with main."
  fi
}

# Wait on PIDs
# Useful when lazily trying to schedule jobs on a physical node or workstation without a scheduler
anywait(){
    if [[ $# -eq 0 ]] || [[ $1 == "--help" ]]; then
        echo "Block on PIDs (anywait)"
        echo "usage: anywait [pid]"
        echo "       anywait [pid1] ... [pidN]"
        echo ""
        echo "Waits on all PIDs to finish / die with 10 second intervals."
    else
        # Borrowed this from StackOverflow:
        # https://stackoverflow.com/a/1127675
        for pid in "$@"; do
            while kill -0 "$pid"; do
                # echo "$pid is still alive, waiting..."
                sleep 10
            done
        done
    fi
}

# My go-to rsync with deletion
ry () {
    rsync -a -H -x -v --numeric-ids --delete --progress -e "ssh -T -o Compression=no -x" $1 $2
}

# My go-to rsync WITHOUT deletion
rys () {
    rsync -a -H -x -v --numeric-ids --ignore-existing --progress -e "ssh -T -o Compression=no -x" $1 $2
}

# I have a custom "vif" to force it to support true color
# without upsetting zsh
# Opens vifm and cds into the directory you :q from.
vicd() {
    local dst="$(command vif --choose-dir - "$@")"
    if [ -z "$dst" ]; then
        echo 'Directory picking cancelled/failed'
        return 1
    fi
    cd "$dst"
}

# SMUX: SSH and then attach to tmux
smux () {
    ssh -t ${@} 'bash -c "tmux a || tmux"'
}

# NVIDIA Bug report
## Handy function when I want to grab logs related to GPU issues quickly
## and send over to Lambda.
if [[ -f /usr/bin/nvidia-bug-report.sh ]]; then
    nvb () {
        dts=$(date +"%Y%m%d-%H%M%S")
        nvb_filename="nvidia-bug-report-${dts}.gz"
        ipmi_filename="ipmitool-${dts}.txt"
        echo "Storing logs in $nvb_filename and $ipmi_filename"
        sudo /usr/bin/nvidia-bug-report.sh --output-file $nvb_filename "$@"
        sudo ipmitool sel elist >> $ipmi_filename
    }
fi

# Diff so fancy
diffs () {
    diff -u $1 $2 | diff-so-fancy | less
}

# Google Drive downloader
# Download locally and rsync to remote no more!
# Thanks Minchul!
dget () {
    drive_url_to_id () {
        local URL=$1
        local FILE_ID=${URL#*google.com/}
        FILE_ID=${FILE_ID#*file\/d\/}
        FILE_ID=${FILE_ID#*file\/}
        FILE_ID=${FILE_ID#*files\/}
        FILE_ID=${FILE_ID#*\?id=}
        FILE_ID=${FILE_ID/\/view/}
        echo $FILE_ID
    }
    download_from_gdrive() {
        curl -H "Authorization: Bearer $2" \
        "https://www.googleapis.com/drive/v3/files/$1?alt=media" -o $3
    }
    # Args:
    if [[ $# -eq 0 ]] || [[ $1 == "--help" ]]; then
        echo "Google Drive Get (dget)"
        echo "usage: dget [url | id]"
        echo "       dget [url | id] [access key]"
        echo "       dget [url | id] [access key] [output file name]"
        echo ""
        echo "       dget https://drive.google.com/file/d/XXXXXX/view"
        echo ""
        echo "To get access key, go to https://developers.google.com/oauthplayground/"
        echo "  and select Drive API v3, then select"
        echo " \`https://www.googleapis.com/auth/drive.readonly\` and then click "
        echo " \`Authorize APIs\`, and get to Step 2. Once in step 2, click the blue button"
        echo " \`Exchange authorization code for tokens\` and you should see your access token"
        echo " in the JSON shown on the right hand side view."
    elif [[ $# -eq 1 ]]; then
        DRIVE_URL=$1
        FILE_ID=$(drive_url_to_id $DRIVE_URL)
        echo "Downloading Google Drive file: $DRIVE_URL"
        echo "ID: $FILE_ID"
        echo ""
        echo "Open https://developers.google.com/oauthplayground/"
        echo "  and select Drive API v3, then select"
        echo " \`https://www.googleapis.com/auth/drive.readonly\` and then click "
        echo " \`Authorize APIs\`, and get to Step 2. Once in step 2, click the blue button"
        echo " \`Exchange authorization code for tokens\` and you should see your access token"
        echo " in the JSON shown on the right hand side view."
        echo ""
        echo "Enter access key: "
        read DRIVE_ACCESS_KEY
        echo ""
        echo "Enter output file name (e.g. file.zip): "
        read DRIVE_FILENAME
        download_from_gdrive $FILE_ID $DRIVE_ACCESS_KEY $DRIVE_FILENAME
    elif [[ $# -eq 1 ]]; then
        DRIVE_URL=$1
        FILE_ID=$(drive_url_to_id $DRIVE_URL)
        DRIVE_ACCESS_KEY=$2
        echo "Enter output file name (e.g. file.zip): "
        read DRIVE_FILENAME
        download_from_gdrive $FILE_ID $DRIVE_ACCESS_KEY $DRIVE_FILENAME
    elif [[ $# -eq 3 ]]; then
        DRIVE_URL=$1
        FILE_ID=$(drive_url_to_id $DRIVE_URL)
        DRIVE_ACCESS_KEY=$2
        DRIVE_FILENAME=$3
        download_from_gdrive $FILE_ID $DRIVE_ACCESS_KEY $DRIVE_FILENAME
    else
        echo "Invalid option: $@"
    fi
}


# TODO: clean these up and make it conditional
# The GPU args only make sense when nvidia-container-toolkit is installed.

# Run in docker
trun () {
    # Arguments:
    # 1: image name
    # 2: command

    UNID="${RANDOM}${RANDOM}${RANDOM}" # random id for temp files

    PREF="$HOME/.tmp/dockerrun/meta/$UNID/" # meta will hold passwd, group, and command script
    LPREF="$HOME/.tmp/dockerrun/run/$UNID/" # run will hold stdout and stderr log
    mkdir -p $PREF
    mkdir -p $LPREF
    SCR="$LPREF/cmd.sh"
    LOG="$LPREF/log.txt"

    IMAGE=$1
    shift

    # Dump chdir and command into script
    THISDIR=$(pwd)
    echo "cd $THISDIR" >> $SCR
    echo "$@" >> $SCR
    chmod +x ${SCR}

    # Clone group and passwd and ro mount so you run rootless
    # Opts borrowed from https://github.com/avolkov1/multinode_containers
    getent group > ${PREF}/group
    getent passwd > ${PREF}/passwd
    USEROPT1="-v${PREF}passwd:/etc/passwd:ro"
    USEROPT2="-v${PREF}group:/etc/group:ro"

    docker run --rm --gpus=all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 -u $(id -u):$(id -g) -e HOME=$HOME -e USER=$USER -v /home:/home $USEROPT1 $USEROPT2 ${IMAGE} ${SCR} 2>&1 | tee ${LOG}
}

# Start interactive terminal in container
tbash () {
    # Arguments:
    # 1: image name
    UNID="${RANDOM}${RANDOM}${RANDOM}" # random id for temp files
    PREF="$HOME/.tmp/dockerrun/meta/$UNID/" # meta will hold passwd and group
    mkdir -p $PREF

    IMAGE=$1
    shift

    # Command is just bash
    SCR="/bin/bash"

    # Clone group and passwd and ro mount so you run rootless
    # Opts borrowed from https://github.com/avolkov1/multinode_containers
    getent group > ${PREF}/group
    getent passwd > ${PREF}/passwd
    USEROPT1="-v${PREF}passwd:/etc/passwd:ro"
    USEROPT2="-v${PREF}group:/etc/group:ro"

    # Opts borrowed from https://github.com/avolkov1/multinode_containers
    docker run -it --rm --gpus=all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 -u $(id -u):$(id -g) -e HOME=$HOME -e USER=$USER -v /home:/home $USEROPT1 $USEROPT2 ${IMAGE} ${SCR}
}