"""
Full reconstruction pipeline:
1. MASt3R-SfM: camera poses + dense point cloud (COLMAP fails on textureless CAD)
2. Mesh extraction: Poisson reconstruction + ball pivoting from the point cloud
3. Preview rendering: multiple viewpoints of point cloud and mesh
4. Export: PLY, OBJ, GLB, COLMAP format, preview PNGs, metadata JSON

Runs on all parts in data/parts/
"""
import os
import sys
import glob
import json
import torch
import numpy as np
import trimesh
import open3d as o3d
from PIL import Image

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "mast3r"))
import mast3r.utils.path_to_dust3r  # noqa: F401
from dust3r.utils.image import load_images
from mast3r.image_pairs import make_pairs
from mast3r.cloud_opt.sparse_ga import sparse_global_alignment
from dust3r.utils.device import to_numpy
from mast3r.model import AsymmetricMASt3R


PARTS_DIR = "data/parts"
OUTPUT_DIR = "data/mast3r_results"
DEVICE = "cuda"
IMAGE_SIZE = 512
MIN_CONF = 0.001  # CAD parts have low MASt3R confidence but valid geometry


_model = None
def get_model():
    global _model
    if _model is None:
        print("Loading MASt3R model...")
        _model = AsymmetricMASt3R.from_pretrained(
            "naver/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric"
        ).to(DEVICE)
        _model.eval()
    return _model


