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
from torch._inductor.runtime.triton_helpers import math as tl_math
import numpy as np
import torch.utils.data
import torch.nn as torch_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_add_arange_copy_cos_div_fill_lift_fresh_mul_repeat_sin_sub_zeros_like_0(
    in_ptr0, out_ptr0, out_ptr1, xnumel, XBLOCK: tl.constexpr):
    xnumel = 48
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x0 = xindex % 3
    tmp0 = tl.load(in_ptr0 + x2, xmask)
    tmp1 = x0
    tmp2 = tl.full([1], 1, tl.int32)
    tmp3 = tmp1 == tmp2
    tmp4 = tl.full([1], 2, tl.int64)
    tmp5 = tmp1 >= tmp4
    tmp6 = tl.load(in_ptr0 + x2, tmp5 & xmask, other=0.0)
    tmp7 = -1 + x0
    tmp8 = tmp7.to(tl.float32)
    tmp9 = tmp6 * tmp8
    tmp10 = 3.141592653589793
    tmp11 = tmp9 * tmp10
    tmp12 = tl_math.sin(tmp11)
    tmp13 = tmp12 / tmp11
    tmp14 = tl.full(tmp13.shape, 0.0, tmp13.dtype)
    tmp15 = tl.where(tmp5, tmp13, tmp14)
    tmp16 = tl.full([1], 1, tl.int64)
    tmp17 = tmp1 < tmp16
    tmp18 = tl.load(in_ptr0 + x2, tmp17 & xmask, other=0.0)
    tmp19 = tmp18 * tmp8
    tmp20 = tmp19 * tmp10
    tmp21 = tl_math.sin(tmp20)
    tmp22 = tmp21 / tmp20
    tmp23 = tl.full(tmp22.shape, 0.0, tmp22.dtype)
    tmp24 = tl.where(tmp17, tmp22, tmp23)
    tmp25 = 0.0
    tmp26 = tl.where(tmp17, tmp24, tmp25)
    tmp27 = tl.where(tmp5, tmp15, tmp26)
    tmp28 = 1.0
    tmp29 = tl.where(tmp3, tmp28, tmp27)
    tmp30 = tmp0 * tmp29
    tmp31 = 6.283185307179586
    tmp32 = tmp8 * tmp31
    tmp33 = 0.3333333333333333
    tmp34 = tmp32 * tmp33
    tmp35 = tl_math.cos(tmp34)
    tmp36 = 0.46
    tmp37 = tmp35 * tmp36
    tmp38 = 0.54
    tmp39 = tmp37 + tmp38
    tmp40 = tmp30 * tmp39
    tmp41 = tmp8 * tmp10
    tmp42 = tl_math.sin(tmp41)
    tmp43 = tmp42 / tmp41
    tmp44 = tmp43.to(tl.int64)
    tmp45 = tl.full(tmp44.shape, 0, tmp44.dtype)
    tmp46 = tl.where(tmp5, tmp44, tmp45)
    tmp47 = tl.where(tmp17, tmp44, tmp45)
    tmp48 = tl.full([1], 0, tl.int64)
    tmp49 = tl.where(tmp17, tmp47, tmp48)
    tmp50 = tl.where(tmp5, tmp46, tmp49)
    tmp51 = tl.where(tmp3, tmp16, tmp50)
    tmp52 = tmp51.to(tl.float32)
    tmp53 = tmp52 - tmp30
    tl.store(out_ptr0 + x2, tmp40, xmask)
    tl.store(out_ptr1 + x2, tmp53, xmask)


