import argparse
import numpy as np
import os
import sys
import yaml

from lib.utils import load_graph_data
from model.pytorch.dcrnn_supervisor import DCRNNSupervisor


def run_dcrnn(args):
    with open(args.config_filename) as f:
        supervisor_config = yaml.load(f)

        graph_pkl_filename = supervisor_config['data'].get('graph_pkl_filename')
        sensor_ids, sensor_id_to_ind, adj_mx = load_graph_data(graph_pkl_filename)

        supervisor = DCRNNSupervisor(random_seed=int(args.random_seed),adj_mx=adj_mx, **supervisor_config)
        mean_score, outputs = supervisor.evaluate('test')
        np.savez_compressed(args.output_filename, **outputs)
        print("MAE : {}".format(mean_score))
        print('Predictions saved as {}.'.format(args.output_filename))


# if __name__ == '__main__':
#     sys.path.append(os.getcwd())
#     parser = argparse.ArgumentParser()
#     parser.add_argument('--use_cpu_only', default=False, type=str, help='Whether to run tensorflow on cpu.')
#     parser.add_argument('--config_filename', default='data/model/pretrained/COV-19/config.yaml', type=str,
#                         help='Config file for pretrained model.')
#     parser.add_argument('--output_filename', default='data/alltest_deepgleam_point_week33_seed1_epo69.npz')
#     args = parser.parse_args()
#     run_dcrnn(args)


if __name__ == '__main__':
    sys.path.append(os.getcwd())
    parser = argparse.ArgumentParser()
    parser.add_argument('--use_cpu_only', default=False, type=str, help='Whether to run tensorflow on cpu.')
    parser.add_argument('--config_filename', default='data/model/pretrained/COV-19/config.yaml', type=str,
                        help='Config file for pretrained model.')
    parser.add_argument('--output_filename', default='test.npz', type=str,
                        help='Output filename.')
    parser.add_argument('--random_seed',default=1,type=int, help='Select random seed.')
    args = parser.parse_args()
    run_dcrnn(args)
