import re
import torch
import random
from tqdm import tqdm
from unidecode import unidecode
from torch.utils.data import Dataset
from transformers import GPT2Model, GPT2LMHeadModel, PreTrainedModel
from samplings import top_p_sampling, top_k_sampling, temperature_sampling
from utils import PATCH_SIZE, PATCH_LENGTH, PATCH_SAMPLING_BATCH_SIZE


class Patchilizer:
    """
    A class for converting music bars to patches and vice versa.
    """

    def __init__(self):
        self.delimiters = ["|:", "::", ":|", "[|", "||", "|]", "|"]
        self.regexPattern = f"({'|'.join(map(re.escape, self.delimiters))})"
        self.pad_token_id = 0
        self.bos_token_id = 1
        self.eos_token_id = 2

    def split_bars(self, body):
        """
        Split a body of music into individual bars.
        """
        bars = re.split(self.regexPattern, "".join(body))
        bars = list(filter(None, bars))
        # remove empty strings
        if bars[0] in self.delimiters:
            bars[1] = bars[0] + bars[1]
            bars = bars[1:]

        bars = [bars[i * 2] + bars[i * 2 + 1] for i in range(len(bars) // 2)]
        return bars

    def bar2patch(self, bar, patch_size=PATCH_SIZE):
        """
        Convert a bar into a patch of specified length.
        """
        patch = [self.bos_token_id] + [ord(c) for c in bar] + [self.eos_token_id]
        patch = patch[:patch_size]
        patch += [self.pad_token_id] * (patch_size - len(patch))
        return patch

    def patch2bar(self, patch):
        """
        Convert a patch into a bar.
        """
        return "".join(
            chr(idx) if idx > self.eos_token_id else ""
            for idx in patch
            if idx != self.eos_token_id
        )

    def encode(
        self,
        abc_code,
        patch_length=PATCH_LENGTH,
        patch_size=PATCH_SIZE,
        add_special_patches=False,
    ):
        """
        Encode music into patches of specified length.
        """
        lines = unidecode(abc_code).split("\n")
        lines = list(filter(None, lines))  # remove empty lines

        body = ""
        patches = []

        for line in lines:
            if len(line) > 1 and (
                (line[0].isalpha() and line[1] == ":") or line.startswith("%%score")
            ):
                if body:
                    bars = self.split_bars(body)
                    patches.extend(
                        self.bar2patch(
                            bar + "\n" if idx == len(bars) - 1 else bar, patch_size
                        )
                        for idx, bar in enumerate(bars)
                    )
                    body = ""

                patches.append(self.bar2patch(line + "\n", patch_size))

            else:
                body += line + "\n"

        if body:
            patches.extend(
                self.bar2patch(bar, patch_size) for bar in self.split_bars(body)
            )

        if add_special_patches:
            bos_patch = [self.bos_token_id] * (patch_size - 1) + [self.eos_token_id]
            eos_patch = [self.bos_token_id] + [self.eos_token_id] * (patch_size - 1)
            patches = [bos_patch] + patches + [eos_patch]

        return patches[:patch_length]

    def decode(self, patches):
        """
        Decode patches into music.
        """
        return "".join(self.patch2bar(patch) for patch in patches)


class PatchLevelDecoder(PreTrainedModel):
    """
    An Patch-level Decoder model for generating patch features in an auto-regressive manner.
    It inherits PreTrainedModel from transformers.
    """

    def __init__(self, config):
        super().__init__(config)
        self.patch_embedding = torch.nn.Linear(PATCH_SIZE * 128, config.n_embd)
        torch.nn.init.normal_(self.patch_embedding.weight, std=0.02)
        self.base = GPT2Model(config)

    def forward(self, patches: torch.Tensor) -> torch.Tensor:
        """
        The forward pass of the patch-level decoder model.
        :param patches: the patches to be encoded
        :return: the encoded patches
        """
        patches = torch.nn.functional.one_hot(patches, num_classes=128).float()
        patches = patches.reshape(len(patches), -1, PATCH_SIZE * 128)
        patches = self.patch_embedding(patches.to(self.device))

        return self.base(inputs_embeds=patches)


class CharLevelDecoder(PreTrainedModel):
    """
    A Char-level Decoder model for generating the characters within each bar patch sequentially.
    It inherits PreTrainedModel from transformers.
    """

    def __init__(self, config):
        super().__init__(config)
        self.pad_token_id = 0
        self.bos_token_id = 1
        self.eos_token_id = 2
        self.base = GPT2LMHeadModel(config)

    def forward(
        self,
        encoded_patches: torch.Tensor,
        target_patches: torch.Tensor,
        patch_sampling_batch_size: int,
    ):
        """
        The forward pass of the char-level decoder model.
        :param encoded_patches: the encoded patches
        :param target_patches: the target patches
        :return: the decoded patches
        """
        # preparing the labels for model training
        target_masks = target_patches == self.pad_token_id
        labels = target_patches.clone().masked_fill_(target_masks, -100)

        # masking the labels for model training
        target_masks = torch.ones_like(labels)
        target_masks = target_masks.masked_fill_(labels == -100, 0)

        # select patches
        if (
            patch_sampling_batch_size != 0
            and patch_sampling_batch_size < target_patches.shape[0]
        ):
            indices = list(range(len(target_patches)))
            random.shuffle(indices)
            selected_indices = sorted(indices[:patch_sampling_batch_size])

            target_patches = target_patches[selected_indices, :]
            target_masks = target_masks[selected_indices, :]
            encoded_patches = encoded_patches[selected_indices, :]
            labels = labels[selected_indices, :]

        # get input embeddings
        inputs_embeds = torch.nn.functional.embedding(
            target_patches, self.base.transformer.wte.weight
        )

        # concatenate the encoded patches with the input embeddings
        inputs_embeds = torch.cat(
            (encoded_patches.unsqueeze(1), inputs_embeds[:, 1:, :]), dim=1
        )

        return self.base(
            inputs_embeds=inputs_embeds, attention_mask=target_masks, labels=labels
        )

    def generate(self, encoded_patch: torch.Tensor, tokens: torch.Tensor):
        """
        The generate function for generating a patch based on the encoded patch and already generated tokens.
        :param encoded_patch: the encoded patch
        :param tokens: already generated tokens in the patch
        :return: the probability distribution of next token
        """
        encoded_patch = encoded_patch.reshape(1, 1, -1)
        tokens = tokens.reshape(1, -1)

        # Get input embeddings
        tokens = torch.nn.functional.embedding(tokens, self.base.transformer.wte.weight)

        # Concatenate the encoded patch with the input embeddings
        tokens = torch.cat((encoded_patch, tokens[:, 1:, :]), dim=1)

        # Get output from model
        outputs = self.base(inputs_embeds=tokens)

        # Get probabilities of next token
        probs = torch.nn.functional.softmax(outputs.logits.squeeze(0)[-1], dim=-1)

        return probs


class TunesFormer(PreTrainedModel):
    """
    TunesFormer is a hierarchical music generation model based on bar patching.
    It includes a patch-level decoder and a character-level decoder.
    It inherits PreTrainedModel from transformers.
    """

    def __init__(self, encoder_config, decoder_config, share_weights=False):
        super().__init__(encoder_config)
        self.pad_token_id = 0
        self.bos_token_id = 1
        self.eos_token_id = 2
        if share_weights:
            max_layers = max(
                encoder_config.num_hidden_layers, decoder_config.num_hidden_layers
            )

            max_context_size = max(encoder_config.max_length, decoder_config.max_length)

            max_position_embeddings = max(
                encoder_config.max_position_embeddings,
                decoder_config.max_position_embeddings,
            )

            encoder_config.num_hidden_layers = max_layers
            encoder_config.max_length = max_context_size
            encoder_config.max_position_embeddings = max_position_embeddings
            decoder_config.num_hidden_layers = max_layers
            decoder_config.max_length = max_context_size
            decoder_config.max_position_embeddings = max_position_embeddings

        self.patch_level_decoder = PatchLevelDecoder(encoder_config)
        self.char_level_decoder = CharLevelDecoder(decoder_config)

        if share_weights:
            self.patch_level_decoder.base = self.char_level_decoder.base.transformer

    def forward(
        self,
        patches: torch.Tensor,
        patch_sampling_batch_size: int = PATCH_SAMPLING_BATCH_SIZE,
    ):
        """
        The forward pass of the TunesFormer model.
        :param patches: the patches to be both encoded and decoded
        :return: the decoded patches
        """
        patches = patches.reshape(len(patches), -1, PATCH_SIZE)
        encoded_patches = self.patch_level_decoder(patches)["last_hidden_state"]

        return self.char_level_decoder(
            encoded_patches.squeeze(0)[:-1, :],
            patches.squeeze(0)[1:, :],
            patch_sampling_batch_size,
        )

    def generate(
        self,
        patches: torch.Tensor,
        tokens: torch.Tensor,
        top_p: float = 1,
        top_k: int = 0,
        temperature: float = 1,
        seed: int = None,
    ):
        """
        The generate function for generating patches based on patches.
        :param patches: the patches to be encoded
        :return: the generated patches
        """
        patches = patches.reshape(len(patches), -1, PATCH_SIZE)
        encoded_patches = self.patch_level_decoder(patches)["last_hidden_state"]

        if tokens == None:
            tokens = torch.tensor([self.bos_token_id], device=self.device)

        generated_patch = []
        random.seed(seed)

        while True:
            if seed != None:
                n_seed = random.randint(0, 1000000)
                random.seed(n_seed)

            else:
                n_seed = None

            prob = (
                self.char_level_decoder.generate(encoded_patches[0][-1], tokens)
                .cpu()
                .detach()
                .numpy()
            )

            prob = top_p_sampling(prob, top_p=top_p, return_probs=True)
            prob = top_k_sampling(prob, top_k=top_k, return_probs=True)

            token = temperature_sampling(prob, temperature=temperature, seed=n_seed)

            generated_patch.append(token)
            if token == self.eos_token_id or len(tokens) >= PATCH_SIZE - 1:
                break

            else:
                tokens = torch.cat(
                    (tokens, torch.tensor([token], device=self.device)), dim=0
                )

        return generated_patch, n_seed


class PatchilizedData(Dataset):
    def __init__(self, items, patchilizer):
        self.texts = []

        for item in tqdm(items):
            text = item["control code"] + "\n".join(
                item["abc notation"].split("\n")[1:]
            )
            input_patch = patchilizer.encode(text, add_special_patches=True)
            input_patch = torch.tensor(input_patch)
            if torch.sum(input_patch) != 0:
                self.texts.append(input_patch)

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        return self.texts[idx]