#!/usr/bin/env python
'''
    Submission script to DICE for src/unfolding_tests/create_unfolding_pull_data.py
'''
from __future__ import print_function
from optparse import OptionParser
from dps.config.xsection import XSectionConfig
from dps.condor import job
from dps.condor.jobtypes.unfolding_pull_job_new import UnfoldingPullJob
from math import log10

def get_tau_values(min, max, spacing):
    taus = []
    logMin = log10(min)
    logMax = log10(max)
    r = int(logMax - logMin)
    for i in range(r*spacing):
        t = 10**(logMin + i/float(spacing))
        taus.append(t)
    return taus

def main():
    parser = OptionParser(__doc__)
    parser.add_option("-c", "--com", dest="com", type=int,
                      help="Centre-of-mass energy to be used for submission")
    parser.add_option("-v", "--variables", dest="variables",
                      help="Variables to be analysed, separated by commas")
    parser.add_option("-i", "--input_directory", dest="input_directory",
                      help="Toy MC input directory", default='data/toy_mc/')
    parser.add_option("-n", dest="n_toy_data",
                      help="Number of toy datasets to analyse", default=3000)
    parser.add_option("--scan_tau", dest="scan_tau", action='store_true',
                      help="Produce unfolding pulls for a range of tau values", default=False)
    parser.add_option("--tau_min", dest="tau_min",
                      help="Minimum tau", default=1E-6)
    parser.add_option("--tau_max", dest="tau_max",
                      help="Maximum tau", default=1e-1)
    parser.add_option("--spacing", dest="spacing",
                      help="Number points between each factor of 10 in tau", default = 5)
    parser.add_option("--doBestTau", dest="do_best_tau", action='store_true',
                      help="In addition to scanning tau, include best choice for tau from config", default = False)
    parser.add_option("-m", "--method", dest="method",
                      help="Unfolding method: TUnfold (default), RooUnfoldSvd, TSVDUnfold, RooUnfoldTUnfold, RooUnfoldInvert, RooUnfoldBinByBin, RooUnfoldBayes",
                      default='TUnfold')
    parser.add_option("-s", "--samples", dest="samples",
                      help="Which input toy samples to run on, separated by commas: madgraph, powhegPythia, amcatnlo",
                      default='powhegPythia')

    (options, _) = parser.parse_args()

    # first the global variables
    centre_of_mass_energy = options.com
    config = XSectionConfig(centre_of_mass_energy)
    method = options.method
    n_toy_data = int( options.n_toy_data )
    output_folder = 'data/pull_data'

    do_best_tau = options.do_best_tau
    tau_range = []
    scan_tau = options.scan_tau
    if scan_tau:
        tau_range.extend( get_tau_values( options.tau_min, options.tau_max, options.spacing ) )

    samples = options.samples.split(',')
    channels = ['combined']
    variables = options.variables.split(',')
    input_file_directory = options.input_directory

    pull_jobs = []

    pull_job = UnfoldingPullJob(
        input_file_directory=input_file_directory,
        method=method,
        channels=channels,
        centre_of_mass=centre_of_mass_energy,
        response=config.unfolding_central,
        samples=samples,
        variables=variables,
        n_toy_data=n_toy_data,
        output_folder=output_folder,
        do_best_tau=do_best_tau,
        tau_values=tau_range
    )

    # One job per variable
    n_sub_jobs = len(samples) * len(variables) * len(channels) * len(tau_range)
    if do_best_tau:
        n_sub_jobs += len(samples) * len(variables) * len(channels)

    condor_job = job.Condor(n_sub_jobs, -1, request_memory=10)
    condor_job.add_job( pull_job )

    condor_job.submit()
    print('All jobs submitted')


if __name__ == '__main__':
    main()
