import numpy as np

try:
    from scipy.spatial import cKDTree as KDTree
except Exception:
    try:
        from sklearn.neighbors import NearestNeighbors as SKNearest
        KDTree = None
    except Exception:
        raise ImportError("Requires scipy.spatial.cKDTree or sklearn.neighbors. Install scipy or scikit-learn.")


def sample_points_from_mesh(vertices: np.ndarray, faces: np.ndarray, n_samples: int) -> np.ndarray:
    """
    Uniformly sample points on mesh surface.

    vertices: (n,3) array or flattened (n*3,)
    faces: (f,3) array of indices or flattened (f*3,)
    n_samples: number of points to sample

    Returns: (n_samples, 3) sampled points (float32)
    """
    v = np.asarray(vertices).reshape(-1, 3).astype(np.float64)
    f = np.asarray(faces).reshape(-1, 3).astype(np.int64)

    v0 = v[f[:, 0], :]
    v1 = v[f[:, 1], :]
    v2 = v[f[:, 2], :]

    # triangle areas
    tri_edges = np.cross(v1 - v0, v2 - v0)
    tri_areas = 0.5 * np.linalg.norm(tri_edges, axis=1)
    area_sum = tri_areas.sum()
    if area_sum == 0:
        # Degenerate mesh: return repeated vertices
        idx = np.random.randint(0, v.shape[0], size=n_samples)
        return v[idx].astype(np.float32)

    # probabilities
    probs = tri_areas / area_sum

    # sample triangle indices according to area
    tri_indices = np.random.choice(len(f), size=n_samples, p=probs)

    # sample barycentric coordinates
    r1 = np.sqrt(np.random.rand(n_samples))
    r2 = np.random.rand(n_samples)
    a = 1.0 - r1
    b = r1 * (1.0 - r2)
    c = r1 * r2

    pts = (a[:, None] * v0[tri_indices] +
           b[:, None] * v1[tri_indices] +
           c[:, None] * v2[tri_indices])

    return pts.astype(np.float32)


def _nn_distances(a_pts: np.ndarray, b_pts: np.ndarray):
    """
    Compute nearest-neighbor Euclidean distances from each point in a_pts to nearest in b_pts.

    Returns distances (not squared).
    """
    if a_pts.shape[0] == 0:
        return np.array([], dtype=np.float32)
    if b_pts.shape[0] == 0:
        # return inf
        return np.full((a_pts.shape[0],), np.inf, dtype=np.float32)

    if KDTree is not None:
        tree = KDTree(b_pts)
        dists, _ = tree.query(a_pts, k=1)
        return dists.astype(np.float32)
    else:
        # fallback to sklearn
        nbrs = SKNearest(n_neighbors=1, algorithm='auto').fit(b_pts)
        dists, _ = nbrs.kneighbors(a_pts)
        return dists[:, 0].astype(np.float32)


def chamfer_distance_from_meshes(pred_vertices: np.ndarray,
                                 pred_faces: np.ndarray,
                                 gt_vertices: np.ndarray,
                                 gt_faces: np.ndarray,
                                 n_samples: int = 100000,
                                 return_raw: bool = False):
    """
    Compute Chamfer distance between predicted mesh and ground-truth mesh.

    pred_vertices/pred_faces: mesh A (prediction)
    gt_vertices/gt_faces: mesh B (ground truth)
    n_samples: number of samples per mesh (default 100k). Lower for speed, e.g. 10k.
    return_raw: if True, also return the sampled point clouds and per-point distances.

    Returns:
      If return_raw is False:
        dict with keys:
          'cd_l2_sq'      : bidirectional mean squared L2 (mean of squared distances)
          'cd_l2'         : bidirectional mean L2 (mean of distances)
          'A_to_B_l2_sq'  : mean squared distances from A->B
          'B_to_A_l2_sq'  : mean squared distances from B->A
          'A_to_B_l2'     : mean distances A->B
          'B_to_A_l2'     : mean distances B->A
      If return_raw is True:
        (metrics_dict, pts_pred, pts_gt, dists_pred_to_gt, dists_gt_to_pred)
    """
    pts_pred = sample_points_from_mesh(pred_vertices, pred_faces, n_samples)
    pts_gt = sample_points_from_mesh(gt_vertices, gt_faces, n_samples)

    d_pred_to_gt = _nn_distances(pts_pred, pts_gt)  # distances from pred samples to nearest gt
    d_gt_to_pred = _nn_distances(pts_gt, pts_pred)  # distances from gt samples to nearest pred

    # L2 (distances) and L2^2 (squared)
    A_to_B_l2 = float(np.mean(d_pred_to_gt))
    B_to_A_l2 = float(np.mean(d_gt_to_pred))
    cd_l2 = 0.5 * (A_to_B_l2 + B_to_A_l2)

    A_to_B_l2_sq = float(np.mean(d_pred_to_gt ** 2))
    B_to_A_l2_sq = float(np.mean(d_gt_to_pred ** 2))
    cd_l2_sq = 0.5 * (A_to_B_l2_sq + B_to_A_l2_sq)

    metrics = {
        'cd_l2_sq': cd_l2_sq,
        'cd_l2': cd_l2,
        'A_to_B_l2_sq': A_to_B_l2_sq,
        'B_to_A_l2_sq': B_to_A_l2_sq,
        'A_to_B_l2': A_to_B_l2,
        'B_to_A_l2': B_to_A_l2,
        'n_samples_per_mesh': n_samples,
    }

    if return_raw:
        return metrics, pts_pred, pts_gt, d_pred_to_gt, d_gt_to_pred
    return metrics


if __name__ == "__main__":
    # Quick example using a simple triangle meshes (triangles)
    # Pred: unit right triangle in XY plane
    pred_verts = np.array([[0, 0, 0],
                           [1, 0, 0],
                           [0, 1, 0]], dtype=np.float32)
    pred_faces = np.array([[0, 1, 2]], dtype=np.int32)

    # GT: slightly translated triangle
    gt_verts = pred_verts + np.array([0.00, 0.00, 0.5], dtype=np.float32)
    gt_faces = pred_faces.copy()

    metrics = chamfer_distance_from_meshes(pred_verts, pred_faces, gt_verts, gt_faces,
                                           n_samples=20000)
    print("Chamfer metrics:", metrics)