import torch
from torch._inductor.select_algorithm import extern_kernels
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 libdevice
from abc import abstractmethod
from torch import nn
from torch.nn.utils import weight_norm
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda
reinterpret_tensor = torch._C._dynamo.guards._reinterpret_tensor


@triton.jit
def triton_per_fused__weight_norm_interface_0(in_out_ptr0, in_ptr0, in_ptr1,
    out_ptr0, xnumel, rnumel, XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 4
    xoffset = tl.program_id(0) * XBLOCK
    xoffset + tl.arange(0, XBLOCK)[:, None]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[None, :]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r0 = rindex
    tmp0 = tl.load(in_ptr0 + r0, None)
    tmp6 = tl.load(in_ptr1 + 0)
    tmp7 = tl.broadcast_to(tmp6, [XBLOCK, RBLOCK])
    tmp1 = tmp0 * tmp0
    tmp2 = tl.broadcast_to(tmp1, [XBLOCK, RBLOCK])
    tmp4 = tl.sum(tmp2, 1)[:, None]
    tmp5 = libdevice.sqrt(tmp4)
    tmp8 = tmp7 / tmp5
    tmp9 = tmp0 * tmp8
    tl.debug_barrier()
    tl.store(in_out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp5, None)
    tl.store(out_ptr0 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp9, None)


@triton.jit
def triton_per_fused_mean_1(in_out_ptr0, in_ptr0, xnumel, rnumel, XBLOCK:
    tl.constexpr):
    RBLOCK: tl.constexpr = 64
    xoffset = tl.program_id(0) * XBLOCK
    xoffset + tl.arange(0, XBLOCK)[:, None]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[None, :]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r0 = rindex
    tmp0 = tl.load(in_ptr0 + r0, None)
    tmp1 = tl.broadcast_to(tmp0, [XBLOCK, RBLOCK])
    tmp3 = tl.sum(tmp1, 1)[:, None]
    tmp4 = 64.0
    tmp5 = tmp3 / tmp4
    tl.debug_barrier()
    tl.store(in_out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp5, None)


def call(args):
    primals_1, primals_2, primals_3 = args
    args.clear()
    assert_size_stride(primals_1, (1, 1), (1, 1))
    assert_size_stride(primals_2, (1, 4), (4, 1))
    assert_size_stride(primals_3, (4, 4, 4, 4), (64, 16, 4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((1, 1), (1, 1), torch.float32)
        buf1 = buf0
        del buf0
        buf2 = empty_strided_cuda((1, 4), (4, 1), torch.float32)
        get_raw_stream(0)
        triton_per_fused__weight_norm_interface_0[grid(1)](buf1, primals_2,
            primals_1, buf2, 1, 4, XBLOCK=1, num_warps=2, num_stages=1)
        buf3 = empty_strided_cuda((64, 1), (1, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(primals_3, (64, 4), (4, 1), 0),
            reinterpret_tensor(buf2, (4, 1), (1, 4), 0), out=buf3)
        buf4 = empty_strided_cuda((), (), torch.float32)
        buf5 = buf4
        del buf4
        triton_per_fused_mean_1[grid(1)](buf5, buf3, 1, 64, XBLOCK=1,
            num_warps=2, num_stages=1)
        del buf3
    return buf5, buf2, primals_1, primals_2, buf1, reinterpret_tensor(primals_3
        , (64, 4), (4, 1), 0)


class HyperNet(nn.Module):
    """This module is responsible for taking the losses from all tasks and return a single loss term.
    We can think of this as our learnable loss criterion

    """

    def __init__(self, main_task, input_dim):
        super().__init__()
        self.main_task = main_task
        self.input_dim = input_dim

    def forward(self, losses, outputs=None, labels=None, data=None):
        """

        :param losses: losses form each task. This should be a tensor of size (batch_size, self.input_dim)
        :param outputs: Optional. Parameters model output.
        :param labels: Optional. Target.
        :param data: Optiona. Parameters model input.
        :return:
        """
        pass

    def _init_weights(self):
        pass

    def get_weights(self):
        """

        :return: list of model parameters
        """
        return list(self.parameters())


class MonoHyperNet(HyperNet):
    """Monotonic Hypernets

    """

    def __init__(self, main_task, input_dim, clamp_bias=False):
        super().__init__(main_task=main_task, input_dim=input_dim)
        self.clamp_bias = clamp_bias

    def get_weights(self):
        """

        :return: list of model parameters
        """
        return list(self.parameters())

    @abstractmethod
    def clamp(self):
        pass


class MonoLinearHyperNetNew(MonoHyperNet):
    """Linear weights, e.g. \\sum_j lpha_j * l_j

    """

    def __init__(self, main_task, input_dim, skip_connection=False,
        clamp_bias=False, init_value=1.0, weight_normalization=True):
        super().__init__(main_task=main_task, input_dim=main_task,
            clamp_bias=clamp_bias)
        self.init_value = init_value
        self.skip_connection = skip_connection
        self.linear = nn.Linear(input_dim, 1, bias=False)
        self._init_weights()
        self.weight_normalization = weight_normalization
        if self.weight_normalization:
            self.linear = weight_norm(self.linear)

    def _init_weights(self):
        self.linear.weight = nn.init.constant_(self.linear.weight, self.
            init_value)

    def clamp(self):
        """make sure parameters are non-negative

        """
        if self.weight_normalization:
            self.linear.weight_v.data.clamp_(0)
            self.linear.weight_g.data.clamp_(0)
        else:
            self.linear.weight.data.clamp_(0)
        if self.linear.bias is not None and self.clamp_bias:
            self.linear.bias.data.clamp_(0)

    def forward(self, input_0):
        primals_1 = self.linear.weight_g
        primals_2 = self.linear.weight_v
        primals_3 = input_0
        output = call([primals_1, primals_2, primals_3])
        return output[0]


Measured latency_ms: 0.037872