import copy

import torch
import numpy as np
import pandas as pd
from tqdm import tqdm
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from datasets import load_dataset

def load_model(model_id='meta-llama/Llama-2-7b-chat', device='cuda'):
    nf4_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type='nf4',
        bnb_4bit_double_quant=True,
        bnb_4bit_compute_dtype=torch.bfloat16,
    )

    model = AutoModelForCausalLM.from_pretrained(
            model_id, 
            torch_dtype=torch.float16,
            load_in_4bit=True,
            quantization_config=nf4_config,
            trust_remote_code=True,
            device_map='auto')
    print(model.__class__)
    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)

    return model, tokenizer


def do_summarization_dataset(model_id='meta-llama/Llama-2-7b-chat', lan='en', count=1000, device='cuda'):
    model, tokenizer = load_model(model_id, device)

    d = load_dataset('GEM/wiki_lingua', lan, split='test').filter(lambda sample, idx: idx < count, with_indices=True)

    model_type = model.config.model_type
    sep_token, sep_token_id = tokenizer.unk_token, tokenizer.unk_token_id
    if model_type == 'llama':
        prompt = (
            f'<s>[INST] <<SYS>> {{system_prompt}} <</SYS>> {{content}} [/INST]'   
            f'{{sep_token}} {{summary}} {{eos_token}}'
        )
    elif model_type == 'chatglm':
        # NOTE: hack sep_token
        sep_token = '🥺'
        sep_token_id = 64482
        prompt = (
            f'[Round 1]\n问：请帮我总结下下面这段文字\n{{content}}\n答：{{sep_token}}{{summary}} {{eos_token}}'
        )
    elif model_type == 'bloom':
        prompt = (
            f'Q: {{system_prompt}} {{content}} A: {{sep_token}}{{summary}} {{eos_token}}'
        )

    def apply_prompt_template(sample):
        txt = prompt.format(
            system_prompt='Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\nWrite a detailed summary of the input.',
            content=sample['source'],
            summary=sample['target'],
            eos_token=tokenizer.eos_token,
            sep_token=sep_token,
        )
        print('==>', txt)
        print('==>', tokenizer(txt).input_ids)
        # print('==>', tokenizer(txt.split(sep_token)[0]).input_ids)
        # print('==>', tokenizer(txt.split(sep_token)[1]).input_ids)
        return {
            'text': txt,
        }                
    
    d = d.map(apply_prompt_template, remove_columns=list(d.features), keep_in_memory=True)
    d = d.map(lambda sample: tokenizer(sample['text']), batched=True, remove_columns=list(d.features))

    nlls = []
    for input_ids in tqdm(d['input_ids'][:count]):
        # print(input_ids)
        sep = input_ids.index(sep_token_id)
        input_ids = input_ids[:sep]   input_ids[sep 1:]
        input_ids = torch.tensor(np.array([input_ids]))
        target_ids = input_ids.clone()
        target_ids[:, :sep] = -100

        with torch.no_grad():
            outputs = model(input_ids, labels=target_ids)
            neg_log_likelihood = outputs.loss
            print(neg_log_likelihood.shape, neg_log_likelihood)
            # print(input_ids.shape, outputs.logits.shape)
            # print(tokenizer.batch_decode(torch.argmax(outputs.logits, dim=2)))
        nlls.append(neg_log_likelihood)

    return torch.exp(torch.stack(nlls).mean().to('cuda'))

print('meta-llama/Llama-2-7b-hf', do_summarization_dataset('meta-llama/Llama-2-7b-hf', lan='zh'))
print('Llama2-Chinese-7b-Chat', do_summarization_dataset('FlagAlpha/Llama2-Chinese-7b-Chat', lan='zh'))
print('THUDM/chatglm2-6b', do_summarization_dataset('THUDM/chatglm2-6b', lan='zh'))