#
# earthquake_tsunami function
#
"""
This function returns a callable object representing an initial water
   displacement generated by a submarine earthqauke.
   
The returned object is a callable okada function that represents
the initial water displacement generated by a submarine earthuake.

The underlying theory is based on:
 Okada, Y. (1985) Surface Deformation due to shear and tensile faults in a half
 space. Bulletin of the Seismological Society of America, 75(4)::1135-1154

To use this function, you first must compile okada_tsunami_fortran.f (which should be in
the same directory). This can be done with:

f2py -c okada_tsunami_fortran.f -m okada_tsunami_fortran

NOTE: Once I found that the above failed, but the following worked:
f2py -c --fcompiler=gnu95 okada_tsunami_fortran.f -m okada_tsunami_fortran

"""
from __future__ import print_function
from builtins import range
from builtins import object
from future.utils import raise_
import numpy

try:
    import okada_tsunami_fortran
except:
    raise Exception('okada_tsunami_fortran does not appear to be compiled.\n \
    You can compile it by running the command:\n \
    "f2py -c okada_tsunami_fortran.f -m okada_tsunami_fortran"\n \
    in the shallow_water directory')
 
from pyproj import Proj
import copy

def earthquake_source(
             source=None,                     # List, numeric array or constant
             filename=None,                    # Okada input file
             domain=None,
             lon_lat_degrees=False,
             lon_before_lat=True,
             utm_zone=None,
             proj4string=None,
             dmax=-1.,
             verbose=False):
    """ Return a function representing the water surface displacement caused by
        a given earthquake slip geometry.  This uses Okada's (1985) functions to
        compute the deformation at the ground surface as a result of the slip.
        
        The earthquake slip geometry is represented with one or more
        rectangular sub-faults, each with a given centroid, depth, length, width,
        strike, dip, and displacement

        The user may either provide this information using 'source' (a numpy
        array), or with a file.

        INPUT PARAMETERS:

        source = numpy array, with rows containing the following parameters for
                 every unit slip surface:

                [centroid_x, centroid_y, centroid_depth,
                 strike, dip, length, width, dis1, dis2,dis3]

                See 'DETAILS' for more information.

        filename = a filename containing the information for 'source'. 
                 The first row is a header (e.g. column labels), and is skipped.

        domain = An anuga 'domain' object. Used to transform coordinates onto
                 the ANUGA coordinate system (which internally has lower-left = (0,0)).
                 To avoid transforming coordinates, use domain=None.
                 The latter would typically be useful in idealised
                 (non-spatial) simulations.
       
        lon_lat_degrees = True/False --- Are centroid_x, centroid_y in degrees (True),
                                 or metres (False)?  'False' is appropriate for 
                                 UTM projected coordinates, or
                                 for non-spatial coordinates in units of metres. 
                                 The latter typically arise in
                                 ANUGA simulations that don't represent a particular 
                                 region of earth.

        lon_before_lat = True / False. If True, then the first column of 'source' is
                         assumed to be longitude (x), and the second is latitude. 
                         Otherwise, the first is latitude (y), and the second is longitude
            

        utm_zone = The utm zone. If utm_zone is not None and lon_lat_degrees=True, then 
                   the coordinates are projected to the UTM zone.
                   Use integers, with negative to represent south. 
                   e.g. UTM zone 55S is represented with utm_zone=-55. 
                        UTM zone 55N would be utm_zone=55

        proj4string = Alternative to utm_zone (more flexible). If proj4string is not None and lon_lat_degrees=True,
                    then the coordinates are projected to the projection defined by this proj4string.
                    Only one of utm_zone OR proj4string may be specified.  

        OUTPUT:
            A callable function(x,y), which outputs the sea floor deformation 
            associated with the earthquake.
            Typically, this is assumed to be identical to the water surface deformation, 
            which should be okay if the earthquake duration is suffuciently short, and the 
            propogation speed of the earthquake is much larger than the ocean gravity wave 
            speed [ = sqrt(g*depth)]

        USAGE:
           tsunami_func = earthquake_source(filename='example_source.txt', lon_lat_degrees=True, utm_zone=-51,verbose=True)

           tsunami_func(3000., 2000.)

            Here 'example_source.txt' looks like this (top few rows including header):
                fault_source
                121.37900  -8.62500  3.5 -285.0 30.0   10.0  10.00   0.01   1.12 0.0
                121.39900  -8.70000  8.5 -285.0 30.0   10.0  10.00  -0.24   1.09 0.0
                121.42000  -8.77600 13.5 -285.0 30.0   10.0  10.00  -0.07   1.62 0.0
                121.44000  -8.85100 18.5 -285.0 30.0   10.0  10.00   0.00   0.00 0.0
                121.46700  -8.60200  3.5 -285.0 30.0   10.0  10.00   0.21   0.87 0.0

            # Alternative usage if the coordinate order is swapped
            tsunami_source = okada_tsunami.earthquake_source(filename='example_source.txt',\
                         verbose=True, lon_lat_degrees=True, lon_before_lat=False, utm_zone=-56, domain=domain)

        DETAILS:
                 centroid_x,centroid_y, are x,y coordinates of the unit slip surface.
                 Their units are either 'm' (if lon_lat_degrees = False), or else they are in decimal degrees (lon/lat)
                 In the latter case, they will be reprojected to the utm_zone, since ANUGA coordinates are always in metres. 
 
                 centroid_depth = depth to the centroid of the fault plain (in km)

                 strike is the fault angle (in degrees), measured from clockwise
                 from north (or equivalently clockwise from any line x==constant, since we
                 assume cartesian coordinates).

                 dip is the slip surface dip (degrees downward from the horizontal). By convention, 
                 the fault always dips to the right of the strike

                 length = the length of the sub-fault (km)
                 width  = the width of the sub-fault (km) DOWN THE DIP (i.e. width=surface_width/cos(dip))

                 dis1, dis2, dis3 are the fault displacement components (m) in the
                 directions 'along strike', 'up-dip', and 'perpendicular to the slip plain'.
                 
                 e.g. for a strike slip earthquake, with slip = 1: dis1 = 1, dis2 = 0, dis3 = 0.
                      for a thrust slip earthquake, with slip = 1: dis1 = 0, dis2 = 1, dis3 = 0
                      for a purely tensional displacement: dis1 = 0, dis2 = 0, dis3 > 0.
    """

    if filename is None and source is None:
        raise Exception('Must have either Okada filename of numeric input for earthquake source')

    # Read from file
    if filename is not None:
        # Check that we haven't also specified a source
        if source is not None:
            raise Exception('Cannot use both filename and numeric input for earthquake source')

        # Open file
        try:
           lns = open(filename).readlines()
        except:
           raise_(Exception, 'Cannot read okada file %s' % filename)
        # Pack into 'source'
        source = []
        for line in lns[1:]:
           source.append([float(x) for x in line.split()])

        # Coerce to array
        source = numpy.array(source)

    # Hack to treat arrays with a single row
    if(len(source.shape)==1):
        source=source.reshape((1, source.shape[0]))
    
    # Switch the order of the spatial coordinates if needed 
    if(lon_before_lat==False):
        source2 = copy.copy(source)
        source2[:,0] = copy.copy(source[:,1])
        source2[:,1] = copy.copy(source[:,0])
        source = source2
 
    # Convert from lon_lat coordinates if needed
    if(lon_lat_degrees==True):
        if (utm_zone is not None) or (proj4string is not None):

            if (utm_zone is not None) and (proj4string is not None):
                raise Exception('Only one of utm_zone or proj4string can not be None')

            # Add a few sanity checks
            lat_abs_max = max(abs(source[:,1]))
            long_abs_max = max(abs(source[:,0]))
            if(long_abs_max > 360):
                raise Exception('Longitude absolute values > 360: Check your inputs')
            if(lat_abs_max > 90):
                raise Exception('Latitude absolute values > 90: Check your inputs')

            if(utm_zone is not None): 
                # Set up function to reproject coordinates
                if verbose: print('project utm_zone = ', utm_zone)
                p = Proj(proj='utm',south=utm_zone<0,zone=abs(utm_zone),ellps='WGS84')
            else:
                if verbose: print('proj4string = ',proj4string) 
                p = Proj(proj4string)

            # Store lon/lat for later printing
            lon_store=copy.copy(source[:,0])
            lat_store=copy.copy(source[:,1])

            for i in range(len(source[:,0])):
                source[i,0], source[i,1] = p(source[i,0], source[i,1])      
        else:
            raise Exception('Need utm_zone for lon_lat coordinates')

    # Sanity checks
    dip=source[:,4]
    mindip=min(dip)
    maxdip=max(dip)
    if((mindip<0.)|(maxdip>90.)):
        raise Exception('ERROR: Dip should be between 0 and 90 degrees')

    depth=source[:,2]
    width=source[:,6]
    topdepth=depth-width/2.0*numpy.sin(dip/180.*numpy.pi)
    if(topdepth.min() < 0.0):
        print('  easting   northing   depth strike    dip length  width  disl1   disl2  disl3 topdepth')
        for i in range(len(source[:,0])):
            if(topdepth[i]<0.0):
                print(source[i,:], topdepth[i])

        raise Exception('ERROR: The tops of some sub-faults are above the ground (printed above)')

    #A few print statements
    if verbose is True:
        print('Earthquake parameters:')
        if lon_lat_degrees==False:
            print('  easting   northing   depth strike    dip length  width  disl1   disl2  disl3')
            for i in range(0,len(source)):
               print('%9.3f  %10.3f  %4.1f  %5.1f   %4.1f  %5.1f  %5.1f  %5.2f  %5.2f %5.2f' % \
                      (source[i,0],source[i,1],source[i,2],source[i,3],source[i,4],source[i,5],\
                       source[i,6],source[i,7],source[i,8], source[i,9]))
        else:
            print('longitude  latitude transformed_easting transformed_northing depth  strike  dip  length  width  disl1 disl2 disl3')
            for i in range(0,len(source)):
                print(' %7.2f  %7.2f  %9.1f  %10.1f  %4.1f  %5.1f   %4.1f  %5.1f  %5.1f  %5.2f  %5.2f %5.2f' % \
                       (lon_store[i], lat_store[i], source[i,0],source[i,1],\
                        source[i,2],source[i,3],source[i,4],source[i,5],\
                        source[i,6],source[i,7],source[i,8], source[i,9]))

    # Offset the coordinates, to translate into the ANUGA coordinate system
    if domain is not None:
        xllcorner = domain.geo_reference.get_xllcorner()
        yllcorner = domain.geo_reference.get_yllcorner()
        for i in range(0,len(source)):
            source[i,0] -= xllcorner  # fault origin (relative)
            source[i,1] -= yllcorner

    # By now, we have converted to 'utm' coordinates         
    return Okada_func(source,verbose=verbose,dmax=-1)

