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])]
    else:
        sel_jobs = jobs.select(job_numbers[0], job_numbers[1])
    lfns = {}
    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



dest1 = '/disk/data3/lhcb/elena/B2KShh/ntuples/jobs'
dest2 = '/disk/data3/lhcb/elena/Lb2Lcmunu/ntuples/jobs'
scripts = '/cvmfs/lhcb.cern.ch/lib/lhcb/LHCBDIRAC/LHCBDIRAC_v8r2p29/scripts'
import os

def download(mylist, dest):
    #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]
            if not os.path.isfile('%s/%s/%s.root' % (dest, k, file_id)):
                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 deleteFromGrid(mylist):
    #os.system('SetupProject.sh LHCbDirac')
    #os.system('lhcb-proxy-init')
    for k in mylist.keys():
        l = len(mylist[k])
        for i,rootfile in enumerate(mylist[k]):
            print '%s\t\t: Deleting %s of %s' % (k, i, l)
            os.system('%s/dirac-dms-remove-files %s' % (scripts, rootfile))


def hadd(mylist, dest):
    #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(302, 304)
    # out = '/afs/cern.ch/work/e/egraveri/Analysis/ganga/Lb2Lcmunu'
    # #pickle.dump(mylist, open(out + '/lfns.pkl', 'wb'))
    # mylist = pickle.load(open(out + '/lfns.pkl', 'rb'))
    # download(mylist, dest2)
    # #deleteFromGrid(mylist)
    # #hadd(mylist)
    #mylist2 = getlist(287, 300)
    out2 = "/home/hep/egraveri/cmtuser/Bender_v27r0p1/Phys/B2KShh/job"
    #pickle.dump(mylist2, open(out2 + '/lfns.pkl', 'wb'))
    mylist2 = pickle.load(open(out2 + '/lfns.pkl', 'rb'))
    #download(mylist2, dest1)   
    hadd(mylist2, dest1)
    #deleteFromGrid(mylist2)