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 math
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_add_embedding_0(in_ptr0, in_ptr1, in_ptr2, in_ptr3,
    in_ptr4, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4 % 4
    x2 = xindex // 16
    x0 = xindex % 4
    x4 = xindex
    tmp0 = tl.load(in_ptr0 + (12 + x1 + 16 * x2), xmask, eviction_policy=
        'evict_last')
    tmp8 = tl.load(in_ptr0 + (8 + x1 + 16 * x2), xmask, eviction_policy=
        'evict_last')
    tmp17 = tl.load(in_ptr0 + (4 + x1 + 16 * x2), xmask, eviction_policy=
        'evict_last')
    tmp26 = tl.load(in_ptr0 + (x1 + 16 * x2), xmask, eviction_policy=
        'evict_last')
    tmp1 = tmp0.to(tl.int64)
    tmp2 = tl.full([XBLOCK], 24, tl.int32)
    tmp3 = tmp1 + tmp2
    tmp4 = tmp1 < 0
    tmp5 = tl.where(tmp4, tmp3, tmp1)
    tl.device_assert((0 <= tmp5) & (tmp5 < 24) | ~xmask,
        'index out of bounds: 0 <= tmp5 < 24')
    tmp7 = tl.load(in_ptr1 + (x0 + 4 * tmp5), xmask)
    tmp9 = tmp8.to(tl.int64)
    tmp10 = tl.full([XBLOCK], 7, tl.int32)
    tmp11 = tmp9 + tmp10
    tmp12 = tmp9 < 0
    tmp13 = tl.where(tmp12, tmp11, tmp9)
    tl.device_assert((0 <= tmp13) & (tmp13 < 7) | ~xmask,
        'index out of bounds: 0 <= tmp13 < 7')
    tmp15 = tl.load(in_ptr2 + (x0 + 4 * tmp13), xmask)
    tmp16 = tmp7 + tmp15
    tmp18 = tmp17.to(tl.int64)
    tmp19 = tl.full([XBLOCK], 32, tl.int32)
    tmp20 = tmp18 + tmp19
    tmp21 = tmp18 < 0
    tmp22 = tl.where(tmp21, tmp20, tmp18)
    tl.device_assert((0 <= tmp22) & (tmp22 < 32) | ~xmask,
        'index out of bounds: 0 <= tmp22 < 32')
    tmp24 = tl.load(in_ptr3 + (x0 + 4 * tmp22), xmask)
    tmp25 = tmp16 + tmp24
    tmp27 = tmp26.to(tl.int64)
    tmp28 = tl.full([XBLOCK], 13, tl.int32)
    tmp29 = tmp27 + tmp28
    tmp30 = tmp27 < 0
    tmp31 = tl.where(tmp30, tmp29, tmp27)
    tl.device_assert((0 <= tmp31) & (tmp31 < 13) | ~xmask,
        'index out of bounds: 0 <= tmp31 < 13')
    tmp33 = tl.load(in_ptr4 + (x0 + 4 * tmp31), xmask)
    tmp34 = tmp25 + tmp33
    tmp35 = 0.0
    tmp36 = tmp34 + tmp35
    tl.store(out_ptr0 + x4, tmp36, xmask)


def call(args):
    arg0_1, arg1_1, arg2_1, arg3_1, arg4_1 = args
    args.clear()
    assert_size_stride(arg0_1, (4, 4, 4, 4), (64, 16, 4, 1))
    assert_size_stride(arg1_1, (24, 4), (4, 1))
    assert_size_stride(arg2_1, (7, 4), (4, 1))
    assert_size_stride(arg3_1, (32, 4), (4, 1))
    assert_size_stride(arg4_1, (13, 4), (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.float32)
        get_raw_stream(0)
        triton_poi_fused_add_embedding_0[grid(256)](arg0_1, arg1_1, arg2_1,
            arg3_1, arg4_1, buf0, 256, XBLOCK=128, num_warps=4, num_stages=1)
        del arg0_1
        del arg1_1
        del arg2_1
        del arg3_1
        del arg4_1
    return buf0,


class FixedEmbedding(nn.Module):

    def __init__(self, c_in, d_model):
        super(FixedEmbedding, self).__init__()
        w = torch.zeros(c_in, d_model).float()
        w.require_grad = False
        position = torch.arange(0, c_in).float().unsqueeze(1)
        div_term = (torch.arange(0, d_model, 2).float() * -(math.log(
            10000.0) / d_model)).exp()
        w[:, 0::2] = torch.sin(position * div_term)
        w[:, 1::2] = torch.cos(position * div_term)
        self.emb = nn.Embedding(c_in, d_model)
        self.emb.weight = nn.Parameter(w, requires_grad=False)

    def forward(self, x):
        return self.emb(x).detach()


class TemporalEmbeddingNew(nn.Module):

    def __init__(self, d_model, embed_type='fixed', freq='h'):
        super(TemporalEmbeddingNew, self).__init__()
        minute_size = 4
        hour_size = 24
        weekday_size = 7
        day_size = 32
        month_size = 13
        Embed = FixedEmbedding if embed_type == 'fixed' else nn.Embedding
        if freq == 't':
            self.minute_embed = Embed(minute_size, d_model)
        self.hour_embed = Embed(hour_size, d_model)
        self.weekday_embed = Embed(weekday_size, d_model)
        self.day_embed = Embed(day_size, d_model)
        self.month_embed = Embed(month_size, d_model)

    def forward(self, input_0):
        arg1_1 = self.hour_embed.emb.weight
        arg2_1 = self.weekday_embed.emb.weight
        arg3_1 = self.day_embed.emb.weight
        arg4_1 = self.month_embed.emb.weight
        arg0_1 = input_0
        output = call([arg0_1, arg1_1, arg2_1, arg3_1, arg4_1])
        return output[0]


Measured latency_ms: 0.036715