import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset
import time

from editable_model import EditableModel
from utils import _last_encoder_state, _logits

class LU(EditableModel):
    """
    Representation lookup approach. Does not require training.
    """

    def __init__(self, model, config, model_constructor, memory=None):
        super().__init__(model, config, model_constructor)

        self.memory = memory

    def forward(self, *inputs, **kwargs):
        if "bert" in self.config.model.name.lower():
            output, encoder_states = self.model(*inputs, **kwargs, output_hidden_states=True)
        else:
            model_output = self.model(*inputs, **kwargs, output_hidden_states=True)
            encoder_states = _last_encoder_state(model_output)
            output = _logits(model_output)

        if self.memory is not None:
            for i, encoder_state in enumerate(encoder_states):
                if "gpt2" in self.config.model.name.lower():
                    # NOTE: broken
                    memory_prefixes, memory_labels = self.memory
                    prefix_means = encoder_state.cumsum(0).detach() / torch.arange(1, encoder_state.shape[0] + 1, device=encoder_state.device).view(-1, 1)
                    dist_mat = (prefix_means.unsqueeze(1) - memory_prefixes.unsqueeze(0)).norm(2, dim=-1)

                    min_dists, min_idxs = dist_mat.min(-1)
                    memory_mask = (min_dists < self.config.lu.threshold)
                    onehot_logits = self.config.lu.onehot_logit * F.one_hot(memory_labels[min_idxs], output.shape[-1]).float()
                    output[i, memory_mask] = onehot_logits[memory_mask]
                elif "bart" in self.config.model.name.lower() or "t5" in self.config.model.name.lower():
                    avg_encoder_state = encoder_state.detach().mean(0)
                    memory_keys, memory_labels = self.memory
                    dists = torch.norm(avg_encoder_state - memory_keys, dim=-1)
                    closest_dist = dists.min()
                    closest_idx = dists.argmin()
                    closest_v = memory_labels[closest_idx]

                    if closest_dist < self.config.lu.threshold:
                        output[i] = torch.zeros((1, kwargs['labels'].shape[1], output.shape[2]), device=output.device)
                        for j, idx in enumerate(closest_v):
                            if j >= output.shape[1]:
                                break
                            output[i, j, idx] = self.config.lu.onehot_logit
                        if "t5" not in self.config.model.name.lower():
                            # T5 does not shift targets in the loss
                            output[i] = output[i].roll(-1, -2)
                else:
                    avg_encoder_state = encoder_state.detach().mean(0)
                    memory_keys, memory_labels = self.memory
                    dists = torch.norm(avg_encoder_state - memory_keys, dim=-1)
                    closest_dist = dists.min()
                    closest_idx = dists.argmin()
                    closest_v = memory_labels[closest_idx]

                    if closest_dist < self.config.lu.threshold:
                        output[i] = self.config.lu.onehot_logit * (2 * closest_v - 1)  # Return onehot_logit or -onehot_logit

        return output

    def edit(self, batch, condition=None):
        edit_model = self.model.eval()
        if "bert" in self.config.model.name.lower():
            _, encoder_states = self.model(**batch, output_hidden_states=True)
        else:
            encoder_states = _last_encoder_state(self.model(**batch, output_hidden_states=True))

        memory_keys = []
        memory_labels = []
        for encoder_state, label in zip(encoder_states, batch["labels"]):
            if "gpt2" in self.config.model.name.lower():
                # NOTE: broken
                avg_encoder_states = (encoder_state.cumsum(0).detach() / torch.arange(1, encoder_state.shape[0] + 1, device=encoder_state.device).view(-1, 1))[-10:, :]
                memory = (avg_encoder_states, label[-10:])
            else:
                avg_encoder_state = encoder_state.detach().mean(0)
                memory_keys.append(avg_encoder_state)
                memory_labels.append(label)

        memory = (torch.stack(memory_keys), torch.stack(memory_labels))
        return LU(self.model.eval(), self.config, self.model_constructor, memory), {}