grond

#!/usr/bin/env python

import math
import sys
import logging
from optparse import OptionParser

from pyrocko import util
from pyrocko.gf import Range

import grond

logger = logging.getLogger('main')

km = 1000.


def d2u(d):
    if isinstance(d, dict):
        return dict((k.replace('-', '_'), v) for (k, v) in d.iteritems())
    else:
        return d.replace('-', '_')


subcommand_descriptions = {
    'init': 'print example configuration',
    'events': 'print available event names for given configuration',
    'check': 'check data and configuration',
    'go': 'run Grond optimization',
    'forward': 'run forward modelling',
    'harvest': 'manually run harvesting',
    'plot': 'plot optimization result',
    'export': 'export results',
    'qc-polarization': 'check sensor orientations with polarization analysis',
}

subcommand_usages = {
    'init': 'init [options]',
    'events': 'events <configfile>',
    'check': 'check <configfile> <eventnames> ... [options]',
    'go': 'go <configfile> <eventnames> ... [options]',
    'forward': (
        'forward <rundir> [options]',
        'forward <configfile> <eventnames> ... [options]'),
    'harvest': 'harvest <rundir> [options]',
    'plot': 'plot <plotnames> <rundir> [options]',
    'export': 'export (best|mean|ensemble|stats) <rundirs> ... [options]',
    'qc-polarization': 'qc-polarization <configfile> <eventname> '
                       '<targetconfigname> [options]',
}

subcommands = subcommand_descriptions.keys()

program_name = 'grond'

usage_tdata = d2u(subcommand_descriptions)
usage_tdata['program_name'] = program_name

usage = '''%(program_name)s <subcommand> [options] [--] <arguments> ...

Subcommands:

    init            %(init)s
    events          %(events)s
    check           %(check)s
    go              %(go)s
    forward         %(forward)s
    harvest         %(harvest)s
    plot            %(plot)s
    export          %(export)s
    qc-polarization %(qc_polarization)s

To get further help and a list of available options for any subcommand run:

    %(program_name)s <subcommand> --help

''' % usage_tdata


def add_common_options(parser):
    parser.add_option(
        '--loglevel',
        action='store',
        dest='loglevel',
        type='choice',
        choices=('critical', 'error', 'warning', 'info', 'debug'),
        default='info',
        help='set logger level to '
             '"critical", "error", "warning", "info", or "debug". '
             'Default is "%default".')


def process_common_options(options):
    util.setup_logging(program_name, options.loglevel)


def cl_parse(command, args, setup=None, details=None):
    usage = subcommand_usages[command]
    descr = subcommand_descriptions[command]

    if isinstance(usage, basestring):
        usage = [usage]

    susage = '%s %s' % (program_name, usage[0])
    for s in usage[1:]:
        susage += '\n%s%s %s' % (' '*7, program_name, s)

    description = descr[0].upper() + descr[1:] + '.'

    if details:
        description = description + '\n\n%s' % details

    parser = OptionParser(usage=susage, description=description)

    if setup:
        setup(parser)

    add_common_options(parser)
    (options, args) = parser.parse_args(args)
    process_common_options(options)
    return parser, options, args


def die(message, err=''):
    if err:
        sys.exit('%s: error: %s \n %s' % (program_name, message, err))
    else:
        sys.exit('%s: error: %s' % (program_name, message))


def help_and_die(parser, message):
    parser.print_help(sys.stderr)
    sys.stderr.write('\n')
    die(message)


def command_init(args):

    dataset_config = grond.DatasetConfig(
        stations_path='stations.txt',
        events_path='events.txt',
        waveform_paths=['data'])

    target_configs = [grond.TargetConfig(
        distance_min=10*km,
        distance_max=1000*km,
        channels=['Z', 'R', 'T'],
        interpolation='multilinear',
        store_id='global_2s',
        inner_misfit_config=grond.InnerMisfitConfig(
            fmin=0.01,
            fmax=0.1))]

    s2 = math.sqrt(2.0)

    problem_config = grond.CMTProblemConfig(
        name_template='cmt_%(event_name)s',
        distance_min=2.*km,
        nbootstrap=100,
        mt_type='deviatoric',
        ranges=dict(
            time=Range(0, 10.0, relative='add'),
            north_shift=Range(-16*km, 16*km),
            east_shift=Range(-16*km, 16*km),
            depth=Range(1*km, 11*km),
            magnitude=Range(4.0, 6.0),
            rmnn=Range(-s2, s2),
            rmee=Range(-s2, s2),
            rmdd=Range(-s2, s2),
            rmne=Range(-1.0, 1.0),
            rmnd=Range(-1.0, 1.0),
            rmed=Range(-1.0, 1.0),
            duration=Range(1.0, 15.0)))

    config = grond.Config(
        dataset_config=dataset_config,
        target_configs=target_configs,
        problem_config=problem_config)

    print config