def reconstruct_part(part, image_paths):
    print(f"\n{'='*60}")
    print(f"MASt3R reconstruction: {part} ({len(image_paths)} images)")
    print(f"{'='*60}")

    out_dir = os.path.join(OUTPUT_DIR, part)
    os.makedirs(out_dir, exist_ok=True)

    model = get_model()

    # Load images
    print("[1/4] Loading images...")
    imgs = load_images(image_paths, size=IMAGE_SIZE, verbose=True)

    # Make pairs
    print("[2/4] Making image pairs (complete scene graph)...")
    pairs = make_pairs(imgs, scene_graph="complete", prefilter=None, symmetrize=True)

    # Run sparse global alignment
    import tempfile
    cache_dir = tempfile.mkdtemp(suffix=f"_mast3r_{part}")
    os.makedirs(cache_dir, exist_ok=True)

    print("[3/4] Running sparse_global_alignment (MASt3R-SfM)...")
    scene = sparse_global_alignment(
        image_paths, pairs, cache_dir, model,
        lr1=0.07, niter1=300,
        lr2=0.01, niter2=300,
        device=DEVICE,
        opt_depth=True,
        shared_intrinsics=False,
        matching_conf_thr=5.0,
    )

    # Extract results
    focals = scene.get_focals().cpu().numpy()
    principal_points = scene.get_principal_points().cpu().numpy()
    cam2world = scene.get_im_poses().cpu().numpy()

    pts3d, depthmaps, confs = scene.get_dense_pts3d(clean_depth=True, subsample=8)
    pts3d = to_numpy(pts3d)
    confs = to_numpy(confs)
    scene_imgs = scene.imgs

    print(f"  Focals: {focals}")
    print(f"  Cam2world shapes: {cam2world.shape}")
    print(f"  Points per view: {[p.shape for p in pts3d]}")

    # --- Build point cloud (use all points; CAD has low conf but valid geometry) ---
    # Note: pts3d[i] is (H*W, 3) flat, confs[i] is (H, W), scene_imgs[i] is (H, W, 3)
    # We use confidence for filtering but must flatten it
    msk = []
    for i in range(len(pts3d)):
        c = confs[i]
        h, w = c.shape
        m = (c > MIN_CONF).reshape(-1)
        msk.append(m)
    total_conf = sum(m.sum() for m in msk)
    total_all = sum(p.shape[0] for p in pts3d)
    if total_conf < total_all * 0.01:
        print(f"  Low confidence ({total_conf}/{total_all} pts) - using all points")
        msk = [np.ones(p.shape[0], dtype=bool) for p in pts3d]

    all_pts = []
    all_cols = []
    for p, m, im in zip(pts3d, msk, scene_imgs):
        if m.sum() == 0:
            continue
        im_flat = im.reshape(-1, 3)
        all_pts.append(p[m])
        all_cols.append(im_flat[m])
    if not all_pts:
        print("  ERROR: No points at all!")
        return None
    pts = np.concatenate(all_pts, axis=0)
    cols = (np.clip(np.concatenate(all_cols, axis=0), 0, 1) * 255).astype(np.uint8)
    print(f"  Total points: {len(pts)}")
    print(f"  Bounds: {pts.min(0)} to {pts.max(0)}")

    # Clean point cloud: remove outliers
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(pts)
    pcd.colors = o3d.utility.Vector3dVector(cols.astype(float) / 255.0)
    # Estimate normals
    pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.05, max_nn=30))
    pcd.orient_normals_consistent_tangent_plane(10)

    # Remove statistical outliers
    pcd_clean, ind = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=2.0)
    print(f"  After outlier removal: {len(pcd_clean.points)} points")

    # Export point cloud PLY
    ply_path = os.path.join(out_dir, "pointcloud.ply")
    o3d.io.write_point_cloud(ply_path, pcd_clean)
    print(f"  Saved point cloud: {ply_path}")

    # Also save as trimesh for rendering
    tm_pc = trimesh.PointCloud(
        np.asarray(pcd_clean.points),
        colors=(np.asarray(pcd_clean.colors) * 255).astype(np.uint8)
    )

    # --- Mesh extraction ---
    print("[4/4] Extracting mesh...")
    mesh_path = extract_mesh(pcd_clean, out_dir, part)

    # --- Export COLMAP format ---
    colmap_dir = os.path.join(out_dir, "colmap_sparse")
    os.makedirs(colmap_dir, exist_ok=True)
    write_colmap_format(colmap_dir, image_paths, focals, principal_points,
                        cam2world, pts3d, msk, scene_imgs, confs)

    # --- Export GLB ---
    try:
        glb_path = os.path.join(out_dir, "scene.glb")
        write_glb(glb_path, scene_imgs, pts3d, msk, focals, cam2world)
        print(f"  Saved GLB: {glb_path}")
    except Exception as e:
        print(f"  GLB export failed: {e}")

    # --- Render previews ---
    render_previews(tm_pc, mesh_path, out_dir, part)

    # --- Save metadata ---
    meta = {
        "part": part,
        "n_images": len(image_paths),
        "focals": focals.tolist(),
        "principal_points": principal_points.tolist(),
        "cam2world": cam2world.tolist(),
        "n_points_raw": int(len(pts)),
        "n_points_clean": int(len(pcd_clean.points)),
        "bounds_min": pts.min(0).tolist(),
        "bounds_max": pts.max(0).tolist(),
        "method": "MASt3R-SfM (sparse_global_alignment)",
        "image_size": IMAGE_SIZE,
        "min_conf": MIN_CONF,
    }
    with open(os.path.join(out_dir, "metadata.json"), "w") as f:
        json.dump(meta, f, indent=2)
    print(f"  Saved metadata.json")

    return out_dir


