import torch
from src.models import GrcResult, ProcessGrcInput, TokenConfidence


def process_grc(model, tokenizer, input: ProcessGrcInput):
    messages = [
        {
            "content": f"You are a Malay language spelling corrector. I will give you some text written in messy Malay (shortened or mistyped). Rewrite it in correct Malay spelling.\n{input.garbled}",
            "role": "user",
        },
        {"content": input.gt, "role": "assistant"},
    ]
    top_k = 5

    inputs = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=False,
        return_tensors="pt",
        return_assistant_tokens_mask=True,
        return_dict=True,
    ).to("cuda")

    attention_mask = inputs.get("attention_mask", None)
    assistant_mask = inputs.get("assistant_masks", None)

    with torch.no_grad():
        outputs = model(inputs["input_ids"], attention_mask=attention_mask)

        confidences = outputs.logits.softmax(dim=-1)

        shifted_confidences = confidences[:, :-1, :]  # align confidence with preds
        shifted_labels = inputs["input_ids"][:, 1:]
        shifted_assistant_mask = assistant_mask[:, 1:]
        boolean_mask = shifted_assistant_mask.bool()
        labels_on_assistant = shifted_labels[boolean_mask]
        confs_on_assistant = shifted_confidences[
            boolean_mask, :
        ]  # last dimension is for vocab

        values, indices = torch.topk(
            confs_on_assistant, k=top_k, dim=-1
        )  # values stores the confidences
        values = values.to(dtype=torch.float16).cpu().numpy()
        tokens_decoded: list[list[str]] = [
            tokenizer.convert_ids_to_tokens(vals) for vals in indices
        ]

        # choose the values belonging to the labels (which are indices on the last dimension)
        correct_labels_confidences = (
            confs_on_assistant.gather(-1, labels_on_assistant.unsqueeze(1))
            .squeeze(1)
            .to(dtype=torch.float16)
            .cpu()
            .numpy()
        )

        # ranks are positions of correct labels in the sorted array
        # torch.nonzero returns indices along the last dimension
        _, sorted_indices = torch.sort(confs_on_assistant, descending=True, dim=-1)
        correct_labels_ranks = (
            torch.nonzero(
                sorted_indices == labels_on_assistant.unsqueeze(1), as_tuple=False
            )[:, 1]
            + 1
        )
        correct_labels_ranks = correct_labels_ranks.cpu().numpy()
        decoded_labels = tokenizer.convert_ids_to_tokens(labels_on_assistant)

        res = GrcResult(
            prediction_map=[
                TokenConfidence(token=tokens[0], confidence=confidences[0])
                for tokens, confidences in zip(tokens_decoded, values)
            ],
            gt_map=[
                TokenConfidence(token=token, confidence=confidence, rank=rank)
                for (token, confidence, rank) in zip(
                    decoded_labels, correct_labels_confidences, correct_labels_ranks
                )
            ],
            topk=[
                [
                    TokenConfidence(token=token, confidence=confidence)
                    for (token, confidence) in zip(rt, rc)
                ]
                for (rt, rc) in zip(tokens_decoded, values)
            ],
        )
    return res