import numpy as np
import torch
from torch import nn
import torch.nn.functional as F

from .utils import generateStrides, generateGrid


class OmniMaxPool2d(nn.MaxPool2d):
  """
  kernel_size: (H, W)
  """

  def __init__(self, kernel_size=(3, 3), stride=1, padding=0, dilation=1,
               return_indices: bool = False, ceil_mode: bool = False):
    super(OmniMaxPool2d, self).__init__(
      kernel_size, stride, padding, dilation, return_indices, ceil_mode)
    if isinstance(kernel_size, int):
      self.kernel_size = (kernel_size, kernel_size)
    if isinstance(stride, int):
      self.stride = (stride, stride)

    self.convert_option = ["mollweide", "adaptive"][0]
    self.strides = None

    self.grid = None
    self.mask = None

  def genSamplingPattern(self, h, w):
    self.strides = generateStrides(h, w, self.convert_option)

    grid = generateGrid(h, w, self.kernel_size, self.stride, self.strides, self.convert_option)

    with torch.no_grad():
      self.grid = torch.FloatTensor(grid)
      self.grid.requires_grad = False

  def forward(self, x):
    # Generate Sampling Pattern
    B, C, H, W = x.shape

    if self.strides is None:
      self.genSamplingPattern(H, W)

    with torch.no_grad():
      grid = self.grid.repeat((B, 1, 1, 1)).to(x.device)  # (B, H*Kh, W*Kw, 2)
      grid.requires_grad = False

    x = F.grid_sample(x, grid, align_corners=False, mode='bilinear')  # (B, in_c, H*Kh, W*Kw)

    x = F.max_pool2d(
      input=x,
      kernel_size=self.kernel_size, stride=self.kernel_size,
      padding=self.padding, dilation=self.dilation,
      ceil_mode=self.ceil_mode, return_indices=self.return_indices
    )

    return x  # (B, out_c, H/stride_h, W/stride_w)
