'''
Created on Jul 24, 2011
@author: Dilum Bandara
@contact: dilumb@engr.colostate.edu
@version: 0.1
@copyright: Copyright(c) 2011, Dilum Bandara
@license: GPL
'''

from Node import Node
from Query import Query
from TreeNode import TreeNode

import sys

#python Simulate.py [num of nodes] [c_n] [file name]
cachingScheme = 'Our' #Our or PoPCache or None
n = int(sys.argv[1]) #No of nodes
c_n = int(sys.argv[2]) #Cache capacity of a node 
fIn = 'QueryTrace' + sys.argv[3] + '.txt'
fOut = 'TreeDistri.txt'

path = []   #These are made global variables for memory efficiency
nodes = {}
keys = {}


def readQueries(fileName):
    '''
    Read query file & return relevant data
    @param fileName: Name of file to extract data from
    @return: Number of queries  
    '''
    numQueries = 0
    totalHops = 0
    fdIn = open(fIn, 'r')
    for line in fdIn: #Extract path information
        tmp = line.split()
        if len(tmp) == 0: print 'Blank line found' #Skip blank lines
        if tmp[0] in keys: keys[tmp[0]][0] += 1 #Track keys
        else: keys[tmp[0]] = [1, 0.0, 0] #[no of queries, frequency, no of cache entries]
        numQueries += 1 #Track no of keys 
        path.append(Query(tmp[0], tmp[1:-1]))
        totalHops += int(tmp[-1])
    fdIn.close()
    return [numQueries, totalHops]

    
def generateNodes():
    '''
    Extract path information & create nodes with relevant data
    @param path: Path information as a list
    @return: None 
    '''
    for i in path:
        for j in range(len(i.path)):
            if i.path[j] not in nodes: #Check whether node is already in list
                nodes[i.path[j]] = Node(i.path[j]) #Add new node
            if j > 0  and (i.path[j] != i.path[j - 1]): #Not the source node & src & dst not same
                if i.key not in nodes[i.path[j]].traffic: #New key
                    nodes[i.path[j]].traffic[i.key] = {i.path[j - 1]: 1}
                else: #Update info for old key
                    if i.path[j - 1] not in nodes[i.path[j]].traffic[i.key]:
                        if len(nodes[i.path[j]].traffic[i.key]) > 0:
                            nodes[i.path[j]].traffic[i.key][i.path[j - 1]] = 1
                        else:
                            nodes[i.path[j]].traffic[i.key] = {i.path[j - 1]: 1}
                    else: #Already there
                        nodes[i.path[j]].traffic[i.key][i.path[j - 1]] += 1
            if j == len(i.path) - 1: #Add key indexed as successor
                if i.key not in nodes[i.path[j]].idx: #If not already in index
                    nodes[i.path[j]].idx.append(i.key)

    
def generateKeys(numQueries, cachingScheme, n, budget):
    '''
    Given a list of keys & frequencies, generate normalized frequencies & number of 
    cache entries for key 
    @param numQueries: Number of queries
    @param cachingScheme: Which caching scheme to use - Only 'Our', 'PoPCache',
    & 'None' is valid
    @param n: Number of nodes   
    @return: Updated list of keys with frequencies & num of cache entries per key 
    ''' 
    for i in keys: keys[i][1] = keys[i][0]/float(numQueries) #Normalize frequency
    
    total_c_k = 0.0
    if cachingScheme == 'Our': #Our optimized cache capacity allocation
        limit_f_k = 0.0     #p in eqn
        limit_c_k = 0.0     # s(n - 1) in eqn
        for i in keys: #For each key
            keys[i][2] = int(round((keys[i][1] * (budget - limit_c_k))/(1 - limit_f_k)))
            #keys[i][2] = int(math.ceil((keys[i][1] * (budget - limit_c_k))/(1 - limit_f_k)))
            if keys[i][2] >= n - 1: 
                keys[i][2] = n - 1 # Cap number of caching entries to (N - 1)
                limit_f_k += keys[i][1] #Sum popularity
                limit_c_k += keys[i][2] #Sum allocated cache capacity
            total_c_k += keys[i][2]
            if total_c_k >= budget: #Stop when budget is reached
                keys[i][2] = keys[i][2] - int(total_c_k - budget)
                break
    elif cachingScheme == 'PoPCache': #PoPCache cache capacity allocation
        for i in keys:
            keys[i][2] = int(round(keys[i][1] * budget))
            #keys[i][2] = int(math.ceil(keys[i][1] * budget))
            if keys[i][2] > n - 1: keys[i][2] = n - 1 # Cap no of cache entries to (N - 1)
            total_c_k += keys[i][2]
            if total_c_k >= budget: #Stop when budget is reached 
                keys[i][2] = keys[i][2] - int(total_c_k - budget)
                break
    elif cachingScheme != 'None':
        print 'Unknown caching scheme' + cachingScheme
        sys.exit(0) 


