# Copyright 2024, Scott Smith.  MIT License (see LICENSE).

from collections import namedtuple
from typing import TYPE_CHECKING

import numpy as np
import pyarrow as pa

if TYPE_CHECKING:
    from .base import LogFile

# GPS channel names that share a common timebase and may need timing correction
GPS_CHANNEL_NAMES = ("GPS Speed", "GPS Latitude", "GPS Longitude", "GPS Altitude")

# None of the algorithms are slow
# fastest: Fukushima 2006, but worse accuracy
# most accurate: Vermeille, also very compact code
# runner up: Osen


# Also considered:

# https://ea4eoz.blogspot.com/2015/11/simple-wgs-84-ecef-conversion-functions.html
# slower algorithm, not great accuracy

# http://wiki.gis.com/wiki/index.php/Geodetic_system
# poor height accuracy

# Olson, D. K., Converting Earth-Centered, Earth-Fixed Coordinates to
# Geodetic Coordinates, IEEE Transactions on Aerospace and Electronic
# Systems, 32 (1996) 473-476.
# difficult to vectorize (poor for numpy/python)

GPS = namedtuple("GPS", ["lat", "long", "alt"])


# convert lat/long/zoom to web mercator. lat/long are degrees
# returns x,y as floats - integer component is which tile to download
def llz2web(lat, long, zoom=0):
    # wikipedia web mercator projection
    mult = 0.25 * (2 << zoom)
    return (
        mult * (1 + long / 180),
        mult * (1 - np.log(np.tan(np.pi / 4 + np.pi / 360 * lat)) / np.pi),
    )


# returns lat/long as floats in degrees
def web2ll(x, y, zoom=0):
    mult = 1 / (0.25 * (2 << zoom))
    return (
        np.arctan(np.exp(np.pi - np.multiply(np.pi * mult, y))) * 360 / np.pi - 90,
        np.multiply(180 * mult, x) - 180,
    )


# lat, long = degrees
# x, y, z, alt = meters
def lla2ecef(lat, lon, alt):
    a = 6378137
    e = 8.181919084261345e-2
    e_sq = e * e

    lat = lat * (np.pi / 180)
    lon = lon * (np.pi / 180)

    clat = np.cos(lat)
    slat = np.sin(lat)

    N = a / np.sqrt(1 - e_sq * slat * slat)

    x = (N + alt) * clat * np.cos(lon)
    y = (N + alt) * clat * np.sin(lon)
    z = ((1 - e_sq) * N + alt) * slat

    return x, y, z


# Karl Osen. Accurate Conversion of Earth-Fixed Earth-Centered Coordinates to Geodetic Coordinates.
# [Research Report] Norwegian University of Science and Technology. 2017. hal-01704943v2
# https://hal.science/hal-01704943v2/document
# pretty accurate, reasonably fast
def ecef2lla_osen(x, y, z):
    invaa = +2.45817225764733181057e-0014  # 1/(a^2)
    l = +3.34718999507065852867e-0003  # (e^2)/2
    p1mee = +9.93305620009858682943e-0001  # 1-(e^2)
    p1meedaa = +2.44171631847341700642e-0014  # (1-(e^2))/(a^2)
    ll4 = +4.48147234524044602618e-0005  # 4*(l^2) = e^4
    ll = +1.12036808631011150655e-0005  # l^2 = (e^4)/4
    invcbrt2 = +7.93700525984099737380e-0001  # 1/(2^(1/3))
    inv3 = +3.33333333333333333333e-0001  # 1/3
    inv6 = +1.66666666666666666667e-0001  # 1/6

    w = x * x + y * y
    m = w * invaa
    w = np.sqrt(w)
    n = z * z * p1meedaa
    mpn = m + n
    p = inv6 * (mpn - ll4)
    P = p * p
    G = m * n * ll
    H = 2 * P * p + G
    p = None
    # if H < Hmin: return -1
    C = np.cbrt(H + G + 2 * np.sqrt(H * G)) * invcbrt2
    G = None
    H = None
    i = -ll - 0.5 * mpn
    beta = inv3 * i - C - P / C
    C = None
    P = None
    k = ll * (ll - mpn)
    mpn = None
    # Compute t
    t = np.sqrt(np.sqrt(beta * beta - k) - 0.5 * (beta + i)) + np.sqrt(np.abs(0.5 * (beta - i))) * (
        2 * (m < n) - 1
    )
    beta = None
    # Use Newton-Raphson's method to compute t correction
    g = 2 * l * (m - n)
    m = None
    n = None
    tt = t * t
    dt = -(tt * (tt + (i + i)) + g * t + k) / (4 * t * (tt + i) + g)
    g = None
    i = None
    tt = None
    # compute latitude (range -PI/2..PI/2)
    u = t + dt + l
    v = t + dt - l
    dt = None
    zu = z * u
    wv = w * v
    # compute altitude
    invuv = 1 / (u * v)
    return GPS(
        np.arctan2(zu, wv) * (180 / np.pi),
        np.arctan2(y, x) * (180 / np.pi),
        np.sqrt(np.square(w - wv * invuv) + np.square(z - zu * p1mee * invuv)) * (1 - 2 * (u < 1)),
    )


