__all__ = ['summarize', 'unite', 'subtract_bg']

import numbers
import os
import sys
import traceback

import ipy_table
import matplotlib
import matplotlib.cm
import matplotlib.colors
import matplotlib.pyplot as plt
import numpy as np
from IPython.core.getipython import get_ipython
from IPython.display import display
from mpl_toolkits.axes_grid import make_axes_locatable
from sastool.classes2 import Curve, Exposure
from sastool.libconfig import qunit
from sastool.misc.easylsq import FixedParameter, nonlinear_odr
from sastool.misc.errorvalue import ErrorValue

from .atsas import datcmp
from .calculation import correlmatrix
from .io import get_different_distances, load_exposure
from .plotting import plotsascurve
from .utils import print_abscissavalue, putlogo, writemarkdown

def _collect_data_for_summarization(headers, raw, reintegrate, qrange):
    ip = get_ipython()
    data1d = []
    data2d = 0
    headersout = []
    if not headers:
        return
    for head in headers:
        try:
            mo = ip.user_ns['mask_override'](head)
        except KeyError:
            mo = None
        ex = None
        last_exception = None
        try:
            ex = load_exposure(head.fsn, raw=raw, processed=not raw)
            assert isinstance(ex, Exposure)
            if mo is not None:
                try:
                    ex.mask = ex.loader.loadmask(mo)
                except FileNotFoundError:
                    print('Could not load mask: %s' % mo)
                    raise FileNotFoundError('Could not load mask: %s' % mo)
        except FileNotFoundError as exc:
            last_exception = sys.exc_info()
        if ex is None:
            print('Could not load {} 2D file for FSN {:d}. Exception: {}'.format(
                ['processed', 'raw'][raw], head.fsn, '\n'.join(traceback.format_exception(*last_exception))))
            ip.user_ns['badfsns'] = set(ip.user_ns['badfsns'])
            ip.user_ns['badfsns'].add(head.fsn)
            continue
        ex.header = head
        curve = None
        if not reintegrate:
            for l in [l_ for l_ in ip.user_ns['_loaders'] if l_.processed != raw]:
                try:
                    curve = l.loadcurve(head.fsn)
                    break
                except FileNotFoundError:
                    continue
            if curve is None:
                print('Cannot load curve for FSN %d: reintegrating.' % head.fsn)
        if curve is None:
            # this happens if reintegrate==True or if reintegrate==False but the curve could not be loaded.
            curve = ex.radial_average(qrange, errorpropagation=3,
                                      abscissa_errorpropagation=3, raw_result=False)
        curve = curve.sanitize()
        data1d.append(curve)

        data1d[-1].save(os.path.join(ip.user_ns['saveto_dir'], 'curve_%05d.txt' % head.fsn))
        mat = np.zeros((len(data1d[-1]), 3))
        mat[:, 0] = data1d[-1].q
        mat[:, 1] = data1d[-1].Intensity
        mat[:, 2] = data1d[-1].Error
        np.savetxt(os.path.join(ip.user_ns['saveto_dir'], 'curve_%s_%05d.dat' % (head.title, head.fsn)), mat)
        del mat
        data2d = data2d + ex
        headersout.append(ex.header)
    data2d /= len(data1d)
    return data1d, data2d, headersout

def _stabilityassessment(headers, data1d, dist, fig_correlmatrices, correlmatrixaxes, std_multiplier,
                         correlmatrix_colormap,
                         correlmatrix_filename, logarithmic_correlmatrix=True, cormaptest=True):
    # calculate and plot correlation matrix
    cmatrix, badidx, rowavg = correlmatrix(data1d, std_multiplier, logarithmic_correlmatrix)
    rowavgmean = rowavg.mean()
    rowavgstd = rowavg.std()
    writemarkdown('#### Assessing sample stability')
    writemarkdown("- Mean of row averages: " + str(rowavgmean))
    writemarkdown("- Std of row averages: " + str(rowavgstd) + ' (%.2f %%)' % (rowavgstd / rowavgmean * 100))

    img = correlmatrixaxes.imshow(cmatrix, interpolation='nearest', cmap=matplotlib.cm.get_cmap(correlmatrix_colormap))
    cax = make_axes_locatable(correlmatrixaxes).append_axes('right', size="5%", pad=0.1)
    fig_correlmatrices.colorbar(img, cax=cax)
    fsns = [h.fsn for h in headers]

    correlmatrixaxes.set_title('%.2f mm' % dist)
    correlmatrixaxes.set_xticks(list(range(len(data1d))))
    correlmatrixaxes.set_xticklabels([str(f) for f in fsns], rotation='vertical')
    correlmatrixaxes.set_yticks(list(range(len(data1d))))
    correlmatrixaxes.set_yticklabels([str(f) for f in fsns])
    np.savez_compressed(correlmatrix_filename,
                        correlmatrix=cmatrix, fsns=np.array(fsns))

    # Report table on sample stability
    tab = [['FSN', 'Date', 'Discrepancy', 'Relative discrepancy ((x-mean(x))/std(x))', 'Quality', 'Quality (cormap)']]
    badfsns = []
    badfsns_datcmp = []
    if cormaptest:
        matC, matp, matpadj, datcmp_ok = datcmp(*data1d)
    else:
        datcmp_ok = [not x for x in badidx]
    for h, bad, discr, dcmp_ok in zip(headers, badidx, rowavg, datcmp_ok):
        tab.append([h.fsn, h.date.isoformat(), discr, (discr - rowavgmean) / rowavgstd,
                    ["\u2713", "\u2718\u2718\u2718\u2718\u2718"][bad],
                    ["\u2713", "\u2718\u2718\u2718\u2718\u2718"][dcmp_ok != 1]])
        if bad:
            badfsns.append(h.fsn)
        if (not dcmp_ok and not np.isnan(dcmp_ok)):
            badfsns_datcmp.append(h.fsn)
    tab = ipy_table.IpyTable(tab)
    tab.apply_theme('basic')
    return badfsns, badfsns_datcmp, tab, rowavg

