# external modules
import yaml
import collections
import numpy as np

# local modules
import constants as c
from geometryUtils import get2dRotmatFromYaw


class LandmarkMap(object):

    def __init__(self, mapFilename):
        with open(mapFilename, 'r') as mf:
            mapYaml = yaml.load(mf)

        if len(mapYaml) == 0:
            # an empty map is useless
            raise ValueError

        # TODO would be nice to do some indexing for speed here
        self.landmarks = []
        for landmarkYaml in mapYaml[c.MAP_LM_KEY]:
            # create a new landmark and add it to the collection
            landmark = Landmark(landmarkYaml[c.LM_FEAT_KEY],
                                landmarkYaml[c.LM_POS_KEY])
            if c.LM_ORNT_KEY in landmarkYaml:
                landmark.setOrientation(landmarkYaml[c.LM_ORNT_KEY])
            self.landmarks.append(landmark)
        self.size = np.array(mapYaml[c.MAP_SIZE_KEY])

    def getLandmarkMatches(self, landmark, N=1):
        """Returns the N best matches in the map to the given landmark's
        descriptor vector.  If N is not provided, it defaults to 1."""
        sortedLandmarks = sorted(self.landmarks,
                                 key=landmark.getFeatureDistance)
        return sortedLandmarks[:N]

    def getClosestLandmarks(self, point, N=1):
        """Returns the N closest landmarks in the map to the given point.
        If N is not provided, it defaults to 1."""
        sortedLandmarks = sorted(self.landmarks,
                                 key=lambda x: x.getDistanceFromPoint(point))
        return sortedLandmarks[:N]

    def getNDims(self):
        return len(self.landmarks[0].position)


class Landmark(object):

    def __init__(self, featureVec, position, orientation=None, label=None):
        self.featureVec = np.array(featureVec)
        self.position = np.array(position)
        if orientation is not None:
            self.orientation = np.array(orientation)
        else:
            self.orientation = None

        # this is not used yet
        self.label = label

    def getColour(self):
        if len(self.featureVec) < 3:
            # default colour channel values are maximum values
            colour = np.ones(3)
            colour[:len(self.featureVec)] = self.featureVec
            return colour
        else:
            return self.featureVec[:3]

    def getSize(self):
        return 0.2
        # if len(self.featureVec) < 4:
        #     # default size is max size
        #     return 1.
        # else:
        #     return self.featureVec[3]

    def setOrientation(self, newOrientation):
        self.orientation = newOrientation

    def getFeatureVec(self):
        return self.featureVec

    def getLabel(self):
        return self.label

    def getRotatedLandmark(self, rotmat):
        newPos = np.dot(rotmat, self.position)
        newOrnt = None
        if self.orientation is not None:
            newOrnt = np.dot(rotmat, self.orientation)
        return Landmark(self.featureVec, newPos, newOrnt, self.label)

    def getFeatureDistance(self, otherLandmark):
        return np.linalg.norm(self.featureVec - otherLandmark.featureVec)

    def getDistanceFromPoint(self, point):
        return np.linalg.norm(self.position - point)

    def getLandmarkInWorldCoords(self, camPos, camOrientation):
        # TODO implement case where cam orientation has roll and pitch as well
        if not isinstance(camOrientation, collections.Iterable):
            rotmat = get2dRotmatFromYaw(-camOrientation)
        # TODO the above assumes the self.position is always 2D
        newPos = (np.dot(rotmat, self.position) * camPos[2]) + camPos[:2]
        newOrnt = None
        if self.orientation is not None:
            newOrnt = np.dot(rotmat, self.orientation)
        return Landmark(self.featureVec, newPos, newOrnt, self.label)

    def getLandmarkInCamCoords(self, camPos, camOrientation):
        # TODO implement case where cam orientation has roll and pitch as well
        if not isinstance(camOrientation, collections.Iterable):
            # use negative angle to rotate into camera space
            rotmat = get2dRotmatFromYaw(camOrientation)
        # TODO the below assumes the self.position is always 2D
        newPos = np.dot(rotmat, (self.position - camPos[:2])) / camPos[2]
        newOrnt = None
        if self.orientation is not None:
            newOrnt = np.dot(rotmat, self.orientation)
        return Landmark(self.featureVec, newPos, newOrnt, self.label)