# Copyright (c) 2021  PaddlePaddle 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 governing permissions and
# limitations under the License.

import json
import copy
import numpy as np
from paddleslim.nas import GPNAS

# 使用GP-NAS参加[CVPR 2021 NAS国际比赛](https://www.cvpr21-nas.com/competition) Track2 demo
# [CVPR 2021 NAS国际比赛Track2 studio地址](https://aistudio.baidu.com/aistudio/competition/detail/71?lang=en)
# [AI studio GP-NAS demo](https://aistudio.baidu.com/aistudio/projectdetail/1824958) 
# demo 基于paddleslim自研NAS算法GP-NAS:Gaussian Process based Neural Architecture Search 
# 基于本demo的改进版可以获得双倍奖金


def preprare_trainning_data(file_name, t_flag):
    ## t_flag ==1 using all trainning data
    ## t_flag ==2 using half trainning data
    with open(file_name, 'r') as f:
        arch_dict = json.load(f)
    Y_all = []
    X_all = []
    for sub_dict in arch_dict.items():
        Y_all.append(sub_dict[1]['acc'] * 100)
        X_all.append(np.array(sub_dict[1]['arch']).T.reshape(4, 16)[2])
    X_all, Y_all = np.array(X_all), np.array(Y_all)
    X_train, Y_train, X_test, Y_test = X_all[0::t_flag], Y_all[
        0::t_flag], X_all[1::t_flag], Y_all[1::t_flag]
    return X_train, Y_train, X_test, Y_test


if __name__ == '__main__':
    stage1_file = './datasets/Track2_stage1_trainning.json'
    stage2_file = './datasets/Track2_stage2_few_show_trainning.json'
    X_train_stage1, Y_train_stage1, X_test_stage1, Y_test_stage1 = preprare_trainning_data(
        stage1_file, 1)
    X_train_stage2, Y_train_stage2, X_test_stage2, Y_test_stage2 = preprare_trainning_data(
        stage2_file, 2)
    gpnas = GPNAS()
    w = gpnas.get_initial_mean(X_test_stage1, Y_test_stage1)
    init_cov = gpnas.get_initial_cov(X_train_stage1)
    error_list = np.array(
        Y_test_stage2.reshape(len(Y_test_stage2), 1) - gpnas.get_predict(
            X_test_stage2))
    print('RMSE trainning on stage1 testing on stage2:',
          np.sqrt(np.dot(error_list.T, error_list) / len(error_list)))
    gpnas.get_posterior_mean(X_train_stage2[0::3], Y_train_stage2[0::3])
    gpnas.get_posterior_mean(X_train_stage2[1::3], Y_train_stage2[1::3])
    gpnas.get_posterior_cov(X_train_stage2[1::3], Y_train_stage2[1::3])
    error_list = np.array(
        Y_test_stage2.reshape(len(Y_test_stage2), 1) - gpnas.get_predict_jiont(
            X_test_stage2, X_train_stage2[::1], Y_train_stage2[::1]))
    print('RMSE using stage1 as prior:',
          np.sqrt(np.dot(error_list.T, error_list) / len(error_list)))

[(50, 'paddleslim.nas.GPNAS', 'GPNAS', False, 'from paddleslim.nas import GPNAS\n'), (31, 'json.load', 'json.load', False, 'import json\n'), (37, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (37, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (57, 'numpy.dot', 'np.dot', True, 'import numpy as np\n'), (65, 'numpy.dot', 'np.dot', True, 'import numpy as np\n'), (36, 'numpy.array', 'np.array', True, 'import numpy as np\n')]