# external modules
import cv2
import numpy as np

# internal modules
from LandmarkMap import LandmarkMap

VIZ_SCALE = 50
TRUE_COLOUR = np.array((0., 1., 0.))
PROB_COLOUR = np.array((1., 1., 1.))
TEXT_OFFSET = 13

# read in simulator config
# read in the map
# create an output image at a given scale
# display objects on the output image according to their parameters
# publish image to a topic


class Visualizer(object):

    def __init__(self, mapFilename, viewWidth, viewHeight):
        self.landmarkMap = LandmarkMap(mapFilename)
        self.mapImage = renderMap(self.landmarkMap)
        self.viewWidth = viewWidth
        self.viewHeight = viewHeight

    def renderSimulation(self, truePos, trueYaw, estProbs, estPoss, estYaws,
                         seenLandmarks):
        simImage = self.mapImage.copy()
        renderPose(simImage, self.viewWidth, self.viewHeight,
                   truePos, trueYaw)
        for prob, pos, yaw in zip(estProbs, estPoss, estYaws):
            # draw the view given the position and yaw
            renderPose(simImage, self.viewWidth, self.viewHeight,
                       pos, yaw, prob)
        for sl in seenLandmarks:
            slPos = sl.getLandmarkInWorldCoords(truePos, trueYaw).position
            slPos = (slPos * VIZ_SCALE).astype(int)
            cv2.rectangle(simImage, tuple(slPos[:2] - 2*np.ones(2, dtype=int)),
                          tuple(slPos[:2] + 2*np.ones(2, dtype=int)),
                          (0.,0.,0.))
        return simImage


def renderPose(mapImage, viewWidth, viewHeight, pos, yaw, prob=None):
    # draw a dot at the center with an arrow showing the pos and yaw
    pos = (pos * VIZ_SCALE).astype(int)
    colour = TRUE_COLOUR
    if prob is not None:
        colour = PROB_COLOUR
        
    cv2.rectangle(mapImage, tuple(pos[:2] - np.ones(2, dtype=int)),
                  tuple(pos[:2] + np.ones(2, dtype=int)),
                  colour, thickness=-1)
    direction = VIZ_SCALE/2 * np.array([np.cos(-yaw), np.sin(-yaw)])
    cv2.arrowedLine(mapImage, tuple(pos[:2]),
                    tuple(pos[:2] + direction.astype(int)),
                    colour, tipLength=0.4)
    
    # if there is a probability, display it over/under the dot
    if prob is not None:
        fontScale = .8
        if yaw > np.pi:
            # top
            textHeight = int(12 * fontScale)
            textPoint = pos[:2] + (-TEXT_OFFSET, textHeight - TEXT_OFFSET)
        else:
            # bottom
            textPoint = pos[:2] + (-TEXT_OFFSET, TEXT_OFFSET)
        probStr = "%.2f" % prob
        cv2.putText(mapImage, probStr, tuple(textPoint),
                    cv2.FONT_HERSHEY_PLAIN, fontScale, colour)

    # draw lines showing the view of the class
    dispWidth = viewWidth * VIZ_SCALE
    dispHeight = viewHeight * VIZ_SCALE
    cornersCwFromBottomRight = np.zeros((4, 2))
    bottomRightCorner = np.array([dispWidth / 2, dispHeight / 2])
    cornersCwFromBottomRight[0, :] = bottomRightCorner
    cornersCwFromBottomRight[1, :] = bottomRightCorner * [-1, 1]  # bottom left
    cornersCwFromBottomRight[2, :] = bottomRightCorner * [-1, -1] # top left
    cornersCwFromBottomRight[3, :] = bottomRightCorner * [1, -1]  # top right
    rotmat = np.array([[np.cos(yaw), np.sin(yaw)],
                       [-np.sin(yaw), np.cos(yaw)]])
    cornersRot = np.dot(rotmat, cornersCwFromBottomRight.T).T.astype(int)
    cornersRot += pos[:2]

    lineColour = colour
    if prob is not None:
        lineColour = lineColour * prob
    for i in range(len(cornersRot) - 1):
        cv2.line(mapImage, tuple(cornersRot[i,: ]), tuple(cornersRot[i+1, :]),
                 lineColour)
    cv2.line(mapImage, tuple(cornersRot[-1, :]), tuple(cornersRot[0, :]),
             lineColour)

        
def renderMap(landmarkMap):
    if len(landmarkMap.size) != 2:
        raise ValueError("This visualizer only works for 2D maps!")
        
    mapImageShape = (VIZ_SCALE * landmarkMap.size).astype(int)
    mapImage = np.zeros((mapImageShape[1], mapImageShape[0], 3))
    # render each object
    for landmark in landmarkMap.landmarks:
        colour = landmark.getColour()
        size = int(VIZ_SCALE * landmark.getSize())
        pos = (landmark.position * VIZ_SCALE).astype(int)
        # negative thickness draws a filled circle
        cv2.circle(mapImage, tuple(pos), size, colour, thickness=-1)

    return mapImage