def traverseTree(nodeSet, key, parentID):
    '''
    Using depth first approach to traverse the tree & build it
    @param nodeSet: Set of child nodes at a given parent node
    @param key: Key 
    @param parentID: ID of the parent
    @return: child nodes
    '''
    children = []
    for i in nodeSet: #For each sender do recursively
        tmpNode = TreeNode(i, nodeSet[i], parentID)
        if key in nodes[i].traffic:
            tmpNode.children = traverseTree(nodes[i].traffic[key], key, i)
        children.append(tmpNode)
    return children

def generateKeyTree():
    '''
    Generate tree for each key starting from successor of key
    @return: Return dictionary of trees for each key 
    '''
    trees = {}
    for i in nodes: #For each node
        for j in nodes[i].idx: #For each key
            tree = TreeNode(i, -1) #Root node
            if j in nodes[i].traffic: #Make sure traffic exists
                tree.children = traverseTree(nodes[i].traffic[j], j, i)
            trees[j] = tree
    return trees

def findAllNodesinTree(tree, allNodes):
    '''
    Find all the nodes in tree
    @param tree: Tree to span
    @param allNodes: All the nodes
    @return: list of [node id, demand] 
    '''
    for i in range(len(tree.children)): #For each children
        allNodes[tree.children[i].id] = tree.children[i].value
        findAllNodesinTree(tree.children[i], allNodes)
    return allNodes

def findAllNodesinTreeLevel(tree, allNodes, level):
    '''
    Find all the nodes in tree & group them accourding to the level they are in
    @param tree: Tree to span
    @param allNodes: All the nodes
    @param level: Level 
    @return: list of [node id, demand] 
    '''
    for i in range(len(tree.children)): #For each children
        if level not in allNodes: 
            allNodes[level] = {tree.children[i].id: tree.children[i].value}
        else:
            allNodes[level][tree.children[i].id] = tree.children[i].value
        findAllNodesinTreeLevel(tree.children[i], allNodes, level + 1)
    return allNodes
    
def findWeightDistriinTree(keyTree):
    '''
    Find weights for branches in each tree corrosponding to a key 
    @param keyTree: key tree for each key 
    @return: list of [node id, demand] 
    '''
    fdOut = open(fOut, 'w')
    for i in keys: #For each key
        if keys[i][0] < 500: continue #Skip keys with lower no of queries
        allNodes = {}
        allNodes = findAllNodesinTreeLevel(keyTree[i], allNodes, 0)
        j = 1  
        sortedNodes = sorted(allNodes[j], key=allNodes.get, reverse = True)
        #for j in range(len(allNodes)):  
        for k in range(len(sortedNodes)):
            allNodes[j][sortedNodes[k]] = allNodes[j][sortedNodes[k]]/float(keys[i][0])
            #allNodes[j][k] = nodes[k].traffic[i]/float(keys[i][0])
            fdOut.write(str(allNodes[j][sortedNodes[k]]) + '\t')
        #for k in allNodes[j]:
        #    allNodes[j][k] = allNodes[j][k]/float(keys[i][0])
            #allNodes[j][k] = nodes[k].traffic[i]/float(keys[i][0])
        #    fdOut.write(str(allNodes[j][k]) + '\t')
        fdOut.write('\n')
    fdOut.close()