def extract_mesh(pcd, out_dir, part):
    """Extract mesh from point cloud using Poisson reconstruction."""
    meshes = {}

    # Method 1: Poisson reconstruction
    try:
        print("  Poisson reconstruction...")
        mesh_poisson, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(
            pcd, depth=9, width=0, scale=1.1, linear_fit=False
        )
        # Remove low-density vertices
        density_threshold = np.quantile(densities, 0.05)
        vertices_to_remove = densities < density_threshold
        mesh_poisson.remove_vertices_by_mask(vertices_to_remove)
        mesh_poisson.compute_vertex_normals()

        poisson_path = os.path.join(out_dir, f"{part}_poisson.ply")
        o3d.io.write_triangle_mesh(poisson_path, mesh_poisson)
        print(f"    Saved Poisson mesh: {poisson_path} ({len(mesh_poisson.vertices)} verts, {len(mesh_poisson.triangles)} faces)")
        meshes["poisson"] = poisson_path

        # Also export as OBJ
        obj_path = os.path.join(out_dir, f"{part}_poisson.obj")
        o3d.io.write_triangle_mesh(obj_path, mesh_poisson)
        meshes["poisson_obj"] = obj_path
    except Exception as e:
        print(f"    Poisson failed: {e}")

    # Method 2: Ball pivoting
    try:
        print("  Ball pivoting...")
        radii = o3d.utility.DoubleVector([0.005, 0.01, 0.02, 0.04])
        mesh_bp = o3d.geometry.TriangleMesh.create_from_point_cloud_ball_pivoting(
            pcd, radii
        )
        mesh_bp.compute_vertex_normals()

        bp_path = os.path.join(out_dir, f"{part}_ballpivoting.ply")
        o3d.io.write_triangle_mesh(bp_path, mesh_bp)
        print(f"    Saved ball-pivoting mesh: {bp_path} ({len(mesh_bp.vertices)} verts, {len(mesh_bp.triangles)} faces)")
        meshes["ballpivoting"] = bp_path
    except Exception as e:
        print(f"    Ball pivoting failed: {e}")

    return meshes


def rotmat_to_qvec(R):
    Rxx, Ryx, Rzx = R[0, 0], R[1, 0], R[2, 0]
    Rxy, Ryy, Rzy = R[0, 1], R[1, 1], R[2, 1]
    Rxz, Ryz, Rzz = R[0, 2], R[1, 2], R[2, 2]
    K = np.array([
        [Rxx - Ryy - Rzz, 0, 0, 0],
        [Ryx + Rxy, Ryy - Rxx - Rzz, 0, 0],
        [Rzx + Rxz, Rzy + Ryz, Rzz - Rxx - Ryy, 0],
        [Ryz - Rzy, Rzx - Rxz, Rxy - Ryx, Rxx + Ryy + Rzz],
    ]) / 3.0
    vals, vecs = np.linalg.eigh(K)
    qvec = vecs[[3, 0, 1, 2], np.argmax(vals)]
    if qvec[0] < 0:
        qvec *= -1
    return qvec


def write_colmap_format(colmap_dir, image_paths, focals, principal_points,
                        cam2world, pts3d, msk, scene_imgs, confs):
    n = len(image_paths)

    with open(os.path.join(colmap_dir, "cameras.txt"), "w") as f:
        for i in range(n):
            h, w = confs[i].shape
            fx = float(focals[i])
            cx, cy = float(principal_points[i, 0]), float(principal_points[i, 1])
            f.write(f"{i+1} PINHOLE {w} {h} {fx} {fx} {cx} {cy}\n")

    with open(os.path.join(colmap_dir, "images.txt"), "w") as f:
        for i in range(n):
            c2w = cam2world[i]
            w2c = np.linalg.inv(c2w)
            R = w2c[:3, :3]
            t = w2c[:3, 3]
            q = rotmat_to_qvec(R)
            fname = os.path.basename(image_paths[i])
            f.write(f"{i+1} {q[0]} {q[1]} {q[2]} {q[3]} {t[0]} {t[1]} {t[2]} 1 {fname}\n")
            f.write("\n")

    with open(os.path.join(colmap_dir, "points3D.txt"), "w") as f:
        pid = 0
        for i in range(n):
            m = msk[i]
            if m.sum() == 0:
                continue
            pts = pts3d[i][m]
            im_flat = scene_imgs[i].reshape(-1, 3)
            cols = (im_flat[m] * 255).astype(np.uint8)
            for p, c in zip(pts, cols):
                pid += 1
                f.write(f"{pid} {p[0]} {p[1]} {p[2]} {c[0]} {c[1]} {c[2]} 0\n")

    transforms = {"camera_model": "PINHOLE", "frames": []}
    for i in range(n):
        h, w = confs[i].shape
        fx = float(focals[i])
        cx, cy = float(principal_points[i, 0]), float(principal_points[i, 1])
        transforms["frames"].append({
            "file_path": f"images/{os.path.basename(image_paths[i])}",
            "camera_model": "PINHOLE",
            "fl_x": fx, "fl_y": fx,
            "cx": cx, "cy": cy,
            "w": int(w), "h": int(h),
            "transform_matrix": cam2world[i].tolist(),
        })
    with open(os.path.join(colmap_dir, "transforms.json"), "w") as f:
        json.dump(transforms, f, indent=2)

    img_dir = os.path.join(colmap_dir, "images")
    os.makedirs(img_dir, exist_ok=True)
    import shutil
    for p in image_paths:
        shutil.copy(p, os.path.join(img_dir, os.path.basename(p)))
    print(f"  Saved COLMAP format to {colmap_dir}")


