import torch
import numpy as np
import _trme_backend

class TRMEMatmul(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, weight, bias=None, block_size=4):
        """
        input: (N, K)
        weight: (M, K)
        bias: (M,) Optional
        """
        if isinstance(bias, int):
            block_size = bias
            bias = None

        # Quantization (A -> [0..3], B -> Int8)
        # Note: We do simple casting here for the prototype.
        # Ideally calling C++ quantization logic.
        A_cpu = input.detach().cpu().numpy()
        A_q = np.clip(np.abs(A_cpu), 0, 3).astype(np.int8)

        B_cpu = weight.detach().cpu().numpy().T
        B_q = np.ascontiguousarray(np.clip(B_cpu, -127, 127).astype(np.int8))

        N, K = A_q.shape
        _, M = B_q.shape

        # Flatten for Rust FFI
        a_bytes = A_q.tobytes()
        b_bytes = B_q.tobytes()

        # Call Rust Backend
        # rsr_gemm returns list[int] (flattened C)
        c_flat = _trme_backend.rsr_gemm(a_bytes, b_bytes, N, K, M, int(block_size))

        # Reshape and Convert to Tensor
        C_q = np.array(c_flat, dtype=np.int32).reshape(N, M)
        output = torch.from_numpy(C_q.astype(np.float32))

        if bias is not None:
             output += bias.cpu()

        ctx.save_for_backward(input, weight, bias)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        input, weight, bias = ctx.saved_tensors
        grad_input = grad_weight = grad_bias = None

        if ctx.needs_input_grad[0]:
            grad_input = grad_output.matmul(weight)
        if ctx.needs_input_grad[1]:
            grad_weight = grad_output.t().matmul(input)
        if bias is not None and ctx.needs_input_grad[2]:
            grad_bias = grad_output.sum(0)

        return grad_input, grad_weight, grad_bias, None

def matmul(input, weight):
    # Basic matmul interface
    return TRMEMatmul.apply(input, weight)

def linear(input, weight, bias=None, block_size=4):
    # nn.Linear interface
    return TRMEMatmul.apply(input, weight, bias, block_size)
