from functools import partial
import math
import inspect
from typing import Optional

import torch
import torch.nn as nn
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    checkpoint_wrapper,
    CheckpointImpl,
    apply_activation_checkpointing,
)
from torch.nn import functional as F

try:
    from flash_attn import flash_attn_func
except ImportError:
    print("Failed to import flash_attn. You won't be able to use the GPT model.")
    flash_attn_func = None
try:
    import xformers.ops as xops
    # from xformers.components.attention.core import scaled_dot_product_attention
except ImportError:
    print("Failed to import xformers.")
    xops = None


QK_LAYERNORM = False
ActivationFunc = nn.GELU
# ActivationFunc = nn.SiLU


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

    def _norm(self, x):
        return x / (torch.sqrt(x.pow(2).mean(-1, keepdim=True)) + self.eps)

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


class LayerNorm(nn.Module):
    """LayerNorm but with an optional bias. PyTorch doesn't support simply bias=False"""

    def __init__(self, ndim):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(ndim))

    def forward(self, input):
        # TODO: actually bias is fixed to None here, can we just use built in layernorm?
        return F.layer_norm(input, self.weight.shape, self.weight, None, 1e-5)


NormFunc = LayerNorm
# NormFunc = RMSNorm


OptimFunc = torch.optim.AdamW

USE_CUSTOM_ATTN = False


def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    """torch.repeat_interleave(x, dim=2, repeats=n_rep)
    https://github.com/facebookresearch/llama/blob/main/llama/model.py"""
    bs, slen, n_kv_heads, head_dim = x.shape
    if n_rep == 1:
        return x
    return (
        x[:, :, :, None, :]
        .expand(bs, slen, n_kv_heads, n_rep, head_dim)
        .reshape(bs, slen, n_kv_heads * n_rep, head_dim)
    )


