"""BA-T demo: two-view bundle adjustment from an image pair.

Upload two images (A, B). BA-T — an iterative transformer layer inspired by
bundle adjustment — refines the pair and predicts:

  * a dense, coloured 3D point cloud (both views, in camera A's frame),
  * the relative camera pose between the two images.

The model runs entirely on the GPU via ZeroGPU. All outputs share one
arbitrary (non-metric) scale within a pair.
"""
import os
import tempfile

import spaces  # MUST come before any torch / CUDA-touching import
import numpy as np
import torch
import gradio as gr

from ba_t.data import load_views, denormalize
from ba_t.inference import load_model, run_pairs

MODEL_ID = "zhangganlin/BA-T"
CKPT_FILENAME = "bat.pth"
DEFAULT_K = 8
DEFAULT_CONF_THRESH = 3.0

# ---------------------------------------------------------------- model load
ckpt_path = os.path.join(tempfile.gettempdir(), CKPT_FILENAME)
if not os.path.exists(ckpt_path):
    from huggingface_hub import hf_hub_download

    ckpt_path = hf_hub_download(MODEL_ID, CKPT_FILENAME)

model = load_model(ckpt_path, device="cuda", verbose=False)
print(f"BA-T loaded from {MODEL_ID}/{CKPT_FILENAME}", flush=True)


# ---------------------------------------------------------------- helpers
def rotation_angle_deg(R):
    return float(np.degrees(np.arccos(np.clip((np.trace(R) - 1) / 2, -1.0, 1.0))))


def transform_points(T, pts):
    return pts @ T[:3, :3].T + T[:3, 3]


def save_ply(path, pts, colors):
    """Binary little-endian PLY with xyz float32 + rgb uint8."""
    pts = np.asarray(pts, dtype=np.float32).reshape(-1, 3)
    colors = np.asarray(colors, dtype=np.uint8).reshape(-1, 3)
    vertex = np.empty(
        len(pts),
        dtype=[
            ("x", "<f4"), ("y", "<f4"), ("z", "<f4"),
            ("red", "u1"), ("green", "u1"), ("blue", "u1"),
        ],
    )
    vertex["x"], vertex["y"], vertex["z"] = pts[:, 0], pts[:, 1], pts[:, 2]
    vertex["red"], vertex["green"], vertex["blue"] = colors[:, 0], colors[:, 1], colors[:, 2]
    header = (
        "ply\nformat binary_little_endian 1.0\n"
        f"element vertex {len(pts)}\n"
        "property float x\nproperty float y\nproperty float z\n"
        "property uchar red\nproperty uchar green\nproperty uchar blue\nend_header\n"
    )
    with open(path, "wb") as f:
        f.write(header.encode("ascii"))
        f.write(vertex.tobytes())


def pose_matrix_str(T):
    return "\n".join("  ".join(f"{v:+9.4f}" for v in row) for row in T)


