import inspect
from dataclasses import dataclass
from typing import Optional

import torch
import torch.nn as nn

from .fused_linear_cross_entropy import LigerFusedLinearCrossEntropyFunction
from .geglu import LigerGELUMulFunction
from .rms_norm import LigerRMSNormFunction
from .rope import LigerRopeFunction
from .swiglu import LigerSiLUMulFunction
from .tiled_mlp import apply_tiled_mlp


# NOTE: Not compile-friendly --> large deviations to the original implementation under compile
class LigerRMSNorm(nn.Module):
    weight: nn.Parameter
    variance_epsilon: float

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return LigerRMSNormFunction.apply(
            hidden_states,
            self.weight,
            self.variance_epsilon,
            0,
            "llama",
            True,
            None,
        )

    def extra_repr(self):
        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"


class LigerLinear(nn.Module):
    weight: nn.Parameter
    bias: nn.Parameter | None

    can_torch_compile = True

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        def _forward(module, x):
            return nn.functional.linear(x, module.weight, module.bias)

        compute_params = [p for p in self.parameters() if p.requires_grad]

        return apply_tiled_mlp(
            fn=_forward,
            mlp_module=self,
            x=input,
            num_shards=None,
            compute_params=compute_params,
        )


class LigerSwiGLUMLP(nn.Module):
    gate_proj: nn.Linear
    up_proj: nn.Linear
    down_proj: nn.Linear

    can_torch_compile = True

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(LigerSiLUMulFunction.apply(self.gate_proj(x), self.up_proj(x)))


class LigerGEGLUMLP(nn.Module):
    gate_proj: nn.Linear
    up_proj: nn.Linear
    down_proj: nn.Linear

    can_torch_compile = True

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(LigerGELUMulFunction.apply(self.gate_proj(x), self.up_proj(x)))


class LigerTiledSwiGLUMLP(nn.Module):
    gate_proj: nn.Linear
    up_proj: nn.Linear
    down_proj: nn.Linear

    can_torch_compile = True

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        def _mlp_forward(module, x):
            """Internal MLP forward function for tiled computation."""
            gate = module.gate_proj(x)
            up = module.up_proj(x)
            return module.down_proj(LigerSiLUMulFunction.apply(gate, up))

        compute_params = [p for p in self.parameters() if p.requires_grad]

        return apply_tiled_mlp(
            fn=_mlp_forward,
            mlp_module=self,
            x=x,
            num_shards=None,
            compute_params=compute_params,
        )


class LigerTiledGEGLUMLP(nn.Module):
    gate_proj: nn.Linear
    up_proj: nn.Linear
    down_proj: nn.Linear

    can_torch_compile = True

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        def _mlp_forward(module, x):
            """Internal MLP forward function for tiled computation."""
            gate = module.gate_proj(x)
            up = module.up_proj(x)
            return module.down_proj(LigerGELUMulFunction.apply(gate, up))

        compute_params = [p for p in self.parameters() if p.requires_grad]

        return apply_tiled_mlp(
            fn=_mlp_forward,
            mlp_module=self,
            x=x,
            num_shards=None,
            compute_params=compute_params,
        )


