#!/usr/bin/env python

import _init_paths
import caffe
from fast_rcnn.train import get_training_roidb, train_net
from fast_rcnn.config import cfg, cfg_from_file
import xmldb
import help_parser as hp
import os
import numpy as np
import shutil

MODEL_DIR = _init_paths.MODEL_DIR

def parse_args():
	inside = sorted(os.listdir(MODEL_DIR))
	nets = [f for f in inside if os.path.isdir(os.path.join(MODEL_DIR,f)) and f != 'models']
	parser = hp.HelpParser(description='Train a Fast R-CNN network')
	parser.add_argument('use_net', help='initialize with pretrained model weights', choices=nets)
	parser.add_argument('out', help='Name of output dir in model directory')
	parser.add_argument('xml', help='Dataset to train on', nargs='+')
	parser.add_argument('-cache', dest='cache', help='Use cache for preloading dataset [False]', action='store_true', default=False)
	parser.add_argument('-gpu', dest='gpu', help='GPU device id to use [0]', default=0, type=int)
	parser.add_argument('-solver', dest='solver', help='Optional solver prototxt', default=None, type=str)
	parser.add_argument('-iters', dest='max_iters', help='number of iterations to train', default=4000, type=int)
	parser.add_argument('-cfg', dest='cfg', help='optional config file', default=None, type=str)
	parser.add_argument('-rand', dest='rand', help='randomize (do not use a fixed seed)', action='store_true')
	parser.add_argument('-keep_int', dest='remove_int', default=True, action='store_true', help='Keep intermediate models [False]')
	parser.add_argument('-nf', dest='flipped', default=True, action='store_false', help='Turns off using horizontally flipped images')

	args = parser.parse_args()

	return args


if __name__ == '__main__':
	args = parse_args()
	
	if args.solver is None:
		args.solver = os.path.join(MODEL_DIR, args.use_net, 'train', 'end2end', 'solver.pt')
	
	cfg.GPU_ID = args.gpu
	caffe.set_mode_gpu()
	caffe.set_device(args.gpu)
	cfg.TRAIN.USE_FLIPPED = args.flipped
	cfg.TRAIN.ASPECT_GROUPING = False
	cfg.TRAIN.HAS_RPN = True
	cfg.TRAIN.IMS_PER_BATCH = 1
	cfg.TRAIN.BBOX_NORMALIZE_TARGETS_PRECOMPUTED = True
	cfg.TRAIN.RPN_POSITIVE_OVERLAP = 0.7
	cfg.TRAIN.RPN_BATCHSIZE = 256
	cfg.TRAIN.PROPOSAL_METHOD = 'gt'
	
	if args.cfg is not None:
		cfg_from_file(args.cfg)
	
	if not args.rand:
		np.random.seed(cfg.RNG_SEED)
		caffe.set_random_seed(cfg.RNG_SEED)
    	
	imdb = xmldb.xmldb(args.xml, load_cache=args.cache)
	roidb = get_training_roidb(imdb)
	
	outdir = os.path.join(MODEL_DIR, args.out)
	
	models = train_net(args.solver, roidb, outdir, pretrained_model=os.path.join(MODEL_DIR, args.use_net, args.use_net + '.caffemodel'), max_iters=args.max_iters)
	if args.remove_int:
		for m in models[:-1]:
			os.remove(m)
	os.rename(models[-1], os.path.join(outdir, args.out + '.caffemodel'))
	shutil.copy(os.path.join(MODEL_DIR, args.use_net, 'test.pt'), os.path.join(outdir, 'test.pt'))
	shutil.copytree(os.path.join(MODEL_DIR, args.use_net, 'train'), os.path.join(outdir, 'train'), symlinks=True)
	
	
