from __future__ import (absolute_import, print_function, division)

from astropy.table import Table
import numpy as np
from astropy.io import fits

# fix needed till astropy.io get's updated - need to check this
fits.column.ASCII2NUMPY['E'] = fits.column.ASCII2NUMPY['F'] = 'f8'

import matplotlib.pyplot as plt
import string as st
import subprocess
from scipy import interpolate
from scipy.integrate import simps
from pydirtygrid.PhotDG import PhotDG
import h5py

__all__ = ['SpecDG']

class SpecDG:
    def __init__(self):
        """
        Read in the mapping file for all spectra

        Returns
        -------
        (self.)plan:   Table
            table containing the identification of each spectrum,
            as well as its parameters

            (self.)seds:   list
                empty list to save the spectra if you want to
        """
        mapfile = '/astro/dust_kg3/klaw/cloudy2/dirtygrid_db/django/param_table6_combined.tsvx'
        self.plan = Table.read(mapfile, format='ascii')
        self.seds = []
        #self.waves = []

    def findFile(self,file_id):
       """
       Returns the full filename for a given GID from the spectrum mapping.

       Parameter
       ---------
       file_id:   string
          identification number of the spectrum to be read;
          'gid' column of the self.plan Table

       Returns
       -------
       filepath+filename:   string
          the absolute path and name of the .fits file
       """

       alphab = []
       for i in range(26): alphab.append(st.uppercase[i])
       alphab = np.array(alphab)
       ind_id = []
       for i in range(len(file_id)):
           tmp_i = np.where(file_id[i] == alphab)
           if len(tmp_i[0] == 1):
               ind_id.append(tmp_i[0][0])
           else:
               break
       prefix = ''.join(alphab[ind_id])
       filepath = '/astro/dust_kg3/klaw/cloudy2/nasa_fits/'+prefix \
           +'/'+str(file_id[-4:-2])+'/'
       files = []
       proc = subprocess.Popen(['ls', filepath], stdout=subprocess.PIPE)
       for line in proc.stdout.readlines(): files.append(line.rstrip())
       filename = [s for s in files if file_id in s]
       if len(filename) == 0:
           return 'Void'
       else:
           filename = filename[0]
           return filepath+filename

    def findGidFromParam(self, grain, geom, sf_type, metal, age, sfr, tau):
        """
        Returns the GID given a set of parameters
        (If you want to plot or something)

        Parameters
        ----------
        grain:   string
        type of grain

        geom:   string
        geometry

        sf_type: string
        star formation type

        metal:   float
        metallicity

        age:   float
        age of the stellar population

        sfr:   float
        star formation rate

        tau:    float
        optical depth

        Returns
        -------
        file_gid:   string
        identification number of the spectrum with the given set
        of parameters
        """
        # ToDo: A better way to do this?
        ind = np.where( (self.plan['grain'] == grain)
                        & (self.plan['geom'] == geom)
                        & (self.plan['sf_type'] == sf_type)
                        & (self.plan['metal'] == metal)
                        & (self.plan['age'] == age)
                        & (self.plan['sfr'] == sfr)
                        & (self.plan['tau'] == tau) )
        # Not all spectra available
        #  -- Exit if non existant - Until interpolation?
        if len(ind[0]) == 0:
            return
        # If valid, then go on
        file_gid = self.plan['gid'][ind[0][0]]
        return file_gid

    def findParamsFromGid(self, thisgid):
        """
        Returns the parameter values for a given GID
        (Used later to update the cube)

        Parameters
        ----------
        thisgid:   string
        identification number of the spectrum

        Returns
        -------
        thisgt, thisgm, thissf, thismt, thissa, thissr, thista:      integers
        indices of the parameter values for the given spectrum,
        in the multidimensional photometry cube
        """
        ind_gid = np.where(self.plan['gid'] == thisgid)
        thisgt = self.plan['grain'][ind_gid[0][0]]      # Grain type
        thisgm = self.plan['geom'][ind_gid[0][0]]      # Geometry
        thissf = self.plan['sf_type'][ind_gid[0][0]]   # Star formation type
        thismt = self.plan['metal'][ind_gid[0][0]]      # Metal
        thissa = self.plan['age'][ind_gid[0][0]]      # Stellar age
        thissr = self.plan['sfr'][ind_gid[0][0]]      # Star formation rate
        thista = self.plan['tau'][ind_gid[0][0]]      # Optical depth
        return thisgt, thisgm, thissf, thismt, thissa, thissr, thista

    def specGet(self, filename):
        """
        Save a spectrum from the filename

        Parameters
        ----------
        filename:   string
        absolute path to the spectrum you want to read

        Returns
        -------
        (self.)seds:   list
        update the list to add the SED
        (self.)waves:   array
        the wavelengths to the SED
        """
        # Read the fits file
        hdulist = fits.open(filename)
        scidata = hdulist[1].data
        self.seds.append(scidata['Flux'])
        self.waves = scidata['Wavelength']

    def specPlot(self, ind=-1):
        """
        Plot SEDs, either giving a specific index of which if more that
        one saved, or all of them

        Parameters
        ----------
        ind:   integer(s)   (Optional)
        indices of the SED you wish to plot
        """
        plt.xscale('log')
        plt.yscale('log')
        plt.ylabel('Luminosity  [$erg$ $s^{-1}$ $\mu m^{-1}$]', fontsize=15)
        plt.xlabel('Wavelength  [$\mu m$]', fontsize=15)
        if ind == -1:
            print('Plotting all SEDs in object')
            for i in range(len(self.seds)):
                plt.plot(self.waves, self.seds[i])
        else:
            for i in ind: plt.plot(self.waves, np.array(self.seds[ind]))
        # ToDo: need a legend
        # ToDo: specify an figure/subplot if desired?
        plt.show()

    def spec2Phot(self, trans_curve, trans_waves, wave0, energy=1):
        """
        Compute the new photometry

        Parameters
        ----------
        trans_curve:   array
        transmission curve of the new filter /!\ Alread read, as an array

        trans_waves:   array
        wavelengths of the new filter

        wave0:   float
        center/effective wavelength of the new filter

        energy:   integer   (Optional)
        specification for integration with energy(=1, default) or photons

        Returns
        -------
        newcube:   array(3, 6, 2, 5, 50, 29, 25) (float)
        new photometry cube
        """
        # Create an instance of the PhotDG class to call functions
        phtemp = PhotDG(silent=1)
        light_speed = 2.998e14      # microns/s
        n_files = len(self.plan['gid'])
        # Create the array to fill
        newcube = np.empty((3,6,2,5,50,29,25))
        # Read in each spectrum
        for i in range(n_files):
            if (i % 1000) == 0:
                print(i)
            filename = self.findFile(self.plan['gid'][i])
            if filename == 'Void':
                continue
            hdu = fits.open(filename)
            scidata = hdu[1].data
            spec = scidata['Flux']                  # SED
            # Wavelength (same for all files)
            if i == 0:
                dgwave = scidata['Wavelength']
            # Interpolate the spectral response on DG wavelength grid
            # Do it in log space
            func = interpolate.interp1d(trans_waves, np.log10(trans_curve),
                                        kind='slinear')
            tr_ind = np.where( (dgwave > np.min(trans_waves))
                               & (dgwave < np.max(trans_waves)))
            newtran = func(dgwave[tr_ind])
            newtran = 10**newtran
            # Frequencies
            nu0 = light_speed/wave0#[indf]
            nu = light_speed/dgwave[tr_ind]
            # Total response curve, for integration      (Used simps)
            if energy == 1:
                tot_resp_curve = simps(newtran*nu0/nu, dgwave[tr_ind])
                tmp_phot = simps(spec[tr_ind]*newtran, dgwave[tr_ind])
                newphot = tmp_phot / tot_resp_curve
            else:      # In photons. Do the Blackbody too?
                tot_resp_curve = simps(newtran*nu0/nu**2, dgwave[tr_ind])
                tmp_phot = simps(spec[tr_ind]*newtran/nu, dgwave[tr_ind])
                newphot = tmp_phot / tot_resp_curve
            # Fill the array
            curgt, curgm, cursf, curmt, cursa, cursr, curta = \
                self.findParamsFromGid(self.plan['gid'][i])
            indgt, indgm, indsf, indmt, indsa, indsr, indta = \
                phtemp.findIndexFromParams(curgt, curgm, cursf, curmt,
                                           cursa, cursr, curta)
            newcube[indgt, indgm, indsf, indmt, indsa, indsr, indta] = newphot

        return newcube
        #ToDo: maybe save sometimes to avoid all recompute if fails somewhere?

