#!/usr/bin/python

#

# Plots the results from the 3D pose graph optimization. It will draw a line

# between consecutive vertices.  The commandline expects two optional filenames:

#

#   ./plot_results.py --initial_poses filename  --optimized_poses filename

#

# The files have the following format:

#   ID x y z q_x q_y q_z q_w

from mpl_toolkits.mplot3d import Axes3D

import matplotlib.pyplot as plot

import numpy

import sys

from optparse import OptionParser

def set_axes_equal(axes):

    ''' Sets the axes of a 3D plot to have equal scale. '''

    x_limits = axes.get_xlim3d()

    y_limits = axes.get_ylim3d()

    z_limits = axes.get_zlim3d()

    x_range = abs(x_limits[1] - x_limits[0])

    x_middle = numpy.mean(x_limits)

    y_range = abs(y_limits[1] - y_limits[0])

    y_middle = numpy.mean(y_limits)

    z_range = abs(z_limits[1] - z_limits[0])

    z_middle = numpy.mean(z_limits)

    length = 0.5 * max([x_range, y_range, z_range])

    axes.set_xlim3d([x_middle - length, x_middle + length])

    axes.set_ylim3d([y_middle - length, y_middle + length])

    axes.set_zlim3d([z_middle - length, z_middle + length])

parser = OptionParser()

parser.add_option("--initial_poses", dest="initial_poses",

                  default="", help="The filename that contains the original poses.")

parser.add_option("--optimized_poses", dest="optimized_poses",

                  default="", help="The filename that contains the optimized poses.")

parser.add_option("-e", "--axes_equal", action="store_true", dest="axes_equal",

                  default="", help="Make the plot axes equal.")

(options, args) = parser.parse_args()

# Read the original and optimized poses files.

poses_original = None

if options.initial_poses != '':

  poses_original = numpy.genfromtxt(options.initial_poses,

                                    usecols = (1, 2, 3))

poses_optimized = None

if options.optimized_poses != '':

  poses_optimized = numpy.genfromtxt(options.optimized_poses,

                                     usecols = (1, 2, 3))

# Plots the results for the specified poses.

figure = plot.figure()

if poses_original is not None:

  axes = plot.subplot(1, 2, 1, projection='3d')

  plot.plot(poses_original[:, 0], poses_original[:, 1], poses_original[:, 2],

            '-', alpha=0.5, color="green")

  plot.title('Original')

  if options.axes_equal:

    axes.set_aspect('equal')

    set_axes_equal(axes)

if poses_optimized is not None:

  axes = plot.subplot(1, 2, 2, projection='3d')

  plot.plot(poses_optimized[:, 0], poses_optimized[:, 1], poses_optimized[:, 2],

            '-', alpha=0.5, color="blue")

  plot.title('Optimized')

  if options.axes_equal:

    axes.set_aspect('equal')

    set_axes_equal(plot.gca())

# Show the plot and wait for the user to close.

plot.show()