from datetime import datetime

import geopandas as gpd
import movingpandas as mpd
import numpy as np
import pandas as pd
from shapely.geometry import Point

far_west_lon = -179.5
far_west_lat = 0
far_west_pt = Point(far_west_lon, far_west_lat)  # 'POINT (-179.5 0)'
far_west_dist_to_am = 180 - abs(far_west_lon)  # 0.5

far_east_lon = 179.5
far_east_lat = 0
far_east_pt = Point(far_east_lon, far_east_lat)  # 'POINT (179.5 0)'
far_east_dist_to_am = 180 - far_east_lon  # 0.5

times = pd.DatetimeIndex(["2023-01-01 01:10:00", "2023-01-01 01:20:00"])

gdf = gpd.GeoDataFrame(
    {
        "t": times,
        "geometry": [far_west_pt, far_east_pt],
    },
    crs="epsg:4326",
).set_index("t")

# Show spurious interpolation
interp_time = datetime(2023, 1, 1, 1, 15)
traj = mpd.Trajectory(gdf, 1)
traj.get_position_at(interp_time).wkt  # 'POINT (0 0)'

from shapely.geometry import LineString, Point

line = LineString([Point(-179.5, 0), Point(179.5, 0)])
line.interpolate(0.5 * line.length) # 'POINT (0 0)'

# Expected interpolation by shifting the points around 90oE
new_lon_loc = 90
shifted_far_west_lon = new_lon_loc + far_west_dist_to_am  # shifted_far_west_lon
new_far_west_pt = Point(shifted_far_west_lon, far_west_lat)  # 'POINT (90.5 0)'
shifted_far_east_lon = new_lon_loc - far_east_dist_to_am  # 89.5
new_far_east_pt = Point(shifted_far_east_lon, far_east_lat)  # 'POINT (89.5 0)'

gdf = gpd.GeoDataFrame(
    {
        "t": times,
        "geometry": [new_far_west_pt, new_far_east_pt],
    },
    crs="epsg:4326",
).set_index("t")
traj = mpd.Trajectory(gdf, 1)
shifted_expected_pt = traj.get_position_at(interp_time)  # 'POINT (90 0)'

expected_lon_rel_am = shifted_expected_pt.x - new_lon_loc  # 0
if expected_lon_rel_am == 0:
    expected_lon = 179.99999  # my data is (-180., 180.)
elif expected_lon_rel_am > 0:
    expected_lon = -180.0 + expected_lon_rel_am
else:
    expected_lon = 180 + expected_lon_rel_am
expected_pt = Point(expected_lon, shifted_expected_pt.y)  # 'POINT (179.99999 0)'

def _shift_line_to_meridian(line: LineString) -> LineString:
    orig_lat = [lat for _, lat in line.coords]
    orig_lon = [lon for lon, _ in line.coords]
    # 180 -> 0 and -180 -> 0
    new_lon = [0.0 if x == 180 or x == -180 else x for x in orig_lon]
    # e.g. 179 -> -1 and -179 -> 1
    new_lon = [math.copysign(180.0 - abs(x), -x) if x != 0 else x for x in new_lon]
    return LineString(list(zip(new_lon, orig_lat)))


def _shift_point_to_antimeridian(point: Point) -> Point:
    orig_lat = point.y
    orig_lon = point.x
    if orig_lon == 0:
        new_lon = -180.0
    elif orig_lon < 0:
        # e.g. -1 -> 179
        new_lon = orig_lon + 180.0
    else:
        # e.g. 1 -> -179
        new_lon = orig_lon - 180.0
    return Point(new_lon, orig_lat)


def _get_position_at(traj: mpd.Trajectory, t: datetime) -> Point:
    """Fix to mpd.Trajectory.get_position_at to work if
    line crosses antimeridian"""
    prev_row = traj.get_row_at(t, "ffill")
    next_row = traj.get_row_at(t, "bfill")
    t_diff = next_row.name - prev_row.name
    t_diff_at = t - prev_row.name
    line = LineString(
        [
            prev_row[traj.get_geom_column_name()],
            next_row[traj.get_geom_column_name()],
        ]
    )
    if t_diff == 0 or line.length == 0:
        return prev_row[traj.get_geom_column_name()]
    # Check if line crosses antimeridian
    if abs(line.bounds[2] - line.bounds[0]) > 180:
        # Shift points by 180. in the line
        shifted_line = _shift_line_to_meridian(line)
        shifted_interp_pt = shifted_line.interpolate(
            (t_diff_at / t_diff) * shifted_line.length
        )
        # Provide corrected location to antimeridian
        return _shift_point_to_antimeridian(shifted_interp_pt)
    else:
        # Default behaviour
        return line.interpolate(t_diff_at / t_diff * line.length)