class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        assert config.n_embd / config.n_head == config.d_head
        assert config.n_head % config.n_kv_head == 0
        # main
        self.n_head = config.n_head
        self.n_kv_head = config.n_kv_head
        self.d_head = config.d_head
        self.n_embd = config.n_embd
        self.c_attn_k = nn.Linear(config.n_embd, config.n_kv_head * config.d_head, bias=config.bias)
        self.c_attn_v = nn.Linear(config.n_embd, config.n_kv_head * config.d_head, bias=config.bias)
        self.c_attn_q = nn.Linear(config.n_embd, config.n_head * config.d_head, bias=config.bias)
        if QK_LAYERNORM:
            self.q_ln = NormFunc(config.d_head)
            self.k_ln = NormFunc(config.d_head)
        # output projection
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
        # regularization
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
        self.dropout = config.dropout

        self.attention_type = config.attention_type
        assert self.attention_type in ("torch", "tao", "xformers")
        self.attention_sliding_window_size = config.attention_sliding_window_size
        if self.attention_sliding_window_size != -1 and self.attention_type != "tao":
            raise ValueError("attention_sliding_window_size variation is only enabled for tao attention")
        # custom attn mask
        if USE_CUSTOM_ATTN:
            if self.attention_type == "xformers":
                # TODO: this should be inf not 100? right now this OOMs anyways
                attn_mask = (
                    torch.ones(config.block_size, config.block_size, dtype=torch.float32).tril(
                        diagonal=0
                    )
                    * 100
                )
                attn_mask[: config.t_text, : config.t_text] = 100
            else:
                attn_mask = torch.ones(config.block_size, config.block_size, dtype=torch.bool).tril(
                    diagonal=0
                )
                attn_mask[: config.t_text, : config.t_text] = True
            self.register_buffer("attn_mask", attn_mask, persistent=False)

        self.use_rotary_pos_emb = config.use_rotary_pos_emb
        if config.use_rotary_pos_emb:
            # pip install "git+https://github.com/Dao-AILab/flash-attention.git#subdirectory=csrc/rotary
            from flash_attn.layers.rotary import RotaryEmbedding

            self.rotary_pos_emb = RotaryEmbedding(config.d_head)

    def forward(self, x):
        B, T, C = x.size()  # b_size, sequ_len, emb_dim (n_embd)

        dropout = self.dropout if self.training else 0.0

        k = self.c_attn_k(x)
        v = self.c_attn_v(x)
        q = self.c_attn_q(x)

        k = k.view(B, T, self.n_kv_head, self.d_head)
        v = v.view(B, T, self.n_kv_head, self.d_head)
        q = q.view(B, T, self.n_head, self.d_head)

        if QK_LAYERNORM:
            # applying per head as per: 2309.14322 Figure E.8
            q = self.q_ln(q)
            k = self.k_ln(k)

        if self.use_rotary_pos_emb:
            k = k.unsqueeze(2)
            v = v.unsqueeze(2)
            kv = torch.cat([k, v], dim=2)  # this is gross, prob want to merge kv projections
            q, kv = self.rotary_pos_emb(q, kv)  # kv: (batch, seqlen, 2, nheads, headdim)
            k = kv[:, :, 0, :, :]
            v = kv[:, :, 1, :, :]

        # repeat k/v heads if n_kv_heads < n_heads
        n_rep = self.n_head // self.n_kv_head

        if self.attention_type == "tao":
            # tao flash 2 needs (B, S, H, D) == (B, T, nh, hs)
            y = flash_attn_func(
                q,
                k,
                v,
                dropout_p=dropout,
                softmax_scale=None,
                causal=True,
                window_size=(self.attention_sliding_window_size, 0),
            )
        elif self.attention_type == "xformers":
            k = repeat_kv(k, n_rep)
            v = repeat_kv(v, n_rep)
            # Input format ``[B, M, H, K]``; B is batch size, M is sequence length, H is number of heads, K is embeding size per head
            if USE_CUSTOM_ATTN:
                y = xops.memory_efficient_attention(
                    q,
                    k,
                    v,
                    attn_bias=self.attn_mask[None][None].repeat(q.shape[0], q.shape[2], 1, 1),
                )
            else:
                y = xops.memory_efficient_attention(q, k, v, attn_bias=xops.LowerTriangularMask())
        else:
            k = repeat_kv(k, n_rep)
            v = repeat_kv(v, n_rep)
            # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
            k = k.transpose(1, 2)  # (B, nh, T, hs)
            v = v.transpose(1, 2)  # (B, nh, T, hs)
            q = q.transpose(1, 2)  # (B, nh, T, hs)
            attn_mask = None if not USE_CUSTOM_ATTN else self.attn_mask
            is_causal = False if USE_CUSTOM_ATTN else True
            with torch.backends.cuda.sdp_kernel(
                enable_flash=True, enable_math=False, enable_mem_efficient=False
            ):
                # TODO: for compile to work we have to remove this context manager
                y = torch.nn.functional.scaled_dot_product_attention(
                    q, k, v, attn_mask=attn_mask, dropout_p=dropout, is_causal=is_causal
                )
            y = y.transpose(1, 2)

        # re-assemble head outputs side by side
        y = y.contiguous().view(B, T, C)

        # output projection
        y = self.resid_dropout(self.c_proj(y))
        return y


class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        embd_inner = 4 * config.n_embd
        if "silu" in str(ActivationFunc.__name__).lower():
            multiple_of = 256
            embd_inner = int(2 * embd_inner / 3)
            embd_inner = multiple_of * ((embd_inner + multiple_of - 1) // multiple_of)
        self.embd_inner = embd_inner
        self.c_fc = nn.Linear(config.n_embd, self.embd_inner, bias=config.bias)
        self.c_proj = nn.Linear(self.embd_inner, config.n_embd, bias=config.bias)
        self.dropout = nn.Dropout(config.dropout)
        self.activation = ActivationFunc()

    def forward(self, x):
        x = self.c_fc(x)
        x = self.activation(x)
        x = self.c_proj(x)
        x = self.dropout(x)
        return x


class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln_1 = NormFunc(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = NormFunc(config.n_embd)
        self.mlp = MLP(config)

    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x


###
# helper methods (shared across models)
###


def init_weights_simple(self, module):
    if isinstance(module, nn.Linear):
        torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
        if module.bias is not None:
            torch.nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Embedding):
        torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)