def command_events(args):
    def setup(parser):
        pass

    parser, options, args = cl_parse('events', args, setup)
    if len(args) != 1:
        help_and_die(parser, 'missing arguments')

    config_path = args[0]
    config = grond.read_config(config_path)

    for event_name in grond.get_event_names(config):
        print event_name


def command_check(args):
    def setup(parser):
        parser.add_option(
            '--target-ids', dest='target_string_ids', metavar='TARGET_IDS',
            help='process only selected targets. TARGET_IDS is a '
                 'comma-separated list of target IDs. Target IDs have the '
                 'form SUPERGROUP.GROUP.NETWORK.STATION.LOCATION.CHANNEL.')

        parser.add_option(
            '--plot', dest='show_plot', action='store_true',
            help='plot sample synthetics and data.')

        parser.add_option(
            '--waveforms', dest='show_waveforms', action='store_true',
            help='show raw, restituted, projected, and processed waveforms')

        parser.add_option(
            '--nrandom', dest='n_random_synthetics', metavar='N', type='int',
            default=10,
            help='set number of random synthetics to forward model (default: '
                 '10). If set to zero, create synthetics for the reference '
                 'solution.')

    parser, options, args = cl_parse('check', args, setup)
    if len(args) < 2:
        help_and_die(parser, 'missing arguments')

    config_path = args[0]
    event_names = args[1:]

    config = grond.read_config(config_path)

    target_string_ids = None
    if options.target_string_ids:
        target_string_ids = options.target_string_ids.split(',')

    grond.check(
        config,
        event_names=event_names,
        target_string_ids=target_string_ids,
        show_plot=options.show_plot,
        show_waveforms=options.show_waveforms,
        n_random_synthetics=options.n_random_synthetics)


def command_go(args):
    def setup(parser):
        parser.add_option(
            '--force', dest='force', action='store_true',
            help='overwrite existing run directory')
        parser.add_option(
            '--status', dest='status', default='state',
            help='status output selection (choices: state, matrix)')
        parser.add_option(
            '--parallel', dest='nparallel', type='int', default=1,
            help='set number of events to process in parallel')

    parser, options, args = cl_parse('go', args, setup)
    if len(args) < 2:
        help_and_die(parser, 'missing arguments')

    config_path = args[0]
    event_names = args[1:]

    config = grond.read_config(config_path)
    if options.status == 'quiet':
        status = ()
    else:
        status = tuple(options.status.split(','))

    grond.go(
        config,
        event_names=event_names,
        force=options.force,
        status=status,
        nparallel=options.nparallel)


def command_forward(args):
    def setup(parser):
        pass

    parser, options, args = cl_parse('forward', args, setup)
    if len(args) < 1:
        help_and_die(parser, 'missing arguments')

    event_names = args[1:]

    if not event_names:
        help_and_die(parser, 'no event names given')

    run_path = args[0]
    grond.forward(
        run_path,
        event_names=event_names)


def command_harvest(args):
    def setup(parser):
        parser.add_option(
            '--force', dest='force', action='store_true',
            help='overwrite existing harvest directory')
        parser.add_option(
            '--neach', dest='neach', type=int, default=10,
            help='take NEACH best samples from each chain (default: 10)')
        parser.add_option(
            '--weed', dest='weed', type=int, default=0,
            help='weed out bootstrap samples with bad global performance. '
                 '0: no weeding (default), '
                 '1: only bootstrap chains where all NEACH best samples '
                 'global misfit is less than the global average misfit of all '
                 'NEACH best in all chains plus one standard deviation are '
                 'included in the harvest ensemble, '
                 '2: same as 1 but additionally individual samples are '
                 'removed if their global misfit is greater than the global '
                 'average misfit of all NEACH best in all chains, '
                 '3: harvesting is done on the global chain only, bootstrap '
                 'chains are excluded')

    parser, options, args = cl_parse('harvest', args, setup)
    if len(args) != 1:
        help_and_die(parser, 'no rundir')

    run_path, = args
    grond.harvest(
        run_path,
        force=options.force,
        nbest=options.neach,
        weed=options.weed)


