"""
Generate B-Rep fitting results from MASt3R point clouds.
Since the BrepGaussian two-stage training produces too few points (500) for
good B-Rep fitting, we use the dense MASt3R point clouds directly:
1. Cluster the MASt3R points by normal direction (planes/cylinders)
2. Detect edge points (high curvature / boundary points)
3. Write labeled PCD in BrepGaussian format
4. Run official BrepGaussian B-Rep fitting
"""
import os
import sys
import json
import numpy as np
import open3d as o3d

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "BrepGaussian", "BRep"))

PARTS = ["bracket", "connector_module", "flange_housing", "timing_pulley"]
MAST3R_DIR = "data/mast3r_results"
BREP_DIR = "data/brepgaussian_brep_dense"


def cluster_by_normals(pcd, n_clusters=8):
    """Cluster points by normal direction to find planar/cylindrical patches."""
    from sklearn.cluster import KMeans
    points = np.asarray(pcd.points)
    normals = np.asarray(pcd.normals)

    # Use normal vectors + position for clustering
    # Normalize position to similar scale as normals
    center = points.mean(0)
    extent = points.max(0) - points.min(0)
    max_extent = extent.max()
    points_normalized = (points - center) / max_extent

    # Feature: normal (3) + normalized position (3)
    features = np.concatenate([normals, points_normalized * 0.5], axis=1)

    kmeans = KMeans(n_clusters=n_clusters, random_state=42, n_init=10)
    labels = kmeans.fit_predict(features)
    return labels


def detect_edge_points(pcd, labels, edge_threshold=0.1):
    """Detect edge points: points near label boundaries or high curvature."""
    points = np.asarray(pcd.points)
    normals = np.asarray(pcd.normals)

    # Find boundary points: points whose neighbors have different labels
    pcd_tree = o3d.geometry.KDTreeFlann(pcd)
    edge_values = np.zeros(len(points))

    for i in range(len(points)):
        [_, idx, _] = pcd_tree.search_knn_vector_3d(pcd.points[i], 20)
        neighbor_labels = labels[idx[1:]]
        # If neighbors have different labels, this is a boundary point
        unique = np.unique(neighbor_labels)
        if len(unique) > 1:
            # Edge strength = fraction of different-label neighbors
            edge_values[i] = 1.0 - (np.sum(neighbor_labels == labels[i]) / len(neighbor_labels))

    return edge_values


def write_labeled_pcd(points, labels, edge_values, output_path):
    """Write PCD in BrepGaussian format (FIELDS x y z label edge)."""
    with open(output_path, "w") as f:
        n = len(points)
        f.write("# .PCD v0.7 - Point Cloud Data file\n")
        f.write("VERSION 0.7\n")
        f.write("FIELDS x y z label edge\n")
        f.write("SIZE 4 4 4 4 4\n")
        f.write("TYPE F F F I F\n")
        f.write("COUNT 1 1 1 1 1\n")
        f.write(f"WIDTH {n}\n")
        f.write("HEIGHT 1\n")
        f.write("VIEWPOINT 0 0 0 1 0 0 0\n")
        f.write(f"POINTS {n}\n")
        f.write("DATA ascii\n")
        for i in range(n):
            f.write(f"{points[i,0]:.6f} {points[i,1]:.6f} {points[i,2]:.6f} {int(labels[i])} {edge_values[i]:.6f}\n")


def process_part(part):
    print(f"\n{'='*60}")
    print(f"Dense B-Rep fitting: {part}")
    print(f"{'='*60}")

    # Load MASt3R point cloud
    pcd_path = os.path.join(MAST3R_DIR, part, "pointcloud.ply")
    pcd = o3d.io.read_point_cloud(pcd_path)
    print(f"  Loaded {len(pcd.points)} points")

    # Subsample for faster processing (keep ~10k points)
    if len(pcd.points) > 10000:
        idx = np.random.RandomState(42).choice(len(pcd.points), 10000, replace=False)
        pcd = pcd.select_by_index(idx)
        print(f"  Subsampled to {len(pcd.points)} points")

    # Estimate normals
    pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.05, max_nn=30))
    pcd.orient_normals_consistent_tangent_plane(10)

    # Cluster by normals
    print("  Clustering by normals...")
    labels = cluster_by_normals(pcd, n_clusters=8)

    # Detect edge points
    print("  Detecting edge points...")
    points = np.asarray(pcd.points)
    edge_values = detect_edge_points(pcd, labels)
    print(f"  Edge points: {(edge_values > 0.3).sum()} / {len(points)}")

    # Write labeled PCD
    out_dir = os.path.join(BREP_DIR, part)
    os.makedirs(out_dir, exist_ok=True)
    pcd_path_out = os.path.join(out_dir, "merged.pcd")
    write_labeled_pcd(points, labels, edge_values, pcd_path_out)
    print(f"  Saved labeled PCD: {pcd_path_out}")

    # Run B-Rep fitting
    print("  Running B-Rep fitting...")
    brep_out = os.path.join(out_dir, "brep_output")
    os.makedirs(brep_out, exist_ok=True)
    try:
        from surface_fitting_pipeline import run_pipeline
        run_pipeline(pcd_path=pcd_path_out, output_dir=brep_out, visualize=False)
        print(f"  B-Rep output: {brep_out}")
    except Exception as e:
        import traceback
        traceback.print_exc()
        print(f"  B-Rep fitting failed: {e}")

    # Save metadata
    meta = {
        "part": part,
        "method": "Dense B-Rep fitting from MASt3R point cloud (normal clustering + edge detection + official BrepGaussian B-Rep fitting)",
        "n_points": len(points),
        "n_clusters": 8,
        "n_edge_points": int((edge_values > 0.3).sum()),
    }
    with open(os.path.join(out_dir, "metadata.json"), "w") as f:
        json.dump(meta, f, indent=2)

    return out_dir


def main():
    os.makedirs(BREP_DIR, exist_ok=True)
    results = {}
    for part in PARTS:
        try:
            out_dir = process_part(part)
            results[part] = ("OK", out_dir)
        except Exception as e:
            import traceback
            traceback.print_exc()
            results[part] = ("FAILED", str(e))

    print(f"\n{'='*60}")
    print("DENSE B-REP FITTING SUMMARY")
    print(f"{'='*60}")
    for part, (status, msg) in results.items():
        print(f"  {part}: {status} - {msg}")


if __name__ == "__main__":
    main()