#!/usr/bin/env python
# -*- coding: iso-8859-1 -*-

# quakepy/pmc/pmc2gmt.py
# $Id: pmc2gmt.py 157 2009-02-16 11:15:52Z fab $
#
# The QuakePy package
# http://www.quakepy.org
#

############################################################################
#    Copyright (C) 2007-2009 by Fabian Euchner and Danijel Schorlemmer     #
#    fabian@fabian-euchner.de                                              #
#                                                                          #
#    This program is free software; you can redistribute it and#or modify  #
#    it under the terms of the GNU General Public License as published by  #
#    the Free Software Foundation; either version 2 of the License, or     #
#    (at your option) any later version.                                   #
#                                                                          #
#    This program is distributed in the hope that it will be useful,       #
#    but 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.                          #
#                                                                          #
#    You should have received a copy of the GNU General Public License     #
#    along with this program; if not, write to the                         #
#    Free Software Foundation, Inc.,                                       #
#    59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             #
############################################################################

"""
The QuakePy package
http://www.quakepy.org

pmc2gmt.py

extract PMC results from QPGrid XML file and write to GMT column format

commandline options:

    -p VALUE   extract magnitude of completeness for given probability
    -m VALUE   extract probability for given magnitude of completeness
    -d VALUE   extract data for given depth layer, in km (default: first)
    -s         extract station information
    -z         input XML is g'zipped
    -n         input XML has no namespace prefix

standard input:  QPGrid XML file
standard output: GMT format

requires package lxml
"""

__version__  = '$Id: pmc2gmt.py 157 2009-02-16 11:15:52Z fab $'
__revision__ = '$Revision: 157 $'
__author__   = "Fabian Euchner <fabian@fabian-euchner.de>, Danijel Schorlemmer <ds@usc.edu>"
__license__  = "GPL"

import sys
import getopt
import gzip

from lxml import etree

QP_NS = 'http://quakepy.org/xmlns/quakepy/1.0'

MM_MAGNITUDE, MM_PROBABILITY, MM_STATION = range(1, 4)