def command_plot(args):
    from grond import plot

    def setup(parser):
        parser.add_option(
            '--save', dest='save', action='store_true', default=False,
            help='save figures to files')

        parser.add_option(
            '--format', '--formats', dest='formats', default='pdf',
            help='comma-separated list of ouptut formats (default: pdf)')

        parser.add_option(
            '--dpi', '--dpi', dest='dpi', type=float, default=120.,
            help='DPI setting for raster formats (default=120)')

    plotnames_avail = plot.available_plotnames()

    details = '''Available <plotnames> are: %s, or "all". Multiple plots are
selected by specifying a comma-separated list.''' % (
        ', '.join('"%s"' % x for x in plotnames_avail))

    parser, options, args = cl_parse('plot', args, setup, details)

    if len(args) != 2:
        help_and_die(parser, 'two arguments required')

    if args[0] == 'all':
        plotnames = plotnames_avail
    else:
        plotnames = args[0].split(',')

    formats = options.formats.split(',')
    dirname = args[1]

    try:
        plot.plot_result(
            dirname, plotnames,
            save=options.save, formats=formats, dpi=options.dpi)

    except grond.GrondError, e:
        die(str(e))


def command_export(args):

    def setup(parser):
        parser.add_option(
            '--type', dest='type', metavar='TYPE',
            choices=('event', 'source', 'vector'),
            help='select type of objects to be exported. Choices: '
                 '"event" (default), "source", "vector".')

        parser.add_option(
            '--parameters', dest='parameters', metavar='PLIST',
            help='select parameters to be exported. PLIST is a '
                 'comma-separated list where each entry has the form '
                 '"<parameter>[.<measure>]". Available measures: "best", '
                 '"mean", "std", "minimum", "percentile16", "median", '
                 '"percentile84", "maximum".')

        parser.add_option(
            '--output', dest='filename', metavar='FILE',
            help='write output to FILE')

    parser, options, args = cl_parse('export', args, setup)
    if len(args) < 2:
        help_and_die(parser, 'arguments required')

    what = args[0]

    dirnames = args[1:]

    what_choices = ('best', 'mean', 'ensemble', 'stats')

    if what not in what_choices:
        help_and_die(
            parser,
            'invalid choice: %s (choose from %s)' % (
                repr(what), ', '.join(repr(x) for x in what_choices)))

    if options.parameters:
        pnames = options.parameters.split(',')
    else:
        pnames = None

    try:
        grond.export(
            what,
            dirnames,
            filename=options.filename,
            type=options.type,
            pnames=pnames)

    except grond.GrondError, e:
        die(str(e))


def command_qc_polarization(args):

    import grond.qc

    def setup(parser):
        pass

    parser, options, args = cl_parse('qc-polarization', args, setup)
    if len(args) != 3:
        help_and_die(parser, 'missing arguments')

    config_path, event_name, target_config_name = args

    config = grond.read_config(config_path)

    ds = config.get_dataset(event_name)

    engine = config.engine_config.get_engine()

    target_config_names_avail = []
    for target_config in config.target_configs:
        name = '%s.%s' % (target_config.super_group, target_config.group or '')
        if name == target_config_name:
            imc = target_config.inner_misfit_config
            fmin = imc.fmin
            fmax = imc.fmax
            ffactor = imc.ffactor

            store = engine.get_store(target_config.store_id)
            timing = '{cake:P|cake:p|cake:P\\|cake:p\\}'

            grond.qc.polarization(
                ds, store, timing, fmin=fmin, fmax=fmax, ffactor=ffactor)

            return

        target_config_names_avail.append(name)

        die('no target_config named "%s" found. Available: %s' % (
            target_config_name, ', '.join(target_config_names_avail)))


if __name__ == '__main__':

    if len(sys.argv) < 2:
        sys.exit('Usage: %s' % usage)

    args = list(sys.argv)
    args.pop(0)
    command = args.pop(0)

    if command in subcommands:
        globals()['command_' + d2u(command)](args)

    elif command in ('--help', '-h', 'help'):
        if command == 'help' and args:
            acommand = args[0]
            if acommand in subcommands:
                globals()['command_' + acommand](['--help'])

        sys.exit('Usage: %s' % usage)

    else:
        die('no such subcommand: %s' % command)