#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Nov 20 09:47:51 2019

@author: hannes
"""

#General imports
import numpy as np
import networkx as nx

def getPath(robot):
    """Get the shortest path for the robot to the end location
    
    Input arguments:
    robot = robot instance whose graph we have to analize
    """

    graphReversed = robot.graph.reverse()
    predecessors = nx.dfs_successors(graphReversed,robot.endNodeCounter)
    
    path = []
    for key, value in predecessors.items():
        path.append(key)
        
    path.append(value[0])
    path.reverse()
    robot.paths.append(path)
    robot.path = robot.graph.subgraph(path)

def leastCostGoalSet(robots, debug, meeting=True):
    """Get the goal set node with the least cost

    Input arguments:
    robots = robots whose graph is to be checked
    debug = bool for debug information prints
    meeting = bool if we need meeting or not
    """

    robot = robots[0]
    if meeting:
        graph = robot.graph
        dictNodes = nx.get_node_attributes(graph,'goalSet')
        goalNodes = list(dictNodes.keys())
        
        goalNodes = np.asarray(goalNodes)
            
        times = []
        for node in goalNodes:
            times.append(graph.nodes[node]['t'])
            
        times = np.asarray(times)

        # TODO: maybe change to argmax
        robot.endLocation = graph.nodes[goalNodes[np.argmin(times)]]['pos']
        robot.endTotalTime = np.min(times)
        robot.endNodeCounter = goalNodes[np.argmin(times)]

        for r in range(1, len(robots)):
            graph = robots[r].graph
            endNode = robots[0].endNodeCounter
            robots[r].endLocation = graph.nodes[endNode]['pos']
            robots[r].endTotalTime = graph.nodes[endNode]['t']
            robots[r].endNodeCounter = endNode
    else:
        for r in range(0, len(robots)):
            graph = robots[r].graph
            endNode = robots[r].nodeCounter-1
            robots[r].endLocation = graph.nodes[endNode]['pos']
            robots[r].endTotalTime = graph.nodes[endNode]['t']
            robots[r].endNodeCounter = endNode

    if debug:
        print('leastCostGoalSet')
        for r in range(0, len(robots)):
            print('Robot ID: %d' %robots[r].ID)
            print('End Node: %d' %robots[r].endNodeCounter)
            print('End Location: ',robots[r].endLocation)
            print('End Time: %.2f\n' %robots[r].endTotalTime)

def checkGoalSet(graph):
    """Check if there are any nodes in the goal set
    
    Input arguments:
    graph = graph which is to be checked
    """

    dictNodes = nx.get_node_attributes(graph,'goalSet')
    goalNodes = list(dictNodes.keys())
    
    if not goalNodes:
        return False

    return True

def updateGoalSet(robots, timeInterval, commRadius, debug):
    """Check if goal set is still valid after rewiring
    
    Input arguments:
    robots = all robot instances
    timeInterval = arrival time interval for no delay constraint
    debug = bool for debug information prints
    """

    dictNodes = nx.get_node_attributes(robots[0].graph,'goalSet')
    goalNodes = list(dictNodes.keys())
    
    if not goalNodes:
        pass    
    goalNodes = np.asarray(goalNodes)
        
    for node in goalNodes:
        times = []
        positions = []
        for r in range(0, len(robots)): 
            times.append(robots[r].graph.nodes[node]['t'])
            positions.append(robots[r].graph.nodes[node]['pos'])

        times = np.asarray(times)
        positions = np.asarray(positions)
        
        timeDiff = np.abs(times - times[0])
        normDist = np.sqrt(np.sum((positions[0] - positions)**2, axis=1))

        if not all(x <= commRadius for x in normDist) and all(y <= timeInterval for y in timeDiff):
            
            if debug:
                print('Update goal set')
                print(timeDiff)
                print('Deleted goalsetNode')
                print(node)
            
            for r in range(0, len(robots)): 
                del robots[r].graph.nodes[node]['goalSet']                
    
    
def calculateGoalSet(robots, commRadius, timeInterval):
    """check if the new nodes are close to each other and if the time is right so that they belong to the goal set
    
    Input arguments:
    robots = all robot instances
    commRadius = communication range of the robots
    timeInterval = timeinterval for meetings to happen
    """
    positions = []
    times = []
    for r in range(0, len(robots)): 
        positions.append(robots[r].graph.nodes[robots[r].vnewIdx]['pos'])
        times.append(robots[r].graph.nodes[robots[r].vnewIdx]['t'])
    
    positions = np.asarray(positions)
    times = np.asarray(times)
    
    timeDiff = np.abs(times - times[0])
    normDist = np.sqrt(np.sum((positions[0] - positions)**2, axis=1))

    if all(x <= commRadius for x in normDist) and all(y <= timeInterval for y in timeDiff):
        for r in range(0, len(robots)): 
            robots[r].graph.nodes[robots[r].vnewIdx]['goalSet'] = True

def updateSuccessorNodes(robot, startNode):
    """Update successor node weight and times as the rewiring changed the graph structure
    
    Input arguments:
    robot = which robot graph that we are analyzing
    startNode = which node got changed
    """
    allSuccessors = nx.dfs_successors(robot.graph,startNode)

    for key, value in allSuccessors.items():
        for v in value:
            
            startTime = robot.graph.nodes[key]['t']
            startPos = robot.graph.nodes[key]['pos']
            succPos = robot.graph.nodes[v]['pos']
            succTotalCost, succEdgeCost, information = cost(startTime,startPos,succPos,robot)
            
            robot.graph[key][v]['weight'] = succEdgeCost
            robot.graph.nodes[v]['t'] = succTotalCost
            robot.graph.nodes[v]['informationGain'] = information             

def rewireGraph(robots, timeInterval, commRadius, debug):
    """Check if there is a node in setVnear which has a smaller cost as child from vnew than previous
    
    Input arguments:
    robots = which robot graph that we are analyzing
    timeInterval = interval for meeting time arrival
    debug = show values
    """

    rewired = False
    
    for r in range(0, len(robots)): 
        robot = robots[r]
        setVnear = robot.setVnear
        timeVnew = robot.totalTime
        posVnew = robot.vnew
            
        nearTotalCost = 0
        nearEdgeCost = 0
        information = 0
        
        for n in range(0,len(setVnear)):
            vnearTime = robot.graph.nodes[setVnear[n]]['t']
            
            vnearInformation = robot.graph.nodes[setVnear[n]]['informationGain']
            vnearUtility = getUtility(vnearTime,vnearInformation)
            
            vnearPos = robot.graph.nodes[setVnear[n]]['pos']
            nearTotalCost, nearEdgeCost, information = cost(timeVnew,posVnew,vnearPos,robot) 
            
            nearUtility = getUtility(nearTotalCost,information)
            
            if (nearTotalCost < vnearTime) or (nearUtility < vnearUtility):
                if list(robot.graph.pred[setVnear[n]]) == []:
                    continue
                predecessor = list(robot.graph.pred[setVnear[n]])[-1]
                robot.graph.remove_edge(predecessor,setVnear[n])
                robot.graph.add_edge(robot.vnewIdx,setVnear[n], weight = nearEdgeCost)
                robot.graph.nodes[setVnear[n]]['t'] = nearTotalCost
                robot.graph.nodes[setVnear[n]]['informationGain'] = information
                updateSuccessorNodes(robot, setVnear[n])
                rewired = True
                    
    if rewired:
        updateGoalSet(robots, timeInterval, commRadius, debug)
            
def cost(nearTime, nearPos, newPos, robot):
    """Calculate cost between two nodes
    
    Input arguments:
    nearTime = time at current near node
    nearPos = position at current near node
    newPos = position of vnew for which we search a less costly parent node
    robot = robot which is updating cost
    """

    normDist = np.sqrt(np.sum((nearPos - newPos)**2))
    nearEdgeCost = normDist/robot.uMax

    information = getInformationGain(robot,newPos)
#    nearEdgeCost = nearEdgeCost/information

    nearTotalCost = nearTime + nearEdgeCost

    return round(nearTotalCost,1), round(nearEdgeCost,1), information 

def extendGraph(robot):
    """Check if there is a node in setVnear which has a smaller cost as parent to vnew than the nearest node
    
    Input arguments:
    robot = which robot graph that we are analyzing
    """
    setVnear = robot.setVnear
    
    vnew = robot.vnew
    vminCost = robot.totalTime

    vminInformation = robot.vnewInformation
    
    vminUtility = getUtility(vminCost,vminInformation)
    
    nearTotalCost = 0
    vminEdgeCost = robot.vnewCost
    
    for n in range(0,len(setVnear)):
        time = robot.graph.nodes[setVnear[n]]['t']
        pos = robot.graph.nodes[setVnear[n]]['pos']
        nearTotalCost, nearEdgeCost, information = cost(time,pos,vnew,robot) 
        
        nearUtility = getUtility(nearTotalCost,information)

        if (nearTotalCost < vminCost) or (nearUtility < vminUtility):        
            robot.nearestNodeIdx = setVnear[n]
            vminCost = nearTotalCost
            vminEdgeCost = nearEdgeCost
            vminUtility = nearUtility
            vminInformation = information
        
    robot.vnewCost = vminEdgeCost
    robot.totalTime = vminCost
    robot.vnewInformation = vminInformation

def buildSetVnear(robot, epsilon, gammaRRT):
    """Examine all nodes and build set with nodes close to vnew with a radius of communication range
    
    Input arguments:
    robot = current robot with its graph
    epsilon = maximum allowed distance traveled
    gammaRRT = parameter for radius of near nodes in RRT star
    """
    setVnear = []
    
    vnew = robot.vnew
    cardinality = robot.graph.number_of_nodes()
    dimension = 2
    
    radius = min(gammaRRT*(pow(np.log(cardinality)/cardinality,1/dimension)), epsilon)
    
    dictNodes = nx.get_node_attributes(robot.graph,'pos')

    nodes = list(dictNodes.values())
    nodes = np.asarray(nodes)

    normDist = np.sqrt(np.sum((nodes - vnew)**2, axis=1))
    
    for n in range(0,len(nodes)):
        if normDist[n] < radius:
            setVnear.append(list(dictNodes.keys())[n])
    
    robot.setVnear = setVnear
    
def steer(robots, epsilon):
    """Steer towards vrand but only as much as allowed by the dynamics
    
    Input arguments:
    robots = robot classes
    epsilon = maximum allowed distance traveled
    """

    #find minTimes for nearest nodes of all robots
    minTimes = []
    for r in range(0, len(robots)): 
        nearestNodeIdx = robots[r].nearestNodeIdx
        graphDict = robots[r].graph.nodes[nearestNodeIdx]
        
        vnearest = list(graphDict.values())
        nearestTime = np.asarray(vnearest[1])
        
        minTimes.append(nearestTime)
    
    #steer towards vrand
    for r in range(0, len(robots)):   
        nearestNodeIdx = robots[r].nearestNodeIdx
        graphDict = robots[r].graph.nodes[nearestNodeIdx]
        
        vnearest = list(graphDict.values())
        
        vrand = robots[r].vrand
        nearestTime = np.asarray(vnearest[1])
        nearestNode = np.asarray(vnearest[0])
    
    
        dist = vrand - nearestNode 
        normDist = np.sqrt(np.sum((nearestNode - vrand)**2))

        s = min(epsilon,normDist)
        travelTime = s/robots[r].uMax

        deltaTcost = travelTime - (nearestTime - min(minTimes))
        
        if deltaTcost > 0:          
            vnew = np.around(nearestNode + robots[r].uMax*deltaTcost*dist/normDist)
            distVnew = np.sqrt(np.sum((nearestNode - vnew)**2))
            travelTimeVnew = distVnew/robots[r].uMax            
        else: 
            vnew = nearestNode
            travelTimeVnew = 0        
        
        totalTimeVnew = travelTimeVnew + nearestTime
        discretization = robots[r].discretization

        if not (0 <= vnew[0] < discretization[0] and 0 <= vnew[1] < discretization[1]):
            if vnew[0] >= discretization[0]:
                vnew[0] = discretization[0]-1
            if vnew[1] >= discretization[1]:
                vnew[1] = discretization[1]-1
            if vnew[0] < 0:
                vnew[0] = 0
            if vnew[1] < 0:
                vnew[1] = 0
        
        information = getInformationGain(robots[r], vnew)
        
        robots[r].vnew = np.around(vnew)
        robots[r].vnewCost = round(travelTimeVnew,1)
        robots[r].totalTime = round(totalTimeVnew,1)
        robots[r].vnewInformation = information
        
def findNearestNode(graph, vrand):
    """Return nearest node index
    
    Input arguments:
    graph = current graph
    vrand = new random node
    """
    dictNodes = nx.get_node_attributes(graph,'pos')

    nodes = list(dictNodes.values())
    nodes = np.asarray(nodes)

    normDist = np.sum((nodes - vrand)**2, axis=1)
    
    return list(dictNodes.keys())[np.argmin(normDist)]

def sampleVrand(discretization, rangeSamples, distribution = 'uniform'):
    """Sample a new random position
    
    Input arguments:
    discretization = grid space needed to check for boundary issues
    rangeSamples = defines uniform range boundaries or mean and std dev for gaussian
    distribution = which distribution to use
    """
    inBoundary = False
    while inBoundary == False:
        if distribution == 'uniform':
            vrand = np.random.uniform([0,0],rangeSamples)
        
        if distribution == 'gaussian':
            vrand = np.random.multivariate_normal(rangeSamples[0],rangeSamples[1])
        if 0 <= vrand[0] < discretization[0] and 0 <= vrand[1] < discretization[1]:
            inBoundary = True
  
    return np.around(vrand)

def getInformationGain(robot, pos):
    """Get expected variance from GP model, either by inferring or lookup
    
    Input arguments:
    robot = robot instance whose GP is to be looked at
    pos = position to check the expected variance for
    """
    
    if robot.optPath:
        ys = robot.expectedVariance[np.int(pos[0]),np.int(pos[1])]
    elif robot.optPoint:
        _, ys = robot.GP.inferGP(robot,pos)
    else:
        return 0
    return ys

# TODO: This doesn't seem to do anything, maybe remove it..
def getUtility(time, informationGain):
    """Utility function which calgulates utility by dividing cost by variance with weight
    
    Input arguments:
    time = cost of node
    informationGain = gain of node
    """

    beta = 4
    if informationGain < 1/beta:
        informationGain = 1/beta
    # return time/(beta*informationGain) 
    return 0
    
def sampleNPoints(robot, discretization, rangeSamples, distribution = 'uniform'):
    """Sample N new random positions and check for highest variance
    
    Input arguments:
    robot = robot instance to access GP
    discretization = grid space needed to check for boundary issues
    rangeSamples = defines uniform range boundaries or mean and std dev for gaussian
    distribution = which distribution to use
    """

    maxVariance = 0
    vrand = np.array([0, 0])
    for _ in range(0,10):
        point = sampleVrand(discretization, rangeSamples, distribution)
        var = getInformationGain(robot, point)
        if var > maxVariance:
            vrand = point
            maxVariance = var
    return vrand

def getInformationGainAlongPath(robot, pos, nearestNodeIdx, epsilon):
    """Calculate information gain along path for each sensor period
    
    Input arguments:
    robot = robot instance to access graph
    pos = vrand from sampling
    nearestNodeIdx = nearest node to vrand for steering
    epsilon = maximum step size
    """

    positionNotReached = True    
    deltaT = robot.deltaT
    discretization = robot.discretization
    
    graphDict = robot.graph.nodes[nearestNodeIdx]
    vnearest = list(graphDict.values())
    
    nearestNode = np.asarray(vnearest[0])
    currentTime = vnearest[1]
    measPos = nearestNode
    var = 0
    
    distance = pos - measPos
    normDist = np.sqrt(np.sum((pos - nearestNode)**2))
    s = min(epsilon,normDist)
    if normDist == 0:
        return 0
    
    vnew = np.around(nearestNode + s*distance/normDist)
    pos = vnew
    
    while positionNotReached:
        
        distance = pos - measPos
        normDist = np.sqrt(np.sum((pos - measPos)**2))
    
        travelTime = normDist/robot.uMax
    
        if travelTime <= deltaT:
            measPos = pos
            positionNotReached = False
        else:
            step = np.around(robot.uMax*deltaT*distance/normDist)
            measPos = measPos + step
            currentTime += deltaT
        
        if not (0 <= measPos[0] < discretization[0] and 0 <= measPos[1] < discretization[1]):
            if measPos[0] >= discretization[0]:
                measPos[0] = discretization[0]-1
            if measPos[1] >= discretization[1]:
                measPos[1] = discretization[1]-1
            if measPos[0] < 0:
                measPos[0] = 0
            if measPos[1] < 0:
                measPos[1] = 0
        
        xRange = np.array([robot.sensingRange, robot.sensingRange +1])
        yRange = np.array([robot.sensingRange, robot.sensingRange +1])
        
        if measPos[0] - robot.sensingRange < 0:
            xRange[0] = measPos[0]
        if measPos[0] + robot.sensingRange > discretization[0]:
            xRange[1] = discretization[0] - measPos[0]
        if measPos[1] - robot.sensingRange < 0:
            yRange[1] = measPos[1]
        if measPos[1] + robot.sensingRange > discretization[1]:
            yRange[1] = discretization[1] - measPos[1]
        
            
        ys = 0
        ym = 0        
        
        if robot.optPath:
            if currentTime % robot.sensorPeriod < 0.1:
                ys = robot.expectedVariance[np.int(measPos[0])-xRange[0]:np.int(measPos[0])-xRange[1],
                                            np.int(measPos[1])-yRange[0]:np.int(measPos[1])-yRange[1]]
                ym = robot.expectedMeasurement[np.int(measPos[0])-xRange[0]:np.int(measPos[0])-xRange[1],
                                            np.int(measPos[1])-yRange[0]:np.int(measPos[1])-yRange[1]]
                ys = np.sum(ys)
                ym = np.sum(ym)
        else:
            if currentTime % robot.sensorPeriod < 0.1:
                ym, ys = robot.GP.inferGP(robot,measPos)
        # TODO: was just ys, tried ys+ym as well and also ym * ys, plus doesn't work to well
        var += ys + 0.5*ym   
    
    return var
