# https://www.researchgate.net/publication/227215135_Transformation_from_Cartesian_to_Geodetic_Coordinates_Accelerated_by_Halley's_Method/link/0912f50af90e6de252000000/download
# "Fukushima 2006"
# fastest, reasonably accurate but not best
def ecef2lla_fukushima2006(x, y, z):
    a = 6378137.0
    finv = 298.257222101
    f = 1.0 / finv
    e2 = (2 - f) * f
    ec2 = 1 - e2
    ec = np.sqrt(ec2)
    # b = a * ec
    c = a * e2
    # PIH = 2 * np.arctan(1.)

    lamb = np.arctan2(y, x)
    s0 = np.abs(z)
    p2 = x * x + y * y
    p = np.sqrt(p2)
    zc = ec * s0
    c0 = ec * p
    c02 = c0 * c0
    s02 = s0 * s0
    a02 = c02 + s02
    a0 = np.sqrt(a02)
    # a03 = a02 * a0
    a03 = a02
    a03 *= a0
    a02 = None
    # s1 = zc * a03 + c * (s02 * s0)
    s02 *= s0
    s02 *= c
    s1 = s02
    s1 += zc * a03
    s02 = None
    # c1 = p * a03 - c * (c02 * c0)
    c02 *= c0
    c02 *= c
    c1 = p * a03 - c02
    c02 = None
    cs0c0 = c * c0 * s0
    # b0 = 1.5 * cs0c0 * ((p*s0 - zc*c0) * a0 - cs0c0)
    zc *= c0
    b0 = cs0c0
    b0 *= 1.5 * ((p * s0 - zc) * a0 - cs0c0)
    a0 = None
    zc = None
    cs0c0 = None
    s1 = s1 * a03 - b0 * s0
    # cc = ec * (c1 * a03 - b0 * c0)
    c1 *= a03
    b0 *= c0
    c1 -= b0
    cc = c1
    cc *= ec
    c1 = None
    a03 = None
    c0 = None
    b0 = None
    phi = np.arctan2(s1, cc)
    s12 = s1 * s1
    cc2 = cc * cc
    h = (p * cc + s0 * s1 - a * np.sqrt(ec2 * s12 + cc2)) / np.sqrt(s12 + cc2)
    s1 = None
    cc = None
    s12 = None
    cc2 = None
    phi = np.copysign(phi, z)

    return GPS(phi * (180 / np.pi), lamb * (180 / np.pi), h)


