#!/bin/bash
#
# Copyright 2021 The Modelbox Project Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language 

CURRDIR=$(pwd)

export MS_ENABLE_GE=1
export MS_GE_TRAIN=1

setup_var() {
    TRAIN_DIR="${CURRDIR}/mnist-train"
    TRAIN_FILE="${TRAIN_DIR}/train.py"
    MINDSPORE_TRAIN_DOWNLOAD_URL="https://obs.dualstack.cn-north-4.myhuaweicloud.com/mindspore-website/notebook/r1.6/tutorials/zh_cn/mindspore_quick_start.py"

    rm -fr ${TRAIN_DIR}
    mkdir ${TRAIN_DIR}
    rm ${TRAIN_DIR}/checkpoint_lenet* -f
    cd ${TRAIN_DIR}
}

download_package() {
    if [ ! -f "${TRAIN_FILE}" ]; then
        wget ${MINDSPORE_TRAIN_DOWNLOAD_URL} -O ${TRAIN_FILE}
        if [ $? -ne 0 ]; then
            echo "download train script failed"
            return 1
        fi
    fi

    return 0
}

setup_export_code() {
    if grep "^export(net.*)" ${TRAIN_FILE}; then
        return;
    fi
    
    echo "
from mindspore import export, load_checkpoint, load_param_into_net
import glob

ckpt = glob.glob('*.ckpt')

net = LeNet5()
param_dict = load_checkpoint(ckpt[0])
load_param_into_net(net, param_dict)
input = np.random.uniform(0.0, 1.0, size=[1, 1, 32, 32]).astype(np.float32)
export(net, Tensor(input), file_name='mnist', file_format='MINDIR')
" >> ${TRAIN_FILE}
}

train() {
    setup_export_code
    if [ $? -ne 0 ]; then
        echo "setup export python script failed."
        return 1
    fi

    python3 ${TRAIN_FILE} --device_target=Ascend
    if [ $? -ne 0 ]; then
            echo "train failed"
            return 1
    fi

    return 0
}

export_lite_model() {
    converter_lite --fmk=MINDIR --modelFile=${TRAIN_DIR}/mnist.mindir --outputFile=${TRAIN_DIR}/mnist
    if [ $? -ne 0 ]; then
            echo "export lite model failed"
            return 1
    fi

    cp ${TRAIN_DIR}/mnist.ms ${CURRDIR}/

    return $?
}


main() {
    setup_var
    if [ $? -ne 0 ]; then
            return 1
    fi

    echo "working directory: ${TRAIN_DIR}"

    download_package
    if [ $? -ne 0 ]; then
            return 1
    fi

    train
    if [ $? -ne 0 ]; then
            return 1
    fi

    export_lite_model
    if [ $? -ne 0 ]; then
            return 1
    fi

    echo "train success."
    echo "model dir: ${TRAIN_DIR}"
    echo "quick start: https://www.mindspore.cn/tutorial/en/r0.5/quick_start/quick_start.html"
}

main