###################################################################
#
# CONVENIENCE FUNCTIONS
#

def rotate_coordinates(coords, angle=0):
    # Rotate coords = [x1, y1] by 'angle' anticlockwise about (0., 0.)
    # Equivalently, write the point given by 'coords' in a new coordinate system, 
    # formed by rotating the axes clockwise about (0., 0.)
    ct=numpy.cos(angle/180.*numpy.pi)    
    st=numpy.sin(angle/180.*numpy.pi)    

    x = ct*coords[0] - st*coords[1]
    y = st*coords[0] + ct*coords[1]

    return [x,y]

def okada_origin_2_slip_centroid(lower_left, eq_depth, eq_length, eq_width,eq_strike, eq_dip):
    # Convenience function to convert from 'okadas' origin for an earthquake
    # slip surface, to the slip surface centroid
    # Note that both are assumed to be in 'real world' coordinates
    # Which are not the same as okada's coordinate system  
    #
    # INPUTS:
    # lower left = [x0, y0] == okada's origin (deepest point at the most negative end of the strike direction) (m)
    # eq_depth = maximum depth of slip surface (km)
    # eq_length, eq_width = rectangular fault dimensions (km)
    # eq_strike, eq_dip  = fault orientation parameters (degrees)

    #import pdb
    #pdb.set_trace()
    new_depth =eq_depth - (eq_width/2.0)*numpy.sin(eq_dip/180.*numpy.pi)
   
    eq_surface_width=eq_width*numpy.cos(eq_dip/180.*numpy.pi)
    ## Rotate the offset to the centroid 
    x_offset, y_offset = rotate_coordinates([eq_length/2.0, eq_surface_width/2.0], -(eq_strike-90.))

    #print x_offset, y_offset
    # Centriod values
    new_x = lower_left[0] + x_offset*1000.
    new_y = lower_left[1] + y_offset*1000.

    return new_x, new_y, new_depth