def write_glb(glb_path, scene_imgs, pts3d, msk, focals, cam2world):
    try:
        from mast3r.demo import _convert_scene_output_to_glb
        _convert_scene_output_to_glb(
            glb_path, scene_imgs, pts3d, msk,
            focals, cam2world,
            as_pointcloud=True, cam_size=0.05,
            transparent_cams=True,
        )
    except Exception:
        all_pts = []; all_cols = []
        for p, m, im in zip(pts3d, msk, scene_imgs):
            if m.sum() == 0:
                continue
            im_flat = im.reshape(-1, 3)
            all_pts.append(p[m])
            all_cols.append((im_flat[m] * 255).astype(np.uint8))
        if all_pts:
            pts = np.concatenate(all_pts, axis=0)
            cols = np.concatenate(all_cols, axis=0)
            pc = trimesh.PointCloud(pts, colors=cols)
            pc.export(glb_path)


def render_previews(pointcloud, mesh_paths, out_dir, part):
    """Render preview images of point cloud and meshes from multiple angles."""
    preview_dir = os.path.join(out_dir, "previews")
    os.makedirs(preview_dir, exist_ok=True)

    angles = [
        ("front", [0, 0, 0]),
        ("side", [0, np.pi/2, 0]),
        ("top", [np.pi/2, 0, 0]),
        ("iso", [np.pi/6, np.pi/4, 0]),
    ]

    # Render point cloud
    for name, euler in angles:
        try:
            sc = trimesh.Scene(pointcloud)
            sc.camera_transform = trimesh.transformations.euler_matrix(*euler)
            png = sc.save_image(resolution=[800, 600])
            path = os.path.join(preview_dir, f"{part}_pc_{name}.png")
            with open(path, "wb") as f:
                f.write(png)
            print(f"    Saved preview: {path}")
        except Exception as e:
            print(f"    PC preview {name} failed: {e}")

    # Render meshes
    for mesh_key, mesh_path in mesh_paths.items():
        try:
            mesh = trimesh.load(mesh_path)
            for name, euler in angles:
                try:
                    sc = trimesh.Scene(mesh)
                    sc.camera_transform = trimesh.transformations.euler_matrix(*euler)
                    png = sc.save_image(resolution=[800, 600])
                    path = os.path.join(preview_dir, f"{part}_{mesh_key}_{name}.png")
                    with open(path, "wb") as f:
                        f.write(png)
                    print(f"    Saved preview: {path}")
                except Exception as e:
                    print(f"    Mesh {mesh_key} preview {name} failed: {e}")
        except Exception as e:
            print(f"    Failed to load {mesh_path} for preview: {e}")


def main():
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    results = {}

    for part in sorted(os.listdir(PARTS_DIR)):
        input_dir = os.path.join(PARTS_DIR, part, "input")
        if not os.path.isdir(input_dir):
            continue
        images = sorted(glob.glob(os.path.join(input_dir, "*.png")))
        if len(images) < 3:
            print(f"Skipping {part}: only {len(images)} images")
            continue
        try:
            out_dir = reconstruct_part(part, images)
            results[part] = ("OK", out_dir)
        except Exception as e:
            import traceback
            traceback.print_exc()
            results[part] = ("FAILED", str(e))

    print(f"\n{'='*60}")
    print("MASt3R RECONSTRUCTION SUMMARY")
    print(f"{'='*60}")
    for part, (status, msg) in results.items():
        print(f"  {part}: {status} - {msg}")
    return results

if __name__ == "__main__":
    main()