from einops.einops import rearrange
import torch
from torch.autograd import Function

import score_computation_cuda
import value_aggregation_cuda


class ScoreComputation(Function):
    @staticmethod
    def forward(ctx, query, key, index):
        is_fp16 = query.dtype == torch.float16
        if is_fp16:
            query = query.float()
            key = key.float()
        x = score_computation_cuda.score_forward(query, key, index)
        ctx.save_for_backward(query, key, index)

        if is_fp16:
            x[0] = x[0].half()
        return x[0]

    @staticmethod
    def backward(ctx, grad_output):
        input1, input2, index = ctx.saved_tensors
        grad_output = grad_output.contiguous()
        x = score_computation_cuda.score_backward(grad_output, input1, input2, index)
        return x[0], x[1], None


score_computation_op = ScoreComputation.apply


class value_aggregation(Function):
    @staticmethod
    def forward(ctx, score, value, index):
        ctx.save_for_backward(score, value, index)
        f = score.shape[2]
        score = rearrange(
            score, "b n f K h -> b (n f) K h"
        )  # [b, N, 4, 4K, H] -> [b, 4N, 4K, H]
        index = rearrange(
            index, "b n f K h -> b (n f) K h"
        )  # [b, N, 4, 4K, H] -> [b, 4N, 4K, H]
        b, N, _, H = score.shape
        D = value.shape[-1]
        # value [b, M, H, D]
        output = score.new_zeros([b, N, H, D]).contiguous()  # b, 4N, H, D

        is_fp16 = score.dtype == torch.float16
        if is_fp16:
            score = score.float()
            value = value.float()
            output = output.float()
        value_aggregation_cuda.value_aggregation_forward(score, value, index, output)
        output = rearrange(output, "b (n f) h d -> b n f h d", f=f)
        if is_fp16:
            output = output.half()
        return output

    @staticmethod
    def backward(ctx, grad_output):
        score, value, index = ctx.saved_tensors
        f = score.shape[2]
        score = rearrange(score, "b n f K h -> b (n f) K h")
        index = rearrange(index, "b n f K h -> b (n f) K h")

        grad_output = grad_output.contiguous()

        grad_score = score.new_zeros(score.shape).contiguous()
        grad_value = value.new_zeros(value.shape).contiguous()

        value_aggregation_cuda.value_aggregation_backward(
            grad_output, score, value, index, grad_score, grad_value
        )
        grad_score = rearrange(grad_score, "b (n f) K h -> b n f K h", f=f)
        return grad_score, grad_value, None


value_aggregation_op = value_aggregation.apply
