import os
from transformers import AutoConfig
from onnxruntime_extensions.onnxprocess import trace_for_onnx, pyfunc_from_model, build_customop_model
from onnxruntime_extensions.onnxprocess import torch_wrapper as torch

# Create a cache directory to store pretrained model.
cache_dir = os.path.join(".", "cache_models")
if not os.path.exists(cache_dir):
    os.makedirs(cache_dir)


model_name_or_path = "gpt2"
device = "cpu"
beam_width = 4
config = AutoConfig.from_pretrained(model_name_or_path, cache_dir=cache_dir)
num_attention_heads = config.n_head
hidden_size = config.n_embd
num_layer = config.n_layer
gpt2_full_model_path = "./gpt2_full.onnx"

# this model was generated by this script
# https://github.com/microsoft/onnxruntime/blob/master/onnxruntime/python/tools/transformers/notebooks/Inference_GPT2-OneStepSearch_OnnxRuntime_CPU.ipynb
onnx_model_path = "gpt2_one_step_search.onnx"
func_one_step = pyfunc_from_model(onnx_model_path)


def get_tokenizer(model_name_or_path, cache_dir):
    from transformers import GPT2Tokenizer  # noqa
    tokenizer = GPT2Tokenizer.from_pretrained(model_name_or_path, cache_dir=cache_dir)

    tokenizer.padding_side = "left"
    tokenizer.pad_token = tokenizer.eos_token
    gpt2_encoder_model_path = './gpt2_tok.onnx'
    build_customop_model('GPT2Tokenizer', gpt2_encoder_model_path, model=tokenizer)
    return tokenizer, pyfunc_from_model(gpt2_encoder_model_path)


def inference_and_dump_full_model(tokenizer, func_tokenizer, input_text, num_tokens_to_produce = 30):

    with trace_for_onnx(input_text, num_tokens_to_produce, names=func_tokenizer.input_names) as tc_sess:
        inputs, num_tokens = tc_sess.get_inputs()
        input_ids, attention_mask = func_tokenizer(inputs, padding=True, padding_side='left')
        attention_mask = attention_mask.type(torch.float)
        position_ids = (attention_mask.long().cumsum(-1) - 1)
        # position_ids.masked_fill_(position_ids < 0, 0)

        # Empty Past State for generating first word
        # batch_size = input_ids.size()[0]
        batch_size = 1
        past_shape = [2, batch_size, num_attention_heads, 0, hidden_size // num_attention_heads]
        empty_past = []
        for _ in range(num_layer):
            empty_past.append(torch.empty(*past_shape).type(torch.float32).to(device))

        beam_select_idx = torch.zeros([1, batch_size]).long()
        input_log_probs = torch.zeros([batch_size, 1])
        input_unfinished_sents = torch.ones([batch_size, 1], dtype=torch.bool)
        prev_step_scores = torch.zeros([batch_size, 1])
        beam_size = beam_width
        prev_step_results = input_ids.clone().detach().to(device)
    
        cfg = torch.control_flow()
        for states in cfg.loop(num_tokens, torch.tensor(True), input_ids, position_ids,
                               attention_mask, beam_select_idx, input_log_probs,
                               input_unfinished_sents, prev_step_results, prev_step_scores, *empty_past):
            step = states[0]
            states[1].symbolic_shape = ['batch_size', 'seq_len']
            states[2].symbolic_shape = ['batch_size', 'seq_len']
            states[3].symbolic_shape = ['batch_size', 'all_seq_len']
            states[4].symbolic_shape = [1, 'batch_size']

            # prev_step_results
            states[7].symbolic_shape = ['batch_size', 'total_seq_len']

            for st_ in states[-num_layer:]:
                st_.symbolic_shape = [2, 'batch_size', num_attention_heads, 'past_seq_len', hidden_size // num_attention_heads]

            prev_attention_mask = states[3]
            outputs = func_one_step(*states[1:])
            last_state = outputs[0].clone().detach().cpu()
            input_ids = last_state.reshape([batch_size * beam_size, -1]).to(device)

            input_unfinished_sents_id = -3
            prev_step_results = outputs[-2].clone().detach().to(device)
            # position_ids = (torch.tensor([context_length + step - 1
            #                                     ]).unsqueeze(0).repeat(batch_size * beam_size, 1).to(device))
            position_ids = torch.zeros([batch_size * beam_size, 1], dtype=torch.int64) + attention_mask.size()[-1]
            factor = (~step.type(torch.bool)).type(torch.int64)
            prev_attention_mask = prev_attention_mask.repeat(factor * (batch_size * beam_size - 1) + 1, 1).to(device)
            attention_mask = torch.cat(
                [
                    prev_attention_mask,
                    torch.ones([batch_size * beam_size, 1], dtype=torch.float),
                ],
                1,
            ).to(device)

            beam_select_idx = outputs[input_unfinished_sents_id - 2].clone().detach().to(device)
            input_log_probs = outputs[input_unfinished_sents_id - 1].clone().detach().to(device)
            input_unfinished_sents = outputs[input_unfinished_sents_id].clone().detach().to(device)
            prev_step_scores = outputs[-1].clone().detach().to(device)

            past = []
            for i in range(num_layer):
                past_i =  outputs[i + 1].clone().detach()
                past.append(past_i.to(device))


            any_unfinished = input_unfinished_sents.any()
            input_ids.symbolic_shape = ['total_batch_size', 'seq_len']
            position_ids.symbolic_shape = ['total_batch_size', 'seq_len']
            attention_mask.symbolic_shape = ['total_batch_size', 'all_seq_len']
            prev_step_results.symbolic_shape = ['total_batch_size', 'step_seq_len']
            for st_ in past:
                st_.symbolic_shape = [2, 'total_batch_size', num_attention_heads, 'all_seq_len', hidden_size // num_attention_heads]
            cfg.flow_output(any_unfinished, input_ids,
                            position_ids, attention_mask, beam_select_idx,
                            input_log_probs, input_unfinished_sents, prev_step_results, prev_step_scores, *past)

        result_id = 6
        all_token_ids = cfg.finalize()[result_id]
        tc_sess.save_as_onnx(gpt2_full_model_path, all_token_ids)

        print(tokenizer.decode(all_token_ids.t[0], skip_special_tokens=True))


def verify_bsfull_model(input_text):
    from onnxruntime_extensions import PyOrtFunction
    gpt2_all = PyOrtFunction.from_model(gpt2_full_model_path)
    outputs = gpt2_all(input_text, 30)
    print(tokenizer.decode(outputs[0], skip_special_tokens=True))


if __name__ == "__main__":
    tokenizer, func_tokenizer = get_tokenizer(model_name_or_path, cache_dir)
    input_text = ['best hotel in bay area.']
    inference_and_dump_full_model(tokenizer, func_tokenizer, input_text)
    verify_bsfull_model(input_text)