# ---------------------------------------------------------------- inference
@spaces.GPU(duration=30)
def bundle_adjust(
    image_a,
    image_b,
    num_refine_iter: int = DEFAULT_K,
    conf_thresh: float = DEFAULT_CONF_THRESH,
):
    """Refine an image pair with BA-T and reconstruct the 3D geometry.

    Args:
        image_a: first image (camera A).
        image_b: second image (camera B, taken from a different viewpoint).
        num_refine_iter: number of BA-layer refinement iterations K (4-24).
        conf_thresh: keep 3D points whose confidence is above this value.

    Returns:
        (ply file of the coloured point cloud in camera A's frame,
         npz file with point maps / confidences / relative poses,
         pose text report)
    """
    if image_a is None or image_b is None:
        raise gr.Error("Please provide both images A and B.")

    view_a, view_b = load_views([image_a, image_b])
    out = run_pairs(model, view_a, [view_b], num_refine_iter=int(num_refine_iter))
    a, b = out["a"][0], out["b"][0]

    T_ba = a["relative_pose"][0].numpy()  # camera A -> camera B
    T_ab = b["relative_pose"][0].numpy()  # camera B -> camera A
    pts_a = a["pts3d"][0].numpy()
    pts_b = b["pts3d"][0].numpy()
    conf_a = a["conf"][0].numpy()
    conf_b = b["conf"][0].numpy()
    img_a = denormalize(view_a["img"][0])
    img_b = denormalize(view_b["img"][0])

    # fused point cloud, both views in the camera frame of A
    keep_a = conf_a > conf_thresh
    keep_b = conf_b > conf_thresh
    if not keep_a.any() and not keep_b.any():
        raise gr.Error(
            "No 3D point passed the confidence threshold. "
            "Try lowering it, or use two images that see the same scene."
        )
    pts = np.concatenate([pts_a[keep_a], transform_points(T_ab, pts_b[keep_b])])
    cols = np.concatenate([img_a[keep_a], img_b[keep_b]])

    with tempfile.NamedTemporaryFile(suffix=".ply", delete=False) as f:
        ply_path = f.name
    save_ply(ply_path, pts, cols)

    with tempfile.NamedTemporaryFile(suffix=".npz", delete=False) as f:
        npz_path = f.name
    np.savez_compressed(
        npz_path,
        pts3d_a=pts_a, pts3d_b=pts_b, conf_a=conf_a, conf_b=conf_b,
        T_b_from_a=T_ba, T_a_from_b=T_ab, img_a=img_a, img_b=img_b,
    )

    # camera centre of B, as a unit direction in camera A's frame (scale is arbitrary)
    t = T_ab[:3, 3] / max(np.linalg.norm(T_ab[:3, 3]), 1e-12)
    report = (
        f"Relative rotation (A -> B): {rotation_angle_deg(T_ba[:3, :3]):.2f} deg\n"
        f"Pose confidence (A->B / B->A): {float(a['relative_pose_conf'][0]):.3f} / "
        f"{float(b['relative_pose_conf'][0]):.3f}\n"
        f"Camera B centre direction (in A's frame): "
        f"[{t[0]:+.3f}, {t[1]:+.3f}, {t[2]:+.3f}]  (unit; scale is arbitrary)\n"
        f"Points kept: {len(pts)} of {pts_a.size // 3 + pts_b.size // 3}\n\n"
        f"T_B<-A (camera A -> camera B):\n{pose_matrix_str(T_ba)}\n\n"
        f"T_A<-B (camera B -> camera A):\n{pose_matrix_str(T_ab)}"
    )
    return ply_path, npz_path, report


# ---------------------------------------------------------------- ui
CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    gr.Markdown(
        """
        # BA-T: An Iterative Transformer for Two-View Bundle Adjustment

        Upload two images of a scene from different viewpoints. BA-T refines the
        pair with a recursive bundle-adjustment transformer layer and predicts a
        dense 3D point cloud and the relative camera pose between the views.
        *Point maps and translation share one arbitrary (non-metric) scale.*

        [Paper](https://arxiv.org/abs/2606.03287) ·
        [Project page](https://ganlinzhang.xyz/BA-T/) ·
        [Code](https://github.com/zhangganlin/BA-T) ·
        [Model](https://huggingface.co/zhangganlin/BA-T)
        """
    )

    with gr.Column(elem_id="col-container"):
        with gr.Row():
            image_a = gr.Image(type="filepath", label="Image A", height=280)
            image_b = gr.Image(type="filepath", label="Image B", height=280)

        with gr.Accordion("Advanced settings", open=False):
            num_refine_iter = gr.Slider(
                minimum=4, maximum=24, step=1, value=DEFAULT_K,
                label="Refinement iterations K",
                info="The BA layer runs K times; trained with 4-24. More = more refined.",
            )
            conf_thresh = gr.Slider(
                minimum=1.0, maximum=10.0, step=0.5, value=DEFAULT_CONF_THRESH,
                label="Point confidence threshold",
                info="Only 3D points above this confidence enter the point cloud.",
            )

        run_btn = gr.Button("Reconstruct", variant="primary")

        point_cloud = gr.Model3D(
            label="3D point cloud (camera A's frame)",
            camera_position=(0, 0, -3.5),
        )
        with gr.Row():
            pose_report = gr.Textbox(label="Relative camera pose", lines=13)
            npz_file = gr.File(label="Raw predictions (.npz)")

        gr.Examples(
            examples=[
                ["examples/teaser_0.jpg", "examples/teaser_1.jpg"],
                ["examples/twoview_desk_0.jpg", "examples/twoview_desk_1.jpg"],
                ["examples/twoview_shelf_0.jpg", "examples/twoview_shelf_1.jpg"],
                ["examples/twoview_corner_0.jpg", "examples/twoview_corner_1.jpg"],
            ],
            inputs=[image_a, image_b],
            outputs=[point_cloud, npz_file, pose_report],
            fn=bundle_adjust,
            cache_examples=True,
            cache_mode="lazy",
        )

    run_btn.click(
        bundle_adjust,
        inputs=[image_a, image_b, num_refine_iter, conf_thresh],
        outputs=[point_cloud, npz_file, pose_report],
        api_name="reconstruct",
    )

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)