def summarize(reintegrate=True, dist_tolerance=3, qranges=None,
              samples=None, raw=False, late_radavg=True, graph_ncols=3,
              std_multiplier=3, graph_extension='png',
              graph_dpi=80, correlmatrix_colormap='coolwarm',
              image_colormap='viridis', correlmatrix_logarithmic=True, cormaptest=True):
    """Summarize scattering patterns and curves for all samples defined
    by the global `allsamplenames`.

    Inputs:
        reintegrate (bool, default=True): if the curves are to be obained
            by reintegrating the patterns. Otherwise 1D curves are loaded.
        dist_tolerance (float, default=3): sample-to-detector distances
            nearer than this are considered the same
        qranges (dict): a dictionary mapping approximate sample-to-detector
            distances (within dist_tolerance) to one-dimensional np.ndarrays
            of the desired q-range of the reintegration.
        samples (list or None): the names of the samples to summarize. If
            None, all samples defined by ``allsamplenames`` are used.
        raw (bool, default=False): if raw images are to be treated instead
            the evaluated ones (default).
        late_radavg (bool, default=True): if the scattering curves are to
            be calculated from the summarized scattering pattern. If False,
            scattering curves are calculated from each pattern and will be
            averaged.
        graph_ncols: the number of columns in graphs (2D patterns,
            correlation matrices)
        std_multiplier: if the absolute value of the relative discrepancy
            is larger than this limit, the exposure is deemed an outlier.
        graph_extension: the extension of the produced hardcopy files.
        graph_dpi: resolution of the graphs
        correlmatrix_colormap: name of the colormap to be used for the
            correlation matrices (resolved by matplotlib.cm.get_cmap())
        image_colormap: name of the colormap to be used for the scattering
            patterns (resolved by matplotlib.cm.get_cmap())
        correlmatrix_logarithmic: if the correlation matrix has to be
            calculated from the logarithm of the intensity.
    """
    if qranges is None:
        qranges = {}
    ip = get_ipython()
    data2d = {}
    data1d = {}
    headers_tosave = {}
    rowavg = {}
    if raw:
        writemarkdown('# Summarizing RAW images.')
        headers = ip.user_ns['_headers']['raw']
        rawpart = '_raw'  # this will be added in the filenames saved
    else:
        writemarkdown('# Summarizing CORRECTED images.')
        headers = ip.user_ns['_headers']['processed']
        rawpart = ''  # nothing will be added in the filenames saved

    if samples is None:
        samples = sorted(ip.user_ns['allsamplenames'])
    for samplename in samples:
        writemarkdown('## ' + samplename)
        headers_sample = [h for h in headers if h.title == samplename]
        data2d[samplename] = {}
        rowavg[samplename] = {}
        data1d[samplename] = {}
        headers_tosave[samplename] = {}
        dists = get_different_distances([h for h in headers if h.title == samplename], dist_tolerance)
        if not dists:
            writemarkdown('No measurements from sample, skipping.')
            continue
        fig_2d = plt.figure()
        fig_curves = plt.figure()
        fig_correlmatrices = plt.figure()
        distaxes = {}
        correlmatrixaxes = {}
        ncols = min(len(dists), graph_ncols)
        nrows = int(np.ceil(len(dists) / ncols))
        onedimaxes = fig_curves.add_axes((0.1, 0.3, 0.8, 0.5))
        onedimstdaxes = fig_curves.add_axes((0.1, 0.1, 0.8, 0.2))
        for distidx, dist in enumerate(dists):
            writemarkdown("### Distance " + str(dist) + " mm")
            headers_narrowed = [h for h in headers_sample if abs(float(h.distance) - dist) < dist_tolerance]
            distaxes[dist] = fig_2d.add_subplot(
                nrows, ncols, distidx + 1)
            correlmatrixaxes[dist] = fig_correlmatrices.add_subplot(
                nrows, ncols, distidx + 1)
            # determine the q-range to be used from the qranges argument.
            try:
                distkey_min = min([np.abs(k - dist)
                                   for k in qranges if np.abs(k - dist) < dist_tolerance])
            except ValueError:
                # no matching key in qranges dict
                qrange = None  # request auto-determination of q-range
            else:
                distkey = [
                    k for k in qranges if np.abs(k - dist) == distkey_min][0]
                qrange = qranges[distkey]

            (data1d[samplename][dist], data2d[samplename][dist], headers_tosave[samplename][dist]) = \
                _collect_data_for_summarization(headers_narrowed, raw, reintegrate, qrange)

            badfsns, badfsns_datcmp, tab, rowavg[samplename][dist] = _stabilityassessment(
                headers_tosave[samplename][dist],
                data1d[samplename][dist], dist,
                fig_correlmatrices,
                correlmatrixaxes[dist], std_multiplier, correlmatrix_colormap,
                os.path.join(ip.user_ns['saveto_dir'], 'correlmatrix_%s_%s' % (
                    samplename,
                    ('%.2f' % dist).replace('.', '_')) +
                             rawpart + '.npz'),
                logarithmic_correlmatrix=correlmatrix_logarithmic,
                cormaptest=cormaptest)

            if 'badfsns' not in ip.user_ns:
                ip.user_ns['badfsns'] = {}
            elif 'badfsns_datcmp' not in ip.user_ns:
                ip.user_ns['badfsns_datcmp'] = {}
            ip.user_ns['badfsns'] = set(ip.user_ns['badfsns']).union(badfsns)
            ip.user_ns['badfsns_datcmp'] = set(ip.user_ns['badfsns_datcmp']).union(badfsns_datcmp)
            display(tab)

            # Plot the image
            try:
                data2d[samplename][dist].imshow(axes=distaxes[dist], show_crosshair=False,
                                                norm=matplotlib.colors.LogNorm(),
                                                cmap=matplotlib.cm.get_cmap(image_colormap))
            except ValueError:
                print('Error plotting 2D image for sample %s, distance %.2f' % (samplename, dist))
            distaxes[dist].set_xlabel('q (' + qunit() + ')')
            distaxes[dist].set_ylabel('q (' + qunit() + ')')
            distaxes[dist].set_title(
                '%.2f mm (%d curve%s)' % (dist, len(headers_tosave[samplename][dist]),
                                          ['', 's'][len(headers_tosave[samplename][dist]) > 1]))

            # Plot the curves
            Istd = np.stack([c.Intensity for c in data1d[samplename][dist]], axis=1)
            for c, h in zip(data1d[samplename][dist], headers_tosave[samplename][dist]):
                color = 'green'
                if h.fsn in badfsns_datcmp:
                    color = 'magenta'
                if h.fsn in badfsns:
                    color = 'red'
                c.loglog(axes=onedimaxes, color=color)
            if Istd.shape[1] > 1:
                onedimstdaxes.loglog(data1d[samplename][dist][0].q, Istd.std(axis=1) / Istd.mean(axis=1) * 100, 'b-')
            if not late_radavg:
                data1d[samplename][dist] = Curve.average(
                    *data1d[samplename][dist])
            else:
                data1d[samplename][dist] = (
                    data2d[samplename][dist].radial_average(
                        qrange,
                        errorpropagation=3,
                        abscissa_errorpropagation=3, raw_result=False))
            data1d[samplename][dist].loglog(
                label='Average', lw=2, color='k', axes=onedimaxes)

            ##Saving image, headers, mask and curve
            # data2d[samplename][dist].write(
            #    os.path.join(ip.user_ns['saveto_dir'],
            #                 samplename + '_'+(
            #                     '%.2f' % dist).replace('.', '_') +
            #                 rawpart + '.npz'), plugin='CREDO Reduced')
            # data2d[samplename][dist].header.write(
            #    os.path.join(ip.user_ns['saveto_dir'],
            ###                 samplename + '_'+(
            #                     '%.2f' % dist).replace('.', '_') +
            #                 rawpart +'.log'), plugin='CREDO Reduced')
            # data2d[samplename][dist].mask.write_to_mat(
            #    os.path.join(ip.user_ns['saveto_dir'],
            #                 data2d[samplename][dist].mask.maskid+'.mat'))
            data1d[samplename][dist].save(os.path.join(ip.user_ns['saveto_dir'],
                                                       samplename + '_' + ('%.2f' % dist).replace('.',
                                                                                '_') + rawpart + '.txt'))

            # Report on qrange and flux
            q_ = data1d[samplename][dist].q
            qmin = q_[q_ > 0].min()
            writemarkdown('#### Q-range & flux')
            writemarkdown(
                '- $q_{min}$: ' + print_abscissavalue(qmin, headers_tosave[samplename][dist][0].wavelength, dist))
            writemarkdown('- $q_{max}$: ' + print_abscissavalue(data1d[samplename][dist].q.max(),
                                                                headers_tosave[samplename][dist][0].wavelength, dist))
            writemarkdown('- Number of $q$ points: ' + str(len(data1d[samplename][dist])))
            meastime = sum([h.exposuretime for h in headers_tosave[samplename][dist]])
            writemarkdown("- from %d exposures, total exposure time %.0f sec <=> %.2f hr" % (
                len(headers_tosave[samplename][dist]),
                meastime, meastime / 3600.))
            try:
                flux = [h.flux for h in headers_tosave[samplename][dist]]
                flux = ErrorValue(np.mean(flux), np.std(flux))
                writemarkdown("- beam flux (photon/sec): %s" % flux)
            except KeyError:
                writemarkdown("- *No information on beam flux: dealing with raw data.*")
        onedimaxes.set_xlabel('')
        onedimaxes.set_ylabel('$d\\Sigma/d\\Omega$ (cm$^{-1}$ sr$^{-1}$)')
        # plt.legend(loc='best')
        onedimaxes.grid(True, which='both')
        onedimaxes.axis('tight')
        onedimaxes.set_title(samplename)
        onedimstdaxes.set_xlabel('q (' + qunit() + ')')
        onedimstdaxes.set_ylabel('Rel.std.dev. of intensity (%)')
        onedimstdaxes.grid(True, which='both')
        onedimstdaxes.set_xlim(*onedimaxes.get_xlim())
        onedimstdaxes.set_xscale(onedimaxes.get_xscale())
        putlogo(fig_curves)
        putlogo(fig_2d)
        fig_2d.tight_layout()
        fig_correlmatrices.suptitle(samplename)
        fig_correlmatrices.tight_layout()
        fig_2d.savefig(
            os.path.join(ip.user_ns['auximages_dir'],
                         'averaging2D_' +
                         samplename + rawpart + '.' + graph_extension),
            dpi=graph_dpi)
        fig_curves.savefig(
            os.path.join(ip.user_ns['auximages_dir'],
                         'averaging1D_' +
                         samplename + rawpart + '.' + graph_extension),
            dpi=graph_dpi)
        putlogo(fig_correlmatrices)
        fig_correlmatrices.savefig(
            os.path.join(ip.user_ns['auximages_dir'],
                         'correlation_' +
                         samplename + rawpart + '.' + graph_extension),
            dpi=graph_dpi)
        writemarkdown("### Collected images from all distances")
        plt.show()
    writemarkdown("Updated badfsns list:")
    writemarkdown('[' + ', '.join(str(f) for f in ip.user_ns['badfsns']) + ']')
    writemarkdown("Updated badfsns list using datcmp:")
    writemarkdown('[' + ', '.join(str(f) for f in ip.user_ns['badfsns_datcmp']) + ']')
    ip.user_ns['_data1d'] = data1d
    ip.user_ns['_data2d'] = data2d
    ip.user_ns['_headers_sample'] = headers_tosave
    ip.user_ns['_rowavg'] = rowavg

def _merge_two_curves(curve1: Curve, curve2: Curve, qmin, qmax, qsep, use_additive_constant=False):
    """Merge two scattering curves

    :param curve1: the first curve (longer distance)
    :type curve1: sastool.classes.curve.GeneralCurve
    :param curve2: the second curve (shorter distance)
    :type curve2: sastool.classes.curve.GeneralCurve
    :param qmin: lower bound of the interval for determining the scaling factor
    :type qmin: float
    :param qmax: upper bound of the interval for determining the scaling factor
    :type qmax: float
    :param qsep: separating (tailoring) point for the merge
    :type qsep: float
    :return: merged_curve, factor, background, stat
    :rtype tuple of a sastool.classes2.curve.Curve and a float
    """
    curve1=curve1.sanitize()
    curve2=curve2.sanitize()
    if len(curve1.trim(qmin, qmax)) > len(curve2.trim(qmin, qmax)):
        curve2_interp = curve2.trim(qmin, qmax)
        curve1_interp = curve1.interpolate(curve2_interp.q)
    else:
        curve1_interp = curve1.trim(qmin, qmax)
        curve2_interp = curve2.interpolate(curve1_interp.q)
    if use_additive_constant:
        bg_init = 0
    else:
        bg_init = FixedParameter(0)
    factor, bg, stat = nonlinear_odr(curve2_interp.Intensity, curve1_interp.Intensity,
                                     curve2_interp.Error, curve1_interp.Error,
                                     lambda x, factor, bg: x * factor + bg, [1.0, bg_init])
    return Curve.merge(curve1 - bg, curve2 * factor, qsep), factor, bg, stat

def _scale_two_exposures(exp1, exp2, qmin, qmax, N=10, use_additive_constant=False):
    qrange = np.linspace(qmin, qmax, N)
    rad1 = exp1.radial_average(qrange=qrange, raw_result=False)
    rad2 = exp2.radial_average(qrange=qrange, raw_result=False)
    if use_additive_constant:
        bg_init = 0
    else:
        bg_init = FixedParameter(0)
    factor, bg, stat = nonlinear_odr(rad2.y, rad1.y, rad2.dy, rad1.dy, lambda x, factor, bg: x * factor + bg,
                                     [1, bg_init])
    return factor, bg

