from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from evo.core import sync
from evo.core.trajectory import PoseTrajectory3D
from evo.tools import plot
from plyfile import PlyData, PlyElement


def plot_trajectory(pred_traj, gt_traj=None, title="", filename="", align=True, correct_scale=True):
    assert isinstance(pred_traj, PoseTrajectory3D)

    if gt_traj is not None:
        assert isinstance(gt_traj, PoseTrajectory3D)
        gt_traj, pred_traj = sync.associate_trajectories(gt_traj, pred_traj)

        if align:
            pred_traj.align(gt_traj, correct_scale=correct_scale)

    plot_collection = plot.PlotCollection("PlotCol")
    fig = plt.figure(figsize=(8, 8))
    plot_mode = plot.PlotMode.xz # ideal for planar movement
    ax = plot.prepare_axis(fig, plot_mode)
    ax.set_title(title)
    if gt_traj is not None:
        plot.traj(ax, plot_mode, gt_traj, '--', 'gray', "Ground Truth")
    plot.traj(ax, plot_mode, pred_traj, '-', 'blue', "Predicted")
    plot_collection.add_figure("traj (error)", fig)
    plot_collection.export(filename, confirm_overwrite=False)
    plt.close(fig=fig)
    print(f"Saved {filename}")

def save_output_for_COLMAP(name: str, traj: PoseTrajectory3D, points: np.ndarray, colors: np.ndarray, fx, fy, cx, cy, H=480, W=640):
    """ Saves the sparse point cloud and camera poses such that it can be opened in COLMAP """

    colmap_dir = Path(name)
    colmap_dir.mkdir(exist_ok=True)
    scale = 10 # for visualization

    # images
    images = ""
    traj = PoseTrajectory3D(poses_se3=list(map(np.linalg.inv, traj.poses_se3)), timestamps=traj.timestamps)
    for idx, (x,y,z), (qw, qx, qy, qz) in zip(range(1,traj.num_poses+1), traj.positions_xyz*scale, traj.orientations_quat_wxyz):
        images += f"{idx} {qw} {qx} {qy} {qz} {x} {y} {z} 1\n\n"
    (colmap_dir / "images.txt").write_text(images)

    # points
    points3D = ""
    colors_uint = (colors * 255).astype(np.uint8).tolist()
    for i, (p,c) in enumerate(zip((points*scale).tolist(), colors_uint), start=1):
        points3D += f"{i} " + ' '.join(map(str, p + c)) + " 0.0 0 0 0 0 0 0\n"
    (colmap_dir / "points3D.txt").write_text(points3D)

    # camera
    (colmap_dir / "cameras.txt").write_text(f"1 PINHOLE {W} {H} {fx} {fy} {cx} {cy}")
    print(f"Saved COLMAP-compatible reconstruction in {colmap_dir.resolve()}")

def save_ply(name: str, points: np.ndarray, colors: np.ndarray):
    points_ply = np.array([(x,y,z,r,g,b) for (x,y,z),(r,g,b) in zip(points, colors)],
                        dtype=[('x', '<f4'), ('y', '<f4'), ('z', '<f4'),('red', 'u1'), ('green', 'u1'),('blue', 'u1')])
    el = PlyElement.describe(points_ply, 'vertex',{'some_property': 'f8'},{'some_property': 'u4'})
    PlyData([el], text=True).write(f"{name}.ply")
    print(f"Saved {name}.ply")