import math  # noqa

import torch  # noqa
from torch.nn.functional import one_hot  # noqa
from torch.nn.modules.utils import _pair, _ntuple
from torch.nn.functional import unfold, fold

from hw2_functional import *  # noqa
from hw2_functional import __all__ as __old_all__

__new_all__ = ['view', 'add', 'conv2d', 'max_pool2d']
__all__ = __old_all__ + __new_all__ # 


#################################################
# conv2d
#################################################

def conv2d(x, w, b=None, padding=0, stride=1, dilation=1, groups=1, ctx=None):
  """A differentiable convolution of 2d tensors.

  Note: Read this following documentation regarding the output's shape.
  https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html#torch.nn.Conv2d

  Backward call:
    backward_fn: conv2d_backward
    args: y, x, w, b, padding, stride, dilation

  Args:
    x (torch.Tensor): The input tensor.
      Has shape `(batch_size, in_channels, height, width)`.
    w (torch.Tensor): The weight tensor.
      Has shape `(out_channels, in_channels, kernel_height, kernel_width)`.
    b (torch.Tensor): The bias tensor. Has shape `(out_channels,)`.
    padding (Tuple[int, int] or int, Optional): The padding in each dimension (height, width).
      Defaults to 0.
    stride (Tuple[int, int] or int, Optional): The stride in each dimension (height, width).
      Defaults to 1.
    dilation (Tuple[int, int] or int, Optional): The dilation in each dimension (height, width).
      Defaults to 1.
    groups (int, Optional): Number of groups. Defaults to 1.
    ctx (List, optional): The autograd context. Defaults to None.

  Returns:
    y (torch.Tensor): The output tensor.
      Has shape `(batch_size, out_channels, out_height, out_width)`.
  """
  assert w.size(0) % groups == 0, \
    f'expected w.size(0)={w.size(0)} to be divisible by groups={groups}'
  assert x.size(1) % groups == 0, \
    f'expected x.size(1)={x.size(1)} to be divisible by groups={groups}'
  assert x.size(1) // groups == w.size(1), \
    f'expected w.size(1)={w.size(1)} to be x.size(1)//groups={x.size(1)}//{groups}'

  # BEGIN SOLUTION
  raise NotImplementedError
  # END SOLUTION


def conv2d_backward(y, x, w, b, padding, stride, dilation, groups):
  """Backward computation of `conv2d`.

  Propagates the gradients of `y` (in `y.grad`) to `x`, `w` and `b` (if `b` is not None),
  and accumulates them in `x.grad`, `w.grad` and `b.grad`, respectively.

  Args:
    y (torch.Tensor): The output tensor.
      Has shape `(batch_size, out_channels, out_height, out_width)`.
    x (torch.Tensor): The input tensor.
      Has shape `(batch_size, in_channels, height, width)`.
    w (torch.Tensor): The weight tensor.
      Has shape `(out_channels, in_channels, kernel_height, kernel_width)`.
    b (torch.Tensor): The bias tensor. Has shape `(out_channels,)`.
    padding (Tuple[int, int] or int, Optional): The padding in each dimension (height, width).
      Defaults to 0.
    stride (Tuple[int, int] or int, Optional): The stride in each dimension (height, width).
      Defaults to 1.
    dilation (Tuple[int, int] or int, Optional): The dilation in each dimension (height, width).
      Defaults to 1.
    groups (int, Optional): Number of groups. Defaults to 1.
  """
  # BEGIN SOLUTION
  raise NotImplementedError
  # END SOLUTION


#################################################
# max_pool2d
#################################################