def unite(samplename, uniqmin=[], uniqmax=[], uniqsep=[], graph_ncols=2, graph_subplotpars={'hspace': 0.3},
          graph_extension='png', graph_dpi=80, additive_constant=False):
    ip = get_ipython()
    if isinstance(uniqmin, numbers.Number):
        uniqmin = [uniqmin]
    if isinstance(uniqmax, numbers.Number):
        uniqmax = [uniqmax]
    if isinstance(uniqsep, numbers.Number):
        uniqsep = [uniqsep]
    data1d = ip.user_ns['_data1d'][samplename]
    print("Uniting measurements of sample %s at different s-d distances" % samplename)
    uniparams = {'qmin': uniqmin, 'qmax': uniqmax, 'qsep': uniqsep}
    for p in uniparams:
        uniparams[p] = uniparams[p] + [None] * \
                                      max(0, len(data1d) - 1 - len(uniparams[p]))
    dists = list(reversed(sorted(data1d.keys())))
    if len(dists) < 2:
        print("Less than two distances found for sample %s; no point of uniting." % samplename)
        return
    united = None
    graph_nrows = int(
        np.ceil((len(dists)) / (graph_ncols * 1.0)))
    fig = plt.figure()
    unitedaxis = fig.add_subplot(graph_nrows, graph_ncols, 1)
    factor = 1.0
    for idx, dist1, dist2, qmin, qmax, qsep in zip(list(range(len(dists) - 1)),
                                                   dists[:-1], dists[1:],
                                                   uniparams['qmin'],
                                                   uniparams['qmax'],
                                                   uniparams['qsep']):
        print("    Scaling together distances %f and %f mm" % (dist1, dist2), flush=True)
        if united is None:
            united = data1d[dist1]
        if qmin is None:
            qmin = data1d[dist2].sanitize().q.min()
            print("        Auto-detected qmin:", qmin, flush=True)
        if qmax is None:
            qmax = data1d[dist1].sanitize().q.max()
            print("        Auto-detected qmax:", qmax, flush=True)
        if qsep is None:
            qsep = 0.5 * (qmin + qmax)
            print("        Auto-detected qsep:", qsep, flush=True)
        ax = fig.add_subplot(graph_nrows, graph_ncols, 2 + idx)
        (factor * data1d[dist1]).loglog(axes=ax, label='%.2f mm' % dist1)
        united, factor1, bg, stat = _merge_two_curves(united,
                                                      data1d[dist2], qmin, qmax, qsep,
                                                      use_additive_constant=additive_constant)
        factor = factor1 * factor
        uniparams['qmin'][idx] = qmin
        uniparams['qmax'][idx] = qmax
        uniparams['qsep'][idx] = qsep
        print("        Scaling factor is", factor.tostring(), flush=True)
        if not additive_constant:
            print("        Additive constant has not been used.", flush=True)
        else:
            print("        Additive constant is:", bg.tostring(), flush=True)
        print("        Reduced Chi^2 of the ODR fit:", stat['Chi2_reduced'], flush=True)
        print("        DoF of the ODR fit:", stat['DoF'], flush=True)
        (data1d[dist2] * factor + bg).loglog(axes=ax, label='%.2f mm' % dist2)
        ax.set_xlabel('q (' + qunit() + ')')
        ax.set_ylabel('$d\\Sigma/d\\Omega$ (cm$^{-1}$ sr$^{-1}$)')
        ax.legend(loc='best')
        # ax.grid(which='both')
        ax.axis('tight')
        ax.set_title('Factor: ' + str(factor))
        lims = ax.axis()
        ax.plot([qmin, qmin], lims[2:], '--r', lw=2)
        ax.plot([qmax, qmax], lims[2:], '--r', lw=2)
        ax.plot([qsep, qsep], lims[2:], '--k')
        ax.grid(True, which='both')
    if '_data1dunited' not in ip.user_ns:
        ip.user_ns['_data1dunited'] = {}
    united.loglog(axes=unitedaxis)
    unitedaxis.set_xlabel('q (' + qunit() + ')')
    unitedaxis.set_ylabel('$d\\Sigma/d\\Omega$ (cm$^{-1}$ sr$^{-1}$)')
    unitedaxis.legend(loc='best')
    unitedaxis.set_title('United scattering of %s' % samplename)
    unitedaxis.grid(True, which='both')
    # unitedaxis.grid(which='both')
    unitedaxis.axis('tight')
    lims = unitedaxis.axis()
    for qs in uniparams['qsep']:
        unitedaxis.plot([qs] * 2, lims[2:], '--r')
    ip.user_ns['_data1dunited'][samplename] = united
    putlogo()
    fig.subplots_adjust(**graph_subplotpars)
    plt.savefig(
        os.path.join(ip.user_ns['auximages_dir'], 'uniting_' + samplename + '.' + graph_extension), dpi=graph_dpi)
    print("    United curve spans the following ranges:")
    print("        q_min: ",
          print_abscissavalue(united.q.min(), ip.user_ns['_headers_sample'][samplename][dists[0]][0].wavelength))
    print("        q_max: ",
          print_abscissavalue(united.q.max(), ip.user_ns['_headers_sample'][samplename][dists[0]][0].wavelength))
    print("        q_max/q_min:", united.q.max() / united.q.min())
    print("        I_min: ", united.Intensity.min(), "cm^{-1}")
    print("        I_max: ", united.Intensity.max(), "cm^{-1}")
    print("        I_max/I_min:", united.Intensity.max() / united.Intensity.min())
    print("        # of points: ", len(united))
    united.save(os.path.join(ip.user_ns['saveto_dir'], 'united_' + samplename + '.txt'))
    plt.show()

def subtract_bg(samplename, bgname, factor=1, distance=None, disttolerance=2,
                subname=None, qrange=(), graph_extension='png', graph_dpi=80):
    """Subtract background from measurements.

    Inputs:
        samplename: the name of the sample
        bgname: the name of the background measurements. Alternatively, it can
            be a numeric value (float or ErrorValue), which will be subtracted.
            If None, this constant will be determined by integrating the
            scattering curve in the range given by qrange.
        factor: the background curve will be multiplied by this
        distance: if None, do the subtraction for all sample-to-detector distances.
            Otherwise give here the value of the sample-to-detector distance.
        qrange: a tuple (qmin, qmax)
        disttolerance: the tolerance in which two distances are considered
            equal.
        subname: the sample name of the background-corrected curve. The default
            is samplename + '-' + bgname
    """
    ip = get_ipython()
    data1d = ip.user_ns['_data1d']
    data2d = ip.user_ns['_data2d']
    if 'subtractedsamplenames' not in ip.user_ns:
        ip.user_ns['subtractedsamplenames'] = set()
    subtractedsamplenames = ip.user_ns['subtractedsamplenames']
    if subname is None:
        if isinstance(bgname, str):
            subname = samplename + '-' + bgname
        else:
            subname = samplename + '-const'
    if distance is None:
        dists = data1d[samplename]
    else:
        dists = [d for d in data1d[samplename] if abs(d - distance) < disttolerance]
    for dist in dists:
        if isinstance(bgname, str):
            if not disttolerance:
                if dist not in data1d[bgname]:
                    print(
                        'Warning: Missing distance %g for background measurement (samplename: %s, background samplename: %s)' % (
                            dist, samplename, bgname))
                    continue
                else:
                    bgdist = dist
            else:
                bgdist = sorted([(d, r) for (d, r) in [(d, np.abs(d - dist)) for d in list(data1d[bgname].keys())] if
                                 r <= disttolerance], key=lambda x: x[1])[0][0]
        if subname not in data1d:
            data1d[subname] = {}
        if subname not in data2d:
            data2d[subname] = {}
        if subname not in ip.user_ns['_headers_sample']:
            ip.user_ns['_headers_sample'][subname] = {}
        data1_s = data1d[samplename][dist]
        data2_s = data2d[samplename][dist]
        if isinstance(bgname, str):
            data1_bg = data1d[bgname][bgdist]
            data2_bg = data2d[bgname][bgdist]
            if factor is None:
                factor = data1_s.trim(*qrange).momentum(0) / data1_bg.trim(*qrange).momentum(0)
        elif bgname is None:
            data1_bg = data1_s.trim(*qrange).momentum(0)
            data2_bg = data1_bg
        else:
            data1_bg = bgname
            data2_bg = bgname
        if factor is None:
            factor = 1
        data1d[subname][dist] = data1_s - factor * data1_bg
        data2d[subname][dist] = data2_s - factor * data2_bg
        data1d[subname][dist].save(
            os.path.join(ip.user_ns['saveto_dir'], subname + '_' + ('%.2f' % dist).replace('.', '_') + '.txt'))
        ip.user_ns['_headers_sample'][subname][dist] = ip.user_ns['_headers_sample'][samplename][
            dist]  # ugly hack, I have no better idea.
        plt.figure()
        plotsascurve(samplename, dist=dist)
        if isinstance(bgname, str):
            plotsascurve(bgname, dist=dist, factor=factor)
        plotsascurve(subname, dist=dist)
        plt.savefig(os.path.join(ip.user_ns['auximages_dir'],
                                 'subtractbg_' + samplename + '.' + graph_extension),
                    dpi=graph_dpi)

        subtractedsamplenames.add(subname)


# -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4

#
# LICENSE
#
# Copyright (C) 2010-2018 GEM Foundation, G. Weatherill, M. Pagani,
# D. Monelli.
#
# The Hazard Modeller's Toolkit is free software: you can redistribute
# it and/or modify it under the terms of the GNU Affero General Public
# License as published by the Free Software Foundation, either version
# 3 of the License, or (at your option) any later version.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>
#
# DISCLAIMER
#
# The software Hazard Modeller's Toolkit (openquake.hmtk) provided herein
# is released as a prototype implementation on behalf of
# scientists and engineers working within the GEM Foundation (Global
# Earthquake Model).
#
# It is distributed for the purpose of open collaboration and in the
# hope that it will be useful to the scientific, engineering, disaster
# risk and software design communities.
#
# The software is NOT distributed as part of GEM’s OpenQuake suite
# (https://www.globalquakemodel.org/tools-products) and must be considered as a
# separate entity. The software provided herein is designed and implemented
# by scientific staff. It is not developed to the design standards, nor
# subject to same level of critical review by professional software
# developers, as GEM’s OpenQuake software suite.
#
# Feedback and contribution to the software is welcome, and can be
# directed to the hazard scientific staff of the GEM Model Facility
# (hazard@globalquakemodel.org).
#
# The Hazard Modeller's Toolkit (openquake.hmtk) is therefore distributed WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
# for more details.
#
# The GEM Foundation, and the authors of the software, assume no
# liability for use of the software.

# -*- coding: utf-8 -*-
"""
"""

import numpy as np
from openquake.hmtk.seismicity.occurrence.utils import input_checks, recurrence_table
from openquake.hmtk.seismicity.occurrence.base import (
    SeismicityOccurrence, OCCURRENCE_METHODS)
from openquake.hmtk.seismicity.occurrence.aki_maximum_likelihood import AkiMaxLikelihood

@OCCURRENCE_METHODS.add(
    'calculate', **{
        'completeness': True,
        'reference_magnitude': 0.0,
        'magnitude_interval': 0.1,
        'Average Type': ['Weighted', 'Harmonic']})
