from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence

import numpy as np

from SLAM_engine import load_slam_config
from SLAM_Filter import VoxelPersistenceFilter, FilterConfig as PersistFilterConfig
from SLAM_Navigation import NavigationExportConfig, NavigationExporter


_LIDAR_FRAME_DT = 0.10  # MID-360 outputs at 10 Hz → 0.10 s per frame


def _voxel_downsample_numpy(points: np.ndarray, voxel: float) -> np.ndarray:
    if points is None or len(points) == 0:
        return np.zeros((0, 3), dtype=np.float32)
    if voxel <= 0:
        return np.asarray(points, dtype=np.float32)
    pts = np.asarray(points, dtype=np.float32)
    keys = np.floor(pts / float(voxel)).astype(np.int32)
    uniq, inv = np.unique(keys, axis=0, return_inverse=True)
    out = np.zeros((len(uniq), 3), dtype=np.float64)
    cnt = np.zeros((len(uniq),), dtype=np.int64)
    np.add.at(out, inv, pts.astype(np.float64, copy=False))
    np.add.at(cnt, inv, 1)
    out /= np.maximum(cnt[:, None], 1)
    return out.astype(np.float32)


def _transform_points(points: np.ndarray, pose: np.ndarray) -> np.ndarray:
    p = np.asarray(points, dtype=np.float32)
    tf = np.asarray(pose, dtype=np.float32)
    if tf.shape != (4, 4):
        return p
    r = tf[:3, :3]
    t = tf[:3, 3]
    return (p @ r.T) + t


def load_replay_npz(path: str | Path) -> tuple[List[np.ndarray], Optional[List[np.ndarray]]]:
    p = Path(path)
    data = np.load(str(p), allow_pickle=True)
    if "frames" not in data:
        raise RuntimeError("Replay file must contain 'frames'.")
    frames_raw = data["frames"]
    frames: List[np.ndarray] = []
    if isinstance(frames_raw, np.ndarray) and frames_raw.dtype == object:
        for fr in frames_raw:
            frames.append(np.asarray(fr, dtype=np.float32))
    else:
        arr = np.asarray(frames_raw, dtype=np.float32)
        if arr.ndim == 3 and arr.shape[2] >= 3:
            for i in range(arr.shape[0]):
                frames.append(arr[i, :, :3].astype(np.float32))
        else:
            raise RuntimeError("Unsupported 'frames' format in replay npz.")

    poses = None
    if "poses" in data:
        poses_raw = np.asarray(data["poses"])
        if poses_raw.ndim == 3 and poses_raw.shape[1:] == (4, 4):
            poses = [poses_raw[i].astype(np.float32) for i in range(poses_raw.shape[0])]
    return frames, poses


def run_replay(
    frames: Sequence[np.ndarray],
    poses: Optional[Sequence[np.ndarray]] = None,
    cfg: Optional[Dict[str, Any]] = None,
    export_nav: bool = False,
    nav_base_name: str = "replay_nav",
) -> Dict[str, Any]:
    config = cfg or load_slam_config()
    filt_cfg = PersistFilterConfig(
        voxel_size=float(config["filter"]["voxel_size"]),
        hit_threshold=int(config["filter"]["hits_threshold"]),
        decay_seconds=float(config["filter"]["persistence"]["decay_seconds"]),
        max_voxels=int(config["filter"]["persistence"]["max_voxels"]),
    )
    filt = VoxelPersistenceFilter(filt_cfg)

    map_cfg = config.get("map", {})
    loc_cfg = config.get("localization", {})
    nav_cfg = NavigationExportConfig.from_dict(config.get("navigation_export", {}))

    maps_dir = config.get("app", {}).get("maps_dir", "DataMap")
    out_dir = Path(maps_dir).expanduser()
    if not out_dir.is_absolute():
        out_dir = (Path(__file__).resolve().parent / out_dir).resolve()
    out_dir.mkdir(parents=True, exist_ok=True)
    nav_exporter = NavigationExporter(nav_cfg, str(out_dir))

    for i, frame in enumerate(frames):
        pts = np.asarray(frame, dtype=np.float32)
        if len(pts) == 0:
            continue
        if poses is not None and i < len(poses):
            pts = _transform_points(pts, np.asarray(poses[i], dtype=np.float32))
        filt.update(pts, now=float(i) * _LIDAR_FRAME_DT)

    stable_world = filt.get_stable_points()
    stable_save = _voxel_downsample_numpy(stable_world, float(map_cfg.get("save_voxel", 0.06)))
    stable_n = int(len(stable_save))

    min_map = int(map_cfg.get("min_points_to_save", 550))
    min_loc = int(loc_cfg.get("min_points_for_localize", 550))
    min_nav = int(nav_cfg.min_points)
    nav_band_n = int(nav_exporter.count_nav_points(stable_save))

    report: Dict[str, Any] = {
        "frames": int(len(frames)),
        "stable_points": stable_n,
        "map_export_pass": bool(stable_n >= min_map),
        "localization_ready": bool(stable_n >= min_loc),
        "nav_band_points": nav_band_n,
        "nav_export_pass": bool(nav_band_n >= min_nav),
        "thresholds": {
            "map": int(min_map),
            "localization": int(min_loc),
            "nav": int(min_nav),
        },
    }

    if poses is not None and len(poses) >= 2:
        path_len = 0.0
        p0 = np.asarray(poses[0], dtype=np.float64)
        p_last = np.asarray(poses[len(poses) - 1], dtype=np.float64)
        prev = p0
        for i in range(1, len(poses)):
            cur = np.asarray(poses[i], dtype=np.float64)
            path_len += float(np.linalg.norm(cur[:3, 3] - prev[:3, 3]))
            prev = cur
        disp = float(np.linalg.norm(p_last[:3, 3] - p0[:3, 3]))
        eff = 0.0 if path_len <= 1e-9 else float(disp / path_len)
        report["motion_kpi"] = {
            "path_length_m": float(path_len),
            "net_displacement_m": float(disp),
            "path_efficiency": float(eff),
        }

    if export_nav and report["nav_export_pass"]:
        nav_out = nav_exporter.export(nav_base_name, stable_save)
        report["nav_export"] = nav_out

    return report


