import numpy as np


def visible_cells(h):  # calculate visible ERA5 surface cells (30 km x 30 km) as function of balloon height
    if h < 0:
        res = 4
    elif h > 40000:
        res = 2601
    else:
        res = h * 2597/40000 + 4
    return res


def transform(lon, lat):  # longitude and latitude (in degree) to x, y, z (in m) in WGS84 coordinate system
    # WGS 84 reference coordinate system parameters
    a = 6378137.0  # major axis [m]
    e2 = 6.69437999014e-3  # eccentricity squared

    lon_rad = np.radians(lon)
    lat_rad = np.radians(lat)
    # convert to cartesian coordinates
    r_n = a / (np.sqrt(1 - e2 * (np.sin(lat_rad) ** 2)))
    x = r_n * np.cos(lat_rad) * np.cos(lon_rad)
    y = r_n * np.cos(lat_rad) * np.sin(lon_rad)
    z = r_n * (1 - e2) * np.sin(lat_rad)
    return x, y, z


def radii(lat, h):  # get radii to transform local velocity vectors in m/s to change in lat, lon (in deg/s)
    # WGS 84 reference coordinate system parameters
    a = 6378137.0  # major axis [m]
    e2 = 6.69437999014e-3  # eccentricity squared

    lat_rad = np.radians(lat)
    # convert to total distance to center of Earth (used for latitudes)
    r_lat = h + a / (np.sqrt(1 - e2 * (np.sin(lat_rad) ** 2)))
    # determine distance to Earth rotation axis (used for longitudes)
    r_lon = np.cos(lat_rad) * r_lat
    return r_lon, r_lat