class BMaxLikelihood(SeismicityOccurrence):
    """ Implements maximum likelihood calculations taking into account time
    variation in completeness"
    """

    def calculate(self, catalogue, config, completeness=None):
        """ Calculates recurrence parameters a_value and b_value, and their
        respective uncertainties

        :param catalogue: Earthquake Catalogue
            An instance of :class:`openquake.hmtk.seismicity.catalogue`
        :param dict config:
            A configuration dictionary; the only parameter that can be
            defined in this case if the type of average to be applied
            in the calculation
        :param list or numpy.ndarray completeness:
            Completeness table
        """

        # Input checks
        cmag, ctime, ref_mag, dmag, config = input_checks(catalogue,
                                                          config,
                                                          completeness)

        # Check the configuration
        if not config['Average Type'] in ['Weighted', 'Harmonic']:
            raise ValueError('Average type not recognised in bMaxLiklihood!')
        return self._b_ml(catalogue, config, cmag, ctime, ref_mag, dmag)

    def _b_ml(self, catalogue, config, cmag, ctime, ref_mag, dmag):
        end_year = float(catalogue.end_year)
        catalogue = catalogue.data
        ival = 0
        mag_eq_tolerance = 1E-5
        aki_ml = AkiMaxLikelihood()

        while ival < np.shape(ctime)[0]:

            id0 = np.abs(ctime - ctime[ival]) < mag_eq_tolerance
            m_c = np.min(cmag[id0])

            print('--- ctime', ctime[ival], ' m_c', m_c)

            # Find events later than cut-off year, and with magnitude
            # greater than or equal to the corresponding completeness
            # magnitude. m_c - mag_eq_tolerance is required to correct
            # floating point differences.
            id1 = np.logical_and(
                catalogue['year'] >= ctime[ival],
                catalogue['magnitude'] >= (m_c - mag_eq_tolerance))
            # Get a- and b- value for the selected events
            temp_rec_table = recurrence_table(catalogue['magnitude'][id1],
                                              dmag,
                                              catalogue['year'][id1],
                                              end_year - ctime[ival] + 1)

            bval, sigma_b = aki_ml._aki_ml(temp_rec_table[:, 0],
                                           temp_rec_table[:, 1], dmag, m_c)

            if ival == 0:
                gr_pars = np.array([np.hstack([bval, sigma_b])])
                neq = np.sum(id1)  # Number of events
            else:
                gr_pars = np.vstack([gr_pars, np.hstack([bval, sigma_b])])
                neq = np.hstack([neq, np.sum(id1)])
            ival = ival + np.sum(id0)

        # Get average GR parameters
        bval, sigma_b = self._average_parameters(
            gr_pars, neq, config['Average Type'])
        aval = self._calculate_a_value(bval,
                                       np.float(np.sum(neq)),
                                       cmag,
                                       ctime,
                                       catalogue['magnitude'],
                                       end_year,
                                       dmag)
        sigma_a = self._calculate_a_value(bval + sigma_b,
                                          np.float(np.sum(neq)),
                                          cmag,
                                          ctime,
                                          catalogue['magnitude'],
                                          end_year,
                                          dmag)
        if not config['reference_magnitude']:
            return bval,\
                sigma_b,\
                aval,\
                sigma_a - aval
        else:
            rate = 10. ** (aval - bval * config['reference_magnitude'])
            sigma_rate = 10. ** (sigma_a -
                                 bval * config['reference_magnitude']) - rate
            return bval,\
                sigma_b,\
                rate,\
                sigma_rate

    def _average_parameters(self, gr_params, neq, average_type='Weighted'):
        """
        Calculates the average of a set of Gutenberg-Richter parameters
        depending on the average type

        :param numpy.ndarray gr_params:
            Gutenberg-Richter parameters [b, sigma_b, a, sigma_a]
        :param numpy.ndarray neq:

        """
        if np.shape(gr_params)[0] != neq.size:
            raise ValueError('Number of weights does not correspond'
                             ' to number of parameters')

        if 'Harmonic' in average_type:
            average_parameters = self._harmonic_mean(gr_params, neq)
        else:
            average_parameters = self._weighted_mean(gr_params, neq)
        bval = average_parameters[0]
        sigma_b = average_parameters[1]
        return bval, sigma_b

    def _calculate_a_value(self, bvalue, nvalue, cmag, cyear, magnitude,
                           end_year, dmag):
        """
        Calculates the a-value using the method of Weichert (1980) and
        McGuire (2004)
        """
        mmin = cmag[0]
        mmax = np.max(magnitude)
        if mmax > np.max(cmag):
            cmag = np.hstack([cmag, mmax + dmag])
        target_mag = (cmag[:-1] + cmag[1:]) / 2.
        nyear = end_year - cyear + 1.
        beta = bvalue * np.log(10.)
        rate_mmin = nvalue * np.sum(np.exp(-beta * target_mag)) /\
            np.sum(nyear * np.exp(-beta * target_mag))
        return np.log10(rate_mmin) + bvalue * mmin

    def _weighted_mean(self, parameters, neq):
        '''Simple weighted mean'''
        weight = neq.astype(float) / np.sum(neq)
        if np.shape(parameters)[0] != weight.size:
            raise ValueError('Parameter vector not same shape as weights')
        else:
            average_value = np.zeros(np.shape(parameters)[1], dtype=float)
            for iloc in range(0, np.shape(parameters)[1]):
                average_value[iloc] = np.sum(parameters[:, iloc] * weight)
        return average_value

    def _harmonic_mean(self, parameters, neq):
        '''Harmonic mean'''
        weight = neq.astype(float) / np.sum(neq)
        if np.shape(parameters)[0] != weight.size:
            raise ValueError('Parameter vector not same shape as weights')

        average_value = np.zeros(np.shape(parameters)[1], dtype=float)
        for iloc in range(0, np.shape(parameters)[1]):
            average_value[iloc] = 1. / np.sum(
                (weight * (1. / parameters[:, iloc])))
        return average_value


#!/usr/bin/env python

import copy
from importlib import import_module
import logging

import numpy as np

from openquake.hazardlib.gsim.base import GMPE
from openquake.hazardlib.gsim.boore_2014 import BooreEtAl2014
from openquake.hazardlib.gsim.campbell_bozorgnia_2014 import CampbellBozorgnia2014
from openquake.hazardlib.imt import PGA, PGV, SA
from openquake.hazardlib import const
from openquake.hazardlib.valid import gsim
from openquake.hazardlib.contexts import RuptureContext

from shakelib.conversions.imt.abrahamson_bhasin_2020 import AbrahamsonBhasin2020
from shakelib.conversions.imc.boore_kishida_2017 import BooreKishida2017
from shakelib.sites import Sites

# Special case GMPEs:
from shakelib.gmpe.nga_east import NGAEast

def set_sites_depth_parameters(sites, gmpe):
    """
    Need to select the appropriate z1pt0 value for different GMPEs.
    Note that these are required site parameters, so even though
    OQ has these equations built into the class in most cases.
    I have submitted an issue to OQ requesting subclasses of these
    methods that do not require the depth parameters in the
    SitesContext to make this easier.

    Args:
        sites:1 An OQ sites context.
        gmpe: An OQ GMPE instance.

    Returns:
        An OQ sites context with the depth parameters set for the
        requested GMPE.
    """
    if gmpe == "[MultiGMPE]":
        return sites

    Sites._addDepthParameters(sites)

    if (
        gmpe == "[AbrahamsonEtAl2014]"
        or gmpe == "[AbrahamsonEtAl2014]\nregion = 'TWN'"
        or gmpe == "[AbrahamsonEtAl2014]\nregion = 'CHN'"
    ):
        sites.z1pt0 = sites.z1pt0_ask14_cal
    if gmpe == "[AbrahamsonEtAl2014]\nregion = 'JPN'":
        sites.z1pt0 = sites.z1pt0_ask14_jpn
    if gmpe == "[ChiouYoungs2014]" or isinstance(gmpe, BooreEtAl2014):
        sites.z1pt0 = sites.z1pt0_cy14_cal
    if isinstance(gmpe, CampbellBozorgnia2014):
        if (
            gmpe == "[CampbellBozorgnia2014JapanSite]"
            or gmpe == "[CampbellBozorgnia2014HighQJapanSite]"
            or gmpe == "[CampbellBozorgnia2014LowQJapanSite]"
        ):
            sites.z2pt5 = sites.z2pt5_cb14_jpn
        else:
            sites.z2pt5 = sites.z2pt5_cb14_cal
    if (
        gmpe == "[ChiouYoungs2008]"
        or gmpe == "[Bradley2013]"
        or gmpe == "[Bradley2013Volc]"
    ):
        sites.z1pt0 = sites.z1pt0_cy08
    if gmpe == "[CampbellBozorgnia2008]":
        sites.z2pt5 = sites.z2pt5_cb07
    if gmpe == "[AbrahamsonSilva2008]":
        sites.z1pt0 = gmpe._compute_median_z1pt0(sites.vs30)

    return sites

def stuff_context(sites, rup, dists):
    """
    Function to fill a rupture context with the contents of all of the
    other contexts.

    Args:
        sites (SiteCollection): A SiteCollection object.

        rup (RuptureContext): A RuptureContext object.

        dists (DistanceContext): A DistanceContext object.

    Returns:
        RuptureContext: A new RuptureContext whose attributes are all of
        the elements of the three inputs.
    """
    ctx = RuptureContext()

    for name in [name for name in vars(sites) if not name.startswith("__")]:
        setattr(ctx, name, getattr(sites, name))
    for name in [name for name in vars(rup) if not name.startswith("__")]:
        setattr(ctx, name, getattr(rup, name))
    for name in [name for name in vars(dists) if not name.startswith("__")]:
        setattr(ctx, name, getattr(dists, name))

    return ctx

def get_gmpe_from_name(name, conf):

    # Only import the NullGMPE when we're testing
    # We'll want to import any other GMPEs we add at the top of this module
    # so that gsim() picks them up; anything in OQ is already included
    if name == "NullGMPE":
        mod = import_module(conf["gmpe_modules"][name][1])
    return gsim(name)

