"""
.. module:: cobaya.grid_tools.gridrun

:Synopsis: script for running a grid of jobs
:Author: Antony Lewis and Jesus Torrado (based on CosmoMC version of the same code)

"""

import hashlib
import os

from cobaya.grid_tools import batchjob_args, jobqueue
from cobaya.grid_tools.gridmanage import run_and_wait
from cobaya.tools import warn_deprecation


def grid_run(args=None):
    warn_deprecation()
    Opts = batchjob_args.BatchArgs(
        prog='cobaya-grid-run',
        desc='Submit jobs to run chains or importance sample',
        notExist=True, notall=True, converge=True)
    jobqueue.addArguments(Opts.parser, combinedJobs=True)
    Opts.parser.add_argument('--noqueue', nargs='?', default=None, type=int, const=1,
                             help='run directly, not using queue, ignoring most other '
                                  'arguments (e.g. for fast tests). Optional argument '
                                  'specifies how many to run in parallel.'
                                  'To use mpi, include mpirun in the --program argument.')
    Opts.parser.add_argument('--subitems', action='store_true',
                             help='include sub-grid items')
    Opts.parser.add_argument('--not_queued', action='store_true')
    Opts.parser.add_argument('--minimize', action='store_true',
                             help='Run minimization jobs')
    Opts.parser.add_argument('--importance_minimize', action='store_true',
                             help='Run minimization jobs for chains '
                                  'that are importance sampled')
    Opts.parser.add_argument('--minimize_failed', action='store_true',
                             help='run where minimization previously failed')
    Opts.parser.add_argument('--checkpoint_run', nargs='?', default=None, const=0,
                             type=float,
                             help='run if stopped and not finished; if optional value '
                                  'given then only run chains with convergence '
                                  'worse than the given value')
    Opts.parser.add_argument('--importance_ready', action='store_true',
                             help='where parent chain has converged and stopped')
    Opts.parser.add_argument('--importance_changed', action='store_true',
                             help='run importance jobs where the parent chain has '
                                  'changed since last run')
    Opts.parser.add_argument('--parent_converge', type=float, default=0,
                             help='minimum R-1 convergence for importance job parent')
    Opts.parser.add_argument('--parent_stopped', action='store_true',
                             help='only run if parent chain is not still running')
    (batch, args) = Opts.parseForBatch(args)

    if args.not_queued:
        assert not args.noqueue, 'Cannot use --noqueue and --not_queued'
        print('Getting queued names...')
        queued = jobqueue.queue_job_names(args.batchPath)

    def notQueued(name):
        return not any(name in job for job in queued)

    variant = ''
    if args.importance_minimize:
        variant = '_minimize'
        if args.importance is None:
            args.importance = []
    if args.minimize:
        args.noimportance = True
        variant = '_minimize'
    if args.importance is None:
        if args.importance_changed or args.importance_ready:
            args.importance = []
        else:
            args.noimportance = True
    isMinimize = args.importance_minimize or args.minimize
    if args.combine_one_job_name:
        print('Combining multiple (hopefully fast) into single job script: ' +
              args.combine_one_job_name)
    yaml_files = []
    if args.noqueue:
        program = jobqueue.get_defaulted("program",
                                         default=jobqueue.default_program,
                                         program=args.program)
    else:
        jobqueue.check_arguments(**args.__dict__)
    processes = set()

    def jobName():
        s = "-".join([os.path.splitext(os.path.basename(ini))[0] for ini in yaml_files])
        if len(yaml_files) < 2 or len(s) < 70:
            return s
        base = os.path.basename(yaml_files[0])
        if len(base) > 70:
            base = base[:70]
        return base + '__' + hashlib.md5(s.encode('utf8')).hexdigest()[:16]

    def submitJob(ini):
        if not args.dryrun:
            print('Submitting...' + ini)
        else:
            print('... ' + ini)
        if args.noqueue:
            if not args.dryrun:
                run_and_wait(processes, program.split(' ') + [ini], args.noqueue)
            return
        yaml_files.append(ini)
        if args.combine_one_job_name:
            return
        if len(yaml_files) >= args.runs_per_job:

            if args.runs_per_job > 1:
                print('--> jobName: ', jobName())
            jobqueue.submitJob(jobName(), yaml_files, **args.__dict__)
            yaml_files.clear()

    for jobItem in Opts.filteredBatchItems(wantSubItems=args.subitems):
        if ((not args.notexist or isMinimize and not jobItem.chainMinimumExists()
             or not isMinimize and not jobItem.chainExists()) and (
                    not args.minimize_failed or not jobItem.chainMinimumConverged())
            and (isMinimize or args.notall is None or not jobItem.allChainExists(
                    args.notall))) \
                and (not isMinimize or getattr(jobItem, 'want_minimize', True)):
            if not args.parent_converge or not jobItem.isImportanceJob or \
                    jobItem.parent.hasConvergeBetterThan(args.parent_converge):
                if args.converge == 0 or not jobItem.hasConvergeBetterThan(
                        args.converge, returnNotExist=True):
                    if args.checkpoint_run is None or jobItem.wantCheckpointContinue(
                            args.checkpoint_run) and jobItem.notRunning():
                        if (not jobItem.isImportanceJob or isMinimize
                                or (args.importance_ready
                                    and jobItem.parent.chainFinished()
                                    or not args.importance_ready
                                    and jobItem.parent.chainExists())
                                and (not args.importance_changed
                                     or jobItem.parentChanged())
                                and (not args.parent_stopped
                                     or jobItem.parent.notRunning())):
                            if not args.not_queued or notQueued(jobItem.name):
                                submitJob(jobItem.yaml_file(variant))

    if args.noqueue:
        if args.noqueue > 1:
            run_and_wait(processes)
    elif yaml_files:
        if args.runs_per_job > 1:
            print('--> jobName: ', jobName())
        jobqueue.submitJob(args.combine_one_job_name or jobName(), yaml_files,
                           sequential=args.combine_one_job_name is not None,
                           **args.__dict__)
