# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
from dataclasses import dataclass
from typing import Optional

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


def find_multiple(n: int, k: int) -> int:
  if n % k == 0:
    return n
  return n + k - (n % k)


@dataclass
class ModelArgs:
  block_size: int = 2048
  vocab_size: int = 32000
  n_layer: int = 32
  n_head: int = 32
  dim: int = 4096
  intermediate_size: int = None
  n_local_heads: int = -1
  head_dim: int = 64
  rope_base: float = 10000
  norm_eps: float = 1e-5
  num_experts: int = 8
  num_activated_experts: int = 2

  def __post_init__(self):
    if self.n_local_heads == -1:
      self.n_local_heads = self.n_head
    if self.intermediate_size is None:
      hidden_dim = 4 * self.dim
      n_hidden = int(2 * hidden_dim / 3)
      self.intermediate_size = find_multiple(n_hidden, 256)
    self.head_dim = self.dim // self.n_head

  @classmethod
  def from_name(cls, name: str):
    if name in transformer_configs:
      return cls(**transformer_configs[name])
    # fuzzy search
    config = [
        config for config in transformer_configs
        if config in str(name).upper() or config in str(name)
    ]
    assert len(config) == 1, name
    return cls(**transformer_configs[config[0]])


transformer_configs = {
    "Mixtral-8x7B-v0.1":
        dict(
            block_size=32768,
            n_layer=32,
            n_head=32,
            n_local_heads=8,
            dim=4096,
            intermediate_size=14336,
            rope_base=1000000.0,
            num_experts=8,
            num_activated_experts=2),
}


class KVCache(nn.Module):

  def __init__(self,
               max_batch_size,
               max_seq_length,
               n_heads,
               head_dim,
               dtype=torch.bfloat16):
    super().__init__()
    cache_shape = (max_batch_size, n_heads, max_seq_length, head_dim)
    self.register_buffer('k_cache', torch.zeros(cache_shape, dtype=dtype))
    self.register_buffer('v_cache', torch.zeros(cache_shape, dtype=dtype))

  def update(self, input_pos, k_val, v_val):
    # input_pos: [S], k_val: [B, H, S, D]
    assert input_pos.shape[0] == k_val.shape[2]

    k_out = self.k_cache
    v_out = self.v_cache
    k_out[:, :, input_pos] = k_val
    v_out[:, :, input_pos] = v_val

    return k_out, v_out


class Transformer(nn.Module):

  def __init__(self, config: ModelArgs) -> None:
    super().__init__()
    self.config = config

    self.tok_embeddings = nn.Embedding(config.vocab_size, config.dim)
    self.layers = nn.ModuleList(
        TransformerBlock(config) for _ in range(config.n_layer))
    self.norm = RMSNorm(config.dim, eps=config.norm_eps)
    self.output = nn.Linear(config.dim, config.vocab_size, bias=False)

    self.freqs_cis: Optional[Tensor] = None
    self.mask_cache: Optional[Tensor] = None
    self.max_batch_size = -1
    self.max_seq_length = -1

  def setup_caches(self, max_batch_size, max_seq_length):
    if self.max_seq_length >= max_seq_length and self.max_batch_size >= max_batch_size:
      return
    head_dim = self.config.dim // self.config.n_head
    max_seq_length = find_multiple(max_seq_length, 8)
    self.max_seq_length = max_seq_length
    self.max_batch_size = max_batch_size
    for b in self.layers:
      b.attention.kv_cache = KVCache(max_batch_size, max_seq_length,
                                     self.config.n_local_heads, head_dim)

    self.freqs_cis = precompute_freqs_cis(self.config.block_size,
                                          self.config.dim // self.config.n_head,
                                          self.config.rope_base)
    self.causal_mask = torch.tril(
        torch.ones(self.max_seq_length, self.max_seq_length, dtype=torch.bool))

  def forward(self, idx: Tensor, input_pos: Optional[Tensor] = None) -> Tensor:
    assert self.freqs_cis is not None, "Caches must be initialized first"
    mask = self.causal_mask[None, None, input_pos]
    freqs_cis = self.freqs_cis[input_pos]
    x = self.tok_embeddings(idx)

    for i, layer in enumerate(self.layers):
      x = layer(x, input_pos, freqs_cis, mask)
    x = self.norm(x)
    logits = self.output(x)
    return logits

  @classmethod
  def from_name(cls, name: str):
    return cls(ModelArgs.from_name(name))


class TransformerBlock(nn.Module):

  def __init__(self, config: ModelArgs) -> None:
    super().__init__()
    self.attention = Attention(config)
    self.block_sparse_moe = MOEFeedForward(config)
    self.ffn_norm = RMSNorm(config.dim, config.norm_eps)
    self.attention_norm = RMSNorm(config.dim, config.norm_eps)

  def forward(self, x: Tensor, input_pos: Tensor, freqs_cis: Tensor,
              mask: Tensor) -> Tensor:
    h = x + self.attention(self.attention_norm(x), freqs_cis, mask, input_pos)
    out = h + self.block_sparse_moe(self.ffn_norm(h))
    return out