class MultiGMPE(GMPE):
    """
    Implements a GMPE that is the combination of multiple GMPEs.

    """

    DEFINED_FOR_TECTONIC_REGION_TYPE = None
    DEFINED_FOR_INTENSITY_MEASURE_TYPES = None
    DEFINED_FOR_INTENSITY_MEASURE_COMPONENT = None
    DEFINED_FOR_STANDARD_DEVIATION_TYPES = set([const.StdDev.TOTAL])
    REQUIRES_SITES_PARAMETERS = None
    REQUIRES_RUPTURE_PARAMETERS = None
    REQUIRES_DISTANCES = None

    def get_mean_and_stddevs(self, sites, rup, dists, imt, stddev_types):
        """
        See superclass `method <http://docs.openquake.org/oq-hazardlib/master/gsim/index.html#openquake.hazardlib.gsim.base.GroundShakingIntensityModel.get_mean_and_stddevs>`__.

        Unlike the superclass method, the stddev list returned by this
        function will have twice as many arrays as are requested in
        stddev_types: The first set will include the standard deviation
        inflation due to the point-source to finite fault conversion (if
        any), and the second set will not include this inflation. In the
        case where a finite rupture is provided (and, thus, no point-source
        to finite rupture adjustments are made) the two sets of stddev
        arrays will be identical. Thus, if::

            stddev_types = [const.StdDev.TOTAL, const.StdDev.INTRA_EVENT,
                            const.StdDev.INTER_EVENT]

        the returned stddev list will contain six arrays: the first three
        will include the point-source inflation, and the second three will
        not.
        """  # noqa

        # ---------------------------------------------------------------------
        # Sort out shapes of the sites and dists elements
        # Need to turn all 2D arrays into 1D arrays because of
        # inconsistencies in how arrays are handled in OpenQuake.
        # ---------------------------------------------------------------------
        shapes = []
        for k, v in sites.__dict__.items():
            if k == "_slots_":
                continue
            if (k != "lons") and (k != "lats"):
                shapes.append(v.shape)
                sites.__dict__[k] = np.reshape(sites.__dict__[k], (-1,))
        for k, v in dists.__dict__.items():
            if k == "_slots_":
                continue
            if (k != "lons") and (k != "lats") and v is not None:
                shapes.append(v.shape)
                dists.__dict__[k] = np.reshape(dists.__dict__[k], (-1,))
        shapeset = set(shapes)
        if len(shapeset) != 1:
            raise Exception("All dists and sites elements must have same shape.")
        else:
            orig_shape = list(shapeset)[0]

        sd_avail = self.DEFINED_FOR_STANDARD_DEVIATION_TYPES
        if not sd_avail.issuperset(set(stddev_types)):
            raise Exception("Requested an unavailable stddev_type.")

        # Evaluate MultiGMPE:
        lnmu, lnsd = self.__get_mean_and_stddevs__(sites, rup, dists, imt, stddev_types)

        # Check for large-distance cutoff/weights
        if hasattr(self, "CUTOFF_DISTANCE"):
            lnmu_large, lnsd_large = self.__get_mean_and_stddevs__(
                sites, rup, dists, imt, stddev_types, large_dist=True
            )
            # Stomp on lnmu and lnsd at large distances
            dist_cutoff = self.CUTOFF_DISTANCE
            lnmu[dists.rjb > dist_cutoff] = lnmu_large[dists.rjb > dist_cutoff]
            for i in range(len(lnsd)):
                lnsd[i][dists.rjb > dist_cutoff] = lnsd_large[i][
                    dists.rjb > dist_cutoff
                ]

        # Undo reshapes of inputs
        for k, v in dists.__dict__.items():
            if k == "_slots_":
                continue
            if (k != "lons") and (k != "lats") and v is not None:
                dists.__dict__[k] = np.reshape(dists.__dict__[k], orig_shape)
        for k, v in sites.__dict__.items():
            if k == "_slots_":
                continue
            if (k != "lons") and (k != "lats"):
                sites.__dict__[k] = np.reshape(sites.__dict__[k], orig_shape)

        # Reshape output
        lnmu = np.reshape(lnmu, orig_shape)
        for i in range(len(lnsd)):
            lnsd[i] = np.reshape(lnsd[i], orig_shape)

        return lnmu, lnsd

    def __get_mean_and_stddevs__(
        self, sites, rup, dists, imt, stddev_types, large_dist=False
    ):

        # ---------------------------------------------------------------------
        # Sort out which set of weights to use
        # ---------------------------------------------------------------------
        if large_dist is False:
            wts = self.WEIGHTS
        else:
            wts = self.WEIGHTS_LARGE_DISTANCE

        # ---------------------------------------------------------------------
        # This is the array to hold the weighted combination of the GMPEs
        # ---------------------------------------------------------------------
        lnmu = np.zeros_like(sites.vs30)
        # ---------------------------------------------------------------------
        # Hold on to the individual means and stddevs so we can compute the
        # combined stddev
        # ---------------------------------------------------------------------
        lnmu_list = []
        lnsd_list = []

        for i, gmpe in enumerate(self.GMPES):
            # -----------------------------------------------------------------
            # Loop over GMPE list
            # -----------------------------------------------------------------

            set_sites_depth_parameters(sites, gmpe)

            # -----------------------------------------------------------------
            # Select the IMT
            # -----------------------------------------------------------------

            gmpe_imts = [
                imt.__name__ for imt in list(gmpe.DEFINED_FOR_INTENSITY_MEASURE_TYPES)
            ]

            if (
                not isinstance(gmpe, MultiGMPE)
                and (imt.string == "PGV")
                and ("PGV" not in gmpe_imts)
            ):
                ab2020 = AbrahamsonBhasin2020(rup.mag)
                timt = SA(ab2020.getTref())
            else:
                timt = imt

            # -----------------------------------------------------------------
            # Grab GMPE_LIMITS in gmpe instance for later as the multigmpe
            # nests downward.
            # -----------------------------------------------------------------
            if hasattr(self, "GMPE_LIMITS"):
                # Remember that GMPE_LIMITS is only present if it is getting
                # loaded from a config... we could change this eventually.
                gmpe.GMPE_LIMITS = self.GMPE_LIMITS

            # -----------------------------------------------------------------
            # Apply GMPE_LIMITS if applicable
            # -----------------------------------------------------------------
            if hasattr(gmpe, "GMPE_LIMITS"):
                gmpes_with_limits = list(gmpe.GMPE_LIMITS.keys())
                gmpe_class_str = str(gmpe).replace("[", "").replace("]", "")
                if gmpe_class_str in gmpes_with_limits:
                    limit_dict = gmpe.GMPE_LIMITS[gmpe_class_str]
                    for k, v in limit_dict.items():
                        if k == "vs30":
                            vs30min = float(v[0])
                            vs30max = float(v[1])
                            sites.vs30 = np.clip(sites.vs30, vs30min, vs30max)
                            Sites._addDepthParameters(sites)

            # -----------------------------------------------------------------
            # Evaluate
            # -----------------------------------------------------------------
            if not isinstance(gmpe, MultiGMPE):
                ctx = stuff_context(sites, rup, dists)
                lmean, lsd = gmpe.get_mean_and_stddevs(
                    ctx, ctx, ctx, timt, stddev_types
                )
            else:
                lmean, lsd = gmpe.get_mean_and_stddevs(
                    sites, rup, dists, timt, stddev_types
                )

            if not isinstance(gmpe, MultiGMPE):
                # -------------------------------------------------------------
                # We may need to inflate the standard deviations to account for
                # the point-source to finite rupture conversion.
                # -------------------------------------------------------------
                lsd_new = self.__inflatePSSigma__(
                    gmpe, lmean, lsd, sites, rup, dists, timt, stddev_types
                )
                for sd in lsd:
                    lsd_new.append(sd)
                lsd = lsd_new

                # -------------------------------------------------------------
                # If IMT is PGV and PGV is not given by the GMPE, then
                # convert from the appropriate PSA
                # -------------------------------------------------------------
                if (imt.string == "PGV") and ("PGV" not in gmpe_imts):
                    lmean, lsd = ab2020.getPGVandSTDDEVS(
                        lmean, lsd, stddev_types, ctx.rrup, ctx.vs30
                    )

                # -------------------------------------------------------------
                # -------------------------------------------------------------
                if self.HAS_SITE[i] is False:
                    lamps = self.__get_site_factors__(
                        sites, rup, dists, timt, default=True
                    )
                    lmean = lmean + lamps

                # -------------------------------------------------------------
                # Convertions due to component definition
                # -------------------------------------------------------------
                imc_in = gmpe.DEFINED_FOR_INTENSITY_MEASURE_COMPONENT
                imc_out = self.DEFINED_FOR_INTENSITY_MEASURE_COMPONENT
                if imc_in != imc_out:
                    bk17 = BooreKishida2017(imc_in, imc_out)
                    lmean = bk17.convertAmps(imt, lmean, dists.rrup, rup.mag)
                    #
                    # The extra sigma from the component conversion appears to
                    # apply to the total sigma, so the question arises as to
                    # how to apportion it between the intra- and inter-event
                    # sigma. Here we assume it all enters as intra-event sigma.
                    #
                    for j, stddev_type in enumerate(stddev_types):
                        if stddev_type == const.StdDev.INTER_EVENT:
                            continue
                        lsd[j] = bk17.convertSigmas(imt, lsd[j])

            # End: if GMPE is not MultiGMPE

            #
            # At this point lsd will have 2 * len(stddev_types) entries, the
            # first group will have the point-source to finite rupture
            # inflation (if any), and the second set will not; in cases where
            # a finite rupture is used, the two sets will be identical
            #

            # -----------------------------------------------------------------
            # Compute weighted mean and collect the elements to compute sd
            # -----------------------------------------------------------------

            lnmu = lnmu + wts[i] * lmean
            lnmu_list.append(lmean)
            lnsd_list = lnsd_list + lsd

        # -----------------------------------------------------------------
        # The mean is a weighted sum of random variables, so the stddev
        # is the weighted sum of of their covariances (effectively). See:
        # https://en.wikipedia.org/wiki/Variance#Weighted_sum_of_variables
        # for an explanation. Also see:
        # http://usgs.github.io/shakemap/manual4_0/tg_processing.html#ground-motion-prediction
        # for a discussion on the way this is implemented here.
        # -------------------------------------------------------------- # noqa

        nwts = len(wts)
        npwts = np.array(wts).reshape((1, -1))
        nsites = len(lnmu)
        # Find the correlation coefficients among the gmpes; if there are
        # fewer than 10 points, just use an approximation (noting that the
        # correlation among GMPEs tends to be quite high).
        if nsites < 10:
            cc = np.full((nwts, nwts), 0.95)
            np.fill_diagonal(cc, 1.0)
        else:
            np.seterr(divide="ignore", invalid="ignore")
            cc = np.reshape(np.corrcoef(lnmu_list), (nwts, nwts))
            np.seterr(divide="warn", invalid="warn")
            cc[np.isnan(cc)] = 1.0

        # Multiply the correlation coefficients by the weights matrix
        # (this is cheaper than multiplying all of elements of each
        # stddev array by their weights since we have to multiply
        # everything by the correlation coefficient matrix anyway))
        cc = ((npwts * npwts.T) * cc).reshape((nwts, nwts, 1))
        nstds = len(stddev_types)
        lnsd_new = []
        for i in range(nstds * 2):
            sdlist = []
            for j in range(nwts):
                sdlist.append(lnsd_list[j * nstds * 2 + i].reshape((1, 1, -1)))
            sdstack = np.hstack(sdlist)
            wcov = (sdstack * np.transpose(sdstack, axes=(1, 0, 2))) * cc
            # This sums the weighted covariance as each point in the output
            lnsd_new.append(np.sqrt(wcov.sum((0, 1))))

        return lnmu, lnsd_new

    @classmethod
    def __from_config__(cls, conf, filter_imt=None):
        """
        Construct a MultiGMPE from a config file.

        Args:
            conf (dict): Dictionary of config options.
            filter_imt (IMT): An optional IMT to filter/reweight the GMPE list.

        Returns:
            MultiGMPE object.

        """
        IMC = getattr(const.IMC, conf["interp"]["component"])
        selected_gmpe = conf["modeling"]["gmpe"]

        logging.debug(f"selected_gmpe: {selected_gmpe}")
        logging.debug(f"IMC: {IMC}")

        # ---------------------------------------------------------------------
        # Allow for selected_gmpe to be found in either conf['gmpe_sets'] or
        # conf['gmpe_modules'], if it is a GMPE set, then all entries must be
        # either a GMPE or a GMPE set (cannot have a GMPE set that is a mix of
        # GMPEs and GMPE sets).
        # ---------------------------------------------------------------------

        if selected_gmpe in conf["gmpe_sets"].keys():
            selected_gmpe_sets = conf["gmpe_sets"][selected_gmpe]["gmpes"]
            gmpe_set_weights = [
                float(w) for w in conf["gmpe_sets"][selected_gmpe]["weights"]
            ]
            logging.debug(f"selected_gmpe_sets: {selected_gmpe_sets}")
            logging.debug(f"gmpe_set_weights: {gmpe_set_weights}")

            # -----------------------------------------------------------------
            # If it is a GMPE set, does it contain GMPEs or GMPE sets?
            # -----------------------------------------------------------------

            set_of_gmpes = all([s in conf["gmpe_modules"] for s in selected_gmpe_sets])
            set_of_sets = all([s in conf["gmpe_sets"] for s in selected_gmpe_sets])

            if set_of_sets is True:
                mgmpes = []
                for s in selected_gmpe_sets:
                    mgmpes.append(
                        cls.__multigmpe_from_gmpe_set__(conf, s, filter_imt=filter_imt)
                    )
                out = MultiGMPE.__from_list__(mgmpes, gmpe_set_weights, imc=IMC)
            elif set_of_gmpes is True:
                out = cls.__multigmpe_from_gmpe_set__(
                    conf, selected_gmpe, filter_imt=filter_imt
                )
            else:
                raise TypeError(
                    "%s must consist exclusively of keys in "
                    "conf['gmpe_modules'] or conf['gmpe_sets']" % selected_gmpe
                )
        elif selected_gmpe in conf["gmpe_modules"].keys():
            modinfo = conf["gmpe_modules"][selected_gmpe]
            # mod = import_module(modinfo[1])
            # tmpclass = getattr(mod, modinfo[0])
            # out = MultiGMPE.__from_list__([tmpclass()], [1.0], imc=IMC)
            out = MultiGMPE.__from_list__(
                [get_gmpe_from_name(modinfo[0], conf)], [1.0], imc=IMC
            )
        else:
            raise TypeError(
                "conf['modeling']['gmpe'] must be a key in "
                "conf['gmpe_modules'] or conf['gmpe_sets']"
            )

        out.DESCRIPTION = selected_gmpe

        # ---------------------------------------------------------------------
        # Deal with GMPE limits
        # ---------------------------------------------------------------------
        gmpe_lims = conf["gmpe_limits"]

        # We need to replace the short name in the dictionary key with module
        # name here since the conf is not available within the MultiGMPE class.
        mods = conf["gmpe_modules"]
        mod_keys = mods.keys()
        new_gmpe_lims = {}
        for k, v in gmpe_lims.items():
            if k in mod_keys:
                new_gmpe_lims[mods[k][0]] = v
            else:
                new_gmpe_lims[k] = v

        out.GMPE_LIMITS = new_gmpe_lims

        return out

    def __multigmpe_from_gmpe_set__(conf, set_name, filter_imt=None):
        """
        Private method for constructing a MultiGMPE from a set_name.

        Args:
            conf (ConfigObj): A ShakeMap config object.
            filter_imt (IMT): An optional IMT to filter/reweight the GMPE list.
            set_name (str): Set name; must correspond to a key in
                conf['set_name'].

        Returns:
            MultiGMPE.

        """
        IMC = getattr(const.IMC, conf["interp"]["component"])

        selected_gmpes = conf["gmpe_sets"][set_name]["gmpes"]
        selected_gmpe_weights = [
            float(w) for w in conf["gmpe_sets"][set_name]["weights"]
        ]

        # Check for large distance GMPEs
        if "weights_large_dist" in conf["gmpe_sets"][set_name].keys():
            if not conf["gmpe_sets"][set_name]["weights_large_dist"]:
                selected_weights_large_dist = None
            else:
                selected_weights_large_dist = [
                    float(w) for w in conf["gmpe_sets"][set_name]["weights_large_dist"]
                ]
        else:
            selected_weights_large_dist = None

        if "dist_cutoff" in conf["gmpe_sets"][set_name].keys():
            if np.isnan(conf["gmpe_sets"][set_name]["dist_cutoff"]):
                selected_dist_cutoff = None
            else:
                selected_dist_cutoff = float(conf["gmpe_sets"][set_name]["dist_cutoff"])
        else:
            selected_dist_cutoff = None

        if "site_gmpes" in conf["gmpe_sets"][set_name].keys():
            if not conf["gmpe_sets"][set_name]["site_gmpes"]:
                selected_site_gmpes = None
            else:
                selected_site_gmpes = conf["gmpe_sets"][set_name]["site_gmpes"]
        else:
            selected_site_gmpes = None

        if "weights_site_gmpes" in conf["gmpe_sets"][set_name].keys():
            if not conf["gmpe_sets"][set_name]["weights_site_gmpes"]:
                selected_weights_site_gmpes = None
            else:
                selected_weights_site_gmpes = conf["gmpe_sets"][set_name][
                    "weights_site_gmpes"
                ]
        else:
            selected_weights_site_gmpes = None

        # ---------------------------------------------------------------------
        # Import GMPE modules and initialize classes into list
        # ---------------------------------------------------------------------
        gmpes = []
        for g in selected_gmpes:
            # This is the old school way of importing the modules; I'm
            # leaving it in here temporarily just for documentation.
            # mod = import_module(conf['gmpe_modules'][g][1])
            # tmpclass = getattr(mod, conf['gmpe_modules'][g][0])
            # gmpes.append(tmpclass())
            gmpe_name = conf["gmpe_modules"][g][0]
            gmpes.append(get_gmpe_from_name(gmpe_name, conf))

        # ---------------------------------------------------------------------
        # Filter out GMPEs not applicable to this period
        # ---------------------------------------------------------------------
        if filter_imt is not None:
            filtered_gmpes, filtered_wts = filter_gmpe_list(
                gmpes, selected_gmpe_weights, filter_imt
            )
        else:
            filtered_gmpes, filtered_wts = gmpes, selected_gmpe_weights

        # ---------------------------------------------------------------------
        # Import site GMPEs
        # ---------------------------------------------------------------------
        if selected_site_gmpes is not None:
            if isinstance(selected_site_gmpes, str):
                selected_site_gmpes = [selected_site_gmpes]
            site_gmpes = []
            for g in selected_site_gmpes:
                # This is the old school way of importing the modules; I'm
                # leaving it in here temporarily just for documentation.
                # mod = import_module(conf['gmpe_modules'][g][1])
                # tmpclass = getattr(mod, conf['gmpe_modules'][g][0])
                # site_gmpes.append(tmpclass())
                gmpe_name = conf["gmpe_modules"][g][0]
                site_gmpes.append(get_gmpe_from_name(gmpe_name, conf))
        else:
            site_gmpes = None

        # ---------------------------------------------------------------------
        # Filter out site GMPEs not applicable to this period
        # ---------------------------------------------------------------------
        if site_gmpes is not None:
            if filter_imt is not None:
                filtered_site_gmpes, filtered_site_wts = filter_gmpe_list(
                    site_gmpes, selected_weights_site_gmpes, filter_imt
                )
            else:
                filtered_site_gmpes = copy.copy(site_gmpes)
                filtered_site_wts = copy.copy(selected_weights_site_gmpes)
        else:
            filtered_site_gmpes = None
            filtered_site_wts = None

        # ---------------------------------------------------------------------
        # Construct MultiGMPE
        # ---------------------------------------------------------------------
        logging.debug(f"    filtered_gmpes: {filtered_gmpes}")
        logging.debug(f"    filtered_wts: {filtered_wts}")

        mgmpe = MultiGMPE.__from_list__(
            filtered_gmpes,
            filtered_wts,
            default_gmpes_for_site=filtered_site_gmpes,
            default_gmpes_for_site_weights=filtered_site_wts,
            imc=IMC,
        )

        # ---------------------------------------------------------------------
        # Append large-distance info if specified
        # ---------------------------------------------------------------------
        if selected_dist_cutoff is not None:
            if filter_imt is not None:
                filtered_gmpes_ld, filtered_wts_ld = filter_gmpe_list(
                    gmpes, selected_weights_large_dist, filter_imt
                )
            else:
                filtered_wts_ld = copy.copy(selected_weights_large_dist)

            mgmpe.CUTOFF_DISTANCE = copy.copy(selected_dist_cutoff)
            mgmpe.WEIGHTS_LARGE_DISTANCE = copy.copy(filtered_wts_ld)

        mgmpe.DESCRIPTION = set_name

        return mgmpe

    @classmethod
    def __from_list__(
        cls,
        gmpes,
        weights,
        imc=const.IMC.GREATER_OF_TWO_HORIZONTAL,
        default_gmpes_for_site=None,
        default_gmpes_for_site_weights=None,
        reference_vs30=760,
    ):
        """
        Construct a MultiGMPE instance from lists of GMPEs and weights.

        Args:
            gmpes (list): List of OpenQuake
                `GMPE <http://docs.openquake.org/oq-hazardlib/master/gsim/index.html#built-in-gsims>`__
                instances.

            weights (list): List of weights; must sum to 1.0.

            imc: Requested intensity measure component. Must be one listed
                `here <http://docs.openquake.org/oq-hazardlib/master/const.html?highlight=imc#openquake.hazardlib.const.IMC>`__.
                The amplitudes returned by the GMPEs will be converted to this
                IMT. Default is 'GREATER_OF_TWO_HORIZONTAL', which is used by
                ShakeMap. See discussion in
                `this section <http://usgs.github.io/shakemap/tg_choice_of_parameters.html#use-of-peak-values-rather-than-mean>`__
                of the ShakeMap manual.

            default_gmpes_for_site (list):
                Optional list of OpenQuake GMPE instance to use as a site term
                for any of the GMPEs that do not have a site term.

                Notes:

                    * We do not check for consistency in the reference rock
                      defintion, so the user nees to be aware of this issue and
                      holds responsibiilty for ensuring compatibility.
                    * We check whether or not a GMPE has a site term by c
                      hecking the REQUIRES_SITES_PARAMETERS slot for vs30.

            default_gmpes_for_site_weights: Weights for default_gmpes_for_site.
                Must sum to one and be same length as default_gmpes_for_site.
                If None, then weights are set to be equal.

            reference_vs30:
                Reference rock Vs30 in m/s. We do not check that this matches
                the reference rock in the GMPEs so this is the responsibility
                of the user.

        """  # noqa

        # ---------------------------------------------------------------------
        # Check that GMPE weights sum to 1.0:
        # ---------------------------------------------------------------------

        if np.abs(np.sum(weights) - 1.0) > 1e-7:
            raise Exception("Weights must sum to one.")

        # ---------------------------------------------------------------------
        # Check that length of GMPE weights equals length of gmpe list
        # ---------------------------------------------------------------------

        if len(weights) != len(gmpes):
            raise Exception("Length of weights must match length of GMPE list.")

        # ---------------------------------------------------------------------
        # Check that gmpes is a list of OQ GMPE instances
        # ---------------------------------------------------------------------

        for g in gmpes:
            if not isinstance(g, GMPE):
                raise Exception(f'"{g}" is a {type(g)} not a GMPE instance.')

        self = cls()
        self.GMPES = gmpes
        self.WEIGHTS = weights

        # ---------------------------------------------------------------------
        # Combine the intensity measure types. This is problematic:
        #   - Logically, we should only include the intersection of the sets
        #     of imts for the different GMPEs.
        #   - In practice, this is not feasible because most GMPEs in CEUS and
        #     subduction zones do not have PGV.
        #   - So instead we will use the union of the imts and then convert
        #     to get the missing imts later in get_mean_and_stddevs.
        # ---------------------------------------------------------------------

        imts = [set(g.DEFINED_FOR_INTENSITY_MEASURE_TYPES) for g in gmpes]
        self.DEFINED_FOR_INTENSITY_MEASURE_TYPES = set.union(*imts)

        # ---------------------------------------------------------------------
        # For VirtualIPE class, we also want to know if ALL of the GMPEs are
        # defined for PGV, in which case we will convert from PGV to MI,
        # otherwise use PGA or Sa.
        # ---------------------------------------------------------------------
        haspgv = [PGV in set(g.DEFINED_FOR_INTENSITY_MEASURE_TYPES) for g in gmpes]
        self.ALL_GMPES_HAVE_PGV = all(haspgv)

        # ---------------------------------------------------------------------
        # Store intensity measure types for conversion in get_mean_and_stddevs.
        # ---------------------------------------------------------------------
        self.IMCs = [g.DEFINED_FOR_INTENSITY_MEASURE_COMPONENT for g in gmpes]

        # ---------------------------------------------------------------------
        # Store the component
        # ---------------------------------------------------------------------
        self.DEFINED_FOR_INTENSITY_MEASURE_COMPONENT = imc

        # ---------------------------------------------------------------------
        # Intersection of GMPE standard deviation types
        # ---------------------------------------------------------------------
        stdlist = [set(g.DEFINED_FOR_STANDARD_DEVIATION_TYPES) for g in gmpes]
        self.DEFINED_FOR_STANDARD_DEVIATION_TYPES = set.intersection(*stdlist)

        # ---------------------------------------------------------------------
        # Need union of site parameters, but it is complicated by the
        # different depth parameter flavors.
        # ---------------------------------------------------------------------
        sitepars = [set(g.REQUIRES_SITES_PARAMETERS) for g in gmpes]
        self.REQUIRES_SITES_PARAMETERS = set.union(*sitepars)

        # ---------------------------------------------------------------------
        # Construct a list of whether or not each GMPE has a site term
        # ---------------------------------------------------------------------
        self.HAS_SITE = ["vs30" in g.REQUIRES_SITES_PARAMETERS for g in gmpes]

        # ---------------------------------------------------------------------
        # Checks and sort out defaults
        # ---------------------------------------------------------------------

        # things to check if default_gmpes_for_site is provided
        if default_gmpes_for_site is not None:
            # check that default_gmpe_for_site are OQ GMPEs or None
            for g in default_gmpes_for_site:
                if not isinstance(g, GMPE):
                    raise Exception(f'"{g}" is not a GMPE instance.')

            # apply default weights if necessary
            if default_gmpes_for_site_weights is None:
                n = len(default_gmpes_for_site)
                default_gmpes_for_site_weights = [1 / n] * n

        # Things to check if one or more GMPE does not have a site term
        if not all(self.HAS_SITE):
            # Raise an exception if no default site is provided
            if default_gmpes_for_site is None:
                raise Exception(
                    "Must provide default_gmpes_for_site if one or"
                    " more GMPE does not have site term."
                )

            # If weights are unspecified, use equal weight
            if default_gmpes_for_site_weights is None:
                default_gmpes_for_site_weights = [
                    1 / len(default_gmpes_for_site)
                ] * len(default_gmpes_for_site)

            # check that length of default_gmpe_for_site matches length of
            # default_gmpe_for_site_weights
            if len(default_gmpes_for_site_weights) != len(default_gmpes_for_site):
                raise Exception(
                    "Length of default_gmpes_for_site_weights "
                    "must match length of default_gmpes_for_site "
                    "list."
                )

            # check weights sum to one if needed
            if not all(self.HAS_SITE):
                if np.sum(default_gmpes_for_site_weights) != 1.0:
                    raise Exception(
                        "default_gmpes_for_site_weights must sum" " to one."
                    )

        # Note: if ALL of the GMPEs do not have a site term (requiring Vs30),
        #       then REQUIRES_SITES_PARAMETERS for the MultiGMPE will not
        #       include Vs30 even though it will be needed to compute the
        #       default site term. So if the site checks have passed to this
        #       point, we should add Vs30 to the set of required site pars:
        self.REQUIRES_SITES_PARAMETERS = set.union(
            set(self.REQUIRES_SITES_PARAMETERS), set(["vs30"])
        )

        self.DEFAULT_GMPES_FOR_SITE = default_gmpes_for_site
        self.DEFAULT_GMPES_FOR_SITE_WEIGHTS = default_gmpes_for_site_weights
        self.REFERENCE_VS30 = reference_vs30

        # ---------------------------------------------------------------------
        # Union of rupture parameters
        # ---------------------------------------------------------------------
        ruppars = [set(g.REQUIRES_RUPTURE_PARAMETERS) for g in gmpes]
        self.REQUIRES_RUPTURE_PARAMETERS = set.union(*ruppars)

        # ---------------------------------------------------------------------
        # Union of distance parameters
        # ---------------------------------------------------------------------
        distpars = [set(g.REQUIRES_DISTANCES) for g in gmpes]
        self.REQUIRES_DISTANCES = set.union(*distpars)

        return self

    def __get_site_factors__(self, sites, rup, dists, imt, default=False):
        """
        Method for computing site amplification factors from the defalut GMPE
        to be applied to GMPEs which do not have a site term.

        **NOTE** Amps are calculated in natural log units and so the ln(amp)
        is returned.

        Args:
            sites (SitesContext): Instance of SitesContext.
            rup (RuptureContext): Instance of RuptureContext.
            dists (DistancesContext): Instance of DistancesContext.
            imt: An instance openquake.hazardlib.imt.
            default (bool): Boolean of whether or not to return the
                amplificaiton factors for the gmpes or default_gmpes_for_site.
                This argument is primarily only intended to be used internally
                for when we just need to access the default amplifications to
                apply to those GMPEs that do not have site terms.

        Returns:
            Site amplifications in natural log units.
        """

        # ---------------------------------------------------------------------
        # Make reference sites context
        # ---------------------------------------------------------------------

        ref_sites = copy.deepcopy(sites)
        ref_sites.vs30 = np.full_like(sites.vs30, self.REFERENCE_VS30)
        # TODO: Should we reset the Sites depth parameters here? Probably.

        # ---------------------------------------------------------------------
        # If default True, construct new MultiGMPE with default GMPE/weights
        # ---------------------------------------------------------------------
        if default is True:
            tmp = MultiGMPE.__from_list__(
                self.DEFAULT_GMPES_FOR_SITE,
                self.DEFAULT_GMPES_FOR_SITE_WEIGHTS,
                self.DEFINED_FOR_INTENSITY_MEASURE_COMPONENT,
            )

        # ---------------------------------------------------------------------
        # If default False, just use self
        # ---------------------------------------------------------------------
        else:
            tmp = self

        lmean, lsd = tmp.get_mean_and_stddevs(
            sites, rup, dists, imt, list(tmp.DEFINED_FOR_STANDARD_DEVIATION_TYPES)
        )
        lmean_ref, lsd = tmp.get_mean_and_stddevs(
            ref_sites, rup, dists, imt, list(tmp.DEFINED_FOR_STANDARD_DEVIATION_TYPES)
        )

        lamps = lmean - lmean_ref

        return lamps

    def __describe__(self):
        """
        Construct a dictionary that describes the MultiGMPE.

        Note: For simplicity, this method ignores issues related to
        GMPEs used for the site term and changes in the GMPE with
        distance. For this level of detail, please see the config files.

        Returns:
            A dictionary representation of the MultiGMPE.
        """
        gmpe_dict = {"gmpes": [], "weights": [], "name": self.DESCRIPTION}

        for i in range(len(self.GMPES)):
            gmpe_dict["weights"].append(self.WEIGHTS[i])
            if isinstance(self.GMPES[i], MultiGMPE):
                gmpe_dict["gmpes"].append(self.GMPES[i].__describe__())
            else:
                gmpe_dict["gmpes"].append(str(self.GMPES[i]))

        return gmpe_dict

    def __inflatePSSigma__(
        self, gmpe, lmean, lsd, sites, rup, dists, imt, stddev_types
    ):
        """
        If the point-source to finite-fault factors are used, we need to
        inflate the intra-event and total standard deviations. We do this
        by standard propagation of error techniques: taking the (numerical)
        derivative of the GMPE (as a function of distance) squared times the
        additional variance from the conversion, added
        to the variance of the GMPE (then taking the square root). We do
        this separately for each of Rrup and Rjb and sum the results.
        If Rrup and Rjb are calculated from a finite rupture model, their
        variance arrays will be "None" and lsd will remain unchanged.
        Otherwise the error inflation will be applied. Normally one or the
        other of Rrup/Rjb will not be used and so that term will be zero; in
        some cases both may be used and both may result in non-zero
        derivatives.

        Args:
            gmpe:
                The GMPE to use for the calculations. Must be a base GMPE and
                not a GMPE set, otherwise no action is taken.
            lmean:
                The mean values returned by the "normal" evaluation of the
                GMPE.
            lsd:
                The standard deviations returned by the "normal" evaluation
                of the GMPE.
            sites:
                The sites context required by the GMPE.
            rup:
                The rupture context required by the GMPE.
            dists:
                The distance context required by the GMPE.
            imt:
                The intensity measure type being evaluated.
            stddev_types:
                The list of stddev types found in lsd.

        Returns:
            list: A list of arrays of inflated standard deviations
            corresponding to the elements of lsd.
        """
        new_sd = []
        delta_distance = 0.01
        delta_var = [0, 0]
        for i, dtype in enumerate(("rrup", "rjb")):
            # Skip dtype if the gmpe does not require it
            if dtype not in gmpe.REQUIRES_DISTANCES:
                continue
            # Skip dtype if it has not been subject to a point-source to
            # finite rupture conversion
            dvar = getattr(dists, dtype + "_var", None)
            if dvar is None:
                continue
            # Add a small amound to the rupture distance (rrup or rjb)
            # and re-evaluate the GMPE
            rup_dist = getattr(dists, dtype)
            rup_dist += delta_distance
            ctx = stuff_context(sites, rup, dists)
            tmean, tsd = gmpe.get_mean_and_stddevs(ctx, ctx, ctx, imt, stddev_types)
            # Find the derivative w.r.t. the rupture distance
            dm_dr = (lmean - tmean) / delta_distance
            # The additional variance is (dm/dr)^2 * dvar
            delta_var[i] = dm_dr ** 2 * dvar
            # Put the rupture distance back to what it was
            rup_dist -= delta_distance
        for i, stdtype in enumerate(stddev_types):
            if stdtype == const.StdDev.INTER_EVENT:
                new_sd.append(lsd[i].copy())
                continue
            new_sd.append(np.sqrt(lsd[i] ** 2 + delta_var[0] + delta_var[1]))
        return new_sd

