import pickle

global jobs

def getlist(*job_numbers):
	for j in job_numbers:
		print j
	if not job_numbers or len(job_numbers) > 2:
		return
	if len(job_numbers) == 1:
		sel_jobs = [jobs(job_numbers[0])]
	lfns = {}
	sel_jobs = jobs.select(job_numbers[0], job_numbers[1])
	print 'Selecting jobs %s' % [j.id for j in sel_jobs]
	for j in sel_jobs:
	    lfns[j.name] = []
	    for sj in j.subjobs:
	        if sj.status == 'completed' or sj.status == 'completing':
	            outfiles = sj.outputfiles
	            rootfile = [f for f in outfiles if '.root' in f.namePattern][0]
	            lfn = rootfile.lfn
	            lfns[j.name].append(lfn)
	        else:
	            print '\tWARNING: no outputs for subjob %s.%s' % (j.id, sj.id)
	#pickle.dump(lfns, open('B2KSph-MC.pkl', 'wb'))
	return lfns



dest = '/disk/users/elena/B2KShh/ntuples/jobs'
scripts = '/cvmfs/lhcb.cern.ch/lib/lhcb/LHCBDIRAC/LHCBDIRAC_v8r2p27/scripts'
import os

def download(mylist):
	#os.system('SetupProject.sh LHCbDirac')
	#os.system('lhcb-proxy-init')
	for k in mylist.keys():
		os.system('mkdir -p %s/%s' % (dest, k))
		l = len(mylist[k])
		for i,rootfile in enumerate(mylist[k]):
			print '%s\t\t: %s of %s' % (k, i, l)
			file_id = rootfile.split('/')[-2]
			file_name = rootfile.split('/')[-1]
			os.system('%s/dirac-dms-get-file %s' % (scripts, rootfile))
			os.system('mv %s %s/%s/%s.root' % (file_name, dest, k, file_id))


def hadd(mylist):
	#os.system('SetupProject Bender v27r0p1')
	paths = ['%s/%s/' % (dest, k) for k in mylist.keys()]
	print 'Found folders:'
	for path in paths:
		target = dest + '/' + path.split('/')[-2] + '.root'
		print '\t %s' % path
		print '\t\t Target is: %s' % target
    	files = os.listdir(path)
    	paths = [path + f for f in files]
    	command = 'hadd -f %s %s' % (target, ' '.join(paths))
    	print command
    	os.system(command)



if __name__ == '__main__':
	#mylist = getlist(243, 256)
	#pickle.dump(mylist, open('lfns.pkl', 'wb'))
	mylist = pickle.load(open('lfns.pkl', 'rb'))
	#download(mylist)
	hadd(mylist)