##########


# -*- coding: utf-8 -*-
from ..._core import bolton_q_sat, ensure_contiguous_state
from sympl import (
    ImplicitTendencyComponent, get_constant, initialize_numpy_arrays_with_properties)
import numpy as np
import logging
try:
    from . import _emanuel_convection
except ImportError as error:
    logging.warning(
        'Import failed. Emanuel Convection is likely not compiled and will not '
        'be available.'
    )
    print(error)

class EmanuelConvection(ImplicitTendencyComponent):
    """
    The Emanuel convection scheme from `[Emanuel and Zivkovic-Rothman]`_

    .. _[Emanuel and Zivkovic-Rothman]:
            http://journals.ametsoc.org/doi/abs/10.1175/1520-0469(1999)056%3C1766%3ADAEOAC%3E2.0.CO%3B2

    """

    input_properties = {
        'air_temperature': {
            'dims': ['*', 'mid_levels'],
            'units': 'degK',
        },
        'specific_humidity': {
            'dims': ['*', 'mid_levels'],
            'units': 'kg/kg',
        },
        'eastward_wind': {
            'dims': ['*', 'mid_levels'],
            'units': 'm s^-1',
        },
        'northward_wind': {
            'dims': ['*', 'mid_levels'],
            'units': 'm s^-1',
        },
        'air_pressure': {
            'dims': ['*', 'mid_levels'],
            'units': 'mbar',
        },
        'air_pressure_on_interface_levels': {
            'dims': ['*', 'interface_levels'],
            'units': 'mbar',
        },
        'cloud_base_mass_flux': {
            'dims': ['*'],
            'units': 'kg m^-2 s^-1',
        },
    }

    diagnostic_properties = {
        'convective_state': {
            'dims': ['*'],
            'units': 'dimensionless',
            'dtype': np.int32,
        },
        'convective_precipitation_rate': {
            'dims': ['*'],
            'units': 'mm day^-1',
        },
        'convective_downdraft_velocity_scale': {
            'dims': ['*'],
            'units': 'm s^-1',
        },
        'convective_downdraft_temperature_scale': {
            'dims': ['*'],
            'units': 'degK',
        },
        'convective_downdraft_specific_humidity_scale': {
            'dims': ['*'],
            'units': 'kg/kg',
        },
        'cloud_base_mass_flux': {
            'dims': ['*'],
            'units': 'kg m^-2 s^-1',
        },
        'atmosphere_convective_available_potential_energy': {
            'dims': ['*'],
            'units': 'J kg^-1',
        },
        'air_temperature_tendency_from_convection': {
            'dims': ['*', 'mid_levels'],
            'units': 'degK day^-1',
        }
    }

    tendency_properties = {
        'air_temperature': {'units': 'degK s^-1'},
        'specific_humidity': {'units': 'kg/kg s^-1'},
        'eastward_wind': {'units': 'm s^-2'},
        'northward_wind': {'units': 'm s^-2'},
    }

    def __init__(self,
                 minimum_convecting_layer=1,
                 autoconversion_water_content_threshold=0.0011,
                 autoconversion_temperature_threshold=-55,
                 entrainment_mixing_coefficient=1.5,
                 downdraft_area_fraction=0.05,
                 precipitation_fraction_outside_cloud=0.12,
                 speed_water_droplets=50.0,
                 speed_snow=5.5,
                 rain_evaporation_coefficient=1.0,
                 snow_evaporation_coefficient=0.8,
                 convective_momentum_transfer_coefficient=0.7,
                 downdraft_surface_velocity_coefficient=10.0,
                 convection_bouyancy_threshold=0.9,
                 mass_flux_relaxation_rate=0.1,
                 mass_flux_damping_rate=0.1,
                 reference_mass_flux_timescale=300.,
                 **kwargs):
        """

        Args:

            minimum_convecting_layer (int, optional):
                The least model level from which convection can be initiated. Normally set
                to :code:`1` if using bulk PBL schemes. Else, it should be set to the first
                model level at which the temperature is defined.

            autoconversion_water_content_threshold (float, optional):
                The amount of water vapour in :math:`kg/kg`
                above which condensation occurs in warm (above freezing point of water)
                clouds. This value linearly reduces to zero between the freezing point and
                the :code:`autoconversion_temperature_threshold`.

            autoconversion_temperature_threshold (float, optional):
                The temperature in :math:`^\circ C` below which
                all water vapour is converted to rain/snow.

            entrainment_mixing_coefficient (float, optional):
                The coefficient of mixing for entrainment of environmental air into
                the cloud.

            downdraft_area_fraction (float, optional):
                The fractional area covered by unsaturated downdrafts.

            precipitation_fraction_outside_cloud (float, optional):
                The fraction of precipitation falling outside the cloud.

            speed_water_droplets (float, optional):
                The speed of descent of water droplets in :math:`Pa/s`.

            speed_snow (float, optional):
                The speed of descent of snow in :math:`Pa/s`.

            rain_evaporation_coefficient (float, optional):
                Coefficient governing the rate of evaporation of rain.

            snow_evaporation_coefficient (float, optional):
                Coefficient governing the rate of evaporation of snow.

            convective_momentum_transfer_coefficient (float, optional):
                Coefficient **between 0 and 1** governing momentum transport
                by clouds. A value of 1 **shuts off** momentum transport.

            downdraft_surface_velocity_coefficient (float, optional):
                Coefficient mulitplying the downdraft mass flux to calculate
                the downdraft velocity scale.

            convection_bouyancy_threshold (float, optional):
                The maximum negative temperature perturbation in :math:`degK` a parcel can
                have below the temperature at its level of free convection.
                If difference is greater, and previous cloud base mass flux
                is zero, there is no convection.

            mass_flux_relaxation_rate (float, optional):
                Coefficient governing the rate of relaxation to subcloud-layer
                quasi-equilibrium.

            mass_flux_damping_rate (float, optional):
                Coefficient which damps the currently calculated mass flux towards
                the value from the previous time step.

            reference_mass_flux_timescale (float, optional):
                Timescale used to calculate the actual damping coefficient along with
                :code:`mass_flux_damping_rate` and the current time step.

        """

        if (convective_momentum_transfer_coefficient < 0 or
                convective_momentum_transfer_coefficient > 1):
            raise ValueError("Momentum transfer coefficient must be between 0 and 1.")

        if (downdraft_area_fraction < 0 or
                downdraft_area_fraction > 1):
            raise ValueError("Downdraft fraction must be between 0 and 1.")

        if (precipitation_fraction_outside_cloud < 0 or
                precipitation_fraction_outside_cloud > 1):
            raise ValueError("Outside cloud precipitation fraction must be between 0 and 1.")

        self._con_mom_txfr = convective_momentum_transfer_coefficient
        self._downdraft_area_frac = downdraft_area_fraction
        self._precip_frac_outside_cloud = precipitation_fraction_outside_cloud
        self._min_conv_layer = minimum_convecting_layer
        self._crit_humidity = autoconversion_water_content_threshold
        self._crit_temp = autoconversion_temperature_threshold
        self._entrain_coeff = entrainment_mixing_coefficient
        self._droplet_speed = speed_water_droplets
        self._snow_speed = speed_snow
        self._rain_evap = rain_evaporation_coefficient
        self._snow_evap = snow_evaporation_coefficient
        self._beta = downdraft_surface_velocity_coefficient
        self._dtmax = convection_bouyancy_threshold
        self._mf_damp = mass_flux_damping_rate
        self._alpha = mass_flux_relaxation_rate
        self._mf_timescale = reference_mass_flux_timescale
        self._ntracers = 0
        self._set_fortran_constants()
        super(EmanuelConvection, self).__init__(**kwargs)

    def _set_fortran_constants(self):
        self._g = get_constant('gravitational_acceleration', 'm/s^2')
        self._Cpd = get_constant('heat_capacity_of_dry_air_at_constant_pressure', 'J/kg/degK')
        self._Cpv = get_constant('heat_capacity_of_vapor_phase', 'J/kg/degK')
        self._Rdair = get_constant('gas_constant_of_dry_air', 'J/kg/degK')
        self._Rcond = get_constant('gas_constant_of_vapor_phase', 'J/kg/degK')
        self._Lv = get_constant('latent_heat_of_condensation', 'J/kg')
        self._rho_condensible = get_constant('density_of_liquid_phase', 'kg/m^3')
        self._Cl = get_constant('specific_enthalpy_of_vapor_phase', 'J/kg')
        _emanuel_convection.init_emanuel_convection(
            self._min_conv_layer,
            self._crit_humidity,
            self._crit_temp, self._entrain_coeff,
            self._downdraft_area_frac,
            self._precip_frac_outside_cloud,
            self._droplet_speed, self._snow_speed,
            self._rain_evap, self._snow_evap,
            self._con_mom_txfr, self._dtmax,
            self._beta, self._alpha, self._mf_damp,
            self._Cpd, self._Cpv, self._Cl,
            self._Rcond, self._Rdair,
            self._Lv, self._g,
            self._rho_condensible, self._mf_timescale)

    @ensure_contiguous_state
    def array_call(self, raw_state, timestep):
        """
        Get convective heating and moistening.

        Args:

            raw_state (dict):
                The state dictionary of numpy arrays satisfying this
                component's input properties.

        Returns:

            tendencies (dict), diagnostics (dict):
                * The heating and moistening tendencies
                * Any diagnostics associated.

        """
        self._set_fortran_constants()

        num_cols, num_levs = raw_state['air_temperature'].shape

        max_conv_level = num_levs - 3

        tendencies = initialize_numpy_arrays_with_properties(
            self.tendency_properties, raw_state, self.input_properties
        )
        diagnostics = initialize_numpy_arrays_with_properties(
            self.diagnostic_properties, raw_state, self.input_properties
        )

        q_sat = bolton_q_sat(
            raw_state['air_temperature'],
            raw_state['air_pressure'] * 100,
            self._Cpd, self._Cpv
        )

        _emanuel_convection.convect(
            num_levs,
            num_cols,
            max_conv_level,
            self._ntracers,
            timestep.total_seconds(),
            raw_state['air_temperature'],
            raw_state['specific_humidity'],
            q_sat,
            raw_state['eastward_wind'],
            raw_state['northward_wind'],
            raw_state['air_pressure'],
            raw_state['air_pressure_on_interface_levels'],
            diagnostics['convective_state'],
            diagnostics['convective_precipitation_rate'],
            diagnostics['convective_downdraft_velocity_scale'],
            diagnostics['convective_downdraft_temperature_scale'],
            diagnostics['convective_downdraft_specific_humidity_scale'],
            raw_state['cloud_base_mass_flux'],
            diagnostics['atmosphere_convective_available_potential_energy'],
            tendencies['air_temperature'],
            tendencies['specific_humidity'],
            tendencies['eastward_wind'],
            tendencies['northward_wind'])

        diagnostics['air_temperature_tendency_from_convection'][:] = (
            tendencies['air_temperature'] * 86400.)
        diagnostics['cloud_base_mass_flux'][:] = raw_state['cloud_base_mass_flux']
        return tendencies, diagnostics


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

