import torch
import torch.nn as nn

from config import (
    DEVICE,
    START_TOKEN,
    END_TOKEN,
    MAX_OUTPUT_LENGTH,
)

from model.transformer import MinecraftTransformer


class BuilderModel(nn.Module):
    """
    High-level wrapper around the MinecraftTransformer.
    """

    def __init__(self, tokenizer):
        super().__init__()

        self.tokenizer = tokenizer

        self.transformer = MinecraftTransformer(
            tokenizer.text_vocab_size,
            tokenizer.output_vocab_size,
        )

    def forward(self, src, tgt):
        """
        Training forward pass.

        src = prompt tokens
        tgt = decoder input tokens (already shifted by train.py)
        """

        return self.transformer(
            src=src,
            tgt=tgt,
        )

    @torch.no_grad()
    def generate(
            self,
            prompt,
            max_length=MAX_OUTPUT_LENGTH,
    ):
        """
        Generate a Minecraft structure from a prompt.
        """

        self.eval()

        device = next(self.parameters()).device

        # Encode prompt
        src = torch.tensor(
            [self.tokenizer.encode_prompt(prompt)],
            dtype=torch.long,
            device=device,
        )

        # Encode with transformer
        memory = self.transformer.encode(src)

        start_id = self.tokenizer.output_to_id[START_TOKEN]
        end_id = self.tokenizer.output_to_id[END_TOKEN]

        generated = [start_id]

        for _ in range(max_length):

            tgt = torch.tensor(
                [generated],
                dtype=torch.long,
                device=device,
            )

            logits = self.transformer.decode(
                tgt,
                memory,
            )

            probs = torch.softmax(logits[0, -1], dim=-1)

            next_token = torch.argmax(probs).item()

            print(
                next_token,
                self.tokenizer.id_to_output[next_token],
                f"{probs[next_token].item():.3f}"
            )

            generated.append(next_token)

            if next_token == end_id:
                break

        blocks = self.tokenizer.decode_blocks(generated)

        # Don't count the <START> token
        token_count = max(0, len(generated) - 1)

        return blocks, token_count

    def save(self, path):
        """
        Save model weights.
        """

        torch.save(
            self.state_dict(),
            path,
        )

    def load(self, path):
        """
        Load model weights.
        """

        self.load_state_dict(
            torch.load(
                path,
                map_location=DEVICE,
            )
        )

    def to_device(self):
        """
        Move model to configured device.
        """

        return self.to(DEVICE)