class Attention(nn.Module):

  def __init__(self, config: ModelArgs):
    super().__init__()
    assert config.dim % config.n_head == 0

    total_head_dim = (config.n_head +
                      2 * config.n_local_heads) * config.head_dim
    # key, query, value projections for all heads, but in a batch
    self.wqkv = nn.Linear(config.dim, total_head_dim, bias=False)
    self.wo = nn.Linear(config.dim, config.dim, bias=False)
    self.kv_cache = None

    self.n_head = config.n_head
    self.head_dim = config.head_dim
    self.n_local_heads = config.n_local_heads
    self.dim = config.dim
    self._register_load_state_dict_pre_hook(self.load_hook)

  def load_hook(self, state_dict, prefix, *args):
    if prefix + "wq.weight" in state_dict:
      wq = state_dict.pop(prefix + "wq.weight")
      wk = state_dict.pop(prefix + "wk.weight")
      wv = state_dict.pop(prefix + "wv.weight")
      state_dict[prefix + "wqkv.weight"] = torch.cat([wq, wk, wv])

  def forward(self,
              x: Tensor,
              freqs_cis: Tensor,
              mask: Tensor,
              input_pos: Optional[Tensor] = None) -> Tensor:
    bsz, seqlen, _ = x.shape

    kv_size = self.n_local_heads * self.head_dim
    q, k, v = self.wqkv(x).split([self.dim, kv_size, kv_size], dim=-1)

    q = q.view(bsz, seqlen, self.n_head, self.head_dim)
    k = k.view(bsz, seqlen, self.n_local_heads, self.head_dim)
    v = v.view(bsz, seqlen, self.n_local_heads, self.head_dim)

    q = apply_rotary_emb(q, freqs_cis)
    k = apply_rotary_emb(k, freqs_cis)

    q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v))

    if self.kv_cache is not None:
      k, v = self.kv_cache.update(input_pos, k, v)

    k = k.repeat_interleave(self.n_head // self.n_local_heads, dim=1)
    v = v.repeat_interleave(self.n_head // self.n_local_heads, dim=1)
    y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0)

    y = y.transpose(1, 2).contiguous().view(bsz, seqlen, self.dim)

    y = self.wo(y)
    return y


class ConditionalFeedForward(nn.Module):

  def __init__(self, config):
    super().__init__()
    self.w1 = nn.Parameter(
        torch.empty(config.num_experts, config.intermediate_size, config.dim))
    self.w2 = nn.Parameter(
        torch.empty(config.num_experts, config.dim, config.intermediate_size))
    self.w3 = nn.Parameter(
        torch.empty(config.num_experts, config.intermediate_size, config.dim))

  def forward(self, x: Tensor, expert_indices: Tensor) -> Tensor:
    w1_weights = self.w1[expert_indices]  # [T, A, D, D]
    w3_weights = self.w3[expert_indices]  # [T, A, D, D]
    w2_weights = self.w2[expert_indices]  # [T, A, D, D]
    x1 = F.silu(torch.einsum('ti,taoi -> tao', x, w1_weights))
    x3 = torch.einsum('ti, taoi -> tao', x, w3_weights)
    expert_outs = torch.einsum('tao, taio -> tai', (x1 * x3), w2_weights)
    return expert_outs


class MOEFeedForward(nn.Module):

  def __init__(self, config) -> None:
    super().__init__()
    self.gate = nn.Linear(config.dim, config.num_experts, bias=False)
    self.cond_ffn = ConditionalFeedForward(config)
    self.dim = config.dim
    self.num_activated_experts = config.num_activated_experts

  def forward(self, x: Tensor) -> Tensor:
    x = x.view(-1, self.dim)
    # T = num_tokens, E = num_experts, D = hidden dim, A = activated experts
    # x: [T, D]
    scores = self.gate(x)  # [T, E]
    expert_weights = F.softmax(scores, dim=-1)
    expert_weights, expert_indices = torch.topk(
        expert_weights, self.num_activated_experts, dim=-1)  # [T, A], [T, A]
    expert_weights /= expert_weights.sum(dim=-1, keepdim=True)  # [T, A]
    expert_outs = self.cond_ffn(x, expert_indices)
    return torch.einsum('tai,ta -> ti', expert_outs, expert_weights)


class RMSNorm(nn.Module):

  def __init__(self, dim: int, eps: float = 1e-5):
    super().__init__()
    self.eps = eps
    self.weight = nn.Parameter(torch.ones(dim))

  def _norm(self, x):
    return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)

  def forward(self, x: Tensor) -> Tensor:
    output = self._norm(x.float()).type_as(x)
    return output * self.weight


def precompute_freqs_cis(seq_len: int,
                         n_elem: int,
                         base: int = 10000) -> Tensor:
  freqs = 1.0 / (
      base**(torch.arange(0, n_elem, 2)[:(n_elem // 2)].float() / n_elem))
  t = torch.arange(seq_len, device=freqs.device)
  freqs = torch.outer(t, freqs)
  freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
  cache = torch.stack([freqs_cis.real, freqs_cis.imag], dim=-1)
  return cache.to(dtype=torch.bfloat16)


def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
  xshaped = x.float().reshape(*x.shape[:-1], -1, 2)
  freqs_cis = freqs_cis.view(1, xshaped.size(1), 1, xshaped.size(3), 2)
  x_out2 = torch.stack(
      [
          xshaped[..., 0] * freqs_cis[..., 0] -
          xshaped[..., 1] * freqs_cis[..., 1],
          xshaped[..., 1] * freqs_cis[..., 0] +
          xshaped[..., 0] * freqs_cis[..., 1],
      ],
      -1,
  )

  x_out2 = x_out2.flatten(3)
  return x_out2.type_as(x)