def main():

    # format string for longitude and latitude
    lonlat_fmt = '%8.4f'

    # format string for magnitudes
    mag_fmt = '%5.2f'

    # format string for probabilities
    prob_fmt = '%e'
    
    # default mode: output values for completeness map
    mapMode = MM_MAGNITUDE

    # default probability level: 0.999
    targetProbability = 0.999

    # default is non-compressed input
    zipMode = False

    # default: no depth layer given (use first in XML document)
    targetDepthLayer = None

    # default namespace: QuakePy
    qpns = QP_NS

    # Read commandline arguments
    cmdParams = sys.argv[1:]
    opts, args = getopt.gnu_getopt( cmdParams,
                                    'p:m:d:sznh',
                                    [ 'probability=', 'magnitude=', 'depthLayer=', 'station',
                                      'compressed', 'nonamespace', 'help' ] )

    for option, parameter in opts:

        if option == '-p' or option == '--probability':
            targetProbability = float( parameter )
            mapMode = MM_MAGNITUDE
        if option == '-m' or option == '--magnitude':
            targetMagnitude = float( parameter )
            mapMode = MM_PROBABILITY
        if option == '-d' or option == '--depthLayer':
            targetDepthLayer = float( parameter )
        if option == '-s' or option == '--station':
            mapMode = MM_STATION
        if option == '-z' or option == '--compressed':
            zipMode = True
        if option == '-n' or option == '--nonamespace':
            qpns = ''

        if option == '-h' or option == '--help':
            PrintHelp()
            sys.exit()

    if zipMode == True:
        xml = gzip.GzipFile( fileobj = sys.stdin ).read()
    else:
        xml = sys.stdin.read()
        
    # check if it is XML
    if not xml.startswith('<?xml'):
        raise IOError, 'pmc2gmt - input stream is not XML'

    tree = etree.fromstring( xml )

    # Extract station list
    if mapMode == MM_STATION:

        if qpns == '':
            search_pattern = 'PMCStations/PMCStation'
        else:
            search_pattern = '{%s}PMCStations/{%s}PMCStation' % ( qpns, qpns )
            
        stations = tree.findall( search_pattern )

        for curr_station in stations:

            lat = curr_station.get( 'latitude' )
            lon = curr_station.get( 'longitude' )
            stationCode  = curr_station.get( 'stationCode' )
            networkCode  = curr_station.get( 'networkCode' )
            locationCode = curr_station.get( 'locationCode' )

            # if location code is not present in XML, output will be 'None' in last column
            print '%s\t%s\t%s\t%s\t%s' % ( lon, lat, stationCode, networkCode, locationCode )

    # Extract either detection probability or completeness magnitude
    else:

        # find right depthLayer
        if qpns == '':
            search_pattern = 'grid/depthLayer'
        else:
            search_pattern = '{%s}grid/{%s}depthLayer' % ( qpns, qpns )

        depth_layers = tree.findall( search_pattern )

        # loop over depthLayers
        depthLayerFound = False
        for dl in depth_layers:

            if depthLayerFound is True:
                sys.exit()

            # get current depth layer 'at' value
            at = dl.get( 'at' )

            if targetDepthLayer is None:
                depthLayerFound = True
            elif float( at ) == targetDepthLayer:
                depthLayerFound = True
            else:
                continue

            # loop over cells
            if qpns == '':
                search_pattern = 'cell'
            else:
                search_pattern = '{%s}cell' % ( qpns )

            cell = dl.findall( search_pattern )

            # loop over lon/lat cells
            for c in cell:

                # get lat, lon
                lat = c.get( 'lat' )
                lon = c.get( 'lon' )

                ## magnitude map, target probability given
                if mapMode == MM_MAGNITUDE:

                    # set format string for output
                    fmt_str = lonlat_fmt + '\t' + lonlat_fmt + '\t' + mag_fmt

                    # get all mp bins
                    if qpns == '':
                        search_pattern = 'PMCData/mp'
                    else:
                        search_pattern = '{%s}PMCData/{%s}mp' % ( qpns, qpns )
                
                    mp = c.findall( search_pattern )

                    # check if it is correct probability
                    mp_found = False

                    for m in mp:

                        prob = m.get( 'probability' )
                        if floatEqual( float(prob), targetProbability ):

                            # correct value in cell, output mp value
                            print fmt_str % ( float( lon ), float( lat ), float( m.text ) )

                            mp_found = True
                            break

                    if mp_found is not True:
                        
                        # Correct MP not found, compute from probability list
                        if qpns == '':
                            search_pattern = 'PMCData/probability'
                        else:
                            search_pattern = '{%s}PMCData/{%s}probability' % ( qpns, qpns )
                            
                        prob = c.findall( search_pattern )

                        # Search for target probability
                        probList = []
                        for p in prob:
                            if float( p.text ) > targetProbability:
                                probList.append( float( p.get( 'magnitude' ) ) )

                        if len( probList ) > 0:
                            print fmt_str % ( float( lon), float( lat ), min( probList ) )
                        else:
                            fmt_str = lonlat_fmt + '\t' + lonlat_fmt + '\t%s'
                            print fmt_str % ( float( lon), float( lat ), 'nan' )

                ## probability map, target magnitude given
                else:

                    # set format string for output
                    fmt_str = lonlat_fmt + '\t' + lonlat_fmt + '\t' + prob_fmt
                    
                    # get all probability bins
                    if qpns == '':
                        search_pattern = 'PMCData/probability'
                    else:
                        search_pattern = '{%s}PMCData/{%s}probability' % ( qpns, qpns )
                        
                    prob = c.findall( search_pattern )

                    # check if it is correct magnitude
                    prob_found = False
                    
                    for p in prob:

                        mag = p.get( 'magnitude' )

                        if floatEqual( float( mag ), targetMagnitude ):

                            # correct value in cell, output
                            print fmt_str % ( float( lon), float( lat ), float( p.text ) )

                            prob_found = True
                            break
                        
                    if prob_found is not True:
                        fmt_str = lonlat_fmt + '\t' + lonlat_fmt + '\t%s'
                        print fmt_str % ( float( lon), float( lat ), 'nan' )
                    

def PrintHelp():
    print 'converts QPGrid XML file to GMT input'
    print 'Usage: pmc2gmt.py [OPTION]'
    print '  Options'
    print '   -p VALUE, --probability=<value>    Target probability for completeness map'
    print '   -m VALUE, --magnitude=<value>      Target magnitude for probability map'
    print '   -d VALUE, --depthLayer=<value>     Depth layer (in km) for which data is extracted'
    print '   -s, --station                      Station locations for station map'
    print '   -z, --compressed                   Compressed input stream (gzip)'
    print '   -n, --nonamespace                  XML input has no namespace prefix'
    print '   -h, --help                         print this information'
    
def floatEqual( f1, f2, epsilon = 1e-12 ):
    """
    checks if two floating point values can be considered equal
    returns True if difference of the two values is smaller or equal to epsilon
    """
    if abs( f1 - f2 ) > epsilon:
        return False
    else:
        return True
    
main()