import torch
from torch.autograd import Function

from pointops._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