#!/usr/bin/env python

"""
    latlon2cartesian.py

    !! These are parameterized differently than stuff listed in textbooks...
"""

import numpy as np

EARTH_RADIUS = 6371
def haversine_distance(point1, point2, radius=EARTH_RADIUS):
    latitude1, longitude1 = np.radians(point1)
    latitude2, longitude2 = np.radians(point2)

    dlongitude = longitude2 - longitude1
    dlatitude = latitude2 - latitude1
    a = np.sin(dlatitude/2)**2 + np.cos(latitude1) * np.cos(latitude2) * np.sin(dlongitude/2)**2
    c = 2 * np.arcsin(np.sqrt(a))
    km = radius * c
    return km


def latlon2cartesian(latlon, d=1):
    lat, lon = np.radians(latlon)
    return np.array([
        d * np.cos(lat) * np.cos(-lon), # x
        d * np.cos(lat) * np.sin(-lon), # y
        d * np.sin(lat),                # z
    ])

def latlon2cartesian_vec(latlon, d=1):
    latlon = np.radians(latlon)
    return np.array([
        d * np.cos(latlon[:,0]) * np.cos(-latlon[:,1]), # x
        d * np.cos(latlon[:,0]) * np.sin(-latlon[:,1]), # y
        d * np.sin(latlon[:,0]),                        # z
    ]).T


def cartesian2latlon(xyz):
    d = (xyz ** 2).sum()
    x, y, z = xyz
    latlon = np.array([
        np.pi / 2 - np.arccos(z / d), # lat
        - np.arctan2(y, x),           # lon
    ])
    return np.degrees(latlon)


if __name__ == "__main__":
    # Some points
    chicago  = np.array([41.8781, -87.6298])
    auckland = np.array([-36.8485, 174.7633])
    beijing  = np.array([39.9042, 116.4074])

    # Unit norm
    assert (latlon2cartesian(chicago) ** 2).sum() == 1, '(chicago) norm != 1'
    assert (latlon2cartesian(auckland) ** 2).sum() == 1, '(auckland) norm != 1'
    assert (latlon2cartesian(beijing) ** 2).sum() == 1, '(beijing) norm != 1'

    # Reversible
    assert np.allclose(cartesian2latlon(latlon2cartesian(chicago)), chicago), '(chicago) reversal failed'
    assert np.allclose(cartesian2latlon(latlon2cartesian(auckland)), auckland), '(auckland) reversal failed'
    assert np.allclose(cartesian2latlon(latlon2cartesian(beijing)), beijing), '(beijing) reversal failed'

    # Can compute distance
    assert np.allclose(
        haversine_distance(chicago, beijing),
        EARTH_RADIUS * np.arccos((latlon2cartesian(chicago) * latlon2cartesian(beijing)).sum()),
    ), '(chicago -> beijing) bad distance'

    assert np.allclose(
        haversine_distance(chicago, auckland),
        EARTH_RADIUS * np.arccos((latlon2cartesian(chicago) * latlon2cartesian(auckland)).sum()),
    ), '(chicago -> auckland) bad distance'

    assert np.allclose(
        haversine_distance(beijing, auckland),
        EARTH_RADIUS * np.arccos((latlon2cartesian(beijing) * latlon2cartesian(auckland)).sum()),
    ), '(beijing -> auckland) bad distance'

    print('pass')