"""Given (x,wave,matrices, slit_profile), extract the flux from each order. For
readability, we keep this separate from the simulator.... but the simulator is
required in order to run this.

To run, create a simulated fits file (e.g. "test_blue.fits") using ghost module then:

blue_high = pymfe.Extractor(pymfe.ghost.Arm('blue', 'high'))

flux,var = blue_high.two_d_extract("test_blue.fits")

plt.plot(blue_high.w_map[0,:], flux[0,:,0])

"""

from __future__ import division, print_function
import numpy as np
import matplotlib.pyplot as plt
try:
    import pyfits
except:
    import astropy.io.fits as pyfits
import pdb
from astropy.modeling import models, fitting
import matplotlib.cm as cm

class Extractor():
    """A class for each arm of the spectrograph. The initialisation function takes a
    single string representing the configuration. For GHOST, it can be "red" or "blue".

    The extraction is defined by 3 key parameters: an "x_map", which is equivalent to
    2dFDR's tramlines and contains a physical x-coordinate for every y (dispersion direction)
    coordinate and order, and a "w_map", which is the wavelength corresponding to every y
    (dispersion direction) coordinate and order.

    sim must include:

    spectral_format_with_matrix()
    make_lenslets(fluxes)

    fluxes (nlenslets x nobj) array
    nl (nlenslets)
    szx (size in x [non-dispersion] direction)
    mode (string,for error messages)
    lenslet_width, im_slit_sz, microns_pix (together define the make_lenslets output)
    """

    def __init__(self,sim,transpose_data=True,badpixmask=[]):
        self.sim = sim
        self.transpose_data=transpose_data
        self.badpixmask = badpixmask
        self.x_map,self.w_map,self.blaze,self.matrices = self.sim.spectral_format_with_matrix()
        #Fill in the slit dimensions in "simulator pixel"s. based on if we are in the
        #high or standard resolution mode.
        self.define_profile(sim.fluxes)

        #Set some default pixel offsets for each lenslet, as used for a square lenslet profile
        ny = self.x_map.shape[1]
        nm = self.x_map.shape[0]
        pix_offset_ix = np.append(np.append([0],np.arange(1,sim.nl).repeat(2)),sim.nl)
        self.square_offsets = np.empty( (2*sim.nl,nm) )
        # The [0,0] component of "matrices" measures the size of a detector pixel in the
        # simulated slit image space. i.e. slitmicrons/detpix.
        for i in range(sim.nl):
            self.square_offsets[:,i] = (pix_offset_ix - sim.nl/2.0) * sim.lenslet_width / self.matrices[i,self.x_map.shape[1]//2,0,0]
        self.sim_offsets = np.empty( (self.sim.im_slit_sz,nm) )
        #Creat an array of slit positions in microns. !!! Add an optional offset to this, i.e. a 1D offset !!!
        im_slit_pix_in_microns = (np.arange(self.sim.im_slit_sz) - self.sim.im_slit_sz/2.0) * self.sim.microns_pix
        for i in range(nm):
            self.sim_offsets[:,i] = im_slit_pix_in_microns / self.matrices[i,self.x_map.shape[1]//2,0,0]
        #To aid in 2D extraction, let's explicitly compute the y offsets corresponding to these x offsets...
        #The "matrices" map pixels back to slit co-ordinates.
        self.slit_tilt = np.zeros( (nm,ny) )
        for i in range(nm):
            for j in range(ny):
                invmat = np.linalg.inv( self.matrices[i,j] )
                #What happens to the +x direction?
                x_dir_map = np.dot(invmat,[1,0])
                self.slit_tilt[i,j] = x_dir_map[1]/x_dir_map[0]

    def define_profile(self,fluxes):
        """ Manually define the slit profile as used in lenslet extraction. As this is
        a low-level function, all lenslets must be defined. e.g. by convention, for the
        star lenslets of the high resolution mode, lenslets 0,1 and 21 through 27 would
        be zero. Also """

        if fluxes.shape[0] != self.sim.nl:
            print("Error: {0:s} resolution mode must have {1:d} lenslets".format(self.sim.mode,self.sim.nl))
        else:
            self.square_profile = np.empty( (fluxes.shape[0]*2, fluxes.shape[1]) )
            self.sim_profile = np.empty( (self.sim.im_slit_sz, fluxes.shape[1]) )
            for i in range(fluxes.shape[1]):
                self.square_profile[:,i] = np.array(fluxes[:,i]).repeat(2)
                im_slit=self.sim.make_lenslets(fluxes=fluxes[:,i])
                self.sim_profile[:,i] = np.sum(im_slit, axis=0)

    def one_d_extract(self, data=[], file='', lenslet_profile='sim', rnoise=3.0):
        """ Extract flux by integrating down columns (the "y" direction), using an
        optimal extraction method.

        Given that some of this code is in common with two_d_extract, the routines could
        easily be merged... however that would make one_d_extract less readable.

        Parameters
        ----------
        data: numpy array (optional)
            Image data, transposed so that dispersion is in the "y" direction. Note that
            this is the transpose of a conventional echellogram. Either data or file
            must be given

        file: string (optional)
            A fits file with conventional row/column directions containing the data to be
            extracted.

        lenslet_profile: 'square' or 'sim'
            Shape of the profile of each fiber as used in the extraction. For a final
            implementation, 'measured' should be a possibility. 'square' assigns each
            pixel uniquely to a single lenslet. For testing only

        badpix: (float array, float array)
            Output of e.g. np.where giving the bad pixel coordinates.

        rnoise: float
            The assumed readout noise.

        WARNING: Binning not implemented yet"""

        if len(data)==0:
            if len(file)==0:
                print("ERROR: Must input data or file")
            else:
                if self.transpose_data:
                    #Transpose the data from the start.
                    data = pyfits.getdata(file).T
                else:
                    data = pyfits.getdata(file)

        ny = self.x_map.shape[1]
        nm = self.x_map.shape[0]
        nx = self.sim.szx

        #Number of "objects"
        no = self.square_profile.shape[1]
        extracted_flux = np.zeros( (nm,ny,no) )
        extracted_var = np.zeros( (nm,ny,no) )

        #Assuming that the data are in photo-electrons, construct a simple model for the
        #pixel inverse variance.
        pixel_inv_var = 1.0/(np.maximum(data,0)/self.sim.gain + rnoise**2)
        pixel_inv_var[self.badpixmask]=0.0

        #Loop through all orders then through all y pixels.
        for i in range(nm):
            print("Extracting order: {0:d}".format(i))
            #Based on the profile we're using, create the local offsets and profile vectors
            if lenslet_profile == 'square':
                offsets = self.square_offsets[:,i]
                profile = self.square_profile
            elif lenslet_profile == 'sim':
                offsets = self.sim_offsets[:,i]
                profile = self.sim_profile
            nx_cutout = 2*int( (np.max(offsets) - np.min(offsets))/2 ) + 2
            phi = np.empty( (nx_cutout,no) )
            for j in range(ny):
                #Check for NaNs
                if self.x_map[i,j] != self.x_map[i,j]:
                    extracted_var[i,j,:] = np.nan
                    continue
                #Create our column cutout for the data and the PSF. !!! Is "round" correct on the next line???
                x_ix = int(np.round(self.x_map[i,j])) - nx_cutout//2 + np.arange(nx_cutout,dtype=int) + nx//2
                for k in range(no):
                    phi[:,k] = np.interp(x_ix - self.x_map[i,j] - nx//2, offsets, profile[:,k])
                    phi[:,k] /= np.sum(phi[:,k])
                #Deal with edge effects...
                ww = np.where( (x_ix >= nx) | (x_ix < 0) )[0]
                x_ix[ww]=0
                phi[ww,:]=0.0

                #Stop here.
#                if i==10:
#                    pdb.set_trace()

                #Cut out our data and inverse variance.
                col_data = data[j,x_ix]
                col_inv_var = pixel_inv_var[j,x_ix]
                #Fill in the "c" matrix and "b" vector from Sharp and Birchall equation 9
                #Simplify things by writing the sum in the computation of "b" as a matrix
                #multiplication. We can do this because we're content to invert the
                #(small) matrix "c" here. Equation 17 from Sharp and Birchall
                #doesn't make a lot of sense... so lets just calculate the variance in the
                #simple explicit way.
                col_inv_var_mat = np.reshape(col_inv_var.repeat(no), (nx_cutout,no) )
                b_mat = phi * col_inv_var_mat
                c_mat = np.dot(phi.T,phi*col_inv_var_mat)
                pixel_weights = np.dot(b_mat,np.linalg.inv(c_mat))
                extracted_flux[i,j,:] = np.dot(col_data,pixel_weights)
                extracted_var[i,j,:] = np.dot(1.0/np.maximum(col_inv_var,1e-12),pixel_weights**2)
                #if ((i % 5)==1) & (j==ny//2):
                #if (i%5==1) & (j==ny//2):
                #if (j==ny//2):
                #    pdb.set_trace()

        return extracted_flux, extracted_var

    def two_d_extract(self, file='', data=[], lenslet_profile='sim', rnoise=3.0, deconvolve=True):
        """ Extract using 2D information. The lenslet model used is a collapsed profile,
        in 1D but where we take into account the slit shear/rotation by interpolating this
        1D slit profile to the nearest two pixels along each row (y-axis in code).

        One key difference to Sharp and Birchall is that c_kj (between equations 8 and 9)
        is the correct normalisation for a (fictitious) 1-pixel wide PSF centered exactly
        on a pixel, but not for a continuum. We normalise correctly for a continuum by
        having one of the \phi functions being one-pixel wide along the slit, and the
        other being unbounded in the dispersion direction.

        Note that the input data has to be the transpose of a conventional echellogram

        TODO:
        1) Neaten the approximate matrix inverse square root

        Parameters
        ----------
        data: numpy array (optional)
            Image data, transposed so that dispersion is in the "y" direction. Note that
            this is the transpose of a conventional echellogram. Either data or file
            must be given

        file: string (optional)
            A fits file with conventional row/column directions containing the data to be
            extracted.

        lenslet_profile: 'square' or 'sim'
            Shape of the profile of each fiber as used in the extraction. For a final
            implementation, 'measured' should be a possibility. 'square' assigns each
            pixel uniquely to a single lenslet. For testing only

        rnoise: float
            The assumed readout noise.

        deconvolve: bool
            Do we deconvolve so that neighboring extracted spectral points
            are statistically independent? This is an approximate deconvolution (a linear
            function of 5 neighboring pixels) so is reasonably robust. """

        if len(data)==0:
            if len(file)==0:
                print("ERROR: Must input data or file")
            else:
                #Transpose the data from the start.
                data = pyfits.getdata(file).T

        ny = self.x_map.shape[1]
        nm = self.x_map.shape[0]
        nx = self.sim.szx

        #Number of "objects"
        no = self.square_profile.shape[1]
        extracted_flux = np.zeros( (nm,ny,no) )
        extracted_var = np.zeros( (nm,ny,no) )
        extracted_covar = np.zeros( (nm,ny-1,no) )

        #Assuming that the data are in photo-electrons, construct a simple model for the
        #pixel inverse variance.
        pixel_inv_var = 1.0/(np.maximum(data,0) + rnoise**2)
        pixel_inv_var[self.badpixmask]=0.0

        #Loop through all orders then through all y pixels.
        for i in range(nm):
            print("Extracting order index: {0:d}".format(i))
            #Based on the profile we're using, create the local offsets and profile vectors
            if lenslet_profile == 'sim':
                offsets = self.sim_offsets[:,i]
                profile = self.sim_profile
            else:
                print("Only sim lenslet profile available for 2D extraction so far...")
                raise userwarning
            nx_cutout = 2*int( (np.max(offsets) - np.min(offsets))/2 ) + 2
            ny_cutout = 2*int(nx_cutout * np.nanmax(np.abs(self.slit_tilt)) / 2) + 3
            for j in range(ny):
                phi = np.zeros( (ny_cutout,nx_cutout,no) )
                phi1d = np.zeros( (ny_cutout,nx_cutout,no) )
                #Check for NaNs
                if self.x_map[i,j] != self.x_map[i,j]:
                    extracted_var[i,j,:] = np.nan
                    continue
                #Create our column cutout for the data and the PSF
                x_ix = int(self.x_map[i,j]) - nx_cutout//2 + np.arange(nx_cutout,dtype=int) + nx//2
                y_ix = j + np.arange(ny_cutout, dtype=int) - ny_cutout//2
                for k in range(no):
                    x_prof = np.interp(x_ix - self.x_map[i,j] - nx//2, offsets, profile[:,k])
                    y_pix = (x_ix - self.x_map[i,j] - nx//2) * self.slit_tilt[i,j] + ny_cutout//2
                    frac_y_pix = y_pix - y_pix.astype(int)
                    subx_ix = np.arange(nx_cutout,dtype=int)
                    phi[y_pix.astype(int),subx_ix,k] = (1-frac_y_pix)*x_prof
                    phi[y_pix.astype(int)+1,subx_ix,k] = frac_y_pix*x_prof
                    phi[:,:,k] /= np.sum(phi[:,:,k])
                    x_prof /= np.sum(x_prof)
                    phi1d[:,:,k] = np.tile(x_prof,ny_cutout).reshape( (ny_cutout, nx_cutout) )
                #Deal with edge effects...
                ww = np.where( (x_ix >= nx) | (x_ix < 0) )[0]
                x_ix[ww]=0
                phi[:,ww,:]=0.0
                phi1d[:,ww,:]=0.0
                ww = np.where( (y_ix >= ny) | (y_ix < 0) )[0]
                y_ix[ww]=0
                phi[ww,:,:]=0.0
                xy = np.meshgrid(y_ix, x_ix, indexing='ij')
                #Cut out our data and inverse variance.
                col_data = data[xy].flatten()
                col_inv_var = pixel_inv_var[xy].flatten()
                #Fill in the "c" matrix and "b" vector from Sharp and Birchall equation 9
                #Simplify things by writing the sum in the computation of "b" as a matrix
                #multiplication. We can do this because we're content to invert the
                #(small) matrix "c" here. Equation 17 from Sharp and Birchall
                #doesn't make a lot of sense... so lets just calculate the variance in the
                #simple explicit way.
                col_inv_var_mat = np.reshape(col_inv_var.repeat(no), (ny_cutout*nx_cutout,no) )
                phi = phi.reshape( (ny_cutout*nx_cutout,no) )
                phi1d = phi1d.reshape( (ny_cutout*nx_cutout,no) )
                b_mat = phi * col_inv_var_mat
                c_mat = np.dot(phi.T,phi1d*col_inv_var_mat)
                pixel_weights = np.dot(b_mat,np.linalg.inv(c_mat))
#                if (j==1000):
#                        pdb.set_trace()
                extracted_flux[i,j,:] = np.dot(col_data,pixel_weights)
                extracted_var[i,j,:] = np.dot(1.0/np.maximum(col_inv_var,1e-12),pixel_weights**2)
                if (j > 0):
                    extracted_covar[i,j-1,:] = np.dot(1.0/np.maximum(col_inv_var,1e-12),pixel_weights* \
                        np.roll(last_pixel_weights,-nx_cutout, axis=0))
                last_pixel_weights = pixel_weights.copy()
#                if (j > 591):
#                    pdb.set_trace()
        if (deconvolve):
            #Create the diagonals of the matrix Q gradually, using the Taylor approximation for
            #the matrix inverse.
            #(Bolton and Schlegel 2009, equation 10)
            #D = diag(C)
            #A = D^{-1/2} (C-D) D^{-1/2}, so C = D^{1/2}(I + A)D^{1/2}
            #Then if Q = (I - 1/2 A + 3/8 A^2) D^{-1/2}
            #... then C^{-1} = QQ, approximately.
            #Note that all of this effort doesn't really seem to achieve much at all in practice...
            #an extremely marginal improvement in resolution... but at least formal pixel-to-pixel
            #data independence is returned.
            extracted_sig = np.sqrt(extracted_var)
            a_diag_p1 = extracted_covar/extracted_sig[:,:-1,:]/extracted_sig[:,1:,:]
#            a_diag_m1 = extracted_covar/extracted_var[:,1:,:]
            Q_diag = np.ones( (nm,ny,no) )
            Q_diag[:,:-1,:] += 3/8.0*a_diag_p1**2
            Q_diag[:,1:,:]  += 3/8.0*a_diag_p1**2
#            Q_diag[:,:-1,:] += 3/8.0*a_diag_p1*a_diag_m1
#            Q_diag[:,1:,:]  += 3/8.0*a_diag_p1*a_diag_m1
            Q_diag /= extracted_sig
            extracted_sqrtsig = np.sqrt(extracted_sig)
            Q_diag_p2 = 3/8.0*a_diag_p1[:,:-1,:]*a_diag_p1[:,1:,:]/extracted_sqrtsig[:,2:,:]/extracted_sqrtsig[:,:-2,:]
#            Q_diag_m2 = 3/8.0*a_diag_m1[:,:-1,:]*a_diag_m1[:,1:,:]/extracted_sig[:,:-2,:]
#            Q_diag_m1 = -0.5*a_diag_m1/extracted_sig[:,:-1,:]
            Q_diag_p1 = -0.5*a_diag_p1/extracted_sqrtsig[:,1:,:]/extracted_sqrtsig[:,:-1,:]
    #The approximation doesn't seem to be quite right, with the ~3% uncertainty on the diagonal of cinv, when there should
    #only be a ~1% uncertainty (obtained by going to the next term in the Taylor expansion). But pretty close...
    #Q = np.diag(Q_diag[0,:,0]) + np.diag(Q_diag_m1[0,:,0],k=-1) + np.diag(Q_diag_p1[0,:,0],k=+1) + np.diag(Q_diag_p2[0,:,0],k=+2) + np.diag(Q_diag_m2[0,:,0],k=-2)
    #cinv_approx = np.dot(Q,Q)
    #cinv = np.diag(extracted_var[0,:,0]) + np.diag(extracted_covar[0,:,0],k=1) + np.diag(extracted_covar[0,:,0],k=-1)
    #cinv = np.linalg.inv(cinv)
            #Now we have a sparse matrix with 5 terms. We need to sum down the rows, ignoring the
            #edge pixels
#            s_vect = Q_diag[:,2:-2,:] + Q_diag_p1[:,1:-2,:] + Q_diag_m1[:,2:-1,:] + Q_diag_p2[:,:-2,:] + Q_diag_m2[:,2:,:]
            s_vect = Q_diag.copy()
            s_vect[:,:-1,:] += Q_diag_p1
            s_vect[:,:-2,:] += Q_diag_p2
            s_vect[:,1:,:] += Q_diag_p1
            s_vect[:,2:,:] += Q_diag_p2
            new_var = 1.0/s_vect**2
            new_flux = extracted_flux*Q_diag/s_vect
            new_flux[:,:-1,:] += extracted_flux[:,1:,:]*Q_diag_p1/s_vect[:,1:,:]
            new_flux[:,:-2,:] += extracted_flux[:,2:,:]*Q_diag_p2/s_vect[:,2:,:]
            new_flux[:,1:,:] += extracted_flux[:,:-1,:]*Q_diag_p1/s_vect[:,:-1,:]
            new_flux[:,2:,:] += extracted_flux[:,:-2,:]*Q_diag_p2/s_vect[:,:-2,:]

            #Fill in the Variance and Flux arrays with NaNs, so that the (not computed) edges
            #are undefined.
 #           new_flux = np.empty_like(extracted_flux)
 #           new_var = np.empty_like(extracted_var)
 #           new_flux[:,:,:]=np.nan
 #           new_var[:,:,:]=np.nan
            #Now fill in the arrays.
 #           new_var[:,2:-2,:] = 1.0/s_vect**2
 #           new_flux[:,2:-2,:] =  extracted_flux[:,2:-2,:]*Q_diag[:,2:-2,:]/s_vect
            #
 #           new_flux[:,2:-2,:] += extracted_flux[:,1:-3,:]*Q_diag_p1[:,1:-2,:]/s_vect
 #           new_flux[:,2:-2,:] += extracted_flux[:,3:-1,:]*Q_diag_p1[:,2:-1,:]/s_vect
 #           new_flux[:,2:-2,:] += extracted_flux[:,:-4,:] *Q_diag_p2[:,:-2,:]/s_vect
 #           new_flux[:,2:-2,:] += extracted_flux[:,4:,:]  *Q_diag_p2[:,2:,:]/s_vect

            return new_flux, new_var
        else:
            return extracted_flux, extracted_var

    def find_lines(self,data,arcfile='lines.txt',outfile='arclines.txt', hw=10,flat_data=[]):
        """Find lines near the locations of input arc lines.

        Parameters
        ----------
        data: numpy array
            data array

        arcfile: string
            file containing lines """

        #First, extract the data
        flux,var = self.one_d_extract(data=data, rnoise=self.sim.rnoise)
        #Read in the lines
        lines = np.loadtxt(arcfile)
        #Only use the first lenslet.
        flux = flux[:,:,0]
        ny = self.x_map.shape[1]
        nm = self.x_map.shape[0]
        nx = self.sim.szx
        lines_out=[]
        if len(flat_data)>0:
            data_to_show = data - 0.05*flat_data
        else:
            data_to_show = data.copy()
        plt.clf()
        plt.imshow( np.arcsinh( (data_to_show-np.median(data_to_show))/1e2) , interpolation='nearest', aspect='auto', cmap=cm.gray)
        for m_ix in range(nm):
            w_ix = np.interp(lines,self.w_map[m_ix,:],np.arange(ny))
            ww = np.where( (w_ix >= hw) & (w_ix < ny-hw) )[0]
            w_ix = w_ix[ww]
            arclines_to_fit = lines[ww]
            for i,ix in enumerate(w_ix):
                x = np.arange(ix-hw,ix+hw,dtype=np.int)
                y = flux[m_ix,x]
                y -= np.min(y) #Rough...
                if np.max(y)< 25*self.sim.rnoise:
                    continue
                g_init = models.Gaussian1D(amplitude=np.max(y), mean=x[np.argmax(y)], stddev=1.5)
                fit_g = fitting.LevMarLSQFitter()
                g = fit_g(g_init, x, y)
                #Wave, ypos, xpos, m, amplitude, fwhm
                xpos = nx//2+np.interp(g.mean.value,np.arange(ny),self.x_map[m_ix])
                ypos = g.mean.value
                plt.plot(xpos,ix,'bx')
                plt.plot(xpos,ypos,'rx') #!!! Maybe around the other way?
                plt.text(xpos+10,ypos,str(arclines_to_fit[i]),color='green',fontsize=10)
                lines_out.append( [arclines_to_fit[i],ypos,xpos,m_ix+self.sim.m_min,g.amplitude.value, g.stddev.value*2.3548] )
        plt.axis([0,nx,ny,0])
        lines_out = np.array(lines_out)
        np.savetxt(outfile,lines_out,fmt='%9.4f %7.2f %7.2f %2d %7.1e %4.1f')


#!/usr/bin/env python

import copy
from importlib import import_module

import numpy as np

from openquake.hazardlib.gsim.base import GMPE
from openquake.hazardlib.imt import PGA, PGV, SA
from openquake.hazardlib import const

from shakelib.conversions.imt.newmark_hall_1982 import NewmarkHall1982
from shakelib.conversions.imc.boore_kishida_2017 import BooreKishida2017
from shakelib.sites import Sites

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

    To do

        * Allow site to be based on a model that isn't a GMPE (e.g.,
          Borcherdt).

    """

    DEFINED_FOR_TECTONIC_REGION_TYPE = None
    DEFINED_FOR_INTENSITY_MEASURE_TYPES = None
    DEFINED_FOR_INTENSITY_MEASURE_COMPONENT = None
    DEFINED_FOR_STANDARD_DEVIATION_TYPES = None
    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>`__.
        """  # noqa

        # 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]

        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

        # ---------------------------------------------------------------------
        # Sort out shapes of sites and dists elements
        # ---------------------------------------------------------------------

        shapes = []
        for k, v in sites.__dict__.items():
            if (k is not 'lons') and (k is not 'lats'):
                shapes.append(v.shape)
        for k, v in dists.__dict__.items():
            if (k is not 'lons') and (k is not 'lats'):
                shapes.append(v.shape)

        shapeset = set(shapes)
        if len(shapeset) != 1:
            raise Exception(
                'All sites and dists elements must have same shape.')
        else:
            orig_shape = list(shapeset)[0]

        # Need to turn all 2D arrays into 1D arrays because of
        # inconsistencies in how arrays are handled in OpenQuake.
        for k, v in dists.__dict__.items():
            if (k is not 'lons') and (k is not 'lats'):
                dists.__dict__[k] = np.reshape(dists.__dict__[k], (-1,))
        for k, v in sites.__dict__.items():
            if (k is not 'lons') and (k is not 'lats'):
                sites.__dict__[k] = np.reshape(sites.__dict__[k], (-1,))

        # ---------------------------------------------------------------------
        # These are arrays to hold the weighted combination of the GMPEs
        # ---------------------------------------------------------------------
        lnmu = np.zeros_like(sites.vs30)
        sd_avail = self.DEFINED_FOR_STANDARD_DEVIATION_TYPES
        if not sd_avail.issuperset(set(stddev_types)):
            raise Exception("Requested an unavailable stddev_type.")

        lnsd2 = [np.zeros_like(sites.vs30) for a in stddev_types]

        for i in range(len(self.GMPES)):
            # -----------------------------------------------------------------
            # Loop over GMPE list
            # -----------------------------------------------------------------

            gmpe = self.GMPES[i]

            sites = MultiGMPE.set_sites_depth_parameters(sites, gmpe)

            # -----------------------------------------------------------------
            # Evaluate GMPEs
            # -----------------------------------------------------------------

            gmpe_imts = [imt.__name__ for imt in
                         gmpe.DEFINED_FOR_INTENSITY_MEASURE_TYPES]
            if (isinstance(imt, PGV)) and ("PGV" not in gmpe_imts):
                # -------------------------------------------------------------
                # If IMT is PGV and PGV is not given by the GMPE, then
                # convert from PSA10.
                # -------------------------------------------------------------
                if self.HAS_SITE[i] is True:
                    psa10, psa10sd = gmpe.get_mean_and_stddevs(
                        sites, rup, dists, SA(1.0), stddev_types)
                else:
                    lamps = self.get_site_factors(
                        sites, rup, dists, SA(1.0), default=True)
                    psa10, psa10sd = gmpe.get_mean_and_stddevs(
                        sites, rup, dists, SA(1.0), stddev_types)
                    psa10 = psa10 + lamps

                lmean, lsd = NewmarkHall1982.psa102pgv(psa10, psa10sd[0])
            else:
                if self.HAS_SITE[i] is True:
                    lmean, lsd = gmpe.get_mean_and_stddevs(
                        sites, rup, dists, imt, stddev_types)
                else:
                    lamps = self.get_site_factors(
                        sites, rup, dists, imt, default=True)
                    lmean, lsd = gmpe.get_mean_and_stddevs(
                        sites, rup, dists, imt, stddev_types)
                    lmean = lmean + lamps

            # -----------------------------------------------------------------
            # Convertions due to component definition
            # -----------------------------------------------------------------

            imc_in = gmpe.DEFINED_FOR_INTENSITY_MEASURE_COMPONENT
            imc_out = self.DEFINED_FOR_INTENSITY_MEASURE_COMPONENT
            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 in range(len(lnsd2)):
                if stddev_types[j] == const.StdDev.INTER_EVENT:
                    continue
                lsd[j] = bk17.convertStddevs(imt, lsd[j], dists.rrup, rup.mag)

            # -----------------------------------------------------------------
            # Compute weighted mean and sd
            # -----------------------------------------------------------------

            lnmu = lnmu + wts[i] * lmean

            # Note: the lnsd2 calculation isn't complete until we drop out of
            # this loop and substract lnmu**2
            for j in range(len(lnsd2)):
                lnsd2[j] = lnsd2[j] + wts[i] * (lmean**2 + lsd[j]**2)

        for j in range(len(lnsd2)):
            lnsd2[j] = lnsd2[j] - lnmu**2

        lnsd = [np.sqrt(a) for a in lnsd2]

        # Undo reshapes of inputs
        for k, v in dists.__dict__.items():
            if (k is not 'lons') and (k is not 'lats'):
                dists.__dict__[k] = np.reshape(dists.__dict__[k], orig_shape)
        for k, v in sites.__dict__.items():
            if (k is not 'lons') and (k is not '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

    @classmethod
    def from_config(cls, conf, filter_imt=None, verbose=False):
        """
        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.
            verbose (bool): Print verbose output for debugging.

        Returns:
            MultiGMPE object.

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

        if verbose is True:
            print('selected_gmpe: %s' % selected_gmpe)
            print('IMC: %s' % 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']]
            if verbose is True:
                print('selected_gmpe_sets: %s' % selected_gmpe_sets)
                print('gmpe_set_weights: %s' % 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, verbose=verbose))
                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,
                    verbose=verbose)
            else:
                raise Exception("%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)
        else:
            raise Exception("conf['modeling']['gmpe'] must be a key in "
                            "conf['gmpe_modules'] or conf['gmpe_sets']")

        out.DESCRIPTION = selected_gmpe
        return out

    def __multigmpe_from_gmpe_set(conf, set_name, filter_imt=None,
                                 verbose=False):
        """
        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:
            mod = import_module(conf['gmpe_modules'][g][1])
            tmpclass = getattr(mod, conf['gmpe_modules'][g][0])
            gmpes.append(tmpclass())

        # ---------------------------------------------------------------------
        # 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:
                mod = import_module(conf['gmpe_modules'][g][1])
                tmpclass = getattr(mod, conf['gmpe_modules'][g][0])
                site_gmpes.append(tmpclass())
        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
        # ---------------------------------------------------------------------
        if verbose is True:
            print('    filtered_gmpes: %s' % filtered_gmpes)
            print('    filtered_wts: %s' % 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)

        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("\"%s\" is not a GMPE instance." % g)

        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 = [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 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 = [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("\"%s\" is not a GMPE instance." % g)

            # 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(
            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 = [g.REQUIRES_RUPTURE_PARAMETERS for g in gmpes]
        self.REQUIRES_RUPTURE_PARAMETERS = set.union(*ruppars)

        # ---------------------------------------------------------------------
        # Union of distance parameters
        # ---------------------------------------------------------------------
        distpars = [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.ones_like(sites.vs30) * self.REFERENCE_VS30

        # ---------------------------------------------------------------------
        # 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

    @staticmethod
    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.
        """

        sites = Sites._addDepthParameters(sites)

        if gmpe == 'AbrahamsonEtAl2014()':
            sites.z1pt0 = sites.z1pt0_ask14_cal
        if gmpe == 'ChiouYoungs2014()':
            # Also BooreEtAl2014() if using subclass with depth parameter
            sites.z1pt0 = sites.z1pt0_cy14_cal
        if gmpe == 'CampbellBozorgnia2014()':
            sites.z2pt5 = sites.z2pt5_cb14_cal
        if gmpe == 'ChiouYoungs2008()':
            sites.z1pt0 = sites.z1pt0_cy08
        if gmpe == 'CampbellBozorgnia2008()':
            sites.z2pt5 = sites.z2pt5_cb07

        return sites

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 imt in
                 get_gmpe_coef_table(g).non_sa_coeffs]
        swts = [w for g, w in zip(gmpes, wts) if imt in
                get_gmpe_coef_table(g).non_sa_coeffs]
    elif imt == PGV():
        sgmpe = []
        swts = []
        for i in range(len(gmpes)):
            if (imt in get_gmpe_coef_table(gmpes[i]).non_sa_coeffs) 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 Exception('No applicable GMPEs from GMPE list for %s' % 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.

    """
    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("GMPE %s does not contain sa_coeffs attribute." % gmpe)


#pylint: disable=dangerous-default-value, unused-import
"""
  This module contains material models containing thermal, fluid, and
  material properties.  These models can be stored to and recalled from
  XML files for archiving.
"""

from collections import ChainMap
import xml.etree.ElementTree as ET

import numpy as np
import scipy.interpolate as inter

from neml import parse, models

class DeformationMaterial:
  """
    Incredibly thin wrapper around a NEML XML file
  """
  def __init__(self, xmlfile, modelname):
    """
      Parameters:
        xmlfile:    file location for the input file
        modelname:  which model to load
    """
    self.xmlfile = xmlfile
    self.modelname = modelname

  def get_neml_model(self):
    """
      Return the actual model for use in a solve
    """
    return parse.parse_xml(self.xmlfile, self.modelname)

class ThermalMaterial:
  """
    Material thermal properties.

    This object needs to provide:
      1) material name
      2) the conductivity, as a function of temperature and its derivative
      3) the diffusivity, as a function of temperature and its derivative
  """
  def get_dict(self):
    """
      Returns the data as a dictionary
    """
    raise NotImplementedError()

  def get_type(self):
    """
      Return the string type for the data
    """
    raise NotImplementedError()

  @classmethod
  def load(cls, fname, modelname):
    """
      Load from a dictionary

      Parameters:
        fname       filename
    """
    tag, typ = find_name(fname, modelname)
    data = load_node(tag)[modelname]

    if typ == "PiecewiseLinearThermalMaterial":
      return PiecewiseLinearThermalMaterial.load(data)
    elif typ == "ConstantThermalMaterial":
      return ConstantThermalMaterial.load(data)
    else:
      raise ValueError("Unknown ThermalMaterial type %s" % typ)

  def save(self, fname, modelname):
    """
      Save the model to an XML file as modelname

      Parameters:
        fname       filename to use
        modelname   (base tag) to use
    """
    root = ET.Element("models")

    save_node(modelname, self.get_dict(), root,
        attrib = {"type": self.get_type()})

    tree = ET.ElementTree(element = root)
    tree.write(fname)

class PiecewiseLinearThermalMaterial(ThermalMaterial):
  """
    Interpolate thermal properties linearly from a table
  """
  def __init__(self, name, temps, cond, diff):
    """
      Properties:
        name:           material name
        temps:          list of temperature points
        cond:           list of conductivity values
        diff:           list of diffusivity values
    """
    if len(temps) != len(cond) or len(temps) != len(diff):
      raise ValueError("The lists of temperatures, conductivity,"
          "and diffusivity values must have equal lengths!")

    self.name = name
    self.temps = np.array(temps)
    self.cond = np.array(cond)
    self.diff = np.array(diff)

    self.fcond, self.dfcond = make_piecewise(self.temps, self.cond)
    self.fdiff, self.dfdiff = make_piecewise(self.temps, self.diff)

  def conductivity(self, T):
    """
      Conductivity as a function of temperature

      Parameters:
        T       temperature
    """
    return self.fcond(T)

  def diffusivity(self, T):
    """
      Diffusivity as a function of temperature

      Parameters:
        T       temperature
    """
    return self.fdiff(T)

  def dconductivity(self, T):
    """
      Derivative of conductivity as a function of temperature

      Parameters:
        T       temperature
    """
    return self.dfcond(T)

  def ddiffusivity(self, T):
    """
      Derivative of diffusivity as a function of temperature

      Parameters:
        T       temperature
    """
    return self.dfdiff(T)

  def get_type(self):
    """
      String type
    """
    return "PiecewiseLinearThermalMaterial"

  def get_dict(self):
    """
      Pickled dictionary
    """
    return {"name": self.name, "temps": string_array(self.temps),
        "cond": string_array(self.cond), "diff": string_array(self.diff)}

  @classmethod
  def load(cls, values):
    """
      Load from a dictionary

      Parameters:
        values  dictionary values
    """
    return cls(values["name"], destring_array(values["temps"]), destring_array(values["cond"]),
        destring_array(values["diff"]))

class ConstantThermalMaterial(ThermalMaterial):
  """
    Constant thermal properties
  """
  def __init__(self, name, k, alpha):
    """
      Properties:
        name:           material name
        k:              conductivity
        alpha:          diffusivity
    """
    self.name = name
    self.cond = k
    self.diff = alpha

  def conductivity(self, T):
    """
      Conductivity as a function of temperature

      Parameters:
        T       temperature
    """
    return T * 0.0 + self.cond

  def diffusivity(self, T):
    """
      Diffusivity as a function of temperature

      Parameters:
        T       temperature
    """
    return T * 0.0 + self.diff

  def dconductivity(self, T):
    """
      Derivative of conductivity as a function of temperature

      Parameters:
        T       temperature
    """
    return T * 0.0

  def ddiffusivity(self, T):
    """
      Derivative of diffusivity as a function of temperature

      Parameters:
        T       temperature
    """
    return T * 0.0

  def get_type(self):
    """
      String type
    """
    return "ConstantThermalMaterial"

  def get_dict(self):
    """
      Pickled dictionary
    """
    return {"name": self.name, "k": str(self.cond),
        "alpha": str(self.diff)}

  @classmethod
  def load(cls, values):
    """
      Load from a dictionary

      Parameters:
        values  dictionary values
    """
    return cls(values["name"], float(values["k"]), float(values["alpha"]))

class FluidMaterial:
  """
    Properties for convective heat transfer.

    This object needs to store:
      1) Fluid name
      2) A map between a ThermalMaterial name and the corresponding
         temperature-dependent film coefficient and its derivative
  """
  def get_dict(self):
    """
      Returns the data as a dictionary
    """
    raise NotImplementedError()

  def get_type(self):
    """
      Return the string type for the data
    """
    raise NotImplementedError()

  @classmethod
  def load(cls, fname, modelname):
    """
      Load a FluidMaterial object from a file

      Parameters:
        fname       file name to load from
    """
    tag, typ = find_name(fname, modelname)
    data = load_node(tag)[modelname]

    if typ == "PiecewiseLinearFluidMaterial":
      return PiecewiseLinearFluidMaterial.load(data)
    elif typ == "ConstantFluidMaterial":
      return ConstantFluidMaterial.load(data)
    else:
      raise ValueError("Unknown FluidMaterial type %s" % typ)

  def save(self, fname, modelname):
    """
      Save the model to an XML file as modelname

      Parameters:
        fname       filename to use
        modelname   (base tag) to use
    """
    root = ET.Element("models")

    save_node(modelname, self.get_dict(), root,
        attrib = {"type": self.get_type()})

    tree = ET.ElementTree(element = root)
    tree.write(fname)

class ConstantFluidMaterial(FluidMaterial):
  """
    Supply a mapping between the material type and a constant
    film coefficient
  """
  def __init__(self, data):
    """
      Dictionary of the form {name: value} mapping
      a material name to the definition of the piecewise linear map

      Parameters:
        data:       the dictionary
    """
    self.data = data

  def get_dict(self):
    return {k: str(v) for k, v in self.data.items()}

  def get_type(self):
    return "ConstantFluidMaterial"

  @classmethod
  def load(cls, values):
    """
      Load from a dictionary

      Parameters:
        values      dictionary data
    """
    data = {k: float(val) for k, val in values.items()}
    return cls(data)

  def coefficient(self, material, T):
    """
      Return the film coefficient for the given material and temperature

      Parameters:
        material:       material name
        T:              temperatures

    """
    if material in self.data:
      return T*0.0 + self.data[material]
    else:
      return T*0.0 + self.data["default"]

  # pylint: disable=unused-argument
  def dcoefficient(self, material, T):
    """
      Return the derivative of the film coefficient with respect to
      temperature for the give material and temperature.

      Parameters:
        material:       material name
        T:              temperatures
    """
    return T * 0.0

class PiecewiseLinearFluidMaterial(FluidMaterial):
  """
    Supply a mapping between the material type and a piecewise linear
    interpolate defining the film coefficient as a function of temperature.
  """
  def __init__(self, data):
    """
      Dictionary of the form {name: (temperatures, values)} mapping
      a material name to the definition of the piecewise linear map

      Parameters:
        data:       the dictionary
    """
    self.data = data

    self.fns = {name: make_piecewise(T, v) for name, (T,v) in data.items()}

  def get_dict(self):
    return {k: {'temp': string_array(T),
      'values': string_array(v)} for k, (T, v) in self.data.items()}

  def get_type(self):
    return "PiecewiseLinearFluidMaterial"

  @classmethod
  def load(cls, values):
    """
      Load from a dictionary

      Parameters:
        values      dictionary data
    """
    data = {k: (destring_array(pair['temp']),
      destring_array(pair['values'])) for k, pair in values.items()}
    return cls(data)

  def coefficient(self, material, T):
    """
      Return the film coefficient for the given material and temperature

      Parameters:
        material:       material name
        T:              temperatures

    """
    if material in self.fns.keys():
      return self.fns[material][0](T)
    else:
      return self.fns["default"][0](T)

  def dcoefficient(self, material, T):
    """
      Return the derivative of the film coefficient with respect to
      temperature for the give material and temperature.

      Parameters:
        material:       material name
        T:              temperatures
    """
    if material in self.fns.keys():
      return self.fns[material][1](T)
    else:
      return self.fns["default"][1](T)

class StructuralMaterial:
  """
  Properties for structural material

  Supply
    1) cycles to failure as a function of temperature and strain range
    2) time to rupture as a function of temperaure and stress
    3) checks creep-fatigue interaction diagram
  """
  def __init__(self, data):
    self.data = data

  def cycles_to_fail(self, pname, temp, erange):
    """
        Returns fatigue cycles to failure at a given temperature and strain range

        Parameters:
          pname:       property name ("nominalFatigue")
          erange:      strain range in mm/mm
          temp:        temperature in K
    """
    pdata = self.data[pname]
    T, a, n, cutoff = [],[],[],[]

    for i in pdata:
      T.append(destring_array(pdata[i]["T"]))
      a.append(destring_array(pdata[i]["a"]))
      n.append(destring_array(pdata[i]["n"]))
      cutoff.append(destring_array(pdata[i]["cutoff"]))

      if np.array(a).shape != np.array(n).shape:
        raise ValueError("\tThe lists of a and n must have equal lengths!")

    inds=np.array(T).argsort(axis=0)
    T = np.array(T)[inds]
    a = np.array(a)[inds]
    n = np.array(n)[inds]
    cutoff = np.array(cutoff)[inds]

    if temp > max(T):
      raise ValueError("\ttemperature is out of range for cycle to failure determination")

    for i in range(np.size(T, axis=0)):
      if temp<=T[i]:
        polysum = 0.0
        if erange<=cutoff[i]:
          erange = cutoff[i][0][0]
        for (b,m) in zip(a[i][0],n[i][0]):
          polysum+=b*np.log10(erange)**m
        break

    return 10**polysum

  def time_to_rupture(self, pname, temp, stress):
    """
        Returns time to rupture at a given temperature and stress

        Parameters:
          pname:       property name ("averageRupture" or "lowerboundRupture")
          stress:      stress in MPa
          temp:        temperature in K
      """
    pdata = self.data[pname]

    a=destring_array(pdata["a"])
    n=destring_array(pdata["n"])
    C=destring_array(pdata["C"])

    if a.shape != n.shape:
      raise ValueError("The lists of a and n must have equal lengths!")

    if stress.shape != temp.shape:
      raise ValueError("Stress and temperature must have the same shape!")

    zeros = stress == 0.0
    not_zeros = np.logical_not(zeros)

    res = np.zeros(stress.shape)
    for (b,m) in zip(a,n):
      res[not_zeros] += b*np.log10(stress[not_zeros])**m
    res[not_zeros] = 10.0**(res[not_zeros]/temp[not_zeros]-C)
    res[zeros] = np.inf

    return res

  def inside_envelope(self, pname, damage_fatigue, damage_creep):
    """
        Returns True if the point lies inside the design envelope and False if not

        Parameters:
          pname:               property name ("cfinteraction")
          damage_fatigue:      fatigue damage fraction
          creep_fatigue:       creep damage fraction
      """

    if damage_fatigue < 0.0 or damage_creep < 0.0:
      raise ValueError("\tout of range: negative damage fraction")

    pdata = destring_array(self.data[pname])

    x_1 = 0.0
    y_1 = 1.0
    x_2 = pdata[0]
    y_2 = pdata[1]
    x_3 = 1.0
    y_3 = 0.0

    if damage_fatigue < x_2:
      return damage_creep <= ((y_2 - y_1) / (x_2 - x_1) * (damage_fatigue - x_1) + y_1)
    return damage_creep <= ((y_3 - y_2) / (x_3 - x_2) * (damage_fatigue - x_2) + y_2)

  @classmethod
  def load(cls, fname, model):
    """
      Load a Structural Material object from a file

      Parameters:
        fname:       file name to load from
        material:    model name
    """
    tag = ET.parse(fname).getroot().find(model)
    return cls(load_node(tag)[model])

  def save(self, fname, modelname):
    """
      Save to a particular file under a particular model name
    """
    root = ET.Element('models')
    save_node(modelname, self.data, root)
    tree = ET.ElementTree(element = root)
    tree.write(fname)

def make_piecewise(x, y):
  """
    Make two piecewise interpolation functions: a piecewise linear
    interpolate between x and y and the corresponding derivative.
  """
  ydiff = np.zeros(y.shape)
  ydiff[:-1] = np.diff(y) / np.diff(x)
  ydiff[-1] = ydiff[-2]

  return inter.interp1d(x, y), inter.interp1d(x, ydiff, kind = "previous")

def find_name(xmlfile, name):
  """
    Find the base tag with name in an XML file

    Parameters:
      xmlfile:        file name
      name:           tag to look for
  """
  root = ET.parse(xmlfile).getroot()
  tag =  root.find(name)

  return tag, tag.attrib["type"]

def save_node(name, entry, node, attrib = {}):
  """
    Save a dictionary to a particular node

    Parameters:
      name:     name of the new node
      entry:    entry of interest
      node:     ET parent node object

    Additional parameters:
      attribs   node attributes
  """
  nnode = ET.SubElement(node, name, attrib)
  if isinstance(entry, dict):
    for k,v in entry.items():
      save_node(k, v, nnode)
  else:
    nnode.text = entry

def load_node(node):
  """
    The actual function that does the loading by walking the XML file

    Parameters:
      node:      xml node object from ET
  """
  if len(node) > 0:
    return {node.tag: dict(ChainMap(*(load_node(child) for child in node)))}
  else:
    return {node.tag: node.text}

def string_array(array):
  """
    Make a numpy array a space separated string
  """
  return " ".join(map(str, array))

def destring_array(string):
  """
    Make an array from a space separated string
  """
  return np.array(list(map(float, string.split(" "))))