# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
import sys
from pathlib import Path

import torch
from torch import Tensor

import litgpt.model
from litgpt.constants import _THUNDER_AVAILABLE
from litgpt.model import LLaMAMLP as OriginalLLaMAMLP
from thunder.core.proxies import TensorProxy
from thunder.core.transforms import get_grad, mean_backward, put_grads
from thunder.extend import OperatorExecutor, register_executor
from thunder.torch import ne, sum, true_divide

if _THUNDER_AVAILABLE:
    import thunder
    import thunder.torch as ltorch

sys.path.append(str(Path(__file__).parent))

import kernels

unsloth_ex = OperatorExecutor("unsloth", version="0.1")
register_executor(unsloth_ex)


"""
====================
 Cross Entropy Loss
====================
"""


def unsloth_cross_entropy_meta(logits: TensorProxy, labels: TensorProxy) -> tuple[TensorProxy, TensorProxy]:
    return (
        TensorProxy(
            shape=(logits.shape[0],),
            # the cross entropy kernel only supports float32
            dtype=thunder.dtypes.float32,
            device=logits.device,
            requires_grad=logits.requires_grad,
        ),
        TensorProxy(shape=(logits.shape[0],), dtype=thunder.dtypes.float32, device=logits.device, requires_grad=False),
    )


unsloth_cross_entropy = unsloth_ex.register_operator(
    "unsloth_cross_entropy", meta=unsloth_cross_entropy_meta, fn=kernels.cross_entropy_loss._cross_entropy_forward_impl
)


def unsloth_cross_entropy_backward_impl(dlosses: Tensor, logits: Tensor, labels: Tensor, logsumexp: Tensor) -> Tensor:
    # clone() because the kernel writes the grads in the logits
    return kernels.cross_entropy_loss._cross_entropy_backward_impl(dlosses, logits.clone(), logsumexp, labels)


def unsloth_cross_entropy_backward_meta(
    dlosses: TensorProxy, logits: TensorProxy, logsumexp: TensorProxy, labels: TensorProxy
) -> TensorProxy:
    return thunder.TensorProxy(like=logits)


unsloth_cross_entropy_backward = unsloth_ex.register_operator(
    "unsloth_cross_entropy_backward", meta=unsloth_cross_entropy_backward_meta, fn=unsloth_cross_entropy_backward_impl
)


def unsloth_cross_entropy_checker(
    logits: TensorProxy,
    labels: TensorProxy,
    weight: TensorProxy | None = None,
    size_average: bool | None = None,
    ignore_index: int = -100,
    reduce: bool | None = None,
    reduction: str = "mean",
    label_smoothing: float = 0.0,
) -> bool:
    return (
        weight is None
        and size_average is None
        and reduce is None
        and reduction in ("none", "mean")
        and ignore_index == -100
        and label_smoothing == 0.0
        and logits.device.type == "cuda"
        and labels.device.type == "cuda"
    )


def cross_entropy_to_unsloth(
    logits: TensorProxy,
    labels: TensorProxy,
    weight: TensorProxy | None = None,
    size_average: bool | None = None,
    ignore_index: int = -100,
    reduce: bool | None = None,
    reduction: str = "mean",
    label_smoothing: float = 0.0,
) -> tuple[TensorProxy, TensorProxy]:
    loss, logsumexp = unsloth_cross_entropy(logits, labels)
    if reduction == "mean":
        # "mean" reduction is not part of the kernel
        # TODO: this doesn't consider that all elements could be masked, causing a division by 0
        n_items = sum(ne(labels, -100))
        loss = true_divide(sum(loss), n_items)
    elif reduction != "none":
        raise NotImplementedError(reduction)
    return loss, logsumexp


def unsloth_cross_entropy_grad(
    logits: TensorProxy,
    labels: TensorProxy,
    weight: TensorProxy | None = None,
    size_average: bool | None = None,
    ignore_index: int = -100,
    reduce: bool | None = None,
    reduction: str = "mean",
    label_smoothing: float = 0.0,
) -> TensorProxy:
    loss, logsumexp = cross_entropy_to_unsloth(**locals())
    grad = get_grad(loss)
    if reduction == "mean":
        grad = mean_backward(logsumexp.ndim, logsumexp.shape, (0,), grad)
    logits_grad = unsloth_cross_entropy_backward(grad, logits, labels, logsumexp)
    put_grads((logits,), (logits_grad,))
    return loss


# registers as cross entropy implementation, including the execution transform and now a grad transform
unsloth_ex.register_implementation(
    ltorch.cross_entropy,
    checker=unsloth_cross_entropy_checker,
    execution_transform=lambda *args: cross_entropy_to_unsloth(*args)[0],
    grad_transform=unsloth_cross_entropy_grad,
)


"""
=========
 RMSNorm
=========

The RMSNorm kernel is not integrated because it's not numerically equal and it doesn't compute the gradient for the
weight, just for the input.
"""


"""
========
 SwiGLU
========
"""


def swiglu(e: torch.Tensor, g: torch.Tensor) -> torch.Tensor:
    return torch.nn.functional.silu(e) * g