def filter_gmpe_list(gmpes, wts, imt):
    """
    Method to remove GMPEs from the GMPE list that are not applicable
    to a specific IMT. Rescales the weights to sum to one.

    Args:
        gmpes (list): List of GMPE instances.
        wts (list): List of floats indicating the weight of the GMPEs.
        imt (IMT): OQ IMT to filter GMPE list for.

    Returns:
        tuple: List of GMPE instances and list of weights.

    """
    if wts is None:
        n = len(gmpes)
        wts = [1 / n] * n

    per_max = [np.max(get_gmpe_sa_periods(g)) for g in gmpes]
    per_min = [np.min(get_gmpe_sa_periods(g)) for g in gmpes]
    if imt == PGA():
        sgmpe = [g for g in gmpes if PGA in g.DEFINED_FOR_INTENSITY_MEASURE_TYPES]
        swts = [
            w
            for g, w in zip(gmpes, wts)
            if PGA in g.DEFINED_FOR_INTENSITY_MEASURE_TYPES
        ]
    elif imt == PGV():
        sgmpe = []
        swts = []
        for i in range(len(gmpes)):
            if (PGV in gmpes[i].DEFINED_FOR_INTENSITY_MEASURE_TYPES) or (
                per_max[i] >= 1.0 and per_min[i] <= 1.0
            ):
                sgmpe.append(gmpes[i])
                swts.append(wts[i])
    else:
        per = imt.period
        sgmpe = []
        swts = []
        for i in range(len(gmpes)):
            if per_max[i] >= per and per_min[i] <= per:
                sgmpe.append(gmpes[i])
                swts.append(wts[i])

    if len(sgmpe) == 0:
        raise KeyError(f"No applicable GMPEs from GMPE list for {str(imt)}")

    # Scale weights to sum to one
    swts = np.array(swts)
    swts = swts / np.sum(swts)

    return sgmpe, swts

