import argparse
import os
import numpy as np
import cv2
import matplotlib.pyplot as plt

_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
# Default paths: repo layout has input/ and rectified/ next to this package folder.
_DEFAULT_DISPARITY = os.path.normpath(
    os.path.join(_SCRIPT_DIR, "..", "input", "disparity.xml")
)
_DEFAULT_Q = os.path.normpath(os.path.join(_SCRIPT_DIR, "..", "rectified", "Q.cvstore"))


def load_disparity_map(disparity_file):
    """Load disparity from .xml (OpenCV FileStorage), .png, or .npy."""
    disparity_file = os.fspath(disparity_file)
    if disparity_file.endswith(".xml"):
        fs = cv2.FileStorage(disparity_file, cv2.FILE_STORAGE_READ)
        disp_map = fs.getNode("disparity").mat().astype(np.float32)
        fs.release()
    elif disparity_file.endswith(".png"):
        disp_map = cv2.imread(disparity_file, cv2.IMREAD_UNCHANGED).astype(np.float32)
    elif disparity_file.endswith(".npy"):
        disp_map = np.load(disparity_file).astype(np.float32)
    else:
        raise ValueError(
            f"Unsupported disparity format: {disparity_file}. Use .xml, .png, or .npy."
        )
    return disp_map


def load_q_matrix(Q):
    """
    Accept a 4x4 ndarray or a file path:
    - .npy: array saved with numpy
    - OpenCV FileStorage (.cvstore, .xml, .yaml, .yml): matrix node named 'Q'
    """
    if isinstance(Q, (str, os.PathLike)):
        path = os.fspath(Q)
        lower = path.lower()
        if lower.endswith(".npy"):
            Q = np.load(path)
        else:
            fs = cv2.FileStorage(path, cv2.FILE_STORAGE_READ)
            if not fs.isOpened():
                raise ValueError(f"Could not open Q file: {path}")
            q_node = fs.getNode("Q")
            if q_node.empty():
                fs.release()
                raise ValueError(
                    f"No matrix 'Q' found in {path}. Expected OpenCV FileStorage node 'Q'."
                )
            Q = q_node.mat()
            fs.release()
    Q = np.asarray(Q, dtype=np.float64)
    if Q.shape != (4, 4):
        raise ValueError(f"Q must be 4x4, got shape {Q.shape}.")
    return Q


def disparity_to_pointcloud(disparity_file, Q, valid_mask=None):
    """
    Build a 3D point cloud from a disparity map and stereo rectification Q matrix.

    Parameters
    ----------
    disparity_file : str or os.PathLike
        Path to disparity (.xml with 'disparity' node, .png, or .npy).
    Q : ndarray (4, 4) or str path
        4x4 reprojection matrix; .npy or OpenCV store (.cvstore / .xml / .yaml) with node 'Q'.
    valid_mask : ndarray of bool, optional
        If None, uses pixels where disparity > min(disparity) (same as original script).

    Returns
    -------
    points_xyz : ndarray, shape (N, 3)
        3D points in rectified camera coordinates.
    """
    disp_map = load_disparity_map(disparity_file)
    Qm = load_q_matrix(Q)
    point_cloud = cv2.reprojectImageTo3D(disp_map, Qm)
    if valid_mask is None:
        valid_mask = disp_map > disp_map.min()
    else:
        valid_mask = np.asarray(valid_mask, dtype=bool)
    points_xyz = point_cloud[valid_mask].reshape(-1, 3)
    return points_xyz


def pointcloud_with_depth_colors(points_xyz):
    """Attach jet colormap RGB in [0,1] from Z depth (same visualization as before)."""
    depth = points_xyz[:, 2]
    dmin, dmax = depth.min(), depth.max()
    if dmax > dmin:
        depth_norm = (depth - dmin) / (dmax - dmin)
    else:
        depth_norm = np.zeros_like(depth)
    colors = plt.get_cmap("jet")(depth_norm)[:, :3]
    return np.hstack((points_xyz, colors))


def save_ply(filename, points):
    with open(filename, "w") as f:
        f.write("ply\nformat ascii 1.0\n")
        f.write(f"element vertex {points.shape[0]}\n")
        f.write("property float x\nproperty float y\nproperty float z\n")
        f.write("property uchar red\nproperty uchar green\nproperty uchar blue\nend_header\n")
        for point in points:
            x, y, z, r, g, b = point
            f.write(f"{x} {y} {z} {int(r * 255)} {int(g * 255)} {int(b * 255)}\n")


def save_ascii_point_cloud(filename, points):
    with open(filename, "w") as f:
        for point in points:
            x, y, z, r, g, b = point
            f.write(f"{x} {y} {z} {int(r * 255)} {int(g * 255)} {int(b * 255)}\n")


def main():
    parser = argparse.ArgumentParser(
        description="Convert disparity map + Q matrix to a point cloud (PLY/TXT)."
    )
    parser.add_argument(
        "--disparity",
        "-d",
        default=_DEFAULT_DISPARITY,
        help=f"Disparity file (.xml, .png, .npy). Default: {_DEFAULT_DISPARITY}",
    )
    parser.add_argument(
        "--q",
        "-q",
        default=_DEFAULT_Q,
        help=f"Q matrix: .npy or OpenCV FileStorage (.cvstore, .xml, .yaml) with node 'Q'. Default: {_DEFAULT_Q}",
    )
    parser.add_argument(
        "--out-dir",
        default="./output",
        help="Directory for PLY and TXT exports (default: ./output).",
    )
    parser.add_argument(
        "--no-save",
        action="store_true",
        help="Only print point count; do not write PLY/TXT.",
    )
    parser.add_argument(
        "--troubleshooting",
        action="store_true",
        help="Also save ASCII .txt point cloud (default: PLY only)",
    )
    args = parser.parse_args()

    points_xyz = disparity_to_pointcloud(args.disparity, args.q)
    print(f"Point cloud shape: {points_xyz.shape} (N x 3)")

    if args.no_save:
        return

    os.makedirs(args.out_dir, exist_ok=True)
    base = os.path.splitext(os.path.basename(os.fspath(args.disparity)))[0]
    point_cloud_color = pointcloud_with_depth_colors(points_xyz)
    out_ply = os.path.join(args.out_dir, f"Pointcloud_{base}.ply")
    save_ply(out_ply, point_cloud_color)
    print(f"Saved: {out_ply}")
    if args.troubleshooting:
        out_txt = os.path.join(args.out_dir, f"Pointcloud_{base}.txt")
        save_ascii_point_cloud(out_txt, point_cloud_color)
        print(f"Saved: {out_txt}")


if __name__ == "__main__":
    main()