#!/usr/bin/env python3
"""Re-convert a few AV2 train logs with optional lens UNDISTORTION (cv2, Brown-Conrady
radial [k1,k2,0,0,k3] from intrinsics.feather). Produces distorted + undistorted variants
through an identical pipeline so undistortion is the only variable. Portrait/crop is
handled separately at eval via the dataset `letterbox` toggle."""
import os, argparse
from pathlib import Path
import numpy as np, torch, cv2, pandas as pd

from mapgs.config import load_config
from mapgs.geometry.cameras import resize_with_intrinsics
from mapgs.data.unified import write_unified_clip
from mapgs.hdmap.hdmap import HDMap
from mapgs.hdmap.anchors import sample_anchors
from mapgs.data.convert.av2_to_unified import (
    FRONT_CAMERAS, _parse_lanes, _ground_field_from_map, _load_annotations, _track_boxes)
from mapgs.config import MapConfig

LOGS = ["00a6ffc1-6ce9-3bc3-a060-6006e9893a1a", "01bb304d-7bd8-35f8-bbef-7086b688e35e"]


def dist_coeffs(log_dir):
    df = pd.read_feather(log_dir / "calibration" / "intrinsics.feather")
    out = {}
    for _, r in df.iterrows():
        out[r["sensor_name"]] = np.array([r["k1"], r["k2"], 0.0, 0.0, r["k3"]], dtype=np.float64)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--undistort", action="store_true")
    ap.add_argument("--out", required=True)
    ap.add_argument("--store-longside", type=int, default=1024)
    ap.add_argument("--max-clips-per-log", type=int, default=1)
    args = ap.parse_args()

    from av2.datasets.sensor.av2_sensor_dataloader import AV2SensorDataLoader
    from av2.map.map_api import ArgoverseStaticMap
    import imageio.v2 as imageio

    cfg = load_config()
    split_dir = Path("/mnt/william/_av2_raw/train")
    loader = AV2SensorDataLoader(data_dir=split_dir, labels_dir=split_dir)
    out = os.path.join(args.out, "train"); os.makedirs(out, exist_ok=True)
    mcfg = MapConfig(**{**cfg.map.__dict__, "n_anchors": cfg.model.tokens.n_map})
    F = cfg.data.num_frames
    clip_count = 0

    for lid in LOGS:
        log_dir = split_dir / lid
        amap = ArgoverseStaticMap.from_map_dir(log_dir / "map", build_raster=True)
        lanes, boundaries = _parse_lanes(amap)
        ground = _ground_field_from_map(amap, lanes, boundaries, cfg.map.ground_spacing)
        ann = _load_annotations(log_dir)
        coeffs = dist_coeffs(log_dir)
        cams = {c: loader.get_log_pinhole_camera(lid, c) for c in FRONT_CAMERAS}
        H0, W0 = cams["ring_front_center"].height_px, cams["ring_front_center"].width_px
        if args.store_longside:
            sc = args.store_longside / max(H0, W0); H0, W0 = int(round(H0 * sc)), int(round(W0 * sc))
        ts_all = loader.get_ordered_log_lidar_timestamps(lid)
        n_clips_log = 0
        for start in range(0, len(ts_all) - F + 1, F):
            if n_clips_log >= args.max_clips_per_log:
                break
            ts_block = ts_all[start:start + F]
            images, c2ws, Ks = [], [], []
            ok = True
            for ts in ts_block:
                city_SE3_ego = loader.get_city_SE3_ego(lid, int(ts))
                img_v, c2w_v, K_v = [], [], []
                for c in FRONT_CAMERAS:
                    fp = loader.get_closest_img_fpath(lid, c, int(ts))
                    if fp is None:
                        ok = False; break
                    arr = np.asarray(imageio.imread(fp))                  # native HxWx3
                    Knat = np.asarray(cams[c].intrinsics.K, dtype=np.float64)
                    if args.undistort:
                        arr = cv2.undistort(arr, Knat, coeffs[c])         # same K, removes radial distortion
                    t = torch.from_numpy(arr).float().permute(2, 0, 1) / 255.0
                    img, Ksc = resize_with_intrinsics(t, torch.from_numpy(Knat).float(), H0, W0)
                    img_v.append(img); K_v.append(Ksc)
                    c2w = city_SE3_ego.compose(cams[c].ego_SE3_cam).transform_matrix
                    c2w_v.append(torch.from_numpy(np.asarray(c2w)).float())
                if not ok:
                    break
                images.append(torch.stack(img_v)); c2ws.append(torch.stack(c2w_v)); Ks.append(torch.stack(K_v))
            if not ok or len(images) < F:
                continue
            images = torch.stack(images); cam2world = torch.stack(c2ws); K = Ks[0]
            anchors = sample_anchors(HDMap(ground, lanes, boundaries), mcfg, seed=clip_count)
            boxes = _track_boxes(ann, ts_block, loader, lid)
            write_unified_clip(out, f"train_{lid[:8]}_{n_clips_log:03d}", "av2", cfg.data.fps,
                               images, K, cam2world, ground, lanes, boundaries, anchors,
                               boxes=boxes, depth=None, scene_scale=1.0, cameras=FRONT_CAMERAS)
            clip_count += 1; n_clips_log += 1
    print(f"wrote {clip_count} clips to {out} (undistort={args.undistort})")


if __name__ == "__main__":
    main()