import os
import vtkAll as vtk
import math
import time
import re
import numpy as np

from ddapp.timercallback import TimerCallback
from ddapp import objectmodel as om
from ddapp.simpletimer import SimpleTimer
from ddapp.utime import getUtime
from ddapp import robotstate

import pickle
import scipy.interpolate


class PlanPlayback(object):

    def __init__(self):
        self.animationCallback = None
        self.animationTimer = None
        self.interpolationMethod = 'slinear'
        self.playbackSpeed = 1.0
        self.jointNameRegex = ''

    @staticmethod
    def getPlanPoses(msgOrList):

        if isinstance(msgOrList, list):
            messages = msgOrList
            allPoseTimes, allPoses = PlanPlayback.getPlanPoses(messages[0])

            for msg in messages[1:]:
                poseTimes, poses = PlanPlayback.getPlanPoses(msg)
                poseTimes += allPoseTimes[-1]
                allPoseTimes = np.hstack((allPoseTimes, poseTimes[1:]))
                allPoses += poses[1:]
            return allPoseTimes, allPoses

        else:
            msg = msgOrList

            poses = []
            poseTimes = []
            for plan in msg.plan:
                pose = robotstate.convertStateMessageToDrakePose(plan)
                poseTimes.append(plan.utime / 1e6)
                poses.append(pose)
            return np.array(poseTimes), poses

    @staticmethod
    def getPlanElapsedTime(msg):
        startTime = msg.plan[0].utime
        endTime = msg.plan[-1].utime
        return (endTime - startTime) / 1e6


    def stopAnimation(self):
        if self.animationTimer:
            self.animationTimer.stop()


    def setInterpolationMethod(method):
        self.interpolationMethod = method


    def playPlan(self, msg, jointController):

        self.playPlans(poseTimes, [msg], jointController)


    def playPlans(self, messages, jointController):

        assert len(messages)

        poseTimes, poses = self.getPlanPoses(messages)
        self.playPoses(poseTimes, poses, jointController)


    def getPoseInterpolatorFromPlan(self, message):
        poseTimes, poses = self.getPlanPoses(message)
        return self.getPoseInterpolator(poseTimes, poses)


    def getPoseInterpolator(self, poseTimes, poses, unwrap_rpy=True):
        if unwrap_rpy:
            poses = np.array(poses, copy=True)
            poses[:,3:6] = np.unwrap(poses[:,3:6],axis=0)

        if self.interpolationMethod in ['slinear', 'quadratic', 'cubic']:
            f = scipy.interpolate.interp1d(poseTimes, poses, axis=0, kind=self.interpolationMethod)
        elif self.interpolationMethod == 'pchip':
            f = scipy.interpolate.pchip(poseTimes, poses, axis=0)
        return f


    def getPlanPoseMeshes(self, messages, jointController, robotModel, numberOfSamples):

        poseTimes, poses = self.getPlanPoses(messages)
        f = self.getPoseInterpolator(poseTimes, poses)
        sampleTimes = np.linspace(poseTimes[0], poseTimes[-1], numberOfSamples)
        meshes = []

        for sampleTime in sampleTimes:

            pose = f(sampleTime)
            jointController.setPose('plan_playback', pose)
            polyData = vtk.vtkPolyData()
            robotModel.model.getModelMesh(polyData)
            meshes.append(polyData)

        return meshes


    def showPoseAtTime(self, time, jointController, poseInterpolator):
        pose = poseInterpolator(time)
        jointController.setPose('plan_playback', pose)


    def playPoses(self, poseTimes, poses, jointController):

        f = self.getPoseInterpolator(poseTimes, poses)

        timer = SimpleTimer()

        def updateAnimation():

            tNow = timer.elapsed() * self.playbackSpeed

            if tNow > poseTimes[-1]:
                pose = poses[-1]
                jointController.setPose('plan_playback', pose)

                if self.animationCallback:
                    self.animationCallback()

                return False

            pose = f(tNow)
            jointController.setPose('plan_playback', pose)

            if self.animationCallback:
                self.animationCallback()

        self.animationTimer = TimerCallback()
        self.animationTimer.targetFps = 60
        self.animationTimer.callback = updateAnimation
        self.animationTimer.start()
        updateAnimation()


    def picklePlan(self, filename, msg):
        poseTimes, poses = self.getPlanPoses(msg)
        pickle.dump((poseTimes, poses), open(filename, 'w'))


    def getMovingJointNames(self, msg):
        poseTimes, poses = self.getPlanPoses(msg)
        diffs = np.diff(poses, axis=0)
        jointIds =  np.unique(np.where(diffs != 0.0)[1])
        jointNames = [robotstate.getDrakePoseJointNames()[jointId] for jointId in jointIds]
        return jointNames


    def plotPlan(self, msg):

        import matplotlib.pyplot as plt

        poseTimes, poses = self.getPlanPoses(msg)

        poses = np.array(poses)

        if self.jointNameRegex:
            jointIds = range(poses.shape[1])
        else:
            diffs = np.diff(poses, axis=0)
            jointIds = np.unique(np.where(diffs != 0.0)[1])

        jointNames = [robotstate.getDrakePoseJointNames()[jointId] for jointId in jointIds]
        jointTrajectories = [poses[:,jointId] for jointId in jointIds]

        seriesNames = []

        sampleResolutionInSeconds = 0.01
        numberOfSamples = (poseTimes[-1] - poseTimes[0]) / sampleResolutionInSeconds
        xnew = np.linspace(poseTimes[0], poseTimes[-1], numberOfSamples)

        fig = plt.figure()
        ax = fig.add_subplot(111)


        for jointId, jointName, jointTrajectory in zip(jointIds, jointNames, jointTrajectories):

            if self.jointNameRegex and not re.match(self.jointNameRegex, jointName):
                continue

            x = poseTimes
            y = jointTrajectory

            y = np.rad2deg(y)

            if self.interpolationMethod in ['slinear', 'quadratic', 'cubic']:
                f = scipy.interpolate.interp1d(x, y, kind=self.interpolationMethod)
            elif self.interpolationMethod == 'pchip':
                f = scipy.interpolate.pchip(x, y)

            ax.plot(x, y, 'ko')
            seriesNames.append(jointName + ' points')

            ax.plot(xnew, f(xnew), '-')
            seriesNames.append(jointName + ' ' + self.interpolationMethod)


        ax.legend(seriesNames, loc='upper right').draggable()
        ax.set_xlabel('time (s)')
        ax.set_ylabel('joint angle (deg)')
        ax.set_title('joint trajectories')
        plt.show()

@@ -4836,49 +4836,8 @@
     
-import matplotlib.pyplot as plt%0A%0A        
 pose
@@ -4873,16 +4873,143 @@
 ses(msg)
+%0A        self.plotPoses(poseTimes, poses)%0A%0A%0A    def plotPoses(self, poseTimes, poses):%0A%0A        import matplotlib.pyplot as plt
 %0A%0A      

add plotPoses method