def get_gmpe_sa_periods(gmpe):
    """
    Method to extract the SA periods defined by a GMPE.

    Args:
        gmpe (GMPE): A GMPE instance.

    Retunrs:
        list: List of periods.

    """
    if gmpe == "[NGAEast]":
        per = gmpe.per_array
    else:
        ctab = get_gmpe_coef_table(gmpe).sa_coeffs
        ilist = list(ctab.keys())
        per = [i.period for i in ilist]
    return per

def get_gmpe_coef_table(gmpe):
    """
    Method for finding the (or "a") GMPE table.

    Notes:

      *  The reason for the complexity here is that there can be multiple
         coefficient tables, and some of them may not have the sa_coeffs
         attribute, which is the main reason for getting the table.
      *  We are also assuming that if there are more than one  coefficient
         table, the range of periods will be the same across all of the
         tables.

    Args:
        gmpe (GMPE): An OQ GMPE instance.

    Returns:
        The associated coefficient table.

    """
    stuff = gmpe.__dir__()
    coef_list = [s for s in stuff if "COEFFS" in s]
    for coef_sel in coef_list:
        cobj = getattr(gmpe, coef_sel)
        if "sa_coeffs" in cobj.__dir__():
            return cobj
    raise Exception(f"GMPE {gmpe} does not contain sa_coeffs attribute.")