def liger_rotary_pos_emb(
    q: torch.Tensor,
    k: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    unsqueeze_dim: int = 1,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Apply standard rotary positional embedding to ``q`` and ``k``."""
    return LigerRopeFunction.apply(q, k, cos, sin, None, unsqueeze_dim)


@dataclass
class CrossEntropyOutput:
    loss: torch.Tensor
    z_loss: Optional[torch.Tensor] = None
    token_accuracy: Optional[torch.Tensor] = None
    predicted_tokens: Optional[torch.Tensor] = None


def liger_fused_linear_cross_entropy(
    input: torch.Tensor,
    weight: torch.Tensor,
    target: torch.Tensor,
    bias: Optional[torch.Tensor] = None,
    ce_weight: Optional[torch.Tensor] = None,
    ignore_index: int = -100,
    lse_square_scale: float = 0.0,
    label_smoothing: float = 0.0,
    reduction: str = "mean",
    softcap: Optional[float] = None,
    return_z_loss: bool = False,
    accum_dtype: Optional[torch.dtype] = None,
    use_token_scaling: bool = False,
    return_token_accuracy: bool = False,
    return_predicted_tokens: bool = False,
):
    loss, z_loss, token_accuracy, predicted_tokens = LigerFusedLinearCrossEntropyFunction.apply(
        input,
        weight,
        target,
        bias,
        ce_weight,
        ignore_index,
        lse_square_scale,
        label_smoothing,
        reduction,
        softcap,
        return_z_loss,
        accum_dtype,
        use_token_scaling,
        return_token_accuracy,
        return_predicted_tokens,
    )
    if not return_z_loss and not return_token_accuracy and not return_predicted_tokens:
        return loss
    return CrossEntropyOutput(
        loss=loss,
        z_loss=z_loss,
        token_accuracy=token_accuracy,
        predicted_tokens=predicted_tokens,
    )


# NOTE: We have this intentional graph break as we encounter issues such as IMAs and Cublas errors.
#       We know that this is an optimized kernel already so there is less ways to
#       fuse it either way; we rely on torch compile to go through the base model to optimize.
@torch.compiler.disable
def LigerForCausalLMLoss(
    logits: None,  # to match transformers signature
    labels: torch.Tensor,
    vocab_size: int,  # to match transformers signature
    num_items_in_batch: Optional[int] = None,
    ignore_index: int = -100,
    shift_labels: Optional[torch.Tensor] = None,
    hidden_states: torch.Tensor | None = None,
    lm_head_weight: torch.Tensor | None = None,
    lm_head_bias: torch.Tensor | None = None,
    **kwargs,
):
    # To match signature we hide these behind the kwargs but we expect a few kwargs to exist
    hidden_size = kwargs.pop("hidden_size", None)
    final_logit_softcapping = kwargs.pop("final_logit_softcapping", None)

    if hidden_size is None or hidden_states is None or lm_head_weight is None:
        raise ValueError(
            f"`LigerForCausalLMLoss` requires the LLM's weight (found `{lm_head_weight is not None}`),"
            f"the last hidden state (found `{hidden_states is not None}`), and the `hidden_size`"
            f"(found `{hidden_size is not None}`). Please make sure to pass the necessary kwargs."
        )

    applicable_params = inspect.signature(liger_fused_linear_cross_entropy).parameters
    kwargs = {k: v for k, v in kwargs.items() if k in applicable_params}

    if shift_labels is None:
        labels = nn.functional.pad(labels, (0, 1), value=ignore_index)
        shift_labels = labels[..., 1:].contiguous()

    hidden_states = hidden_states.view(-1, hidden_size)
    shift_labels = shift_labels.view(-1).to(hidden_states.device)

    reduction = "sum" if num_items_in_batch is not None else "mean"
    loss = liger_fused_linear_cross_entropy(
        hidden_states,
        lm_head_weight,
        shift_labels,
        bias=lm_head_bias,
        reduction=reduction,
        ignore_index=ignore_index,
        softcap=final_logit_softcapping,
        return_token_accuracy=False,
        return_predicted_tokens=False,
        **kwargs,
    )

    if reduction == "sum":
        loss = loss / num_items_in_batch

    return loss


# Add torch compile support for functions - in this case it's to allow this to be used
# but the function itself will not be compiled (see the note at the function)
liger_rotary_pos_emb.can_torch_compile = True
LigerForCausalLMLoss.can_torch_compile = True


__all__ = [
    "LigerRMSNorm",
    "LigerLinear",
    "LigerSwiGLUMLP",
    "LigerGEGLUMLP",
    "LigerTiledSwiGLUMLP",
    "LigerTiledGEGLUMLP",
    "liger_rotary_pos_emb",
    "LigerForCausalLMLoss",
]