#
# LICENSE
#
# Copyright (C) 2015-2019 GEM Foundation, G. Weatherill, M. Pagani
#
# 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.

"""
Python tools for calculating activity rates on a grid from a source model
"""
import numpy as np
from openquake.hazardlib.sourceconverter import SourceConverter
from openquake.hazardlib import nrml
from openquake.hazardlib.source.complex_fault import ComplexFaultSource
from openquake.hazardlib.source.characteristic import CharacteristicFaultSource
from openquake.hazardlib.source.simple_fault import SimpleFaultSource
from openquake.hazardlib.source.area import AreaSource
from openquake.hazardlib.source.point import PointSource
from openquake.hazardlib.geo.mesh import Mesh
from openquake.hazardlib.geo.polygon import Polygon

class RateGrid(object):
    """
    Class for calculation of activity rate grids

    :param float xspc:
        Longitude spacing of grid
    :param float yspc:
        Latitude spacing of grid
    :param float zspc:
        Depth spacing (km) of grid
    :param np.ndarray xlim:
        Longitude cell bounds
    :param np.ndarray ylim:
        Latitude cell bounds
    :param np.ndarray zlim:
        Depth cell bounds
    :param int nx:
        Number of longitude cells
    :param int ny:
        Number of latitude cells
    :param int nz:
        Number of depth cells
    :param list source_model:
        Seismic source mode
    :param np.ndarray rates:
        Activity rates
    :param float area_discretisation:
        Discretisation step (km) of area sources
    """

    def __init__(self, limits, sources, area_discretisation=10.):
        """
        Instantiate class with grid configurations
        :param list limits:
             Grid configuration [west, east, xspc, south, north, yspc,
                                 upper, lower, zspc]
        """
        self.xspc = limits[2]
        self.yspc = limits[5]
        self.zspc = limits[8]
        self.xlim = np.arange(limits[0], limits[1] + self.xspc, self.xspc)
        self.ylim = np.arange(limits[3], limits[4] + self.yspc, self.yspc)
        self.zlim = np.arange(limits[6], limits[7] + self.zspc, self.zspc)
        self.nx = len(self.xlim)
        self.ny = len(self.ylim)
        self.nz = len(self.zlim)
        self.source_model = sources
        self.rates = np.zeros([self.nx - 1, self.ny - 1, self.nz - 1],
                              dtype=float)
        self.area_discretisation = area_discretisation

    @classmethod
    def from_model_files(cls, limits, input_model, investigation_time=1.0,
                         simple_mesh_spacing=1.0, complex_mesh_spacing=5.0,
                         mfd_width=0.1, area_discretisation=10.0):
        """
        Reads the hazard model from a file

        :param list limits:
             Grid configuration [west, east, xspc, south, north, yspc,
                                 upper, lower, zspc]
        :param str input_model:
            Path to input source model
        :param float investigation_time:
            Investigation time of Poisson model
        :param float simple_mesh_spacing:
            Rupture mesh spacing of simple fault (km)
        :param float complex_mesh_spacing:
            Rupture mesh spacing of complex fault (km)
        :param float mfd_width:
            Spacing (in magnitude units) of MFD
        :param float area_discretisation:
            Spacing of discretisation of area source (km)
        """
        converter = SourceConverter(investigation_time,
                                    simple_mesh_spacing,
                                    complex_mesh_spacing,
                                    mfd_width,
                                    area_discretisation)
        sources = []
        for grp in nrml.to_python(input_model, converter):
            sources.extend(grp.sources)
        return cls(limits, sources, area_discretisation)

    def number_sources(self):
        """
        Returns the number of sources
        """
        return len(self.source_model)

    def get_rates(self, mmin, mmax=np.inf):
        """
        Returns the cumulative rates greater than Mmin

        :param float mmin:
            Minimum magnitude
        """
        nsrcs = self.number_sources()
        for iloc, source in enumerate(self.source_model):
            print("Source Number %s of %s, Name = %s, Typology = %s" % (
                iloc + 1,
                nsrcs,
                source.name,
                source.__class__.__name__))
            if isinstance(source, CharacteristicFaultSource):
                self._get_fault_rates(source, mmin, mmax)
            elif isinstance(source, ComplexFaultSource):
                self._get_fault_rates(source, mmin, mmax)
            elif isinstance(source, SimpleFaultSource):
                self._get_fault_rates(source, mmin, mmax)
            elif isinstance(source, AreaSource):
                self._get_area_rates(source, mmin, mmax)
            elif isinstance(source, PointSource):
                self._get_point_rates(source, mmin, mmax)
            else:
                print("Source type %s not recognised - skipping!" % source)
                continue

    def _get_point_location(self, location):
        """
        Returns the location in the output grid corresponding to the cell in
        which the epicentre lays

        :param location:
            Source hypocentre as instance of :class:
            openquake.hazardlib.geo.point.Point
        :returns:
            xloc - Location of longitude cell
            yloc - Location of latitude cell
        """
        if (location.longitude < self.xlim[0]) or\
                (location.longitude > self.xlim[-1]):
            return None, None
        xloc = int(((location.longitude - self.xlim[0]) / self.xspc) + 1E-7)
        if (location.latitude < self.ylim[0]) or\
                (location.latitude > self.ylim[-1]):
            return None, None
        yloc = int(((location.latitude - self.ylim[0]) / self.yspc) + 1E-7)
        return xloc, yloc

    def _get_point_rates(self, source, mmin, mmax=np.inf):
        """
        Adds the rates for a point source

        :param source:
            Point source as instance of :class:
            openquake.hazardlib.source.point.PointSource
        :param float mmin:
            Minimum Magnitude
        :param float mmax:
            Maximum Magnitude
        """
        xloc, yloc = self._get_point_location(source.location)
        if (xloc is None) or (yloc is None):
            return
        # Get annual rates
        annual_rate = source.get_annual_occurrence_rates()
        mags = np.array([val[0] for val in annual_rate])
        annual_rate = np.array([val[1] for val in annual_rate])
        idx = np.logical_and(mags >= mmin, mags < mmax)
        annual_rate = np.sum(annual_rate[idx])
        for hypo_depth in source.hypocenter_distribution.data:
            zloc = int((hypo_depth[1] - self.zlim[0]) / self.zspc)
            if (zloc < 0) or (zloc >= (self.nz - 1)):
                continue
            else:
                self.rates[xloc, yloc, zloc] += float(hypo_depth[0]) * \
                    annual_rate

    def _get_area_rates(self, source, mmin, mmax=np.inf):
        """
        Adds the rates from the area source by discretising the source
        to a set of point sources

        :param source:
            Area source as instance of :class:
            openquake.hazardlib.source.area.AreaSource
        """
        points = list(source)
        for point in points:
            self._get_point_rates(point, mmin, mmax)

    def _get_fault_rates(self, source, mmin, mmax=np.inf):
        """
        Adds the rates for a simple or complex fault source

        :param source:
            Fault source as instance of :class:
            openquake.hazardlib.source.simple_fault.SimpleFaultSource or
            openquake.hazardlib.source.complex_fault.ComplexFaultSource
        """
        for rupt in list(source.iter_ruptures()):
            valid_rupt = (rupt.mag >= mmin) and (rupt.mag < mmax)
            if not valid_rupt:
                continue
            grd = np.column_stack([rupt.surface.mesh.lons.flatten(),
                                   rupt.surface.mesh.lats.flatten(),
                                   rupt.surface.mesh.depths.flatten()])
            npts = np.shape(grd)[0]
            counter = np.histogramdd(grd,
                                     bins=[self.xlim, self.ylim, self.zlim]
                                     )[0]
            point_rate = rupt.occurrence_rate / float(npts)
            self.rates += (point_rate * counter)