def compare_reports(current: Dict[str, Any], baseline: Dict[str, Any], tol_ratio: float = 0.15) -> Dict[str, Any]:
    checks: List[Dict[str, Any]] = []

    def add_numeric(name: str):
        cur = float(current.get(name, 0.0))
        base = float(baseline.get(name, 0.0))
        if abs(base) < 1e-9:
            ok = abs(cur - base) < 1e-9
            ratio = 0.0
        else:
            ratio = abs(cur - base) / abs(base)
            ok = ratio <= float(tol_ratio)
        checks.append({"metric": name, "ok": bool(ok), "current": cur, "baseline": base, "ratio": ratio})

    for key in ("stable_points", "nav_band_points"):
        add_numeric(key)

    cur_motion = current.get("motion_kpi", {}) if isinstance(current, dict) else {}
    base_motion = baseline.get("motion_kpi", {}) if isinstance(baseline, dict) else {}
    if isinstance(cur_motion, dict) and isinstance(base_motion, dict):
        for key in ("path_length_m", "net_displacement_m", "path_efficiency"):
            if key in cur_motion and key in base_motion:
                cur = float(cur_motion.get(key, 0.0))
                base = float(base_motion.get(key, 0.0))
                if abs(base) < 1e-9:
                    ok = abs(cur - base) < 1e-9
                    ratio = 0.0
                else:
                    ratio = abs(cur - base) / abs(base)
                    ok = ratio <= float(tol_ratio)
                checks.append(
                    {
                        "metric": f"motion_kpi.{key}",
                        "ok": bool(ok),
                        "current": cur,
                        "baseline": base,
                        "ratio": ratio,
                    }
                )

    for key in ("map_export_pass", "localization_ready", "nav_export_pass"):
        cur = bool(current.get(key, False))
        base = bool(baseline.get(key, False))
        checks.append({"metric": key, "ok": bool(cur == base), "current": cur, "baseline": base})

    passed = all(bool(c.get("ok", False)) for c in checks)
    return {"passed": bool(passed), "checks": checks}


def main():
    ap = argparse.ArgumentParser(description="Offline SLAM replay + regression report")
    ap.add_argument("--input", required=True, help="Replay .npz containing frames (and optional poses)")
    ap.add_argument("--out", default="", help="Optional output report JSON path")
    ap.add_argument("--baseline", default="", help="Optional baseline report JSON to compare against")
    ap.add_argument("--tol-ratio", type=float, default=0.15, help="Relative tolerance for numeric metrics")
    ap.add_argument("--export-nav", action="store_true", help="Also export nav map if replay passes")
    args = ap.parse_args()

    frames, poses = load_replay_npz(args.input)
    report = run_replay(frames, poses=poses, export_nav=bool(args.export_nav))

    out: Dict[str, Any] = {"report": report}
    if args.baseline:
        base = json.loads(Path(args.baseline).read_text(encoding="utf-8"))
        base_rep = base.get("report", base) if isinstance(base, dict) else {}
        out["comparison"] = compare_reports(report, base_rep, tol_ratio=float(args.tol_ratio))

    text = json.dumps(out, indent=2)
    print(text)
    if args.out:
        Path(args.out).write_text(text, encoding="utf-8")


if __name__ == "__main__":
    main()