#!/usr/bin/env python3
import os
import sys
import argparse
import logging
from astropy.io.registry import identify_format

import numpy as np
from astropy.io import ascii, fits
from astropy.io.votable import parse_single_table
from astropy.table import Table
from astropy.coordinates import Angle, SkyCoord
from astropy.wcs import WCS
from astropy.time import Time
import astropy.units as u
from mskpy.photometry import airmass_app

parser = argparse.ArgumentParser(
    'lmi-standard-phot.py',
    description='Extract standard stars from LMI photometry catalogs.',
    epilog='Use after lmi-add-cat.py.')

parser.add_argument('files', nargs='+', help='files to search for photometry')
parser.add_argument('--catalogs', default=os.path.expanduser('~/data/catalogs'),
                    help='directory holding the catalog files.')
parser.add_argument('--dmax', type=u.Quantity, default='2 arcsec',
                    help='maximum distance for catalog matches')
parser.add_argument('--keep-all', action='store_true',
                    help='when there are multiple matches, save all measurements, otherwise only save the brightest')
parser.add_argument('--logfile', default='lmi-standard-phot.log')
parser.add_argument('-o', default='standard-phot.txt',
                    help='output file name')
parser.add_argument('-f', action='store_true',
                    help='force overwrite of output file')

args = parser.parse_args()

if os.path.exists(args.o) and not args.f:
    raise SystemExit('Refusing to overwrite output file.')

HB_FILTERS = ['OH', 'NH', 'UC', 'CN', 'C3', 'CO+', 'BC',
              'C2', 'GC', 'H2O+', 'RC']
elevation = 2380 * u.m  # DCT elevation

def setup_logging(logfile):
    logger = logging.Logger('lmi-standard-phot.py')
    logger.setLevel(logging.DEBUG)

    # This test allows logging to work when it is run multiple times
    # from ipython
    if len(logger.handlers) == 0:
        formatter = logging.Formatter('%(levelname)s: %(message)s')

        console = logging.StreamHandler(sys.stdout)
        console.setLevel(logging.DEBUG)
        console.setFormatter(formatter)
        logger.addHandler(console)

        logfile = logging.FileHandler(logfile)
        logfile.setLevel(logging.INFO)
        logfile.setFormatter(formatter)
        logger.addHandler(logfile)

    logger.info('#' * 70)
    logger.info(Time.now().iso + 'Z')
    logger.info('Command line: ' + ' '.join(sys.argv[1:]))
    for handler in logger.handlers:
        if hasattr(handler, 'baseFilename'):
            logger.info('Logging to ' + handler.baseFilename)

    return logger

def landolt09():
    table = (parse_single_table(args.catalogs + "/landolt09-sources.xml")
             .to_table(use_names_over_ids=True))
    cat = dict()
    cat['names'] = table['Star'].data.data.astype(str)
    ra = Angle((table['RAh'].data, table['RAm'].data, table['RAs'].data),
               unit=u.hourangle)
    dec = Angle((table['DEd'].data, table['DEm'].data, table['DEs'].data),
                unit=u.deg)
    dec[table['DE-'].data == b'-'] *= -1
    cat['coords'] = SkyCoord(ra, dec)
    cat['V'] = table['Vmag'].data.data
    cat['R'] = table['Vmag'].data.data - table['V-R'].data.data
    cat['R_err'] = np.sqrt(table['e_Vmag'].data.data**2
                           + table['e_V-R'].data.data**2)
    cat['color'] = table['V-R'].data.data
    cat['color_err'] = table['e_V-R'].data.data
    return cat

def smith02():
    table = (parse_single_table("/home/msk/data/catalogs/smith02-standards.xml")
             .to_table(use_names_over_ids=True))
    cat = dict()
    cat['names'] = table['Name'].data.data.astype(str)
    ra = [s for s in table['RAJ2000']]
    dec = [s for s in table['DEJ2000']]
    cat['coords'] = SkyCoord(ra, dec, unit=[u.hourangle, u.deg])
    cat['SDSS-R'] = table["r'mag"].data.data
    cat['SDSS-R_err'] = table["e_r'mag"].data.data
    cat['color'] = table["g'-r'"].data.data
    cat['color_err'] = table["e_g'-r'"].data.data
    cat['SDSS-G'] = table["r'mag"].data.data + cat['color']
    cat['SDSS-G_err'] = np.sqrt(
        table["e_r'mag"].data.data**2 + cat['color_err']**2)
    return cat

def farnham00():
    table = ascii.read('/home/msk/data/catalogs/hb-standards.txt')
    cat = dict()
    cat['names'] = ['HD {}'.format(n) for n in table['HD']]
    cat['coords'] = SkyCoord(table['RA'], table['Dec'],
                             unit=[u.hourangle, u.deg])
    for filt in HB_FILTERS:
        cat[filt] = table[filt].data.data
        cat[filt + '_err'] = np.repeat(0.01, len(table))
    cat['color'] = table['B-V'].data.data
    cat['color_err'] = np.repeat(0.01, len(table))
    return cat

logger = setup_logging(args.logfile)

catalogs = {
    'Farnham et al. 2000': farnham00(),
    'Landolt 2009': landolt09(),
    'Smith et al 2002': smith02()
}

columns = ['file', 'catalog', 'object', 'date', 'za', 'airmass', 'filter',
           'color', 'dist', 'm', 'm_err', 'm_inst', 'm_inst_err']
formats = {
    'airmass': '{:.4f}',
    'color': '{:.4f}',
    'dist': '{:3f}',
    'm': '{:.4f}',
    'm_err': '{:.4f}',
    'm_inst': '{:.4f}',
    'm_inst_err': '{:.4f}'
}

rows = []
for f in sorted(args.files):
    with fits.open(f, mode='readonly') as hdu:
        h = hdu[0].header
        if h['OBSTYPE'] != 'OBJECT':
            logger.info(f + ': wrong OBSTYPE')
            continue

        if 'CAT' not in hdu:
            logger.info(f + ': missing photometry catalog')
            continue

        date = h['DATE-OBS'].replace('T', ' ')
        za = h['ZA']
        am = airmass_app(h['ZA'] * u.deg, elevation)
        filt = h['filters']
        exptime = h['exptime']

        phot = Table(hdu['CAT'].data)
        phot = phot[phot['krflux'] > 0]  # clean out bad data
        w = WCS(hdu[0])
        radec = np.array(w.pixel_to_world_values(
            list(zip(phot['x'], phot['y']))))
        coords = SkyCoord(*radec.T, unit='deg')
        m = -2.5 * np.log10(phot['krflux'] / exptime)
        merr = 1.0857 * phot['krfluxerr'] / phot['krflux']

        n = 0
        for name, cat in catalogs.items():
            if filt not in cat:
                continue
            indices, dist, d3 = coords.match_to_catalog_sky(cat['coords'])
            i = dist < args.dmax
            indices, dist, d3 = indices[i], dist[i], d3[i]
            for i in set(indices):
                n += 1
                matches = np.flatnonzero((indices == i))
                brightest = matches[m[matches].argmin()]
                for j in matches:
                    if not args.keep_all and j != brightest:
                        continue
                    rows.append((f, name, cat['names'][i], date, za,
                                 am, filt, cat['color'][i], dist[j].arcsec,
                                 cat[filt][i], cat[filt + '_err'][i],
                                 m[j], merr[j]))

        logger.info('{}: {} standards found'.format(f, n))

if len(rows) == 0:
    logger.info('No standards found in any file.')
else:
    tab = Table(rows=rows, names=columns)
    tab.meta['comments'] = [Time.now().iso]
    for c, f in formats.items():
        tab[c].format = f

    tab.write(args.o, format='ascii.fixed_width_two_line', overwrite=args.f)

import numpy,import astropy,from astropy