# Utility routines to accompany the Cassini Jupiter wind velocity dataset.
#
# Roland Young, AOPP, May 2017
#
# These utilites
#    Estimate the random error in the velocity measurements
#    Convert planetocentric to planetographic latitude
#    Convert velocities and errors from spherical to oblate spheroidal geometry
#
# Array arguments passed to these routines are assumed to be NumPy arrays.

import numpy as np

def estimate_random_error(latitude, pair, separation, radius_mean):
    # Estimates the random error in the velocity measurements.
    # Uses Equations 1-2 from the readme.
    # Input:
    #   latitude:    Planetocentric latitude (degrees)
    #   pair:        Image pair for each vector
    #   separation:  Separation between the two images in the pair (s)
    #   radius_mean: Mean radius of the planet (m)
    # Output:
    #   u_err: Estimated random error in zonal velocity u (m.s-1)
    #   v_err: Estimated random error in meridional velocity v (m.s-1)
    # Usage:
    #   u_err, v_err = g14s.estimate_random_error(
    #       latitude, pair, separation, radius_mean)
    # Estimated random error in displacement in pixels
    # Based on combining the three sources of random error in the readme
    err_px = 0.4585
    # Resolution of images (degrees / pixel)
    resolution = 0.05
    u_err = err_px * np.radians(resolution) / separation[pair-1] * radius_mean \
        * np.cos(np.radians(latitude))
    v_err = err_px * np.radians(resolution) / separation[pair-1] * radius_mean
    return u_err, v_err

def convert_latc_to_latg(latitude, radius_ratio):
    # Converts planetocentric latitude to planetographic latitude.
    # Uses Equation 6 from the readme.
    # Input: 
    #   latitude:     Planetocentric latitude (degrees)
    #   radius_ratio: Equatorial radius / Polar radius
    # Output:
    #   latitude_g:   Planetographic latitude (degrees)
    # Usage:
    #   latitude_g = g14s.convert_latc_to_latg(latitude, radius_ratio)
    return np.degrees(np.arctan(radius_ratio**2 * np.tan(np.radians(latitude))))

def convert_uv_to_oblate(latitude, u, v, 
                         radius_equator, radius_mean, radius_pole,
                         u_err=None, v_err=None):
    # Converts velocities and (optionally) velocity errors 
    # from spherical geometry to oblate spheroidal geometry. 
    # Uses Equations 3-10 from the readme.
    # Input:
    #   latitude:        Planetocentric latitude (degrees)
    #   u:               Zonal velocity (m.s-1)
    #   v:               Meridional velocity (m.s-1)
    #   radius_equator:  Equatorial radius (m)
    #   radius_mean:     Mean radius [radius used in spherical geometry] (m)
    #   radius_pole:     Polar radius (m)
    #   u_err [optional] Zonal velocity error (m.s-1)
    #   v_err [optional] Meridional velocity error (m.s-1)
    # Output:
    #   u_oblate: Zonal velocity in oblate spheroidal geometry (m.s-1)
    #   v_oblate: Meridional velocity in oblate spheroidal geometry (m.s-1)
    #   u_err_oblate [optional]: Zonal velocity error in 
    #                            oblate spheroidal geometry (m.s-1)
    #   v_err_oblate [optional]: Meridional velocity error in 
    #                            oblate spheroidal geometry (m.s-1)
    # Usage (velocities only): 
    #   u_oblate, v_oblate = g14s.convert_uv_to_oblate(
    #           latitude, u, v, radius_equator, radius_mean, radius_pole)
    # Usage (velocities and both velocity errors):
    #   u_oblate, v_oblate, u_err_oblate, v_err_oblate = \
    #       g14s.convert_uv_to_oblate(
    #           latitude, u, v, radius_equator, radius_mean, radius_pole,
    #           u_err=u_err, v_err=v_err)
    radius_ratio = radius_equator / radius_pole
    tan_lat2 = np.tan(np.radians(latitude))**2
    little_r = radius_equator / np.sqrt(1.0 + radius_ratio**2*tan_lat2)
    big_r = radius_equator / radius_ratio**2 \
        * ((1.0 + radius_ratio**4*tan_lat2) / 
           (1.0 + radius_ratio**2*tan_lat2))**1.5
    u_conversion = little_r / (radius_mean * np.cos(np.radians(latitude)))
    v_conversion = big_r    /  radius_mean
    u_oblate = u * u_conversion
    v_oblate = v * v_conversion
    if (u_err is not None):
        u_err_oblate = u_err * u_conversion
    if (v_err is not None):
        v_err_oblate = v_err * v_conversion
    if (u_err is not None) & (v_err is not None):
        return u_oblate, v_oblate, u_err_oblate, v_err_oblate
    elif (u_err is not None):
        return u_oblate, v_oblate, u_err_oblate
    elif (v_err is not None):
        return u_oblate, v_oblate, v_err_oblate
    else:
        return u_oblate, v_oblate

if __name__ == '__main__':
    pass