import argparse
from trainer import train_led_graph as led


def parse_config():
	parser = argparse.ArgumentParser()
	parser.add_argument("--cuda", default=True)
	parser.add_argument("--learning_rate", type=int, default=0.002)
	parser.add_argument("--max_epochs", type=int, default=128)

	parser.add_argument('--cfg', default='led_augment')
	parser.add_argument('--gpu', type=int, default=0, help='Specify which GPU to use.')
	parser.add_argument('--train', type=int, default=1, help='Whether train or evaluate.')

	parser.add_argument("--info", type=str, default='graph', help='Name of the experiment. '
															 'It will be used in file creation.')

	# Graph variant knobs.
	parser.add_argument('--top_n', type=int, default=5,
						help='Number of sparse neighbors per agent (max 10 for NBA with A=11).')
	parser.add_argument('--residual_on', type=str, default='eps', choices=['eps', 'y0'],
						help='Where to apply the graph residual: directly on epsilon, or '
							 'on the implied y_0 estimate (then re-projected to epsilon).')
	parser.add_argument('--use_sigma', action='store_true',
						help='If set, pass the initializer variance_estimation to the graph '
							 'as per-agent uncertainty (modulates node features and edges).')
	parser.add_argument('--use_v6_graph', action='store_true',
						help='If set, use MoFlow V6-style RAG-scoring graph (FutureInteractionGraphV6) '
							 'instead of the default hand-crafted distance-based graph.')
	parser.add_argument('--uncertainty_weight', type=float, default=1.0,
						help='Weight for uncertainty NLL loss. Set to 0 for nosigma ablation.')
	parser.add_argument('--edge_mode', type=str, default='full',
						choices=['full', 'dist_only', 'relpos_only', 'heading_only', 'vel_only', 'full_relvel'],
						help='Edge feature mode for RelTrajEncoder ablation.')
	parser.add_argument('--resume_epoch', type=int, default=0,
						help='Resume from this checkpoint epoch (0=start fresh).')
	parser.add_argument('--neighbor_mode', type=str, default='rag',
						choices=['rag', 'l2', 'semantic'],
						help='Neighbor selection: rag (semantic+geo), l2 (closest), or semantic (learned only)')
	return parser.parse_args()


def main(config):
	t = led.Trainer(config)
	if config.train == 1:
		t.fit()
	else:
		t.test_single_model()


if __name__ == "__main__":
	config = parse_config()
	main(config)