# Computing geodetic coordinates from geocentric coordinates
# H. Vermeille, 2003/2004
# http://users.auth.gr/kvek/78_Vermeille.pdf
def ecef2lla_vermeille2003(x, y, z):
    a = 6378137.0
    e = 8.181919084261345e-2

    p = (x * x + y * y) * (1 / (a * a))
    q = ((1 - e * e) / (a * a)) * z * z
    r = (p + q - e**4) * (1 / 6)
    s = (e**4 / 4) * p * q / (r**3)
    p = None
    t = np.cbrt(1 + s + np.sqrt(s * (2 + s)))
    s = None
    u = r * (1 + t + 1 / t)
    r = None
    t = None
    v = np.sqrt(u * u + e**4 * q)
    u += v  # precalc
    w = (e**2 / 2) * (u - q) / v
    q = None
    k = np.sqrt(u + w * w) - w
    D = k * np.sqrt(x * x + y * y) / (k + e**2)
    rtDDzz = np.sqrt(D * D + z * z)
    return GPS(
        (180 / np.pi) * 2 * np.arctan2(z, D + rtDDzz),
        (180 / np.pi) * np.arctan2(y, x),
        (k + e**2 - 1) / k * rtDDzz,
    )


def find_crossing_idx(
    XYZ: np.ndarray, marker: np.ndarray  # coordinates to look up in (X, Y, Z), meters
):  # (lat, long), degrees

    if isinstance(marker, tuple):
        marker = np.array(marker)
    if len(marker.shape) == 1:
        # force it to be a 2d shape to make the rest of the code simpler
        return find_crossing_idx(XYZ, marker.reshape((1, len(marker))))[0]

    # very similar to gps lap insert, but we can assume XYZ is a
    # reference (as opposed to GPS lap insert where we aren't sure if
    # the trajectory of the GPS is correct or not - think pit stops,
    # going off track, etc).  As a result, only one pass is needed.
    # Also we do not need to filter based on minspeed, so no timecodes
    # are used in this function.

    lat = marker[:, 0].reshape((len(marker), 1))
    lon = marker[:, 1].reshape((len(marker), 1))
    SO = np.stack(lla2ecef(lat, lon, 0), axis=2)
    SD = np.stack(lla2ecef(lat, lon, 1000), axis=2) - SO

    O = XYZ[:, :3]
    D = O[1:] - O[:-1]
    O = O[:-1] - SO

    SN = np.sum(SD * SD, axis=2, keepdims=True) * D - np.sum(SD * D, axis=2, keepdims=True) * SD
    t = np.clip(
        -np.sum(SN * O, axis=2, keepdims=True) / np.sum(SN * D, axis=2, keepdims=True),
        0,
        1,
    )

    # XXX This won't work with rally stages (anything not a circuit)
    distsq = np.sum(np.square(O + t * D), axis=2)
    minidx = np.argmin(distsq, axis=1)
    colrange = np.arange(t.shape[0])
    return np.column_stack([minidx + t[colrange, minidx, 0], np.sqrt(distsq[colrange, minidx])])


def find_crossing_dist(
    XYZD: np.ndarray,  # coordinates to look up in (X, Y, Z, Distance), meters
    marker: tuple[float, float],
):  # (lat, long) tuple, degrees
    idx, _ = find_crossing_idx(XYZD, np.array(marker))
    if idx + 1 >= len(XYZD):
        return XYZD[int(idx), 3]
    scale, idx = np.modf(idx)
    return XYZD[int(idx), 3] + scale * (XYZD[int(idx) + 1, 3] - XYZD[int(idx), 3])


