import torch
import math


def rotate_3x3_kernel_adaptive_matrixcompute(weights, num_experts, kernel_theta):
    """
    Args:
        weights: tensor, shape = [num_experts, Cout, Cin, k, k]
        num_experts: number of experts
        kernel_theta: a tensor vector of float number with the size of num_experts
    Return:
        weights: rotated weights tensor, shape = [num_experts, Cout, Cin, 3, 3]
    """
    assert(weights.shape[3] == 3)
    assert(weights.shape[4] == 3)
    assert(max(kernel_theta) <= 45.0)
    assert(min(kernel_theta) >= - 45.0)

    is_clockwise = kernel_theta < 0
    kernel_theta = torch.abs(kernel_theta)

    x_vector = torch.cos(kernel_theta / 180. * math.pi)
    y_vector = torch.sin(kernel_theta / 180. * math.pi)

    a_vector = x_vector - y_vector
    b_vector = torch.mul(x_vector, y_vector)
    c_vector = x_vector + y_vector

    alpha = torch.zeros(num_experts * 9, num_experts * 9, device=weights.device)
    for idx in range(num_experts):
        x, y = x_vector[idx], y_vector[idx]
        a, b, c = a_vector[idx], b_vector[idx], c_vector[idx]
        if is_clockwise[idx] == False:
            # anti-clockwise (theta > 0)
            sub_alpha = torch.tensor([[ a, 1-a, 0.,  0.,     0., 0.,  0.,  0., 0.],           # w'(-1, 1)
                                      [0., x-b,  b,  0., 1-c+b, y-b,  0.,  0., 0.],           # w'( 0, 1)
                                      [0.,  0.,  a,  0.,    0., 1-a,  0.,  0., 0.],           # w'( 1, 1)
                                      [ b, y-b, 0., x-b, 1-c+b,  0.,  0.,  0., 0.],           # w'(-1, 0)
                                      [0.,  0., 0.,  0.,    1.,  0.,  0.,  0., 0.],           # w'( 0, 0)
                                      [0.,  0., 0.,  0., 1-c+b, x-b,  0., y-b,  b],           # w'( 1, 0)
                                      [0.,  0., 0., 1-a,    0.,  0.,   a,  0., 0.],           # w'(-1,-1)
                                      [0.,  0., 0., y-b, 1-c+b,  0.,   b, x-b, 0.],           # w'( 0,-1)
                                      [0.,  0., 0.,  0.,    0.,  0.,  0., 1-a,  a]],          # w'( 1,-1)
                                     device=weights.device)
        else:
            # clockwise (theta < 0)
            sub_alpha = torch.tensor([[ a,  0., 0., 1-a,     0., 0.,  0.,  0., 0.],           # w'(-1, 1)
                                      [ b, x-b, 0., y-b, 1-c+b,  0.,  0.,  0., 0.],           # w'( 0, 1)
                                      [0., 1-a,  a,  0.,    0.,  0.,  0.,  0., 0.],           # w'( 1, 1)
                                      [0.,  0., 0., x-b, 1-c+b,  0.,   b, y-b, 0.],           # w'(-1, 0)
                                      [0.,  0., 0.,  0.,    1.,  0.,  0.,  0., 0.],           # w'( 0, 0)
                                      [0., y-b,  b,  0., 1-c+b, x-b,  0.,  0., 0.],           # w'( 1, 0)
                                      [0.,  0., 0.,  0.,    0.,  0.,   a, 1-a, 0.],           # w'(-1,-1)
                                      [0.,  0., 0.,  0., 1-c+b, y-b,  0., x-b,  b],           # w'( 0,-1)
                                      [0.,  0., 0.,  0.,    0., 1-a,  0.,  0.,  a]],          # w'( 1,-1)
                                     device=weights.device)
        alpha[idx * 9:(idx+1) * 9, idx * 9:(idx+1) * 9] = sub_alpha

    _, Cout, Cin, _, _ = weights.shape  # [num_experts, Cout, Cin, 3, 3]
    weights = weights.transpose(0, 1).transpose(1, 2)
    # ---> [Cout, num_experts, Cin, 3, 3]
    # ---> [Cout, Cin, num_experts, 3, 3]
    weights = weights.contiguous().view(Cout * Cin, num_experts * 9)
    # ---> [Cout * Cin, num_experts * 9]
    weights = weights.transpose(0, 1)
    # ---> [num_experts * 9, Cout * Cin]
    weights = torch.mm(alpha, weights)
    # [num_experts * 9, num_experts * 9] x [num_experts * 9, Cout * Cin] ---> [num_experts * 9, Cout * Cin]
    weights = weights.transpose(0, 1).view(Cout, Cin, num_experts, 3, 3)
    # ---> [Cout * Cin, num_experts * 9]
    # ---> [Cout, Cin, num_experts, 3, 3]
    weights = weights.transpose(1, 2).transpose(0, 1)
    # ---> [Cout, num_experts, Cin, 3, 3]
    # ---> [num_experts, Cout, Cin, 3, 3]

    return weights