from torch import nn
import torch
from modules.GRUCell_with_hyp_bias import GRUCell_with_hyp_bias


class MemoryUpdater(nn.Module):
  def update_memory(self, unique_node_ids, unique_messages, timestamps):
    pass


class GRUMemoryUpdater(MemoryUpdater):
  def __init__(self, memory, message_dimension, memory_dimension, device,manifold,c):
    super(GRUMemoryUpdater, self).__init__()
    self.memory = memory
    self.layer_norm = torch.nn.LayerNorm(memory_dimension)
    self.message_dimension = message_dimension
    self.device = device
    self.manifold = manifold
    self.c =c
    if self.manifold == "Euclidean":
        #original version
        self.memory_updater = nn.GRUCell(input_size=message_dimension,
                                         hidden_size=memory_dimension)
    elif self.manifold == "Hyperboloid":
        self.memory_updater = GRUCell_with_hyp_bias(manifold=manifold, c=c,
                                                    input_size=message_dimension,
                                                    hidden_size=memory_dimension)
    elif self.manifold == "PoincareBall":
        self.memory_updater = GRUCell_with_hyp_bias(manifold=manifold, c=c,
                                                    input_size=message_dimension,
                                                    hidden_size=memory_dimension)


  def update_memory(self, unique_node_ids, unique_messages, timestamps):
    if len(unique_node_ids) <= 0:
      return

    assert (self.memory.get_last_update(unique_node_ids) <= timestamps).all().item(), "Trying to " \
                                                                                     "update memory to time in the past"

    memory = self.memory.get_memory(unique_node_ids)
    self.memory.last_update[unique_node_ids] = timestamps

    updated_memory = self.memory_updater(unique_messages, memory)

    self.memory.set_memory(unique_node_ids, updated_memory)

  def get_updated_memory(self, unique_node_ids, unique_messages, timestamps):
    if len(unique_node_ids) <= 0:
      return self.memory.memory.data.clone(), self.memory.last_update.data.clone()

    assert (self.memory.get_last_update(unique_node_ids) <= timestamps).all().item(), "Trying to " \
                                                                                     "update memory to time in the past"

    updated_memory = self.memory.memory.data.clone()
    updated_memory[unique_node_ids] = self.memory_updater(unique_messages, updated_memory[unique_node_ids])

    updated_last_update = self.memory.last_update.data.clone()
    updated_last_update[unique_node_ids] = timestamps

    return updated_memory, updated_last_update
