import configargparse

PARSER = configargparse.ArgParser(default_config_files=['configs/MountainCar-v0.config'])
PARSER.add('-c', '--config_path', required=False, is_config_file=True, help='config file path')
PARSER.add('--exp_name', required=True, help='name of experiment')
PARSER.add('--env_name', required=True, help='name of environment')
PARSER.add('--max_episodes', required=True, type=int, help='max number of episodes')
PARSER.add('--buffer_size', required=True, type=int, help='buffer size for caching data')
PARSER.add_argument('--layer_norm', dest='layer_norm', action='store_true', help='apply layer normalization')
PARSER.add_argument('--no_layer_norm', dest='layer_norm', action='store_false', help='turn off layer normalization')
PARSER.add_argument('--is_stateful', dest='is_stateful', action='store_true', help='maintain states when acting in the env')
PARSER.add_argument('--double_q', dest='double_q', action='store_true', help='if double-Q is used')
PARSER.add_argument('--dueling_q', dest='dueling_q', action='store_true', help='if dueling-Q is used')
PARSER.add_argument('--use_swa', dest='use_swa', action='store_true', help='whehter use Stochastic Weight Averaging')
PARSER.add_argument('--picky_swa', dest='picky_swa', action='store_true', help='whehter use Picky Stochastic Weight Averaging')
PARSER.add_argument('--render_mode', dest='render_mode', action='store_true')
PARSER.add_argument('--no_render_mode', dest='render_mode', action='store_false')
PARSER.add('--n_actions', required=True, type=int, help='number of discrete actions agent can take')
PARSER.add('--a_width', required=True, type=int, help='width of action vector')
PARSER.add('--o_size', required=True, type=int, help='observation size')
PARSER.add('--z_size', required=True, type=int, help='hidden state size')
PARSER.add('--seed', required=True, type=int, help='random seed for a single experiment')
PARSER.add('--data_path', required=True, type=str, help='path to the training data', default='./')
PARSER.add('--prior_data_path', required=True, type=str, help='path to the human data for the PPL', default='./')
PARSER.add('--zoo_data_path', required=True, type=str, help='path to the zoo agent data for the PPL', default='./')
PARSER.add('--human_data_path', required=True, type=str, help='path to the human data for the state model', default='./')

PARSER.add('--vae_batch_size', required=True, type=int, help='batch size for vae train')
PARSER.add('--vae_n_samples', default=1, type=int, help='number of samples draw from the models')
PARSER.add('--vae_learning_rate', required=True, type=float, help='vae learning rate')
PARSER.add('--vae_weight_decay', required=True, type=float, help='vae decoupled weight decay factor')
PARSER.add('--vae_kl_weight', required=True, type=float, help='vae kl-d weight')
PARSER.add('--vae_kl_regularize_weight', required=True, type=float, help='vae kl-d regularizer weight')
PARSER.add('--vae_num_epoch', required=True, type=int, help='vae num epoch for training')
PARSER.add('--vae_optimizer', type=str, help='Adam, RMSprop, SGD.', default='Adam')
PARSER.add('--sample_average', type=str, help="sample average method during acting, use 'states' or 'EFE_values'", default='states')
PARSER.add('--test_weight_avg', type=str, help="model weight averaging during test, 'ema' or 'winPolyak'", default=None)
PARSER.add('--vae_dropout_rate', default=0, type=float, help='dropout rate applied to VAE modules')
PARSER.add('--l2_reg', default=1e-3, type=float, help='L2 regularization factor')

PARSER.add('--planner_lookahead', type=int, default=30, help='num of time steps to repeat an action')
PARSER.add('--planner_plan_depth', type=int, default=3, help='depth in the planning horizon')
PARSER.add('--planner_n_samples', type=int, default=100, help='num of samples to draw during the state rollout')
PARSER.add('--planner_rho', type=float, default=0.1, help='risk aversion or greediness of the agent')
PARSER.add('--planner_gamma', type=float, default=1.0, help='temperature parameter used in policy selection')
PARSER.add_argument('--planner_full_efe', dest='planner_full_efe', action='store_true')
PARSER.add_argument('--planner_no_full_efe', dest='planner_full_efe', action='store_false')
PARSER.set_defaults(render_mode=True)
PARSER.set_defaults(layer_norm=False)
PARSER.set_defaults(is_stateful=False)
PARSER.set_defaults(double_q=False)
PARSER.set_defaults(dueling_q=False)
PARSER.set_defaults(use_swa=False)
PARSER.set_defaults(picky_swa=False)