"""Unit-sphere geometry and plate kinematics. Vectors are (..., 3); lat/lon in degrees."""
from __future__ import annotations

import numpy as np


def latlon_to_xyz(lat, lon):
    la, lo = np.radians(lat), np.radians(lon)
    return np.stack([np.cos(la) * np.cos(lo), np.cos(la) * np.sin(lo), np.sin(la)], axis=-1)


def xyz_to_latlon(p):
    p = p / np.linalg.norm(p, axis=-1, keepdims=True)
    return np.degrees(np.arcsin(np.clip(p[..., 2], -1, 1))), np.degrees(np.arctan2(p[..., 1], p[..., 0]))


def east_north(p):
    """Local unit east/north vectors. At the poles east is taken as +y (finite, arbitrary)."""
    p = np.asarray(p, dtype=np.float64)
    e = np.stack([-p[..., 1], p[..., 0], np.zeros_like(p[..., 0])], axis=-1)  # z × p
    ne = np.linalg.norm(e, axis=-1, keepdims=True)
    polar = ne[..., 0] < 1e-12
    e = np.where(polar[..., None], np.array([0.0, 1.0, 0.0]), e / np.maximum(ne, 1e-300))
    n = np.cross(p, e)
    return e, n


def tangent_dir(p, q):
    """Unit tangent at p pointing along the great circle toward q."""
    t = q - np.sum(p * q, axis=-1, keepdims=True) * p
    return t / np.maximum(np.linalg.norm(t, axis=-1, keepdims=True), 1e-15)


def gc_dist_km(p, q, radius_km):
    return radius_km * np.arccos(np.clip(np.sum(p * q, axis=-1), -1.0, 1.0))


def motion_to_omega(lat, lon, azimuth_deg, speed_cm_yr, radius_km):
    """Angular velocity (rad/yr) that moves the point (lat, lon) along azimuth at speed."""
    p = latlon_to_xyz(np.array([lat]), np.array([lon]))
    e, n = east_north(p)
    az = np.radians(azimuth_deg)
    v = (np.sin(az) * e[0] + np.cos(az) * n[0]) * (speed_cm_yr / 100.0)  # m/yr
    return np.cross(p[0], v) / (radius_km * 1000.0)


def velocity(p, omega, radius_km):
    """Surface velocity (m/yr) of points p on plates with angular velocity omega (broadcast)."""
    return np.cross(omega, p) * (radius_km * 1000.0)


def rotate_about(p, v, ang):
    """Rotate tangent vectors v about normals p by ang radians (counter-clockwise seen from outside)."""
    ang = np.asarray(ang, dtype=np.float64)[..., None]
    return v * np.cos(ang) + np.cross(p, v) * np.sin(ang)


def azimuth_deg(center, p):
    """Bearing (deg clockwise from north, [0, 360)) from center (3,) to points p (..., 3)."""
    e, n = east_north(center[None])
    t = tangent_dir(np.broadcast_to(center, p.shape), p)
    return np.degrees(np.arctan2(t @ e[0], t @ n[0])) % 360.0


def great_circle_point(p0, t0, dist_km, radius_km):
    """Point reached from p0 moving dist_km along unit tangent t0."""
    a = dist_km / radius_km
    return np.cos(a) * p0 + np.sin(a) * t0