"""Recover distance from the box bottom row and score it against the truth.

This is the quantitative form of "do box bottoms sit at the ground contact
point". Looking at a contact sheet catches a box that is obviously wrong; it
does not catch a box that is twenty pixels low, which is a systematic depth bias
worth several metres at range.

It is also, deliberately, phase 2's ground-plane estimator run early: a pixel is
back-projected to a ray and intersected with the ground plane from
calib/<sensor_id>.yaml, using nothing but focal length, principal point, camera
height and the ground normal. Running it against gt_distance_m today says
whether the calibration handoff actually works before anything is built on it.

Stop signs are excluded by default: they are pole-mounted, so the bottom of the
box is not a ground contact and the estimator does not apply to them.

Usage:
    python -m src.data.check_ground_plane
    python -m src.data.check_ground_plane --min-visibility 0.7 --min-horizon-margin 20
"""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import pandas as pd

from src.common import calib, paths, schema

BUCKETS = [(0, 10), (10, 20), (20, 30), (30, 45), (45, 1000)]


def horizon_row(calibration: calib.Calibration) -> float:
    """The image row whose back-projected ray runs parallel to the ground.

    Distance estimates diverge as a box bottom approaches this row, so it is the
    natural place to put a validity margin.
    """
    normal = calibration.ground_normal_cam()
    # Walk down the centre column until the ray turns parallel to the ground.
    rows = np.linspace(0.0, float(calibration.height), 4001)
    columns = np.full_like(rows, calibration.cx)
    alignment = normal @ calibration.ray(columns, rows)
    crossings = np.flatnonzero(np.sign(alignment[:-1]) != np.sign(alignment[1:]))
    return float(rows[crossings[0]]) if len(crossings) else float(calibration.cy)


def estimate(calibration: calib.Calibration, u: np.ndarray, v: np.ndarray) -> np.ndarray:
    """Ray-plane intersection. Returns forward depth, matching gt_distance_m."""
    normal = calibration.ground_normal_cam()
    rays = calibration.ray(u, v)
    scale = -calibration.cam_height_m / (normal @ rays)
    return scale * rays[2]


def report(frame: pd.DataFrame, calibration: calib.Calibration,
           min_visibility: float, min_margin: float) -> None:
    objects = frame[frame["gt_distance_m"].notna() & (frame["truncation"] < 0.01)]
    objects = objects[objects["class"].isin(schema.GROUND_CONTACT_CLASSES)]
    if objects.empty:
        print("  no ground-contact objects with distance")
        return

    u = ((objects["x1"] + objects["x2"]) / 2).to_numpy()
    v = objects["y2"].to_numpy()
    truth = objects["gt_distance_m"].to_numpy()
    predicted = estimate(calibration, u, v)

    margin = v - horizon_row(calibration)
    visibility = objects["visibility"].to_numpy()
    usable = (
        np.isfinite(predicted) & (predicted > 0) & (margin >= min_margin)
        & (np.isnan(visibility) | (visibility >= min_visibility))
    )

    print(f"  horizon row {horizon_row(calibration):.1f} px, "
          f"{usable.sum()}/{len(objects)} objects pass the validity gate")
    print(f"  {'range':>10}  {'n':>5}  {'m/px':>6}  {'MAE':>8}  {'bias':>8}")
    for low, high in BUCKETS:
        in_bucket = usable & (truth >= low) & (truth < high)
        if in_bucket.sum() < 5:
            continue
        error = predicted[in_bucket] - truth[in_bucket]
        sensitivity = np.median(truth[in_bucket] ** 2
                                / (calibration.fy * calibration.cam_height_m))
        label = f"{low}-{high} m" if high < 1000 else f"{low}+ m"
        print(f"  {label:>10}  {in_bucket.sum():5d}  {sensitivity:6.2f}  "
              f"{np.abs(error).mean():7.2f}m  {np.median(error):+7.2f}m")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--unified-root", type=Path, default=None)
    parser.add_argument("--min-visibility", type=float, default=0.0)
    parser.add_argument("--min-horizon-margin", type=float, default=0.0,
                        help="pixels a box bottom must sit below the horizon row")
    args = parser.parse_args()

    root = paths.unified_root(args.unified_root)
    manifest = schema.objects_only(schema.read_manifest(root / "manifest.parquet"))
    calibrations = calib.load_all(root)

    for sensor_id, calibration in sorted(calibrations.items()):
        rows = manifest[manifest["sensor_id"] == sensor_id]
        print(f"\n{sensor_id}  (cam_height {calibration.cam_height_m:.3f} m, "
              f"fy {calibration.fy:.0f})")
        if rows.empty:
            print("  no objects")
            continue
        report(rows, calibration, args.min_visibility, args.min_horizon_margin)


if __name__ == "__main__":
    main()