
import torch
import torch.nn as nn
import torch.nn.functional as F
import math

try:
    from . import _C
except ImportError:
    pass

def to_sylvester(A):
    i = torch.arange(32, device=A.device).view(1, 32)
    d = torch.arange(32, device=A.device).view(32, 1)
    rows_idx = i.expand(32, 32)
    cols_idx = (i + d) % 32
    return A[..., rows_idx, cols_idx].contiguous()

def from_sylvester(syl_blocks):
    i = torch.arange(32, device=syl_blocks.device).view(32, 1)
    j = torch.arange(32, device=syl_blocks.device).view(1, 32)
    diags_idx = (j - i) % 32
    elems_idx = i.expand(32, 32)
    return syl_blocks[..., diags_idx, elems_idx]

def _matmul_core(A, B):
    is_2d = A.dim() == 2
    if is_2d:
        A = A.unsqueeze(0)
        B = B.unsqueeze(0)
        
    Batch, M, K = A.shape
    _, K2, N = B.shape
    
    pad_m = (32 - M % 32) % 32
    pad_k = (32 - K % 32) % 32
    pad_n = (32 - N % 32) % 32
    
    A_padded = F.pad(A, (0, pad_k, 0, pad_m)) 
    B_padded = F.pad(B, (0, pad_n, 0, pad_k))
    
    M_blocks = A_padded.shape[1] // 32
    K_blocks = A_padded.shape[2] // 32
    N_blocks = B_padded.shape[2] // 32
    
    A_blocked = A_padded.view(Batch, M_blocks, 32, K_blocks, 32).permute(0, 1, 3, 2, 4)
    B_blocked = B_padded.view(Batch, K_blocks, 32, N_blocks, 32).permute(0, 1, 3, 2, 4)
    
    V_A = to_sylvester(A_blocked)
    V_B = to_sylvester(B_blocked)
    
    V_C = _C.forward(V_A, V_B)
    
    C_blocked = from_sylvester(V_C)
    C_padded = C_blocked.permute(0, 1, 3, 2, 4).reshape(Batch, M_blocks * 32, N_blocks * 32)
    
    C_out = C_padded[:, :M, :N]
    
    res = C_out.squeeze(0) if is_2d else C_out
    return res.clone()

class SIDRCFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, A, B):
        ctx.save_for_backward(A, B)
        return _matmul_core(A, B)

    @staticmethod
    def backward(ctx, grad_output):
        A, B = ctx.saved_tensors
        grad_A = grad_B = None
        if ctx.needs_input_grad[0]:
            grad_A = _matmul_core(grad_output, B.transpose(-1, -2))
        if ctx.needs_input_grad[1]:
            grad_B = _matmul_core(A.transpose(-1, -2), grad_output)
        return grad_A, grad_B

def matmul(A, B):
    return SIDRCFunction.apply(A, B)

class SylvesterLinear(nn.Module):
    def __init__(self, in_features, out_features, bias=True, dtype=torch.float16):
        super().__init__()
        self.weight = nn.Parameter(torch.empty((out_features, in_features), dtype=dtype))
        if bias:
            self.bias = nn.Parameter(torch.empty(out_features, dtype=dtype))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        if self.bias is not None:
            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
            nn.init.uniform_(self.bias, -bound, bound)

    def forward(self, input):
        out = matmul(input, self.weight.t())
        if self.bias is not None:
            out = out + self.bias
        return out
