import bitblas
import torch
import torch.nn as nn

from dataclasses import dataclass
from typing import Literal
from bitblas.cache import OperatorCache
from torch.nn import functional as F


def gelu_approx(x):
    return F.gelu(x, approximate="tanh")


@dataclass
class LinearWeights:
    weight: torch.Tensor
    bias: torch.Tensor


class Linear(nn.Module):
    """
    Linear layer with support for bitblas quantization.
    If dtype is torch.int8, it uses bitblas for quantization.
    Otherwise, it uses a standard nn.Linear layer.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        dtype: torch.dtype = None,
        group_size: int = 128,
    ):
        super().__init__()

        if dtype == torch.int8:
            self.linear = bitblas.Linear(
                in_features=in_features,
                out_features=out_features,
                bias=bias,
                with_zeros=True,
                zeros_mode="original",
                with_scaling=True,
                A_dtype="float16",
                W_dtype="uint4",
                accum_dtype="float16",
                out_dtype="float16",
                fast_decoding=True,
                enable_tuning=True,
                group_size=group_size,
            )
        else:
            self.linear = nn.Linear(
                in_features=in_features,
                out_features=out_features,
                bias=bias,
                dtype=torch.float16,
            )

    def forward(self, x):
        return self.linear(x)

    @property
    def weight(self) -> torch.Tensor:
        try:
            return self.linear.weight
        except AttributeError:
            return self.linear.qweight

    @property
    def bias(self) -> torch.Tensor:
        return self.linear.bias


def linear(x: torch.Tensor, w: LinearWeights) -> torch.Tensor:
    return F.linear(x, w.weight, w.bias)


@dataclass
class LayerNormWeights:
    weight: torch.Tensor
    bias: torch.Tensor


def layer_norm(x: torch.Tensor, w: LayerNormWeights) -> torch.Tensor:
    return F.layer_norm(x, w.bias.shape, w.weight, w.bias)


@dataclass
class MLPWeights:
    fc1: LinearWeights
    fc2: LinearWeights
    act: Literal["gelu_approx"] = "gelu_approx"


def mlp(x: torch.Tensor, w: MLPWeights) -> torch.Tensor:

    x = w.fc1(x)
    x = gelu_approx(x)
    x = w.fc2(x)
    return x


@dataclass
class AttentionWeights:
    qkv: LinearWeights
    proj: LinearWeights


def attn(x: torch.Tensor, w: AttentionWeights, n_heads: int) -> torch.Tensor:
    bsz, q_len, d_model = x.shape
    head_dim = d_model // n_heads

    q, k, v = [
        t.view(bsz, q_len, n_heads, head_dim).transpose(1, 2)
        for t in linear(x, w.qkv).chunk(3, dim=-1)
    ]
    out = F.scaled_dot_product_attention(q, k, v)
    out = out.transpose(1, 2).reshape(bsz, q_len, d_model)
    out = linear(out, w.proj)
    return out