"""Kinematic table-tennis ball, ported from ``KinematicPingPongCommand`` (training).

Same drag, fitted bounce regression, net/landing logic, incoming-velocity solver
and virtual-paddle contact model as the Isaac Lab environment, in numpy.
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field

import numpy as np

from .mathutil import quat_apply, quat_from_euler_xyz, quat_mul, quat_to_matrix

# ---- constants (tracking_pingpong_online_closd_ball_env_cfg_FB + KinematicPingPongCommandCfg) ----
BALL_RADIUS = 0.02
GRAVITY = 9.81
DRAG_K = 0.09213392087790545          # 1/m, v <- v / (1 + k|v|dt)
BALL_MAX_SPEED = 12.0
TABLE_HEIGHT = 0.76
TABLE_X_RANGE = (0.50, 3.24)
TABLE_HALF_WIDTH = 0.76
NET_X = 1.87
NET_HEIGHT = 0.91
RETURN_X_RANGE = (1.87, 3.24)
KILL_Z = 0.15
KILL_X_RANGE = (-0.8, 3.6)
KILL_HALF_Y = 1.5
PADDLE_CONTACT_DISTANCE = 0.01
PADDLE_COOLDOWN_S = 0.12
PADDLE_RESTITUTION = 0.8
PADDLE_TANGENT_KEEP = 1.0
PADDLE_VELOCITY_GAIN = 1.0
PADDLE_SEPARATION = 0.04
PADDLE_NORMAL_LOCAL = np.array([0.0, 1.0, 0.0])
MIN_CLOSING_SPEED = 0.05
LANDING_QUALITY_STD = 0.45

# fitted bounce regression: vx' = a0 + a1 vx + a2 vz ; vy' = b0 + b1 vy ; vz' = c0 |vz| + c1 |vz|^2
_BOUNCE_A = (-0.74010409, 0.88783120, -0.28152663)
_BOUNCE_B = (-0.01392620, 0.82615466)
_BOUNCE_C = (0.91228877, -0.00550677)


def drag_step(v: np.ndarray, dt: float) -> np.ndarray:
    speed = float(np.linalg.norm(v))
    v = v / max(1.0 + DRAG_K * speed * dt, 1e-9)
    speed = float(np.linalg.norm(v))
    if speed > BALL_MAX_SPEED:
        v = v * (BALL_MAX_SPEED / speed)
    return v


def bounce_velocity(v: np.ndarray) -> np.ndarray:
    vz_in = max(-float(v[2]), 0.0)
    return np.array(
        [
            _BOUNCE_A[0] + _BOUNCE_A[1] * v[0] + _BOUNCE_A[2] * v[2],
            _BOUNCE_B[0] + _BOUNCE_B[1] * v[1],
            _BOUNCE_C[0] * vz_in + _BOUNCE_C[1] * vz_in * vz_in,
        ],
        dtype=np.float64,
    )


def _on_table(x: float, y: float) -> bool:
    return TABLE_X_RANGE[0] <= x <= TABLE_X_RANGE[1] and abs(y) <= TABLE_HALF_WIDTH


def integrate(p: np.ndarray, v: np.ndarray, dt: float) -> tuple[np.ndarray, np.ndarray, bool, np.ndarray | None]:
    """One ball step: drag, gravity, position, table bounce.  Returns (p, v, bounced, contact_xy)."""
    last_p = p.copy()
    prev_v = v.copy()
    v = drag_step(v, dt)
    v = v + np.array([0.0, 0.0, -GRAVITY * dt])
    p = p + v * dt
    z_table = TABLE_HEIGHT + BALL_RADIUS
    bounced = False
    contact = None
    if last_p[2] > z_table >= p[2] and v[2] < 0.0:
        alpha = float(np.clip((z_table - last_p[2]) / (p[2] - last_p[2]), 0.0, 1.0))
        contact_xy = last_p[:2] + alpha * (p[:2] - last_p[:2])
        if _on_table(float(contact_xy[0]), float(contact_xy[1])):
            t1 = alpha * dt
            vc = drag_step(prev_v, t1) + np.array([0.0, 0.0, -GRAVITY * t1])
            pc = np.array([contact_xy[0], contact_xy[1], z_table])
            vb = bounce_velocity(vc)
            t2 = (1.0 - alpha) * dt
            v = drag_step(vb, t2) + np.array([0.0, 0.0, -GRAVITY * t2])
            p = pc + v * t2
            bounced = True
            contact = pc
    return p, v, bounced, contact


def simulate(p: np.ndarray, v: np.ndarray, duration: float, dt: float) -> tuple[np.ndarray, np.ndarray, int]:
    steps = max(1, int(round(duration / dt)))
    h = duration / steps
    bounces = 0
    for _ in range(steps):
        p, v, b, _ = integrate(p, v, h)
        bounces += int(b)
    return p, v, bounces


@dataclass
class Crossing:
    found: bool
    time: float = math.nan
    pos: np.ndarray = field(default_factory=lambda: np.full(3, np.nan))
    vel: np.ndarray = field(default_factory=lambda: np.full(3, np.nan))
    bounces: int = 0


def predict_plane_crossing(p: np.ndarray, v: np.ndarray, plane_x: float, dt: float, max_time: float = 3.0) -> Crossing:
    """Roll the ball (travelling toward -x) until it crosses ``x = plane_x``."""
    t = 0.0
    bounces = 0
    while t < max_time:
        last_p, last_v = p.copy(), v.copy()
        p, v, b, _ = integrate(p, v, dt)
        bounces += int(b)
        t += dt
        if last_p[0] > plane_x >= p[0]:
            alpha = float(np.clip((last_p[0] - plane_x) / max(last_p[0] - p[0], 1e-9), 0.0, 1.0))
            pos = last_p + alpha * (p - last_p)
            vel = last_v + alpha * (v - last_v)
            return Crossing(True, t - dt + alpha * dt, pos, vel, bounces)
        if p[2] < KILL_Z:
            break
    return Crossing(False)


def solve_incoming_velocity(start: np.ndarray, target: np.ndarray, t_hit: float, rng: np.random.Generator, dt: float, iters: int = 2) -> np.ndarray:
    """Training ``_solve_incoming_velocity_to_hit``: bounce-waypoint guess + shooting corrections."""
    ratio = rng.uniform(0.45, 0.70)
    t1 = max(ratio * t_hit, 0.08)
    bounce = np.array([rng.uniform(0.75, 1.45), target[1] + rng.uniform(-0.10, 0.10), TABLE_HEIGHT + BALL_RADIUS])
    g_vec = np.array([0.0, 0.0, -GRAVITY])
    v = (bounce - start - 0.5 * g_vec * t1 * t1) / t1
    for _ in range(iters):
        pred, _, _ = simulate(start.copy(), v.copy(), t_hit, dt)
        v = v + (target - pred) / max(t_hit, 1e-6)
    return v


@dataclass
class IncomingBall:
    start: np.ndarray
    vel: np.ndarray
    crossing: Crossing


def sample_incoming_ball(
    target: np.ndarray,
    t_hit: float,
    plane_x: float,
    hit_box: tuple[tuple[float, float], tuple[float, float], tuple[float, float]],
    rng: np.random.Generator,
    dt: float,
    attempts: int = 4,
    pos_tol: float = 0.04,
    time_tol: float = 0.04,
) -> IncomingBall | None:
    """Training ``_reset_online_ball`` (synthetic branch): best of ``attempts`` candidates."""
    (x_lo, x_hi), (y_lo, y_hi), (z_lo, z_hi) = hit_box
    best: tuple[int, float, IncomingBall] | None = None
    for _ in range(attempts):
        start = np.array([rng.uniform(1.67, 3.24), float(np.clip(target[1] + rng.uniform(-0.25, 0.25), -0.76, 0.76)), rng.uniform(0.90, 1.50)])
        v = solve_incoming_velocity(start, target, t_hit, rng, dt)
        cr = predict_plane_crossing(start.copy(), v.copy(), plane_x, dt)
        if not cr.found:
            continue
        in_range = x_lo - 1e-5 <= cr.pos[0] <= x_hi + 1e-5 and y_lo - 1e-5 <= cr.pos[1] <= y_hi + 1e-5 and z_lo - 1e-5 <= cr.pos[2] <= z_hi + 1e-5
        one_bounce = cr.bounces == 1
        pos_err = float(np.linalg.norm(cr.pos - target))
        time_err = abs(cr.time - t_hit)
        if in_range and one_bounce and pos_err <= pos_tol and time_err <= time_tol:
            tier = 0
        elif in_range and one_bounce:
            tier = 1
        elif one_bounce:
            tier = 2
        else:
            tier = 3
        score = pos_err + 0.25 * time_err
        cand = IncomingBall(start, v, cr)
        if best is None or (tier, score) < (best[0], best[1]):
            best = (tier, score, cand)
    return best[2] if best is not None else None


@dataclass
class PaddleState:
    pos: np.ndarray       # virtual paddle centre, world
    quat: np.ndarray      # wxyz, world
    lin_vel: np.ndarray   # world


def virtual_paddle(body_pos: np.ndarray, body_quat: np.ndarray, body_lin_vel: np.ndarray, body_ang_vel: np.ndarray,
                   offset_xyz: np.ndarray, offset_rpy: np.ndarray) -> PaddleState:
    """Training ``_get_virtual_paddle_state_for_ball``."""
    r = quat_apply(body_quat, offset_xyz)
    pos = body_pos + r
    quat = quat_mul(body_quat, quat_from_euler_xyz(*offset_rpy))
    vel = body_lin_vel + np.cross(body_ang_vel, r)
    return PaddleState(pos, quat, vel)


@dataclass
class ReturnPrediction:
    over_net: bool
    on_table: bool
    landing_xy: np.ndarray | None
    landing_error: float
    landing_quality: float


def predict_return(p: np.ndarray, v: np.ndarray, target_landing_xy: np.ndarray, dt: float, max_time: float = 3.0) -> ReturnPrediction:
    """Roll the outgoing ball: legal net crossing, first table bounce inside the return box, landing error."""
    over_net = False
    t = 0.0
    while t < max_time:
        last_p = p.copy()
        p, v, bounced, contact = integrate(p, v, dt)
        t += dt
        if last_p[0] < NET_X <= p[0]:
            alpha = (NET_X - last_p[0]) / max(p[0] - last_p[0], 1e-9)
            cross = last_p + alpha * (p - last_p)
            if abs(cross[1]) <= TABLE_HALF_WIDTH and cross[2] > NET_HEIGHT:
                over_net = True
        if bounced and contact is not None:
            on_table = over_net and RETURN_X_RANGE[0] <= contact[0] <= RETURN_X_RANGE[1] and abs(contact[1]) <= TABLE_HALF_WIDTH
            err = float(np.linalg.norm(contact[:2] - target_landing_xy))
            return ReturnPrediction(over_net, on_table, contact[:2].copy(), err, math.exp(-(err / LANDING_QUALITY_STD) ** 2))
        if p[2] < KILL_Z or p[0] > KILL_X_RANGE[1] or abs(p[1]) > KILL_HALF_Y:
            break
    return ReturnPrediction(over_net, False, None, math.nan, 0.0)


class KinematicBall:
    """Live ball state advanced once per control step, with paddle contact detection."""

    def __init__(self, half_width: float, half_height: float):
        self.half_width = float(half_width)
        self.half_height = float(half_height)
        self.pos = np.zeros(3)
        self.vel = np.zeros(3)
        self.active = False
        self.hit = False
        self.over_net = False
        self.landed = False
        self.landing_xy: np.ndarray | None = None
        self.cooldown = 0.0
        self._last_paddle: PaddleState | None = None
        self.hit_paddle_speed = 0.0
        self.outgoing_vel = np.zeros(3)

    def launch(self, pos: np.ndarray, vel: np.ndarray) -> None:
        self.pos = np.asarray(pos, dtype=np.float64).copy()
        self.vel = np.asarray(vel, dtype=np.float64).copy()
        self.active = True
        self.hit = False
        self.over_net = False
        self.landed = False
        self.landing_xy = None
        self.cooldown = 0.0
        self._last_paddle = None

    def deactivate(self) -> None:
        self.active = False

    def step(self, dt: float, paddle: PaddleState | None, forehand: bool) -> dict[str, bool]:
        events = {"hit": False, "landed": False}
        if not self.active:
            self._last_paddle = paddle
            return events
        self.cooldown = max(0.0, self.cooldown - dt)
        last_pos = self.pos.copy()
        self.pos, self.vel, bounced, contact = integrate(self.pos, self.vel, dt)
        if self.hit and last_pos[0] < NET_X <= self.pos[0]:
            alpha = (NET_X - last_pos[0]) / max(self.pos[0] - last_pos[0], 1e-9)
            cross = last_pos + alpha * (self.pos - last_pos)
            if abs(cross[1]) <= TABLE_HALF_WIDTH and cross[2] > NET_HEIGHT:
                self.over_net = True
        if bounced and contact is not None and self.hit and not self.landed:
            self.landed = True
            self.landing_xy = contact[:2].copy()
            events["landed"] = True
        if paddle is not None and not self.hit and self.cooldown <= 0.0:
            if self._try_paddle_contact(last_pos, paddle, forehand):
                events["hit"] = True
        if self.pos[2] < KILL_Z or not (KILL_X_RANGE[0] <= self.pos[0] <= KILL_X_RANGE[1]) or abs(self.pos[1]) > KILL_HALF_Y:
            self.active = False
        self._last_paddle = paddle
        return events

    def _face_ok(self, u: float, w: float) -> bool:
        a = self.half_width + BALL_RADIUS
        b = self.half_height + BALL_RADIUS
        return (u / a) ** 2 + (w / b) ** 2 <= 1.0

    def _try_paddle_contact(self, last_pos: np.ndarray, paddle: PaddleState, forehand: bool) -> bool:
        rot = quat_to_matrix(paddle.quat)
        local = rot.T @ (self.pos - paddle.pos)
        d = float(local[1])
        u, w = float(local[0]), float(local[2])
        contact = abs(d) <= BALL_RADIUS + PADDLE_CONTACT_DISTANCE and self._face_ok(u, w)
        if not contact and self._last_paddle is not None:
            rot_last = quat_to_matrix(self._last_paddle.quat)
            local_last = rot_last.T @ (last_pos - self._last_paddle.pos)
            d_last = float(local_last[1])
            if (d_last < 0.0) != (d < 0.0) and abs(d - d_last) > 1e-9:
                alpha = float(np.clip(-d_last / (d - d_last), 0.0, 1.0))
                ui = local_last[0] + alpha * (u - local_last[0])
                wi = local_last[2] + alpha * (w - local_last[2])
                contact = self._face_ok(float(ui), float(wi))
        if not contact:
            return False
        sign = 1.0 if forehand else -1.0
        normal = rot @ (PADDLE_NORMAL_LOCAL * sign)
        rel = self.vel - paddle.lin_vel
        if float(np.dot(rel, normal)) > 0.0:
            normal = -normal
        vn = float(np.dot(rel, normal))
        if vn > -MIN_CLOSING_SPEED:
            return False
        vt = rel - vn * normal
        self.vel = PADDLE_VELOCITY_GAIN * paddle.lin_vel + (-PADDLE_RESTITUTION * vn * normal + PADDLE_TANGENT_KEEP * vt)
        tangent_w = rot @ np.array([u, 0.0, w])
        self.pos = paddle.pos + tangent_w + PADDLE_SEPARATION * normal
        self.hit = True
        self.cooldown = PADDLE_COOLDOWN_S
        self.hit_paddle_speed = float(np.linalg.norm(paddle.lin_vel))
        self.outgoing_vel = self.vel.copy()
        return True