import networkx as nx
import numpy as np
import matplotlib.pyplot as plt
import utils

class Memory:
    def __init__(self,memory_size=1000,memory_word_size=32):
        self.Gmemory = nx.DiGraph()
        self.StateAttributesDict={}
        self.memory_size = memory_size
        self.memory_word_size = memory_word_size
        self.gamma = 0.9
        self.lr = 0.01   
             
        self._write_content_similarity = utils.CosineSimilarity(1,self.memory_word_size,name="write_content_similarity")
        self._TH = 0.99

    def plotMemory(self):
        nx.draw_networkx(self.Gmemory)

    
    def node_feature_list(self):
        F = [] #这个矩阵能不能在add pair的时候加入？
        # print("len()",len(self.Gmemory))
        for nodeid in self.Gmemory.nodes():
            # print("i",nodeid)
            # print("feature",self.Gmemory.node[nodeid]['feature'])
            F.append(self.Gmemory.nodes[nodeid]['feature'])
            
        return F

        
    def _get_location(self,cur_state,next_state):
        
        if self.Gmemory.number_of_nodes() == 0:
            print("an empty memory graph")
            cur_node = 0
            next_node =1
        else:
            cur_similarity = self._write_content_similarity.onevsall(cur_state,self.node_feature_list())
            next_similarity = self._write_content_similarity.onevsall(next_state,self.node_feature_list())
        
            if max(cur_similarity)<self._TH:
                if max(next_similarity)<self._TH:
                    #print("add new node",max(similarity))
                    cur_node = self.Gmemory.number_of_nodes()
                    next_node = self.Gmemory.number_of_nodes()+1      
                else:
                    cur_node = self.Gmemory.number_of_nodes()
                    next_node = next_similarity.index(max(next_similarity))
            else:
                if max(next_similarity)<self._TH:
                    cur_node = cur_similarity.index(max(cur_similarity))
                    next_node = self.Gmemory.number_of_nodes()
                else:
                    cur_node = cur_similarity.index(max(cur_similarity))
                    next_node = next_similarity.index(max(next_similarity))
            
        return cur_node,next_node
    def add_single_node(self,cur_node,cur_state,act,r):
        self.Gmemory.add_node(cur_node,feature=cur_state,reward=r)
        #self.StateAttributesDict.append(('cur_node':cur_state))

    def check_node_exist(self,cur_node):
        if cur_node in self.Gmemory:
            return True
        else:
            return False

    def update_edge(self,cur_node,act,reward,next_node):
        #不能只遵循Q原则更新边
        old_rew = self.Gmemory[cur_node][next_node]['rew']
        old_visits = self.Gmemory[cur_node][next_node]['visits']
        old_q = self.Gmemory[cur_node][next_node]['q']
        candidate_neighs_of_next_node = list(self.Gmemory[next_node])
        if len(candidate_neighs_of_next_node) ==0 :
            # 如果该节点是终点
            tar = reward # R作为目标
        else:
            weight_list = []
            for neighs_i in candidate_neighs_of_next_node:
                weight_list.append(self.Gmemory[next_node][neighs_i]['q'])
            weight_max = np.max(weight_list)
            tar = reward + self.gamma*weight_max
        delta_q = tar -old_q
        new_q = old_q + self.lr* delta_q
        self.Gmemory.add_edge(cur_node,next_node,label=act,rew=reward,visits=old_visits+1,q=new_q)

    def add_pair(self,cur_node,cur_state,act,r,next_node,next_state):
        #一次要加入两个，暂时只使用一个
        #self.Gmemory.add_node(loc,feature=state,reward=r)#如果之前有也直接替换，不考虑做中和吗？
        # 暂时还不能在这里加入，因为node 列表有可能被替代，而这个list不会，所以二者不等长
        #self.StateAttributesDict.append(state)
        # 每次用node_featrue_list从头捋一遍吧
        
        #if [cur_node,next_node] in self.Gmemory.edges():
        if self.Gmemory.has_edge(cur_node,next_node):
            #print("update edge *****************************")
            self.update_edge(cur_node,act,r,next_node)
        else:
            #增加节点和边
            #print("add node")
            if self.check_node_exist(cur_node):
                if self.check_node_exist(next_node):
                    pass
                else:
                    self.Gmemory.add_node(next_node,feature=next_state,reward=r)
            else:
                if self.check_node_exist(next_node):
                    self.Gmemory.add_node(cur_node,feature=cur_state,reward=0)
                else:
                    self.Gmemory.add_node(cur_node,feature=cur_state,reward=0)
                    self.Gmemory.add_node(next_node,feature=next_state,reward=r)
            self.Gmemory.add_edge(cur_node,next_node,label=act,rew=r,visits=1,q=r)


    def pairwriter(self,transition):
        [cur_state,act,r,next_state]=transition
        cur_node,next_node = self._get_location(cur_state,next_state)
        if cur_node == next_node:
            self.add_single_node(cur_node,cur_state,act,r)
        else:
            self.add_pair(cur_node,cur_state,act,r,next_node,next_state)