from __future__ import annotations

import argparse

from translate_align_pointcloud_to_cad import translate_align_by_front_tail


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Standalone train point-cloud translation alignment tool"
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    translate_align = subparsers.add_parser(
        "translate-align",
        help="Align a whole point cloud to a standard STEP CAD by front/tail centroids",
    )
    translate_align.add_argument("--cad", "-c", default="datasets/CAD/M02.stp", help="Reference STEP CAD file")
    translate_align.add_argument(
        "--pointcloud",
        "-p",
        default="datasets/shift_test_output/",
        help="Input PLY point cloud or a directory containing shifted_whole_pointcloud.ply",
    )
    translate_align.add_argument("--output-dir", "-o", default="results", help="Output directory")
    translate_align.add_argument("--axis", choices=["x", "y", "z"], default="x", help="Train length axis")
    translate_align.add_argument("--crop-ratio", type=float, default=0.01, help="Front/tail crop ratio")
    translate_align.add_argument("--mesh-deflection", type=float, default=0.1, help="CAD meshing deflection")
    translate_align.add_argument("--cad-sample-points", type=int, default=30000, help="Points sampled from CAD mesh")
    translate_align.add_argument("--local-sample-points", type=int, default=1000, help="Points sampled from each end")
    translate_align.add_argument("--scale-pointcloud", type=float, default=1.0, help="Scale factor applied to point cloud")
    translate_align.add_argument("--no-visualize", action="store_true", help="Disable Open3D windows")
    return parser


def main() -> None:
    parser = build_parser()
    args = parser.parse_args()

    if args.command == "translate-align":
        result = translate_align_by_front_tail(
            cad_path=args.cad,
            pointcloud_path=args.pointcloud,
            output_dir=args.output_dir,
            axis=args.axis,
            crop_ratio=args.crop_ratio,
            mesh_deflection=args.mesh_deflection,
            cad_sample_points=args.cad_sample_points,
            local_sample_points=args.local_sample_points,
            scale_pointcloud=args.scale_pointcloud,
            visualize=not args.no_visualize,
        )

        print("\n========== Result ==========")
        print("translation:", result["translation"])
        print("front-tail translation difference:", result["front_tail_translation_difference"])
        print("aligned path:", result["aligned_path"])
        print("matrix path:", result["matrix_path"])
    else:
        parser.error(f"Unsupported command: {args.command}")


if __name__ == "__main__":
    main()