import datetime
from functools import partial
import inspect
import importlib
import math
import os
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
from torch.nn.attention.flex_attention import flex_attention
from einops import rearrange

flex_attention = torch.compile(flex_attention)


def is_ddp():
    return int(os.environ.get("RANK", -1)) != -1


def is_master():
    if is_ddp():
        return int(os.environ["RANK"]) == 0
    return True


def print_with_time(content):
    """Print the content with the current time."""
    print(f"[{datetime.datetime.now().strftime('%Y-%m-%d_%H:%M:%S')}]: {content}")


def print_with_time_master(content):
    if is_master():
        print_with_time(content)


FA_V_MAJOR = None
try:
    from flash_attn_interface import flash_attn_varlen_func, flash_attn_func

    FA_V_MAJOR = 3
except ImportError:
    print_with_time_master("Failed to import flash_attn_3.")
    try:
        from flash_attn import flash_attn_varlen_func, flash_attn_func

        FA_V_MAJOR = 2
    except ImportError:
        print_with_time_master("Failed to import flash_attn. You won't be able to use the GPT model.")
        flash_attn_func = None
        flash_attn_varlen_func = None


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)


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, train_config, layer_idx: int):
        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.layer_idx = layer_idx
        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.rope_theta = config.rope_theta
        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 config.use_qk_norm:
            self.q_ln = LayerNorm(config.d_head)
            self.k_ln = LayerNorm(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(0.0)
        self.resid_dropout = nn.Dropout(train_config.dropout)
        self.dropout = train_config.dropout
        self.use_qk_norm = config.use_qk_norm

        self.attention_type = train_config.attention_type
        assert self.attention_type in ("torch", "tao")
        if self.layer_idx % config.global_every_n_layers != 0:
            self.attention_sliding_window_size = config.attention_sliding_window_size
        else:
            self.attention_sliding_window_size = -1
        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")

        self.use_rotary_pos_emb = config.use_rotary_pos_emb

    def init_weights(self, init_std: float):
        for linear in (self.c_attn_k, self.c_attn_v, self.c_attn_q):
            nn.init.trunc_normal_(linear.weight, mean=0.0, std=0.02)
        nn.init.trunc_normal_(self.c_proj.weight, mean=0.0, std=init_std)

    def forward(self, x, rope_module=None, seq_lens=None, block_mask=None, causal=True):
        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 self.use_qk_norm:
            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 = rope_module(q, kv, seq_lens=seq_lens)  # 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

        fa_dtype_in = q.dtype
        q, k, v = q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16)
        if self.attention_type == "tao" and block_mask is None:
            cu_seqlen = torch.tensor([0] + seq_lens, device=q.device).cumsum(dim=0).to(torch.int32)
            # tao flash 2 needs (B, S, H, D) == (B, T, nh, hs)
            if B == 1:
                if FA_V_MAJOR == 3:
                    seqused_q, seqused_k = None, None
                    try:  # old flash 3 api
                        fa_args = (
                            q[0],
                            k[0],
                            v[0],
                            cu_seqlen,
                            cu_seqlen,
                            max(cu_seqlen),
                            max(cu_seqlen),
                        )
                        y = flash_attn_varlen_func(
                            *fa_args,
                            # dropout_p=dropout,
                            # softmax_scale=None,
                            causal=causal,
                            window_size=(
                                self.attention_sliding_window_size,
                                self.attention_sliding_window_size,
                            ),
                        )
                    except Exception as e:  # new flash 3 api
                        fa_args = (
                            q[0],
                            k[0],
                            v[0],
                            cu_seqlen,
                            cu_seqlen,
                            seqused_q,
                            seqused_k,
                            max(cu_seqlen),
                            max(cu_seqlen),
                        )
                        y = flash_attn_varlen_func(
                            *fa_args,
                            # dropout_p=dropout,
                            # softmax_scale=None,
                            causal=causal,
                            window_size=(
                                self.attention_sliding_window_size,
                                self.attention_sliding_window_size,
                            ),
                        )
                    if isinstance(y, tuple):
                        y = y[0]
                else:
                    fa_args = (
                        q[0],
                        k[0],
                        v[0],
                        cu_seqlen,
                        cu_seqlen,
                        max(cu_seqlen),
                        max(cu_seqlen),
                    )
                    y = flash_attn_varlen_func(
                        *fa_args,
                        # dropout_p=dropout,
                        # softmax_scale=None,
                        causal=causal,
                        window_size=(
                            self.attention_sliding_window_size,
                            self.attention_sliding_window_size,
                        ),
                    )
                    if isinstance(y, tuple):
                        y = y[0]
            else:
                # B is not 1. not packing
                y = flash_attn_func(
                    q,
                    k,
                    v,
                    causal=causal,
                    window_size=(
                        self.attention_sliding_window_size,
                        self.attention_sliding_window_size,
                    ),
                )
                if isinstance(y, tuple):
                    y = y[0]
        else:
            # tao flash 2 needs (B, S, H, D) == (B, T, nh, hs)
            q = rearrange(q, "b n h d -> b h n d")
            k = rearrange(k, "b n h d -> b h n d")
            v = rearrange(v, "b n h d -> b h n d")
            y = flex_attention(
                q,
                k,
                v,
                enable_gqa=True,
                block_mask=block_mask,
            )
            y = rearrange(y, "b h n d -> b n h d")
        y = y.to(fa_dtype_in)

        # 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 Relu2(nn.Module):
    def forward(self, x):
        return torch.pow(F.relu(x), 2)