class ThunderLLaMAMLP(OriginalLLaMAMLP):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x_fc_1 = self.fc_1(x)
        x_fc_2 = self.fc_2(x)
        x = swiglu(x_fc_1, x_fc_2)
        return self.proj(x)


litgpt.model.LLaMAMLP = ThunderLLaMAMLP


def swiglu_forward_meta(e: TensorProxy, g: TensorProxy) -> TensorProxy:
    return TensorProxy(like=e)


litgpt_swiglu = unsloth_ex.register_operator("litgpt_swiglu", meta=swiglu_forward_meta, fn=swiglu, replaces=swiglu)


unsloth_swiglu_forward = unsloth_ex.register_operator(
    "unsloth_swiglu_forward", meta=swiglu_forward_meta, fn=lambda *args: kernels.swiglu_fg_kernel(*args)
)


def unsloth_swiglu_backward_meta(DW: TensorProxy, e: TensorProxy, g: TensorProxy) -> tuple[TensorProxy, TensorProxy]:
    return TensorProxy(like=g), TensorProxy(like=e)


def unsloth_swiglu_backward_fn(DW: Tensor, e: Tensor, g: Tensor) -> tuple[Tensor, tuple]:
    B, T, n_embd = e.shape
    e = e.view(-1, n_embd)
    g = g.view(-1, n_embd)
    DW, e, g = kernels.swiglu_DWf_DW_dfg_kernel(DW, e, g)
    e = e.view(B, T, n_embd)
    g = g.view(B, T, n_embd)
    return g, e


unsloth_swiglu_backward = unsloth_ex.register_operator(
    "unsloth_swiglu_backward", meta=unsloth_swiglu_backward_meta, fn=unsloth_swiglu_backward_fn
)


def swiglu_to_unsloth_checker(e: TensorProxy, g: TensorProxy) -> bool:
    return e.device.type == "cuda" and g.device.type == "cuda"


def unsloth_swiglu_grad(e: TensorProxy, g: TensorProxy) -> TensorProxy:
    h = unsloth_swiglu_forward(**locals())
    grad = get_grad(h)
    e_grad, g_grad = unsloth_swiglu_backward(grad, e, g)
    put_grads((e, g), (e_grad, g_grad))
    return h


unsloth_ex.register_implementation(
    litgpt_swiglu,
    checker=swiglu_to_unsloth_checker,
    execution_transform=unsloth_swiglu_forward,
    grad_transform=unsloth_swiglu_grad,
)


"""
======
 RoPE
======
"""


def apply_rope_meta(x: TensorProxy, cos: TensorProxy, sin: TensorProxy) -> TensorProxy:
    return TensorProxy(like=x)


apply_rope = unsloth_ex.register_operator(
    "litgpt_apply_rope", like=apply_rope_meta, fn=litgpt.model.apply_rope, replaces=litgpt.model.apply_rope
)


def unsloth_apply_rope_meta(
    Q: TensorProxy, cos: TensorProxy, sin: TensorProxy
) -> tuple[TensorProxy, TensorProxy, TensorProxy, int, int, int]:
    batch, n_heads, seq_len, head_dim = Q.shape
    assert seq_len <= cos.shape[-2]
    BLOCK_SIZE, num_warps = kernels.calculate_settings(head_dim // 2)
    div, mod = divmod(n_heads, kernels.rope_embedding.ROPE_GROUP_SIZE)
    n_groups = div + (mod != 0)
    return TensorProxy(like=Q), cos, sin, n_groups, BLOCK_SIZE, num_warps


unsloth_apply_rope = unsloth_ex.register_operator(
    "unsloth_apply_rope", meta=unsloth_apply_rope_meta, fn=kernels._rope_embedding_forward_impl
)


def unsloth_apply_rope_backward_meta(
    dY: TensorProxy, cos: TensorProxy, sin: TensorProxy, n_groups: int, BLOCK_SIZE: int, num_warps: int
) -> TensorProxy:
    return TensorProxy(like=dY)


unsloth_apply_rope_backward = unsloth_ex.register_operator(
    "unsloth_apply_rope_backward", meta=unsloth_apply_rope_backward_meta, fn=kernels._rope_embedding_backward_impl
)


def apply_rope_to_unsloth_checker(x: TensorProxy, cos: TensorProxy, sin: TensorProxy) -> bool:
    return len(x.shape) == 4 and x.device.type == "cuda" and cos.device.type == "cuda" and sin.device.type == "cuda"


def unsloth_apply_rope_grad(x: TensorProxy, cos: TensorProxy, sin: TensorProxy) -> TensorProxy:
    Q, cos, sin, n_groups, BLOCK_SIZE, num_warps = unsloth_apply_rope(x, cos, sin)
    dY = get_grad(Q)
    dX = unsloth_apply_rope_backward(dY, cos, sin, n_groups, BLOCK_SIZE, num_warps)
    put_grads((x,), (dX,))
    return Q


unsloth_ex.register_implementation(
    apply_rope,
    checker=apply_rope_to_unsloth_checker,
    execution_transform=lambda *args: unsloth_apply_rope(*args)[0],
    grad_transform=unsloth_apply_rope_grad,
)