text: # Copyright 2017-present, Facebook, Inc.
  # All rights reserved.
  #
  # This source code is licensed under the license found in the
  # LICENSE file in the root directory of this source tree.
  
  import torch
  import torch.nn as nn
  import torch.optim as optim
  
  import numpy as np
  import quadprog
  
  from .common import MLP, ResNet18, ResNet18_TinyImagenet
  from sklearn.decomposition import PCA
  from scipy.linalg import orth as orth
  
  # Auxiliary functions useful for GEM's inner optimization.
  
  
  def accuracy(output, target, topk=(1, )):
      """Computes the precision@k for the specified values of k"""
      maxk = max(topk)
      batch_size = target.size(0)
  
      _, pred = output.topk(maxk, 1, True, True)
      pred = pred.t()
      correct = pred.eq(target.view(1, -1).expand_as(pred))
  
      res = []
      for k in topk:
          correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
          res.append(correct_k.mul_(100.0 / batch_size))
      return res
  
  
  def compute_offsets(task, nc_per_task, is_cifar):
      """
          Compute offsets for cifar to determine which
          outputs to select for a given task.
      """
      if is_cifar:
          offset1 = task * nc_per_task
          offset2 = (task + 1) * nc_per_task
      else:
          offset1 = 0
          offset2 = nc_per_task
      return offset1, offset2
  
  
  def store_grad(pp, grads, grad_dims, tid):
      """
          This stores parameter gradients of past tasks.
          pp: parameters
          grads: gradients
          grad_dims: list with number of parameters per layers
          tid: task id
      """
      # store the gradients
      grads[:, tid].fill_(0.0)
      cnt = 0
      for param in pp():
          if param.grad is not None:
              beg = 0 if cnt == 0 else sum(grad_dims[:cnt])
              en = sum(grad_dims[:cnt + 1])
              grads[beg:en, tid].copy_(param.grad.data.view(-1))
          cnt += 1
  
  
  def overwrite_grad(pp, newgrad, grad_dims):
      """
          This is used to overwrite the gradients with a new gradient
          vector, whenever violations occur.
          pp: parameters
          newgrad: corrected gradient
          grad_dims: list storing number of parameters at each layer
      """
      cnt = 0
      for param in pp():
          if param.grad is not None:
              beg = 0 if cnt == 0 else sum(grad_dims[:cnt])
              en = sum(grad_dims[:cnt + 1])
              this_grad = newgrad[beg:en].contiguous().view(
                  param.grad.data.size())
              param.grad.data.copy_(this_grad)
          cnt += 1
  
  
  def project2cone2(gradient, memories, margin=0.5, eps=1e-3):
      """
          Solves the GEM dual QP described in the paper given a proposed
          gradient "gradient", and a memory of task gradients "memories".
          Overwrites "gradient" with the final projected update.
  
          input:  gradient, p-vector
          input:  memories, (t * p)-vector
          output: x, p-vector
      """
      memories_np = memories.cpu().t().double().numpy()
      gradient_np = gradient.cpu().contiguous().view(-1).double().numpy()
      t = memories_np.shape[0]
      P = np.dot(memories_np, memories_np.transpose())
      P = 0.5 * (P + P.transpose()) + np.eye(t) * eps
      q = np.dot(memories_np, gradient_np) * -1
      G = np.eye(t)
      h = np.zeros(t) + margin
      v = quadprog.solve_qp(P, q, G, h)[0]
      x = np.dot(v, memories_np) + gradient_np
      gradient.copy_(torch.Tensor(x).view(-1, 1))
  
  
  def project5cone5(gradient, memories, beg, en, margin=0.5, eps=1e-3):
      """
          First we do orthognality of memories and then find the null space of
          these memories gradients
      """
      memories_np = memories[beg:en].cpu().t().double().numpy()
      gradient_np = gradient[beg:en].cpu().contiguous().view(-1).double().numpy()
      memories_np_mean = np.mean(memories_np, axis=0)
      memories_np_sum = np.sum(memories_np, axis=0)
      memories_np_del_mean = memories_np - memories_np_mean.reshape(1, -1)
      # memories_np_pca = PCA(n_components=len(memories_np))
      # memories_np_pca.fit(memories_np_del_mean)
      if len(memories_np) == 1:
          x = gradient_np - np.min([
              (memories_np_sum.transpose().dot(gradient_np) /
               memories_np_sum.transpose().dot(memories_np_sum)), -margin
          ]) * memories_np_sum
      else:
          memories_np_pca = orth(memories_np_del_mean.transpose()).transpose()
  
          # memories_np_orth = memories_np_pca.components_
          memories_np_orth = memories_np_pca
          Pg = gradient_np - memories_np_orth.transpose().dot(
              memories_np_orth.dot(gradient_np))
          Pg_bar = memories_np_sum - memories_np_orth.transpose().dot(
              memories_np_orth.dot(memories_np_sum))
          if memories_np_sum.transpose().dot(Pg) > 0:
              x = Pg
          else:
              x = gradient_np - np.min([
                  memories_np_sum.transpose().dot(Pg) /
                  memories_np_sum.transpose().dot(Pg_bar), -margin
              ]) * memories_np_sum - memories_np_orth.transpose().dot(
                  memories_np_orth.dot(gradient_np)) + memories_np_sum.transpose(
                  ).dot(Pg) / memories_np_sum.transpose().dot(
                      Pg_bar) * memories_np_orth.transpose().dot(
                          memories_np_orth.dot(memories_np_sum))
  
      print("task length: {}".format(len(memories_np)))
      gradient[beg:en].copy_(torch.Tensor(x).view(-1, 1))
  
  
  def agem(gradient, memories):
      memories_np = memories.cpu().t().double().numpy()
      gradient_np = gradient.cpu().contiguous().view(-1).double().numpy()
      x = gradient_np - np.dot(memories_np, gradient_np) / np.dot(
          memories_np, memories_np.transpose()) * memories_np.reshape(-1)
      gradient.copy_(torch.Tensor(x).view(-1, 1))
  
  
  class Net(nn.Module):
      def __init__(self, n_inputs, n_outputs, n_tasks, args):
          super(Net, self).__init__()
          nl, nh = args.n_layers, args.n_hiddens
          self.margin = args.memory_strength
          self.is_cifar = (args.data_file == 'cifar100.pt') or (
              args.data_file == 'cifar10.pt')
          self.is_imagenet = (args.data_file == 'tiny-imagenet-200.pt')
          if self.is_cifar:
              self.net = ResNet18(n_outputs)
          elif self.is_imagenet:
              self.net = ResNet18_TinyImagenet(n_outputs)
          else:
              self.net = MLP([n_inputs] + [nh] * nl + [n_outputs])
  
          self.ce = nn.CrossEntropyLoss()
          self.n_outputs = n_outputs
  
          self.opt = optim.SGD(self.parameters(), args.lr)
  
          self.n_memories = args.n_memories
          self.gpu = args.cuda
  
          # allocate episodic memory
          self.memory_data = torch.FloatTensor(n_tasks, self.n_memories,
                                               n_inputs)
          self.memory_labs = torch.LongTensor(n_tasks, self.n_memories)
          if args.cuda:
              self.memory_data = self.memory_data.cuda()
              self.memory_labs = self.memory_labs.cuda()
              self.net.cuda()
  
          # allocate temporary synaptic memory
          self.grad_dims = []
          self.block_grad_dims = []
  
          for name, module in self.net.named_children():
              block_params = 0
              for param in module.parameters():
                  block_params += param.data.numel()
                  self.grad_dims.append(param.data.numel())
              self.block_grad_dims.append(block_params)
  
          self.grads = torch.Tensor(sum(self.grad_dims), n_tasks)
          if args.cuda:
              self.grads = self.grads.cuda()
  
          # allocate counters
          self.observed_tasks = []
          self.old_task = -1
          self.mem_cnt = 0
          if self.is_cifar or self.is_imagenet:
              self.nc_per_task = int(n_outputs / n_tasks)
          else:
              self.nc_per_task = n_outputs
  
      def forward(self, x, t):
          output = self.net(x)
          if self.is_cifar or self.is_imagenet:
              # make sure we predict classes within the current task
              offset1 = int(t * self.nc_per_task)
              offset2 = int((t + 1) * self.nc_per_task)
              if offset1 > 0:
                  output[:, :offset1].data.fill_(-10e10)
              if offset2 < self.n_outputs:
                  output[:, offset2:self.n_outputs].data.fill_(-10e10)
          return output
  
      def observe(self, x, t, y, ep):
          # update memory
          if t != self.old_task:
              self.observed_tasks.append(t)
              self.old_task = t
  
          # Update ring buffer storing examples from current task
          bsz = y.data.size(0)
          endcnt = min(self.mem_cnt + bsz, self.n_memories)
          effbsz = endcnt - self.mem_cnt
          self.memory_data[t, self.mem_cnt:endcnt].copy_(x.data[:effbsz])
          if bsz == 1:
              self.memory_labs[t, self.mem_cnt] = y.data[0]
          else:
              self.memory_labs[t, self.mem_cnt:endcnt].copy_(y.data[:effbsz])
          self.mem_cnt += effbsz
          if self.mem_cnt == self.n_memories:
              self.mem_cnt = 0
  
          # compute gradient on previous tasks
          if len(self.observed_tasks) > 1:
              for tt in range(len(self.observed_tasks) - 1):
                  self.zero_grad()
                  # fwd/bwd on the examples in the memory
                  past_task = self.observed_tasks[tt]
  
                  offset1, offset2 = compute_offsets(
                      past_task, self.nc_per_task,
                      (self.is_cifar or self.is_imagenet))
                  output = self.forward(self.memory_data[past_task], past_task)
                  ptloss = self.ce(output[:, offset1:offset2],
                                   self.memory_labs[past_task] - offset1)
                  # top1, top5 = accuracy(
                  #     output[:, offset1:offset2],
                  #     self.memory_labs[past_task] - offset1,
                  #     topk=(1, 5))
                  # with open('log.log', 'a') as f:
                  #     f.write(
                  #         "now training task {}, memory {} loss: {} Prec@1 {}\n".
                  #         format(t, tt, ptloss.item(), top1.item()))
                  ptloss.backward()
                  store_grad(self.parameters, self.grads, self.grad_dims,
                             past_task)
  
          # now compute the grad on the current minibatch
          self.zero_grad()
  
          offset1, offset2 = compute_offsets(t, self.nc_per_task,
                                             (self.is_cifar or self.is_imagenet))
          output = self.forward(x, t)
          loss = self.ce(output[:, offset1:offset2], y - offset1)
  
          # top1, top5 = accuracy(
          #     output[:, offset1:offset2], y - offset1, topk=(1, 5))
          # with open('log.log', 'a') as f:
          #     f.write("now training task {}, new_task {} loss: {} Prec@1 {}\n".
          #             format(t, t, loss.item(), top1.item()))
          loss.backward()
  
          # check if gradient violates constraints
          if len(self.observed_tasks) > 1:
              # copy gradient
              store_grad(self.parameters, self.grads, self.grad_dims, t)
              indx = torch.cuda.LongTensor(self.observed_tasks[:-1]) if self.gpu \
                  else torch.LongTensor(self.observed_tasks[:-1])
  
              for cnt in range(len(self.block_grad_dims)):
                  beg = 0 if cnt == 0 else sum(self.block_grad_dims[:cnt])
                  en = sum(self.block_grad_dims[:cnt + 1])
                  if beg == en:
                      continue
                  dotp = torch.mm(self.grads[:, t].unsqueeze(0),
                                  self.grads.index_select(1, indx))
  
                  if (dotp < 0).sum() != 0:
                      project5cone5(
                          self.grads[:, t].unsqueeze(1),
                          self.grads.index_select(1, indx),
                          beg,
                          en,
                          margin=self.margin)
              # copy gradients back
              overwrite_grad(self.parameters, self.grads[:, t], self.grad_dims)
          self.opt.step()
  