import numpy as np

from helpers import box_center, box_wh, clip_box


def _fit_motion(times, centers, max_speed, max_accel):
    design = np.column_stack(
        (np.ones_like(times), times, 0.5 * times * times)
    ).astype(np.float64)
    coefficients, *_ = np.linalg.lstsq(design, centers, rcond=None)
    residuals = np.linalg.norm(centers - design @ coefficients, axis=1)
    median = float(np.median(residuals))
    mad = float(np.median(np.abs(residuals - median)))
    limit = median + max(1.0, 3.0 * 1.4826 * mad)
    keep = residuals <= limit
    if np.count_nonzero(keep) >= 4:
        coefficients, *_ = np.linalg.lstsq(
            design[keep], centers[keep], rcond=None
        )
        residuals = np.linalg.norm(
            centers[keep] - design[keep] @ coefficients, axis=1
        )

    velocity = coefficients[1].astype(np.float64)
    acceleration = coefficients[2].astype(np.float64)
    speed = float(np.linalg.norm(velocity))
    accel = float(np.linalg.norm(acceleration))
    if speed > float(max_speed):
        velocity *= float(max_speed) / max(speed, 1e-6)
    if accel > float(max_accel):
        acceleration *= float(max_accel) / max(accel, 1e-6)
    rms = float(np.sqrt(np.mean(residuals * residuals))) if residuals.size else 0.0
    return coefficients[0], velocity, acceleration, rms, keep


def predict_ballistic(
    observations,
    now_ts,
    frame_w,
    frame_h,
    *,
    lookback=10,
    min_observations=5,
    min_span_sec=0.12,
    max_horizon_sec=0.55,
    max_speed=900.0,
    max_accel=1200.0,
    max_size_rate=1.2,
    max_uncertainty=120.0,
):
    recent = []
    for observation in list(observations or [])[-max(2, int(lookback)):]:
        try:
            ts = float(observation["ts"])
            center = np.asarray(observation["center"], dtype=np.float64)
            box = np.asarray(observation["box"], dtype=np.float64)
        except (KeyError, TypeError, ValueError):
            continue
        if center.shape != (2,) or box.shape != (4,) or not np.all(np.isfinite(center)):
            continue
        if recent and ts <= recent[-1][0]:
            continue
        width, height = box_wh(box)
        if width <= 1.0 or height <= 1.0:
            continue
        recent.append((ts, center, np.array([width, height], dtype=np.float64)))

    if len(recent) < int(min_observations):
        return None
    last_ts = recent[-1][0]
    first_ts = recent[0][0]
    if last_ts - first_ts < float(min_span_sec):
        return None

    times = np.asarray([row[0] - last_ts for row in recent], dtype=np.float64)
    centers = np.asarray([row[1] for row in recent], dtype=np.float64)
    sizes = np.asarray([row[2] for row in recent], dtype=np.float64)
    origin, velocity, acceleration, rms, keep = _fit_motion(
        times, centers, max_speed, max_accel
    )

    horizon = float(np.clip(
        max(0.0, float(now_ts) - last_ts),
        0.0,
        float(max_horizon_sec),
    ))
    predicted_center = (
        origin
        + velocity * horizon
        + 0.5 * acceleration * horizon * horizon
    )

    size_design = np.column_stack((np.ones_like(times), times))
    size_keep = keep if np.count_nonzero(keep) >= 3 else np.ones(len(times), dtype=bool)
    log_sizes = np.log(np.maximum(sizes, 2.0))
    size_coefficients, *_ = np.linalg.lstsq(
        size_design[size_keep],
        log_sizes[size_keep],
        rcond=None,
    )
    size_rate = np.clip(
        size_coefficients[1],
        -float(max_size_rate),
        float(max_size_rate),
    )
    predicted_size = np.exp(size_coefficients[0] + size_rate * horizon)
    last_size = sizes[-1]
    predicted_size = np.clip(predicted_size, 0.65 * last_size, 1.80 * last_size)

    speed = float(np.linalg.norm(velocity))
    accel = float(np.linalg.norm(acceleration))
    uncertainty = float(np.clip(
        6.0 + rms + 0.08 * speed * horizon + 0.12 * accel * horizon * horizon,
        6.0,
        float(max_uncertainty),
    ))
    confidence = float(np.clip(
        np.exp(-rms / max(4.0, float(np.linalg.norm(last_size))))
        * (1.0 - 0.55 * horizon / max(float(max_horizon_sec), 1e-3)),
        0.0,
        1.0,
    ))

    cx, cy = predicted_center
    width, height = predicted_size
    box = clip_box(
        [cx - 0.5 * width, cy - 0.5 * height,
         cx + 0.5 * width, cy + 0.5 * height],
        frame_w,
        frame_h,
    )
    return {
        "box": box,
        "center": box_center(box),
        "velocity": velocity.astype(np.float32),
        "acceleration": acceleration.astype(np.float32),
        "horizon": horizon,
        "uncertainty": uncertainty,
        "confidence": confidence,
        "fit_rms": rms,
    }