#!/usr/bin/env modbuild

PERFORMANCE="yes"
PROFILE="no"
SP="no"

pbuild::prep() {

    local BRANCH
    if [[ "$V" =~ beta$ ]]; then
        BRANCH="ver${V_MAJOR}.${V_MINOR}"
    elif [[ "$V" =~ profile$ ]]; then
        BRANCH="${V_PKG}"
        PROFILE="yes"
    elif [[ "$V" =~ sp$ ]]; then
        BRANCH="ver${V_MAJOR}.${V_MINOR}"
        SP="yes"
    else
        BRANCH="${V_PKG}"
    fi

    if [[ -d "$SRC_DIR/.git" ]]; then
        cd "$SRC_DIR"
        git pull 
    else
        git clone --depth=1 -b "$BRANCH" https://github.com/3dem/relion.git "$SRC_DIR" || return $?
    fi
}


pbuild::pre_configure() {
    # Section will be switched on for -profile version
    #-----------------------------------------------------------
    if [[ "$PROFILE" == "yes" ]]; then
        # this needs to go into the config.yml file for >5.0-1beta
        if [[ "$V_MAJOR" -lt "5" ]]; then
            : pbuild::add_configure_args "-DCMAKE_BUILD_TYPE=BENCHMARKING"
        fi
    fi
    #------------------------------------------------------------

    # Section will be switched on for single-precision  version 
    # Safe memory usage for multibody refinements
    #------------------------------------------------------------
    if [[ "$SP" == "yes" ]]; then
        # this needs to go into the config.yml file for >5.0-1beta
        if [[ "$V_MAJOR" -lt "5" ]]; then
            : pbuild::add_configure_args "-DDoublePrec_CPU=OFF"
        fi
    fi
    #------------------------------------------------------------

    # Set performance oriented compiler flags, specific for a given
    # architecture.
    # -------------------------------------------------------------
    if [[ "$PERFORMANCE" == "yes" && "$PROFILE" == "no" ]]; then
        # set compiler flags per arch
        echo "> setting performance system compiler flags"
        case "$(uname -m)" in
            "aarch64")
                export CC=gcc
                export CFLAGS='-mcpu=neoverse-v2 -O3 -pipe -Wno-error'
                export CXX=g++
                export CXXFLAGS='-mcpu=neoverse-v2 -O3 -pipe -Wno-error'
                export FCFLAGS='-mcpu=neoverse-v2 -O3 -pipe'
                export FFLAGS='-mcpu=neoverse-v2 -O3 -pipe'
                export LDFLAGS='-Wl,--as-needed'
                ;;
            "x86_64")
                export CC=gcc
                export CFLAGS='-march=x86-64-v3 -mtune=znver1 -O3 -pipe -Wno-error'
                export CXX=g++
                export CXXFLAGS='-march=x86-64-v3 -mtune=znver1 -O3 -pipe -Wno-error'
                export FCFLAGS='-march=x86-64-v3 -mtune=znver1 -O3 -pipe'
                export FFLAGS='-march=x86-64-v3 -mtune=znver1 -O3 -pipe'
                export LDFLAGS='-Wl,--as-needed'
                ;;
            *)
                echo ">> Unknown system arch found $(uname -m)! Using basic flags only!"
                export CC=gcc
                export CFLAGS='-O3 -pipe -Wno-error'
                export CXX=g++
                export CXXFLAGS='-O3 -pipe -Wno-error'
                export FCFLAGS='-O3 -pipe'
                export FFLAGS='-O3 -pipe'
                export LDFLAGS='-Wl,--as-needed'
                ;;
        esac
        echo ">> ${CC} ${CFLAGS}"
        echo ">> ${CXX} ${CXXFLAGS}"
    fi
    # -------------------------------------------------------------


    # The following section will only work for versions >= 5.0-beta 
    # -------------------------------------------------------------
    if [[ "$V_MAJOR" -ge "5" ]]; then
        #download blush weights, from zenodo
        echo "> downloading blush weights"
        mkdir -p "$PREFIX/torch/hub/checkpoints/relion_blush"
        wget -q https://zenodo.org/records/10072731/files/blush_v1.0.ckpt.gz -P "$PREFIX/torch/hub/checkpoints/relion_blush" \
        && gunzip "$PREFIX/torch/hub/checkpoints/relion_blush/blush_v1.0.ckpt.gz"
        mv "$PREFIX/torch/hub/checkpoints/relion_blush/blush_v1.0.ckpt" "$PREFIX/torch/hub/checkpoints/relion_blush/v1.0.ckpt" 
        touch  "$PREFIX/torch/hub/checkpoints/relion_blush/v1.0_installed.txt"
        
        #cp classranker that was downloaded on another machine before from BUILDBLOCK_DIR
        echo "> downloading classranker"
        mkdir -p "$PREFIX/torch/hub/checkpoints/relion_class_ranker"
        wget -q ftp://ftp.mrc-lmb.cam.ac.uk/pub/dari/classranker_v1.0.ckpt.gz -P "$PREFIX/torch/hub/checkpoints/relion_class_ranker" \
        && gunzip "$PREFIX/torch/hub/checkpoints/relion_class_ranker/classranker_v1.0.ckpt.gz"
        mv "$PREFIX/torch/hub/checkpoints/relion_class_ranker/classranker_v1.0.ckpt" "$PREFIX/torch/hub/checkpoints/relion_class_ranker/v1.0.ckpt"
        touch  "$PREFIX/torch/hub/checkpoints/relion_class_ranker/v1.0_installed.txt"
        
        # fltk stuff 
        echo "> downloading FLTK"
        mkdir -p "$SRC_DIR/external/fltk"
        wget -q ftp://ftp.mrc-lmb.cam.ac.uk/pub/scheres/fltk-1.3.5-source.tar.gz -P "$SRC_DIR/external/fltk/"

        # Install Miniconda
        echo "> installing miniconda"
        mkdir -p  "$PREFIX/miniconda"
        if [[ "$(uname -m)" =~ "aarch64" ]]; then
            wget -q https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-aarch64.sh -O "$PREFIX/miniconda/miniconda.sh"
        else
            wget -q https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O "$PREFIX/miniconda/miniconda.sh"
        fi
        bash "$PREFIX/miniconda/miniconda.sh" -b -u -p "$PREFIX/miniconda/"

        # Setup conda env; switch into relion repo directory to capture local
        # `.[vis]` pip package dependencies
        echo "> building conda env"
        pushd "$SRC_DIR"
        "$PREFIX/miniconda/condabin/conda" env create -q -f environment.yml
        "$PREFIX/miniconda/condabin/conda" env list
        popd
    fi
    
    #---------------------------------------------------------------
}