class MLP(nn.Module):
    def __init__(self, config, train_config):
        super().__init__()
        embd_inner = 4 * config.n_embd
        if config.activation_f == "silu":
            self.activation = nn.SiLU()
        elif config.activation_f == "gelu":
            self.activation = nn.GELU()
        elif config.activation_f == "relu2":
            self.activation = Relu2()
        else:
            raise NotImplementedError(f"activation function `{config.activation_f}` not supported")
        self.activation_func_str = config.activation_f
        if self.activation_func_str == "silu":
            multiple_of = 256
            embd_inner = int(2 * embd_inner / 3)
            embd_inner = multiple_of * ((embd_inner + multiple_of - 1) // multiple_of)
            self.c_fc_2 = nn.Linear(config.n_embd, embd_inner, bias=config.bias)
        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(train_config.dropout)

    def init_weights(self, init_std: float):
        nn.init.trunc_normal_(self.c_fc.weight, mean=0.0, std=0.02)
        for linear in [self.c_proj] + [self.c_fc_2] if self.activation_func_str == "silu" else []:
            nn.init.trunc_normal_(linear.weight, mean=0.0, std=init_std)

    def forward(self, x):
        if self.activation_func_str == "silu":
            x = self.c_proj(self.activation(self.c_fc(x)) * self.c_fc_2(x))
        else:
            x = self.c_proj(self.activation(self.c_fc(x)))
        x = self.dropout(x)
        return x


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

        self.weight_init_std = 0.02 / (2 * (layer_id + 1)) ** 0.5

    def init_weights(self):
        self.attn.init_weights(self.weight_init_std)
        self.mlp.init_weights(self.weight_init_std)

    def forward(self, x, rope_module=None, seq_lens=None, block_mask=None, causal=True):
        x = x + self.attn(
            self.ln_1(x),
            rope_module=rope_module,
            seq_lens=seq_lens,
            block_mask=block_mask,
            causal=causal,
        )
        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_with_time_master(
        f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters"
    )
    print_with_time_master(
        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(torch.optim.AdamW).parameters
    use_fused = fused_available and device_type == "cuda" and use_fused
    extra_args = dict(fused=True) if use_fused else dict()
    optimizer = torch.optim.AdamW(
        optim_groups, lr=learning_rate, betas=betas, eps=1e-10, **extra_args
    )  # eps just incase: https://github.com/pytorch/pytorch/issues/26218
    print_with_time_master(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_with_time_master("applying fsdp activation checkpointing...")

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