import torch
import torch.nn as nn
import torch.nn.functional as F


class _SpikeAdaptiveMaxPoolNd(nn.Module):

    def __init__(self, output_size):
        super(_SpikeAdaptiveMaxPoolNd, self).__init__()
        self.output_size = output_size
        self.return_indices = True

    def reset_state(self):
        pass


class AdaptiveMaxPool2d(_SpikeAdaptiveMaxPoolNd):
    """Simple port of PyTorch AdaptiveMaxPool2d with small adjustment for spiking operations.
    
    Currently pooling only supports operations on floating point numbers, thus it casts the uint8 spikes to floats back and forth.
    The trace of the 'maximum' spike is also returned. In case of multiple spikes within pooling window, returns first spike of 
    the window (top left corner).
    """

    def forward(self, x, trace):
        x = x
        x, idx = F.adaptive_max_pool2d(x, self.output_size, self.return_indices
            )
        trace = trace[idx]
        return x > 0, trace


def get_inputs():
    return [torch.rand([4, 4, 4, 4]), torch.rand([16, 4, 4, 4])]


def get_init_inputs():
    return [[], {'output_size': 4}]

import torch
import triton
import triton.language as tl
from torch._inductor.runtime.triton_heuristics import grid
from torch._C import _cuda_getCurrentRawStream as get_raw_stream
import torch.nn as nn
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda


@triton.jit
def triton_poi_fused_adaptive_max_pool2d_gt_0(in_ptr0, out_ptr0, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + x0, xmask)
    tmp1 = 0.0
    tmp2 = tmp0 > tmp1
    tl.store(out_ptr0 + x0, tmp2, xmask)


@triton.jit
def triton_poi_fused_index_1(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl
    .constexpr):
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    tl.full([XBLOCK], True, tl.int1)
    x3 = xindex // 64
    x4 = xindex % 1024
    x5 = xindex
    tl.load(in_ptr0 + x3, None, eviction_policy='evict_last')
    tmp1 = tl.load(in_ptr1 + x4, None, eviction_policy='evict_last')
    tl.store(out_ptr0 + x5, tmp1, None)


def call(args):
    arg0_1, arg1_1 = args
    args.clear()
    assert_size_stride(arg0_1, (4, 4, 4, 4), (64, 16, 4, 1))
    assert_size_stride(arg1_1, (16, 4, 4, 4), (64, 16, 4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 4, 4, 4), (64, 16, 4, 1), torch.bool)
        get_raw_stream(0)
        triton_poi_fused_adaptive_max_pool2d_gt_0[grid(256)](arg0_1, buf0, 
            256, XBLOCK=128, num_warps=4, num_stages=1)
        buf1 = empty_strided_cuda((4, 4, 4, 4, 4, 4, 4), (4096, 1024, 256, 
            64, 16, 4, 1), torch.float32)
        triton_poi_fused_index_1[grid(16384)](arg0_1, arg1_1, buf1, 16384,
            XBLOCK=256, num_warps=4, num_stages=1)
        del arg0_1
        del arg1_1
    return buf0, buf1


class _SpikeAdaptiveMaxPoolNd(nn.Module):

    def __init__(self, output_size):
        super(_SpikeAdaptiveMaxPoolNd, self).__init__()
        self.output_size = output_size
        self.return_indices = True

    def reset_state(self):
        pass


class AdaptiveMaxPool2dNew(_SpikeAdaptiveMaxPoolNd):
    """Simple port of PyTorch AdaptiveMaxPool2d with small adjustment for spiking operations.
    
    Currently pooling only supports operations on floating point numbers, thus it casts the uint8 spikes to floats back and forth.
    The trace of the 'maximum' spike is also returned. In case of multiple spikes within pooling window, returns first spike of 
    the window (top left corner).
    """

    def forward(self, input_0, input_1):
        arg0_1 = input_0
        arg1_1 = input_1
        output = call([arg0_1, arg1_1])
        return output[0], output[1]
