import torch
from torch.autograd import Function

from pn2utl._C import attention_relation_step_forward_cuda, attention_relation_step_backward_cuda, \
    attention_fusion_step_forward_cuda, attention_fusion_step_backward_cuda


class AttentionRelationStep(Function):
    @staticmethod
    def forward(ctx, query, key, weight, index_target, index_refer):
        """
        input - query: (n, g, c), key: (n, g, c), weight: (c)  1_c for scatter attention,
                index_target: (m), index_refer: (m)
        output - relation: (M, g)
        """

        assert query.is_contiguous() \
               and key.is_contiguous() \
               and index_target.is_contiguous() \
               and index_refer.is_contiguous() \
               and weight.is_contiguous()

        assert index_target.shape[0] == index_refer.shape[0]

        _, g, c = query.shape
        m = index_target.shape[0]
        output = torch.cuda.FloatTensor(m, g).zero_()
        attention_relation_step_forward_cuda(m, g, c, query, key, weight,
                                             index_target.int(), index_refer.int(), output)
        ctx.save_for_backward(query, key, weight, index_target, index_refer)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        query, key, weight, index_target, index_refer = ctx.saved_tensors
        n, g, c = query.shape
        m = index_target.shape[0]
        grad_query = torch.cuda.FloatTensor(n, g, c).zero_()
        grad_key = torch.cuda.FloatTensor(n, g, c).zero_()
        grad_weight = torch.cuda.FloatTensor(c).zero_()
        attention_relation_step_backward_cuda(m, g, c,
                                              query, grad_query,
                                              key, grad_key,
                                              weight, grad_weight,
                                              index_target.int(), index_refer.int(),
                                              grad_output)
        return grad_query, grad_key, None, None, None


class AttentionFusionStep(Function):
    @staticmethod
    def forward(ctx, weight, value, index_target, index_refer):
        """
        input - weight: (m, g), value: (n, g, c)
                index_target: (m), index_value: (m)
        output - output: (n, g, c)
        """

        assert weight.is_contiguous() \
               and value.is_contiguous() \
               and index_target.is_contiguous() \
               and index_refer.is_contiguous() \
               and weight.is_contiguous()

        assert index_target.shape[0] == index_refer.shape[0]

        n, g, c = value.shape
        m = index_refer.shape[0]
        output = torch.cuda.FloatTensor(n, g, c).zero_()
        attention_fusion_step_forward_cuda(m, g, c, weight, value, index_target.int(), index_refer.int(), output)
        ctx.save_for_backward(weight, value, index_target, index_refer)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        """
        input: grad_output: (n, g, c)
        output: grad_weight: (m, g), grad_value: (n, g, c), none, none
        """
        weight, value, index_target, index_refer = ctx.saved_tensors
        n, g, c = value.shape
        m = index_target.shape[0]
        grad_weight = torch.cuda.FloatTensor(m, g).zero_()
        grad_value = torch.cuda.FloatTensor(n, g, c).zero_()
        attention_fusion_step_backward_cuda(m, g, c,
                                            weight, grad_weight,
                                            value, grad_value,
                                            index_target.int(), index_refer.int(),
                                            grad_output)
        return grad_weight, grad_value, None, None


attention_relation_step = AttentionRelationStep.apply
attention_fusion_step = AttentionFusionStep.apply
