"""Entry point."""

import argparse
import time

import torch

import graphnas.trainer as trainer
import graphnas.utils.tensor_utils as utils
import warnings

warnings.filterwarnings('ignore')
import os



def build_args():
    parser = argparse.ArgumentParser(description='GraphNAS')
    register_default_args(parser)
    args = parser.parse_args()

    return args


def register_default_args(parser):
    parser.add_argument('--mode', type=str, default='train',
                        choices=['train', 'derive'],
                        help='train: Training GraphNAS, derive: Deriving Architectures')
    parser.add_argument('--random_seed', type=int, default=123)
    parser.add_argument("--cuda", type=bool, default=True, required=False,
                        help="run in cuda mode")
    parser.add_argument('--save_epoch', type=int, default=2)
    parser.add_argument('--max_save_num', type=int, default=5)
    # controller

    parser.add_argument('--shared_initial_step', type=int, default=0)
    parser.add_argument('--batch_size', type=int, default=64)
    parser.add_argument('--entropy_mode', type=str, default='reward', choices=['reward', 'regularizer'])
    parser.add_argument('--entropy_coeff', type=float, default=1e-4)
    parser.add_argument('--shared_rnn_max_length', type=int, default=35)
    parser.add_argument('--load_path', type=str, default='')
    parser.add_argument('--layers_of_child_model', type=int, default=1)
    parser.add_argument('--search_mode', type=str, default='Zeng')
    parser.add_argument('--num_hops', type=str, default=8)          ###########
    parser.add_argument('--format', type=str, default='Zeng')
    parser.add_argument('--max_epoch', type=int, default=5)          ###########大循环

    parser.add_argument('--ema_baseline_decay', type=float, default=0.95)
    parser.add_argument('--discount', type=float, default=1.0)
    parser.add_argument('--controller_max_step', type=int, default=50,  #####  多少个模型
                        help='step for controller parameters')
    parser.add_argument('--controller_optim', type=str, default='adam')
    parser.add_argument('--controller_lr', type=float, default=3.5e-4,
                        help="will be ignored if --controller_lr_cosine=True")
    parser.add_argument('--controller_grad_clip', type=float, default=0)
    parser.add_argument('--tanh_c', type=float, default=2.5)
    parser.add_argument('--softmax_temperature', type=float, default=5.0)
    parser.add_argument('--derive_num_sample', type=int, default=20)   ###########
    parser.add_argument('--derive_finally', type=bool, default=True)
    parser.add_argument('--derive_from_history', type=bool, default=True)

    parser.add_argument('--num_granularity', type=int, default=0)

    # child model
    parser.add_argument("--dataset", type=str, default="Citeseer", required=False,
                        help="The input dataset.")
    parser.add_argument("--epochs", type=int, default=100,                  ###########
                        help="number of training epochs")
    parser.add_argument("--retrain_epochs", type=int, default=100,
                        help="number of training epochs")
    parser.add_argument("--multi_label", type=bool, default=False,
                        help="multi_label or single_label task")
    parser.add_argument("--residual", action="store_false",
                        help="use residual connection")
    parser.add_argument("--in-drop", type=float, default=0.0,
                        help="input feature dropout")
    parser.add_argument("--lr", type=float, default=0.2,
                        help="learning rate")
    parser.add_argument("--param_file", type=str, default="cora_test.pkl",
                        help="learning rate")
    parser.add_argument("--optim_file", type=str, default="opt_cora_test.pkl",
                        help="optimizer save path")
    parser.add_argument('--weight_decay', type=float, default=1e-4)
    parser.add_argument('--max_param', type=float, default=5E6)
    parser.add_argument('--supervised', type=bool, default=True)
    parser.add_argument('--submanager_log_file', type=str, default=f"hop_8_gran_16.txt") #sub_manager_logger_file_{time.time()}.txt
    parser.add_argument('--device', type=str, default='1')

    parser.add_argument('--normalization', type=str, default='AugNormAdj',
                    choices=['NormLap', 'Lap', 'RWalkLap', 'FirstOrderGCN',
                             'AugNormAdj', 'NormAdj', 'RWalk', 'AugRWalk',
                             'NoNorm', 'LowPass'],
                    help='Normalization method for the adjacency matrix.')

def main(args):  # pylint:disable=redefined-outer-name

    # os.environ['CUDA_VISIBLE_DEVICE'] = args.device
    if args.cuda and not torch.cuda.is_available():  # cuda is not available
        args.cuda = False
        print("CUDA NOT AVAILABLE!")
    # args.max_epoch = 1
    # args.controller_max_step = 1
    # args.derive_num_sample = 1
    torch.manual_seed(args.random_seed)
    if args.cuda:
        # print("Device Num", torch.cuda.current_device())
        torch.cuda.manual_seed(args.random_seed)

    utils.makedirs(args.dataset)

    trnr = trainer.Trainer(args)

    if args.mode == 'train':
        print(args)
        trnr.train()
    elif args.mode == 'derive':
        trnr.derive()
    else:
        raise Exception(f"[!] Mode not found: {args.mode}")


if __name__ == "__main__":
    args = build_args()
    main(args)
