# -*- coding: utf-8 -*-
"""
Returns the radius of the earth in meters at a given longitude, 
latitude point. Uses the WGS-84 earth ellipsoid. 

Created on 2010-12-20 23:48:46 Z
author: bjones
"""

import numpy as np

def R_earth_m(lon=None, lat=None):
    """
    Input: longitude (deg E), latitude (deg N) 
    Output: radius of earth in meters at corresponding point on ellipsoid
    """
    
    # WGS-84 semi-major and semi-minor axes
    a = 6378137.0 # m
    b = 6356752.314245 # m
    
    if (lon is None) or (lat is None):
        # return best known single estimate (this is R3; could use R1 or R2 as well)
        return (a**2*b)**(1./3.) # geometric mean radius, sphere of equivalent volume
        
    # Convert to physicist's notation spherical coordinates (e.g. Arfken, 1985)
    # phi = azimuthal angle, 0 <= phi < 2 pi
    # theta = polar angle, 0 <= theta <= pi
    # --------------------------------------
    # Conversion notes:
    # phi is longitude converted to radians.
    # theta is colatitude: theta = pi/2 - lat (after lat converted to radians).
    
    # physicist's spherical coordinates
    phi = lon*np.pi/180.
    theta = np.pi/2. - lat*np.pi/180.
    
    """
    The equation of an ellipsoid is 
        x^2/a^2 + y^2/a^2 + z^2/b^2 = 1, 
    with two "a" axes because the earth is fat around the equator.
    Now use x = r sin(theta) cos(phi), y = r sin(theta) sin(phi), z = r cos(theta), 
    and we easily obtain the next equation.
    """
    
    inv_r_squared = (np.sin(theta)*np.cos(phi)/a)**2 + \
                    (np.sin(theta)*np.sin(phi)/a)**2 + \
                    (np.cos(theta)/b)**2
    
    return 1.0/np.sqrt(inv_r_squared)
    
if __name__ == '__main__':
    
    # WGS-84 semi-major and semi-minor axes
    a = 6378137.0 # m
    b = 6356752.314245 # m
    
    print 'R3',R_earth_m(),'m' # default
    
    print 'north',R_earth_m(lat=90,lon=0),'m'
    assert np.allclose(b,R_earth_m(lat=90,lon=0))
    
    print 'south',R_earth_m(lat=-90,lon=0),'m'
    assert np.allclose(b,R_earth_m(lat=-90,lon=0))
    
    print 'east',R_earth_m(lat=0,lon=90),'m'
    assert np.allclose(a,R_earth_m(lat=0,lon=90))
    
    print 'west',R_earth_m(lat=0,lon=-90),'m'
    assert np.allclose(a,R_earth_m(lat=0,lon=-90))
    
    print 'dateline',R_earth_m(lat=0,lon=180),'m'
    assert np.allclose(a,R_earth_m(lat=0,lon=180))
    
    #original definition of meter occurs at this latitude (48.276841):
    original_m = 2.0*np.pi*R_earth_m(lat=48.276841,lon=0)/4.0
    print '10million?', original_m,'m'
    assert np.allclose(1e7,original_m)