#!/usr/bin/env python
'''
 This module takes an integer and translates it into a workload
 e.g. takes 1 gives (QCD_20to30_BCtoE,central) depending on the action wanted
 It will also parse all parameters defined in the condor job description file
 under 'arguments'.
 
 This module is not meant to be used in a manual workflow.
 It is in to be used with the condor script only
 
 Usage:
 job_mapper --operation=test|single|analysis [--noop, --energy, --process, --sample, --mode, --cmssw, --return_*]
 
 mandatory:
 --operation=test|single|single-sample-analysis|analysis
     test: set of default parameters just for testing
     single: run just one job based on the parameters (need to specify energy, mode, sample & cores)
     single-sample-analysis: runs all analysis modes for a given sample and centre-of-mass energy (need to specify energy, cores & sample)
     analysis: run all jobs for a given centre-of-mass energy (need to specify energy & cores)
     
 optional parameters:
 --cores: number of cores to be used per condor job
 --noop: no operation flag. Will create condor job template, but won't submit the jobs
 --energy=<centre-of-mass energy> 
 --process <condor process number> 
 --sample=<valid sample> (TTJet, etc) 
 --mode=<valid mode of operation> (central, JES_up, etc)
 --cmssw=53X|73X|74X|75X
 
 defining return values:
 return_<parameter>: script will return the specified parameter (e.g. return_mode returns the mode)
 return_n_jobs: return the number of jobs
 
'''
from optparse import OptionParser
import sys
import os
from imp import load_source
dirname, _ = os.path.split(os.path.abspath(__file__))
analysis_info = load_source( 'analysis_info', dirname + '/../python/analysis_info.py' )

# same in 7 and 8 TeV
analysis_modes = analysis_info.analysisModes
# don't need 'Test' here
analysis_modes.remove('Test')

samples_7TeV = analysis_info.datasets_7TeV.keys()
samples_8TeV = analysis_info.datasets_8TeV.keys()

def build_matrix(energy, chosen_sample = None):
    samples = None
    if energy == 7:
        samples = samples_7TeV
    if energy == 8:
        samples = samples_8TeV
    result = []
    if chosen_sample:
        if not chosen_sample in samples:
            sys.exit("ERROR Trying to run unknown sample '%s'" % options.sample)
        for mode in analysis_modes:
            result.append((chosen_sample, mode))
        return result, [chosen_sample]
    else:
        for sample in samples:
            for mode in analysis_modes:
                result.append((sample, mode))
    return result, samples

def main(options, args = []):
    if options.operation == 'test':
        # set test defaults
        if not options.energy:
            options.energy = 7
        if not options.cores:
            options.cores = 1
        if not options.mode:
            options.mode = 'central'
        if not options.sample:
            options.sample = 'TTJet'
    # first the static options    
    if options.return_cmssw:
        return options.cmssw
    if options.return_energy:
        return options.energy
    if options.return_cores:
        return options.cores
    if options.return_process:
        return options.process
    if options.return_operation:
        return options.operation
    # now the dynamic options
    matrix, samples = None, None
    if options.operation == 'single-sample-analysis':
        matrix, samples = build_matrix(options.energy, options.sample)
    else:
        matrix, samples = build_matrix(options.energy)
    if options.return_n_jobs:
        if options.operation == 'test' or options.operation == 'single':
            return 1
        else:
            return len(matrix)
    
    if options.return_sample:
        if options.sample:
            if not options.sample in samples:
                sys.exit("ERROR Trying to run unknown sample '%s'" % options.sample)
            return options.sample
        else:
            return matrix[options.process][0]
    
    if options.return_mode:
        if options.mode:
            if not options.mode in analysis_modes:
                sys.exit("Trying to run unknown mode: %s" % options.mode)
            return options.mode
        else:
            return matrix[options.process][1]
        
    if options.return_noop:
        if options.noop:
            return 1
        else:
            return 0

def parse_args(parameters = []):
    parser = OptionParser( __doc__ )
    parser.add_option( '--process', dest = 'process', help = 'process',
                       type = int  )
    parser.add_option( '--cores', dest = 'cores', help = 'cores',
                       type = int  )
    parser.add_option( '--sample', dest = 'sample', help = 'sample' )
    parser.add_option( '--mode', dest = 'mode', help = 'mode' )
    parser.add_option( '--cmssw', dest = 'cmssw', 
                       help = 'cmssw version (53X|73X|74X)', default = '53X' )
    parser.add_option( '--operation', dest = 'operation', 
                       help = 'mode of operation: analysis|single|test', default = 'test' )
    parser.add_option( '--energy', dest = 'energy', 
                       help = 'centre-of-mass energy (7|8|13|14 TeV)',
                       type = int )
    parser.add_option("--noop", action = "store_true", dest = "noop",
                          help = "Returns the noop flag" )
    parser.add_option("--return_sample", action = "store_true", dest = "return_sample",
                          help = "Returns the sample from the map based on process" )
    parser.add_option("--return_mode", action = "store_true", dest = "return_mode",
                          help = "Returns the analysis mode from the map based on process" )
    parser.add_option("--return_cmssw", action = "store_true", dest = "return_cmssw",
                          help = "Returns the CMSSW version" )
    parser.add_option("--return_energy", action = "store_true", dest = "return_energy",
                          help = "Returns the centre-of-mass energy" )
    parser.add_option("--return_cores", action = "store_true", dest = "return_cores",
                          help = "Returns the number of requested cores" )
    parser.add_option("--return_process", action = "store_true", dest = "return_process",
                          help = "Returns the number of requested cores" )
    parser.add_option("--return_operation", action = "store_true", dest = "return_operation",
                          help = "Returns the operation" )
    parser.add_option("--return_n_jobs", action = "store_true", dest = "return_n_jobs",
                          help = "Returns the operation" )
    parser.add_option("--return_noop", action = "store_true", dest = "return_noop",
                          help = "Returns the noop flag" )
    
    options, args = None, None
    if len(parameters) > 0:
        ( options, args ) = parser.parse_args(parameters)
    else:
        ( options, args ) = parser.parse_args()
    return options, args

if __name__ == '__main__':
    options, args = parse_args()
    print main(options, args)
    