@triton.jit
def triton_poi_fused_div_1(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 48
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = xindex // 3
    tmp0 = tl.load(in_ptr0 + x2, xmask)
    tmp1 = tl.load(in_ptr0 + 3 * x1, xmask, eviction_policy='evict_last')
    tmp2 = tl.load(in_ptr0 + (1 + 3 * x1), xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr0 + (2 + 3 * x1), xmask, eviction_policy='evict_last')
    tmp3 = tmp1 + tmp2
    tmp5 = tmp3 + tmp4
    tmp6 = tmp0 / tmp5
    tl.store(out_ptr0 + x2, tmp6, xmask)


@triton.jit
def triton_poi_fused_arange_repeat_2(out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 48
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 3
    x2 = xindex
    tmp0 = -1 + x0
    tl.store(out_ptr0 + x2, tmp0, xmask)


@triton.jit
def triton_poi_fused_add_arange_cos_div_mul_repeat_sum_3(in_ptr0, in_ptr1,
    out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + 3 * x0, xmask, eviction_policy='evict_last')
    tmp8 = tl.load(in_ptr1 + 3 * x0, xmask, eviction_policy='evict_last')
    tmp11 = tl.load(in_ptr0 + (1 + 3 * x0), xmask, eviction_policy='evict_last'
        )
    tmp17 = tl.load(in_ptr1 + (1 + 3 * x0), xmask, eviction_policy='evict_last'
        )
    tmp21 = tl.load(in_ptr0 + (2 + 3 * x0), xmask, eviction_policy='evict_last'
        )
    tmp27 = tl.load(in_ptr1 + (2 + 3 * x0), xmask, eviction_policy='evict_last'
        )
    tmp1 = -2.0943951023931953
    tmp2 = tl_math.cos(tmp1)
    tmp3 = 0.46
    tmp4 = tmp2 * tmp3
    tmp5 = 0.54
    tmp6 = tmp4 + tmp5
    tmp7 = tmp0 * tmp6
    tmp9 = tmp8.to(tl.float32)
    tmp10 = tmp7 * tmp9
    tmp12 = 0.0
    tmp13 = tl_math.cos(tmp12)
    tmp14 = tmp13 * tmp3
    tmp15 = tmp14 + tmp5
    tmp16 = tmp11 * tmp15
    tmp18 = tmp17.to(tl.float32)
    tmp19 = tmp16 * tmp18
    tmp20 = tmp10 + tmp19
    tmp22 = 2.0943951023931953
    tmp23 = tl_math.cos(tmp22)
    tmp24 = tmp23 * tmp3
    tmp25 = tmp24 + tmp5
    tmp26 = tmp21 * tmp25
    tmp28 = tmp27.to(tl.float32)
    tmp29 = tmp26 * tmp28
    tmp30 = tmp20 + tmp29
    tl.store(out_ptr0 + x0, tmp30, xmask)


@triton.jit
def triton_poi_fused_add_arange_cos_div_mul_repeat_4(in_out_ptr0, in_ptr0,
    xnumel, XBLOCK: tl.constexpr):
    xnumel = 48
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x0 = xindex % 3
    x1 = xindex // 3
    tmp0 = tl.load(in_out_ptr0 + x2, xmask)
    tmp13 = tl.load(in_ptr0 + x1, xmask, eviction_policy='evict_last')
    tmp1 = -1 + x0
    tmp2 = tmp1.to(tl.float32)
    tmp3 = 6.283185307179586
    tmp4 = tmp2 * tmp3
    tmp5 = 0.3333333333333333
    tmp6 = tmp4 * tmp5
    tmp7 = tl_math.cos(tmp6)
    tmp8 = 0.46
    tmp9 = tmp7 * tmp8
    tmp10 = 0.54
    tmp11 = tmp9 + tmp10
    tmp12 = tmp0 * tmp11
    tmp14 = tmp12 / tmp13
    tl.store(in_out_ptr0 + x2, tmp14, xmask)


def call(args):
    arg0_1, = args
    args.clear()
    assert_size_stride(arg0_1, (4, 4, 3), (12, 3, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 4, 3), (12, 3, 1), torch.float32)
        buf2 = empty_strided_cuda((4, 4, 3), (12, 3, 1), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_add_arange_copy_cos_div_fill_lift_fresh_mul_repeat_sin_sub_zeros_like_0[
            grid(48)](arg0_1, buf0, buf2, 48, XBLOCK=64, num_warps=1,
            num_stages=1)
        del arg0_1
        buf1 = empty_strided_cuda((4, 4, 3), (12, 3, 1), torch.float32)
        triton_poi_fused_div_1[grid(48)](buf0, buf1, 48, XBLOCK=64,
            num_warps=1, num_stages=1)
        del buf0
        buf3 = empty_strided_cuda((4, 4, 3), (12, 3, 1), torch.int64)
        triton_poi_fused_arange_repeat_2[grid(48)](buf3, 48, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf4 = torch.ops.aten.pow.Scalar(-1, buf3)
        del buf3
        buf5 = buf4
        del buf4
        buf6 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        triton_poi_fused_add_arange_cos_div_mul_repeat_sum_3[grid(16)](buf2,
            buf5, buf6, 16, XBLOCK=16, num_warps=1, num_stages=1)
        del buf5
        buf7 = buf2
        del buf2
        triton_poi_fused_add_arange_cos_div_mul_repeat_4[grid(48)](buf7,
            buf6, 48, XBLOCK=64, num_warps=1, num_stages=1)
        del buf6
    return buf1, buf7


class SincFilterNew(torch_nn.Module):
    """ SincFilter
        Given the cut-off-frequency, produce the low-pass and high-pass
        windowed-sinc-filters.
        If input cut-off-frequency is (batchsize=1, signal_length, 1),
        output filter coef is (batchsize=1, signal_length, filter_order).
        For each time step in [1, signal_length), we calculate one
        filter for low-pass sinc filter and another for high-pass filter.
        
        Example:
        import scipy
        import scipy.signal
        import numpy as np
        
        filter_order = 31
        cut_f = 0.2
        sinc_layer = SincFilter(filter_order)
        lp_coef, hp_coef = sinc_layer(torch.ones(1, 10, 1) * cut_f)
        
        w, h1 = scipy.signal.freqz(lp_coef[0, 0, :].numpy(), [1])
        w, h2 = scipy.signal.freqz(hp_coef[0, 0, :].numpy(), [1])
        plt.plot(w, 20*np.log10(np.abs(h1)))
        plt.plot(w, 20*np.log10(np.abs(h2)))
        plt.plot([cut_f * np.pi, cut_f * np.pi], [-100, 0])
    """

    def __init__(self, filter_order):
        super(SincFilterNew, self).__init__()
        self.half_k = (filter_order - 1) // 2
        self.order = self.half_k * 2 + 1

    def hamming_w(self, n_index):
        """ prepare hamming window for each time step
        n_index (batchsize=1, signal_length, filter_order)
            For each time step, n_index will be [-(M-1)/2, ... 0, (M-1)/2]
            n_index[0, 0, :] = [-(M-1)/2, ... 0, (M-1)/2]
            n_index[0, 1, :] = [-(M-1)/2, ... 0, (M-1)/2]
            ...
        output  (batchsize=1, signal_length, filter_order)
            output[0, 0, :] = hamming_window
            output[0, 1, :] = hamming_window
            ...
        """
        return 0.54 + 0.46 * torch.cos(2 * np.pi * n_index / self.order)

    def sinc(self, x):
        """ Normalized sinc-filter sin( pi * x) / pi * x
        https://en.wikipedia.org/wiki/Sinc_function
        
        Assume x (batchsize, signal_length, filter_order) and 
        x[0, 0, :] = [-half_order, - half_order+1, ... 0, ..., half_order]
        x[:, :, self.half_order] -> time index = 0, sinc(0)=1
        """
        y = torch.zeros_like(x)
        y[:, :, 0:self.half_k] = torch.sin(np.pi * x[:, :, 0:self.half_k]) / (
            np.pi * x[:, :, 0:self.half_k])
        y[:, :, self.half_k + 1:] = torch.sin(np.pi * x[:, :, self.half_k + 1:]
            ) / (np.pi * x[:, :, self.half_k + 1:])
        y[:, :, self.half_k] = 1
        return y

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


Measured latency_ms: 0.074383