# Copyright (c) Together
# Apache 2.0 - Author: Michael Poli
# Adapted for the minimal Evo1 HF port.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor


def grab_first_if_tuple(x):
    if x.__class__.__name__ == "tuple":
        return x[0]
    return x


class RMSNorm(torch.nn.Module):
    def __init__(self, config):
        super().__init__()
        self.eps = config.eps
        self.hidden_size = config.hidden_size
        self.scale = torch.nn.Parameter(torch.ones(self.hidden_size))
        self.register_parameter("scale", self.scale)
        self.use_flash_rmsnorm = config.get("use_flash_rmsnorm", False)
        if self.use_flash_rmsnorm:
            from flash_attn.ops.rms_norm import rms_norm as rmsnorm_func
            self.rmsnorm_func = rmsnorm_func

    def forward(self, x):
        if self.use_flash_rmsnorm:
            return self.rmsnorm_func(x, self.scale, self.eps)
        y = x / (x.norm(2, dim=-1, keepdim=True) * self.hidden_size ** (-1.0 / 2) + self.eps)
        return self.scale * y


class ParallelGatedMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        multiple_of = config.get("inner_size_multiple_of", 64)
        self.act_type = config.get("mlp_activation", "silu")
        if self.act_type == "gelu":
            self.act = F.gelu
        elif self.act_type == "silu":
            self.act = F.silu
        else:
            raise NotImplementedError(f"Unknown mlp_activation: {self.act_type}")

        self.multiple_of = multiple_of * config.model_parallel_size
        inner_size = int(2 * config.hidden_size * 4 / 3)
        inner_size = self.multiple_of * ((inner_size + self.multiple_of - 1) // self.multiple_of)
        if config.get("inner_mlp_size", None) is not None:
            inner_size = config.inner_mlp_size

        self.l1 = nn.Linear(config.hidden_size, inner_size, bias=False)
        self.l2 = nn.Linear(config.hidden_size, inner_size, bias=False)
        self.l3 = nn.Linear(inner_size, config.hidden_size, bias=False)

    def forward(self, z):
        z1, z2 = self.l1(z), self.l2(z)
        z1, z2 = grab_first_if_tuple(z1), grab_first_if_tuple(z2)
        y = self.l3(self.act(z1) * z2)
        return grab_first_if_tuple(y)


class VocabParallelEmbedding(nn.Embedding):
    """Single-process variant of the original VocabParallelEmbedding.

    The original supports tensor-parallel embedding sharding. We keep the
    naming so existing checkpoints load directly, but drop all distributed
    paths since this minimal port runs on a single device.
    """

    def __init__(self, config):
        vocab_size = config.vocab_size
        padding_idx = config.get("padding_idx", None)
        super().__init__(vocab_size, embedding_dim=config.hidden_size, padding_idx=padding_idx)

    def embed(self, x: Tensor) -> Tensor:
        return self.forward(x)

    def unembed(self, u: Tensor) -> Tensor:
        return u @ self.weight.T