def assignCache(cachingScheme, keyTree):
    '''
    Assign cache content to nodes
    @param cachingScheme: What caching scheme to use 
    @param keyTree: key tree for each key 
    '''
    if cachingScheme == 'None': return #Do nothing
    elif cachingScheme != 'Our' and cachingScheme != 'PoPCache':
        print 'Unknown caching scheme' + cachingScheme
        sys.exit(0) 

    for i in keys: #For each key
        if keys[i][2] == 0: #Nothing to assign
            continue
        if cachingScheme == 'Our':
            allNodes = {}
            allNodes = findAllNodesinTree(keyTree[i], allNodes)
            #Sort based on demand
            sortedNodes = sorted(allNodes, key=allNodes.get, reverse = True) #TODO check sorting
            
            if len(sortedNodes) < keys[i][2]: 
                print 'Not enough nodes to assigned', len(sortedNodes), keys[i][2]
                numCache = len(sortedNodes)
            else: numCache = keys[i][2]
           
            for j in range(numCache): #For required no of nodes add key to index
                nodes[sortedNodes[j]].idx.append(i) 
        elif cachingScheme == 'PoPCache':
            allNodes = {}
            allNodes = findAllNodesinTreeLevel(keyTree[i], allNodes, 0)
            #Try all nodes starting from lowest level until desired no of nodes are found
            numCache = 0
            allocatedSofar = 0
            for j in range(len(allNodes)):  
                numCache += len(allNodes[j])
                if numCache >= keys[i][2]: #Found enough nodes 
                    #Sort based on demand
                    sortedNodes = sorted(allNodes[j], key=allNodes.get, reverse = True) #TODO check sorting
                    for k in range(keys[i][2] - allocatedSofar):
                        nodes[sortedNodes[k]].idx.append(i) #Add to index
                    break
                #Need more cache copies
                for k in allNodes[j]: #For each node at current level
                    nodes[k].idx.append(i) #add key to index
                    allocatedSofar += 1  
            if numCache < keys[i][2]:
                print 'Not enough nodes to assign', numCache, keys[i][2]               
        

def query():
    '''
    Replay queries & find new hop count
    @return: List of hops for each query
    '''
    hops = []
    for i in range(len(path)): #For each key
        for j in range(len(path[i].path) - 1): #For each node along path
            if path[i].path[j] == path[i].path[j + 1]: #Source & destination is same
                hops.append(0) #Source destination is same node
                break
            elif path[i].key in nodes[path[i].path[j]].idx:
                hops.append(0) #Indexed in same node due to caching
                break
            elif path[i].key in nodes[path[i].path[j + 1]].idx:
                hops.append(j + 1) #at least 1 hop + ...
                break
    return hops
            
    
if __name__ == '__main__':
    results = readQueries(fIn)
    numQueries = results[0] 
    #print 'No of keys: ' + str(len(keys))
    #print 'No of queries: ' + str(numQueries)
    generateNodes()
    generateKeys(numQueries, cachingScheme, n, n * c_n)
    keyTree = generateKeyTree()
    #findWeightDistriinTree(keyTree)
    #sys.exit(0)
    assignCache(cachingScheme, keyTree)
    hops = query()
    #Calculate stat
    hops_total = 0.0
    hops_ave = 0.0
    hops_min = 999
    hops_max = 0
    for i in hops: 
        hops_total += i
        if i < hops_min: hops_min = i 
        if i > hops_max: hops_max = i
    hops_ave = hops_total/len(hops)
    #print 'Hops [ave, min, max]:', hops_ave, hops_min, hops_max
    idx_total = 0.0
    idx_ave = 0.0
    idx_min = 99999
    idx_max = 0
    for i in nodes:
        idx_total += len(nodes[i].idx)
        if len(nodes[i].idx) < idx_min: idx_min = len(nodes[i].idx)
        if len(nodes[i].idx) > idx_max: idx_max = len(nodes[i].idx)
    idx_ave = idx_total/len(nodes)
    #print 'Index [ave, min, max]:', idx_ave, idx_min, idx_max
    print len(keys), numQueries, results[1]/float(numQueries), hops_ave, hops_min, hops_max, idx_ave, idx_min, idx_max