def find_laps(
    XYZ: np.ndarray,  # coordinates to look up in (X, Y, Z), meters
    timecodes: np.ndarray,  # time for above coordinates, ms
    marker: tuple[float, float],
):  # (lat, long) tuple, degrees
    # gps lap insert.  We assume the start finish "line" is a
    # plane containing the vector that goes through the GPS
    # coordinates sf lat/long from altitude 0 to 1000.  The normal
    # of the plane is generally in line with the direction of
    # travel, given the above constraint.

    # O, D = vehicle vector (O=origin, D=direction, [0]=O, [1]=O+D)
    # SO, SD = start finish origin, direction (plane must contain SO and SO+SD poitns)
    # SN = start finish plane normal

    # D = a*SD + SN
    # 0 = SD . SN
    # combine to get: 0 = SD . (D - a*SD)
    #                  a * (SD . SD) = SD . D
    # plug back into first eq:
    # SN = D - (SD . D) / (SD . SD) * SD
    # or to avoid division, and because length doesn't matter:
    # SN = (SD . SD) * D - (SD. D) * SD

    # now determine intersection with plane SO,SN from vector O,O+D:
    # SN . (O + tD - SO) = 0
    # t * (D . SN) + SN . (O - SO) = 0
    # t = -SN.(O-SO) / D.SN

    SO = np.array(lla2ecef(*marker, 0.0)).reshape((1, 3))
    SD = np.array(lla2ecef(*marker, 1000)).reshape((1, 3)) - SO

    O = XYZ - SO

    D = O[1:] - O[:-1]
    O = O[:-1]

    # VBox seems to need 30, maybe my friend is using an old map description
    marker_size = 30  # meters, how far you can be from the marker to count as a lap

    # Precalculate in which time periods we were traveling at least 4 m/s (~10mph)
    minspeed = np.sum(D * D, axis=1) > np.square((timecodes[1:] - timecodes[:-1]) * (4 / 1000))

    SN = (
        np.sum(SD * SD, axis=1).reshape((len(SD), 1)) * D
        - np.sum(SD * D, axis=1).reshape((len(D), 1)) * SD
    )
    t = np.maximum(-np.sum(SN * O, axis=1) / np.sum(SN * D, axis=1), 0)
    # This only works because the track is considered at altitude 0
    dist = np.sum(np.square(O + t.reshape((len(t), 1)) * D), axis=1)
    pick = (t[1:] <= 1) & (t[:-1] > 1) & (dist[1:] < marker_size**2)

    # Now that we have a decent candidate selection of lap
    # crossings, generate a single normal vector for the
    # start/finish line to use for all lap crossings, to make the
    # lap times more accurate/consistent.  Weight the crossings by
    # velocity and add them together.  As it happens, SN is
    # already weighted by velocity...
    SN = np.sum(SN[1:][pick & minspeed[1:]], axis=0).reshape((1, 3))
    # recompute t, dist, pick
    t = np.maximum(-np.sum(SN * O, axis=1) / np.sum(SN * D, axis=1), 0)
    dist = np.sum(np.square(O + t.reshape((len(t), 1)) * D), axis=1)
    pick = (t[1:] <= 1) & (t[:-1] > 1) & (dist[1:] < marker_size**2)

    lap_markers = [0]
    for idx in np.nonzero(pick)[0] + 1:
        if timecodes[idx] <= lap_markers[-1]:
            continue
        if not minspeed[idx]:
            idx = np.argmax(minspeed[idx:]) + idx
        lap_markers.append(timecodes[idx] + t[idx] * (timecodes[idx + 1] - timecodes[idx]))
    return lap_markers[1:]


ecef2lla = ecef2lla_vermeille2003

if __name__ == "__main__":

    def perf_test():
        import time

        samples = 10000000
        lat = -90.0 + 180.0 * np.random.rand(samples, 1)
        long = -180.0 + 360.0 * np.random.rand(samples, 1)
        alt = -11e3 + (20e3) * np.random.rand(
            samples, 1
        )  # From approximately the bottom of the Mariana trench, to the top of the Everest

        print("generating x,y,z")
        x, y, z = lla2ecef(lat, long, alt)
        algos = [
            ("osen", ecef2lla_osen),
            ("fukushima2006", ecef2lla_fukushima2006),
            ("vermeille2003", ecef2lla_vermeille2003),
        ]
        stats: dict[str, list[float]] = {name: [] for name, algo in algos}
        for _ in range(5):
            for name, algo in algos:
                start = time.time()
                ilat, ilong, ialt = algo(x, y, z)
                duration = time.time() - start
                stats[name].append(duration)
                print("algorithm %s took %.3f" % (name, duration))
                print(
                    "  avg",
                    np.sqrt(np.sum((ilat - lat) ** 2)) / len(ilat),
                    np.sqrt(np.sum((ilong - long) ** 2)) / len(ilong),
                    np.sqrt(np.sum((ialt - alt) ** 2)) / len(ialt),
                )
                print(
                    "  max",
                    np.max(np.abs(ilat - lat)),
                    np.max(np.abs(ilong - long)),
                    np.max(np.abs(ialt - alt)),
                )
        for name, stat in stats.items():
            print(name, ", ".join(["%.3f" % s for s in stat]))

    perf_test()


