import torch
import argparse

from contextlib import nullcontext
from transformers import AutoModelForCausalLM, AutoTokenizer


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_path", type=str, default="gpt2")
    parser.add_argument("--tokenizer_path", type=str, default="gpt2")
    parser.add_argument("--use-flash-attn", action="store_true")
    parser.add_argument("--no-grad", action="store_true")

    args = parser.parse_args()
    return args


def prepare_data(tokenizer, packing: bool, device: torch.device):
    texts = [
        "Hello, how are you?",
        "When is the next holiday?",
        "China is a great country.",
    ]
    encoded = tokenizer(texts)

    # Convert to tensor
    if packing:
        res = torch.zeros((1, sum(len(x) for x in encoded["input_ids"])), dtype=torch.long)
        attention_mask = torch.full((1, 1, res.size(1), res.size(1)), dtype=torch.bfloat16, fill_value=float("-inf"))

        offset = 0
        for i, (input_ids, attn_mask) in enumerate(zip(encoded["input_ids"], encoded["attention_mask"])):
            res[0, offset: offset + len(input_ids)] = torch.tensor(input_ids)
            attention_mask[0, 0, offset: offset + len(attn_mask), offset: offset + len(attn_mask)] = 0.
            offset += len(attn_mask)
    else:
        max_length = max(len(x) for x in encoded["input_ids"])
        res = torch.zeros((len(encoded["input_ids"]), max_length), dtype=torch.long)
        attention_mask = torch.zeros((len(encoded["input_ids"]), max_length), dtype=torch.long)
        for i, (input_ids, attn_mask) in enumerate(zip(encoded["input_ids"], encoded["attention_mask"])):
            res[i, : len(input_ids)] = torch.tensor(input_ids)
            attention_mask[i, : len(attn_mask)] = torch.tensor(attn_mask)
    return res.to(device), attention_mask.to(device)


def main(args):
    model = AutoModelForCausalLM.from_pretrained(
        args.model_path, use_flash_attention_2=args.use_flash_attn, device_map="cuda", torch_dtype=torch.bfloat16
    )
    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path)
    device = torch.device("cuda")

    context = torch.no_grad if args.no_grad else nullcontext
    with context():
        model.eval()
        input_ids_no_pack, attention_mask_no_pack = prepare_data(tokenizer, packing=False, device=device)
        input_ids_pack, attention_mask_pack = prepare_data(tokenizer, packing=True, device=device)

        logits_no_pack = model(input_ids=input_ids_no_pack, attention_mask=attention_mask_no_pack).logits
        logits_pack = model(input_ids=input_ids_pack, attention_mask=attention_mask_pack).logits

        logits_no_pack_flatten = torch.zeros_like(logits_pack)
        offset = 0
        for i in range(logits_no_pack.shape[0]):
            length = attention_mask_no_pack[i].sum().item()
            logits_no_pack_flatten[0, offset: offset + length] = logits_no_pack[i, :length]
            offset += length

        print((logits_no_pack_flatten - logits_pack).sum())


if __name__ == "__main__":
    args = parse_args()
    main(args)