import glob
import json

import torch
from tqdm import tqdm
from transformers import AutoTokenizer
from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions

from models import GPT2LMHeadModel, GPT2LMHeadModelBackground, GPT2OdinModel
from prepare_dataset import _prepare_dialogue

MODELS = {
    "gpt2": GPT2LMHeadModel,
    "gpt2background": GPT2LMHeadModelBackground,
    "gpt2odin": GPT2OdinModel,
}


def _prepare_gpt2_results(
    model_output: CausalLMOutputWithCrossAttentions,
    response: torch.LongTensor,
    *args,
    **kwargs,
):
    scores: torch.FloatTensor = torch.log_softmax(model_output.logits, dim=-1)

    positions = torch.LongTensor(range(response.shape[0]))

    gpt2_target_score = scores[positions, response]
    gpt2_target_score = gpt2_target_score.cpu().detach().tolist()

    gpt2_argmax_score = torch.amax(scores, dim=-1)
    gpt2_argmax_score = gpt2_argmax_score.cpu().detach().tolist()

    return {
        "gpt2_target_score": gpt2_target_score,
        "gpt2_argmax_score": gpt2_argmax_score,
    }


def _prepare_gpt2background_results(
    model_output: CausalLMOutputWithCrossAttentions,
    response: torch.LongTensor,
    *args,
    **kwargs,
):
    scores: torch.FloatTensor = torch.log_softmax(model_output.logits, dim=-1)

    positions = torch.LongTensor(range(response.shape[0]))

    gpt2background_target_score = scores[positions, response]
    gpt2background_target_score = gpt2background_target_score.cpu().detach().tolist()

    gpt2background_argmax_score = torch.amax(scores, dim=-1)
    gpt2background_argmax_score = gpt2background_argmax_score.cpu().detach().tolist()

    return {
        "gpt2background_target_score": gpt2background_target_score,
        "gpt2background_argmax_score": gpt2background_argmax_score,
    }


def _prepare_gpt2odin_results(model_output: tuple, *args, **kwargs):
    lm_logits, h_prod, g_logits = model_output

    h_prod = torch.log_softmax(h_prod, dim=-1)
    h_prod = torch.amax(h_prod, dim=-1)
    h_prod = h_prod.cpu().detach().tolist()

    g_logits = torch.log(g_logits)
    g_logits = g_logits.cpu().detach().tolist()

    return {"h_max": h_prod, "g": g_logits}


MODEL_EVAL_FUNC = {
    "gpt2": _prepare_gpt2_results,
    "gpt2background": _prepare_gpt2background_results,
    "gpt2odin": _prepare_gpt2odin_results,
}


def _iterate_dbdc4_files():
    dbdc4_eval_files = glob.glob("./data/DBDC4_eval_20200314/en/*.json")
    dbdc4_eval_files_len = len(dbdc4_eval_files)

    print(f"Found {dbdc4_eval_files_len} DBDC4 eval files")

    for file in tqdm(dbdc4_eval_files):
        tqdm.write(f"Evaluating: {file}")
        with open(file, "rt") as fin:
            dialogue = json.load(fin)

        d_id = file

        for t_idx, turn in enumerate(dialogue["turns"]):
            if len(turn["annotations"]) == 0:
                continue

            context = " ".join(
                [
                    context_turn["utterance"]
                    for context_turn in dialogue["turns"][:t_idx]
                ]
            ).strip()
            response = turn["utterance"].strip()

            annotation_id = turn["annotation-id"]
            annotations = [
                annotation["breakdown"] for annotation in turn["annotations"]
            ]

            yield dict(
                d_id=d_id,
                context=context,
                response=response,
                annotation_id=annotation_id,
                annotations=annotations,
            )


def eval_model(model_name: str, model_checkpoint: str, cuda: bool, use_amp: bool):
    gpt2_tokenizer = AutoTokenizer.from_pretrained("distilgpt2")

    if cuda:
        device = torch.device("cuda")
    else:
        device = torch.device("cpu")

    # model: torch.nn.Module = MODELS[model_name](use_amp).to(device)

    model = torch.load(model_checkpoint, map_location=device)
    model.use_amp = use_amp
    model.eval()

    eval_results = dict()

    with torch.no_grad():
        for dialogue in _iterate_dbdc4_files():
            context = dialogue["context"]
            response = dialogue["response"]

            input_ids, attention_mask, tokenized_response = _prepare_dialogue(
                gpt2_tokenizer, context, response
            )

            tokenized_response = torch.LongTensor(tokenized_response).to(device)
            input_ids = torch.LongTensor(input_ids).to(device)
            attention_mask = torch.LongTensor(attention_mask).to(device)

            model_output = model(
                **dict(input_ids=input_ids, attention_mask=attention_mask)
            )

            scores = MODEL_EVAL_FUNC[model_name](
                model_output=model_output, response=tokenized_response
            )

            dialogue.update(scores)

            eval_results[dialogue["annotation_id"]] = dialogue

    with open(f"./data/{model_name}_eval.json", "wt") as fout:
        json.dump(obj=eval_results, fp=fout)