class RatePolygon(RateGrid):
    """
    Calculates the rate of events within a polygon

    :param limits:
        Polygon as instance of :class: openquake.hazardlib.geo.polygon.Polygon
    :param float upper_depth:
        Upper seismic depth of the polygon (km)
    :param float lower_depth:
        Lower seismic depth of the polygon (km)
    :param list source_model:
        List of seismic sources
    :param float rates:
        Activity rate of polygon
    :param float area_discretisation:
        Discretisation spacing (km) of the area source
    """

    def __init__(self, limits, sources, area_discretisation=10.):
        """
        Instantiate class with grid configurations
        :param dict limits:
             Configuration as dictionary containing:
             * polygon - OpenQuake Polygon
             * uppper_depth - upper seismogenic depth (km)
             * lower_depth - lower seismogenic depth (km)
        """
        assert isinstance(limits["polygon"], Polygon)
        self.limits = limits["polygon"]
        self.upper_depth = limits["upper_depth"]
        self.lower_depth = limits["lower_depth"]
        self.source_model = sources
        self.rates = 0.0
        self.area_discretisation = area_discretisation

    def _get_point_rates(self, source, mmin, mmax=np.inf):
        """
        Adds the rates for a point source

        :param source:
            Point source as instance of :class:
            openquake.hazardlib.source.point.PointSource
        :param float mmin:
            Minimum Magnitude
        :param float mmax:
            Maximum Magnitude
        """
        src_mesh = Mesh.from_points_list([source.location])
        in_poly = self.limits.intersects(src_mesh)[0]
        if not in_poly:
            return
        else:
            for (mag, rate) in source.get_annual_occurrence_rates():
                if (mag < mmin) or (mag > mmax):
                    return
                else:
                    for (prob, depth) in source.hypocenter_distribution.data:
                        if (depth < self.upper_depth) or\
                                (depth > self.lower_depth):
                            continue
                        else:
                            self.rates += (prob * rate)

    def _get_fault_rates(self, source, mmin, mmax=np.inf):
        """
        Adds the rates for a simple or complex fault source

        :param source:
            Fault source as instance of :class:
            openquake.hazardlib.source.simple_fault.SimpleFaultSource or
            openquake.hazardlib.source.complex_fault.ComplexFaultSource
        """
        for rup in list(source.iter_ruptures()):
            if (rup.mag < mmin) or (rup.mag > mmax):
                # Magnitude outside search range
                continue
            depths = rup.surface.mesh.depths.flatten()
            # Generate simple mesh from surface
            rupt_mesh = Mesh(rup.surface.mesh.lons.flatten(),
                             rup.surface.mesh.lats.flatten(),
                             depths)
            # Mesh points in polygon
            in_poly = self.limits.intersects(rupt_mesh)
            in_depth = np.logical_and(depths >= self.upper_depth,
                                      depths <= self.lower_depth)
            idx = np.logical_and(in_poly, in_depth)
            if np.any(idx):
                node_rate = rup.occurrence_rate / float(len(depths))
                self.rates += (node_rate * np.sum(idx))