def fix_gps_timing_gaps(log: "LogFile", expected_dt_ms: float = 40.0) -> "LogFile":
    """Detect and correct large timing gaps in GPS channels and lap boundaries.

    Some AIM data loggers produce GPS data with spurious timestamp jumps
    (e.g., 65533ms gaps that should be ~40ms). This is likely caused by a
    16-bit overflow bug in the logger firmware. This function detects such
    gaps and corrects the timecodes by removing the excess time.

    The fix is applied in-place to the LogFile's channels dict and laps table.

    Parameters
    ----------
    log : LogFile
        The loaded log file with channels dict.
    expected_dt_ms : float, default=40.0
        Expected time delta between GPS samples in milliseconds.
        Default is 40ms (25 Hz GPS).

    Returns
    -------
    LogFile
        The same LogFile object with corrected GPS timecodes and lap boundaries.
    """
    # Find the first GPS channel that exists
    gps_channel_name = None
    for name in GPS_CHANNEL_NAMES:
        if name in log.channels:
            gps_channel_name = name
            break

    if gps_channel_name is None:
        return log

    # Get the GPS timecodes
    gps_table = log.channels[gps_channel_name]
    gps_time = gps_table.column("timecodes").to_numpy()

    if len(gps_time) < 2:
        return log

    # Detect gaps
    dt = np.diff(gps_time)
    gap_threshold = expected_dt_ms * 10  # 400ms default

    # Find indices where gaps are too large
    gap_indices = np.where(dt > gap_threshold)[0]

    if len(gap_indices) == 0:
        return log

    # Build list of (gap_time, correction) pairs
    gap_corrections = []
    for gap_idx in gap_indices:
        gap_time = gps_time[gap_idx]
        gap_size = dt[gap_idx]
        correction = gap_size - expected_dt_ms
        gap_corrections.append((gap_time, correction))

    # Fix GPS channel timecodes
    gps_time_fixed = gps_time.astype(np.float64)
    for gap_time, correction in gap_corrections:
        gps_time_fixed[gps_time > gap_time] -= correction

    gps_time_fixed = gps_time_fixed.astype(np.int64)

    # Update all GPS channels with corrected timecodes
    for name in GPS_CHANNEL_NAMES:
        if name not in log.channels:
            continue

        table = log.channels[name]
        value_column = table.column(name)
        field = table.schema.field(name)
        metadata = field.metadata

        new_table = pa.table(
            {
                "timecodes": pa.array(gps_time_fixed, type=pa.int64()),
                name: value_column,
            }
        )

        if metadata:
            new_field = new_table.schema.field(name).with_metadata(metadata)
            new_schema = pa.schema([new_table.schema.field("timecodes"), new_field])
            new_table = new_table.cast(new_schema)

        log.channels[name] = new_table

    # Fix lap boundaries (start_time, end_time)
    if log.laps is not None and len(log.laps) > 0:
        start_times = log.laps.column("start_time").to_numpy().astype(np.float64)
        end_times = log.laps.column("end_time").to_numpy().astype(np.float64)

        for gap_time, correction in gap_corrections:
            start_times[start_times > gap_time] -= correction
            end_times[end_times > gap_time] -= correction

        new_laps_data = {}
        for col_name in log.laps.column_names:
            if col_name == "start_time":
                new_laps_data[col_name] = pa.array(start_times.astype(np.int64), type=pa.int64())
            elif col_name == "end_time":
                new_laps_data[col_name] = pa.array(end_times.astype(np.int64), type=pa.int64())
            else:
                new_laps_data[col_name] = log.laps.column(col_name)

        log.laps = pa.table(new_laps_data)

    return log