def max_pool2d(x, kernel_size, padding=0, stride=1, dilation=1, ctx=None):
  """A differentiable convolution of 2d tensors.

  Note: Read this following documentation regarding the output's shape.
  https://pytorch.org/docs/stable/generated/torch.nn.MaxPool2d.html#torch.nn.MaxPool2d

  Backward call:
    backward_fn: max_pool2d_backward
    args: y, x, padding, stride, dilation

  Args:
    x (torch.Tensor): The input tensor. Has shape `(batch_size, in_channels, height, width)`.
    kernel_size (Tuple[int, int] or int): The kernel size in each dimension (height, width).
    padding (Tuple[int, int] or int, Optional): The padding in each dimension (height, width).
      Defaults to 0.
    stride (Tuple[int, int] or int, Optional): The stride in each dimension (height, width).
      Defaults to 1.
    dilation (Tuple[int, int] or int, Optional): The dilation in each dimension (height, width).
      Defaults to 1.
    ctx (List, optional): The autograd context. Defaults to None.

  Returns:
    y (torch.Tensor): The output tensor.
      Has shape `(batch_size, in_channels, out_height, out_width)`.
  """
  # BEGIN SOLUTION
  raise NotImplementedError
  # END SOLUTION


def max_pool2d_backward(y, x, index, kernel_size, padding, stride, dilation):
  """Backward computation of `max_pool2d`.

  Propagates the gradients of `y` (in `y.grad`) to `x` and accumulates it in `x.grad`.

  Args:
    y (torch.Tensor): The output tensor.
      Has shape `(batch_size, in_channels, out_height, out_width)`.
    x (torch.Tensor): The input tensor.
      Has shape `(batch_size, in_channels, height, width)`.
    index (torch.Tensor): Auxilary tensor with indices of the maximum elements. You are
      not restricted to a specific format.
    kernel_size (Tuple[int, int] or int): The kernel size in each dimension (height, width).
    padding (Tuple[int, int] or int, Optional): The padding in each dimension (height, width).
      Defaults to 0.
    stride (Tuple[int, int] or int, Optional): The stride in each dimension (height, width).
      Defaults to 1.
    dilation (Tuple[int, int] or int, Optional): The dilation in each dimension (height, width).
      Defaults to 1.
  """
  # BEGIN SOLUTION
  raise NotImplementedError
  # END SOLUTION


#################################################
# view
#################################################

def view(x, size, ctx=None):
  """A differentiable view function.

  Backward call:
    backward_fn: view_backward
    args: y, x

  Args:
    x (torch.Tensor): The input tensor.
    size (torch.Size): The new size (shape).
    ctx (List, optional): The autograd context. Defaults to None.

  Returns:
    y (torch.Tensor): The output tensor. Has shape `size`.
  """
  # BEGIN SOLUTION
  raise NotImplementedError
  # END SOLUTION


def view_backward(y, x):
  """Backward computation of `view`.

  Propagates the gradients of `y` (in `y.grad`) to `x`, and accumulates them in `x.grad`.

  Args:
    y (torch.Tensor): The output tensor.
    x (torch.Tensor): The input tensor.
  """
  # BEGIN SOLUTION
  raise NotImplementedError
  # END SOLUTION


#################################################
# add
#################################################

def add(a, b, ctx=None):
  """A differentiable addition of two tensors.

  Backward call:
    backward_fn: add_backward
    args: y, a, b

  Args:
    a (torch.Tensor): The first input tensor.
    b (torch.Tensor): The second input tensor. Should have the same shape as `a`.
    ctx (List, optional): The autograd context. Defaults to None.

  Returns:
    y (torch.Tensor): The output tensor. The sum of `a + b`.
  """
  assert a.size() == b.size(), 'tensors should have the same size'
  # BEGIN SOLUTION
  raise NotImplementedError
  # END SOLUTION
  

def add_backward(y, a, b):
  """Backward computation of `add`.

  Propagates the gradients of `y` (in `y.grad`) to `a` and `b`, and accumulates them in `a.grad`,
  `b.grad`, respectively.

  Args:
    y (torch.Tensor): The output tensor.
    a (torch.Tensor): The first input tensor.
    b (torch.Tensor): The second input tensor.
  """
  # BEGIN SOLUTION
  raise NotImplementedError
  # END SOLUTION