########################################################################

#
# Okada class
#

"""
   This is a callable class representing the initial water displacment 
   generated by an earthquake.

"""

class Okada_func(object):

   def __init__(self, source,verbose=False,alp=0.5,dmax=-1.):
      self.elon  = source[:,0]
      self.elat  = source[:,1]
      self.edep  = source[:,2]
      self.strk  = source[:,3]
      self.dip   = source[:,4]
      self.lnt = source[:,5]
      self.wdt = source[:,6]
      self.disl1 = source[:,7]
      self.disl2 = source[:,8]
      self.alp     = alp
      self.dmax    = dmax
      self.verbose = verbose
     
      #print self.elon, self.elat, self.edep, self.strk, self.dip, self.lnt, self.wdt, self.disl1, self.disl2, self.alp, self.dmax, self.verbose
 
   def __call__(self, x, y):
      """Make Okada_func a callable object.
      
      If called as a function, this object returns z values representing
      the initial distribution of water heights at the points (x,y)
      produced by a submarine mass failure.
      """
       
      edsp,ndsp,zdsp = okada_tsunami_fortran.fault_disp(self.alp,self.elon,self.elat,self.edep,self.strk,self.dip,\
                                             self.lnt,self.wdt,self.disl1,self.disl2,x,y,self.dmax)
      return zdsp
 