def get_init_fn(input_dim: int, init_depth: Optional[int] = None):
    """LM layer initialization."""
    std = 1 / math.sqrt(input_dim)
    if init_depth is not None:
        std = std / math.sqrt(2 * init_depth)
    return partial(torch.nn.init.trunc_normal_, mean=0.0, std=std, a=-3 * std, b=3 * std)


def _needs_weight_decay(param_name, param, is_fsdp):
    if not is_fsdp:
        return param.dim() >= 2
    # at this point it should be an FSDP module
    assert param_name.startswith("_fsdp_wrapped_module")
    patterns = ("ln_", "bias")
    for pattern in patterns:
        if pattern in param_name:
            return False
    return True


def configure_optimizers(
    self, weight_decay, learning_rate, betas, device_type, use_fused=True, is_fsdp=False
):
    # start with all of the candidate parameters
    param_dict = {pn: p for pn, p in self.named_parameters()}
    # filter out those that do not require grad
    param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
    # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.
    # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.
    decay_params = [p for n, p in param_dict.items() if _needs_weight_decay(n, p, is_fsdp=is_fsdp)]
    nodecay_params = [p for n, p in param_dict.items() if not _needs_weight_decay(n, p, is_fsdp=is_fsdp)]

    optim_groups = [
        {"params": decay_params, "weight_decay": weight_decay},
        {"params": nodecay_params, "weight_decay": 0.0},
    ]
    num_decay_params = sum(p.numel() for p in decay_params)
    num_nodecay_params = sum(p.numel() for p in nodecay_params)
    print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")
    print(
        f"num non-decayed parameter tensors: {len(nodecay_params)},"
        f" with {num_nodecay_params:,} parameters"
    )
    # Create AdamW optimizer and use the fused version if it is available
    fused_available = "fused" in inspect.signature(OptimFunc).parameters
    use_fused = fused_available and device_type == "cuda" and use_fused
    extra_args = dict(fused=True) if use_fused else dict()
    optimizer = OptimFunc(
        optim_groups, lr=learning_rate, betas=betas, eps=1e-07, **extra_args
    )  # eps just incase: https://github.com/pytorch/pytorch/issues/26218
    print(f"using fused Optimizer: {use_fused}")
    return optimizer


def estimate_mfu(self, fwdbwd_per_iter, dt):
    """estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS"""
    # first estimate the number of flops we do per iteration.
    # see PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311
    N = self.get_num_params()
    cfg = self.config
    L, H, Q, T = cfg.n_layer, cfg.n_head, cfg.n_embd // cfg.n_head, cfg.block_size
    flops_per_token = 6 * N + 12 * L * H * Q * T
    flops_per_fwdbwd = flops_per_token * T
    flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter
    # express our flops throughput as ratio of A100 bfloat16 peak flops
    flops_achieved = flops_per_iter * (1.0 / dt)  # per second
    flops_promised = 312e12  # A100 GPU bfloat16 peak flops is 312 TFLOPS
    mfu = flops_achieved / flops_promised
    return mfu


def estimate_mfu_no_model(n_params, n_layer, n_head, n_embd, block_size, fwdbwd_per_iter, dt):
    """estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS"""
    # first estimate the number of flops we do per iteration.
    # see PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311
    L, H, Q, T = n_layer, n_head, n_embd // n_head, block_size
    flops_per_token = 6 * n_params + 12 * L * H * Q * T
    flops_per_fwdbwd = flops_per_token * T
    flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter
    # express our flops throughput as ratio of A100 bfloat16 peak flops
    flops_achieved = flops_per_iter * (1.0 / dt)  # per second
    flops_promised = 312e12  # A100 GPU bfloat16 peak flops is 312 TFLOPS
    mfu = flops_achieved / flops_promised
    return mfu


non_reentrant_wrapper = partial(
    checkpoint_wrapper,
    checkpoint_impl=CheckpointImpl.NO_REENTRANT,
)


def check_fn(submodule):
    return isinstance(submodule, Block)


def apply_fsdp_checkpointing(model):
    """apply activation checkpointing to model
    returns None as model is updated directly
    """
    print("applying fsdp activation checkpointing...")

    apply_activation_checkpointing(model, checkpoint_wrapper_fn=non_reentrant_wrapper, check_fn=check_fn)