import gradio as gr
import torch
import random
from unidecode import unidecode
import re
from samplings import top_p_sampling, top_k_sampling, temperature_sampling
from transformers import GPT2Model, GPT2LMHeadModel, PreTrainedModel

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
PATCH_LENGTH = 128      # Patch Length
PATCH_SIZE = 32       # Patch Size

PATCH_NUM_LAYERS = 9         # Number of layers in the encoder
CHAR_NUM_LAYERS = 3          # Number of layers in the decoder

NUM_EPOCHS = 32                 # Number of epochs to train for (if early stopping doesn't intervene)
LEARNING_RATE = 5e-5            # Learning rate for the optimizer
PATCH_SAMPLING_BATCH_SIZE = 0   # Batch size for patch during training, 0 for full context
LOAD_FROM_CHECKPOINT = False     # Whether to load weights from a checkpoint
SHARE_WEIGHTS = False            # Whether to share weights between the encoder and decoder

description = """
<div>

<a style="display:inline-block" href='https://github.com/sander-wood/tunesformer'><img src='https://img.shields.io/github/stars/sander-wood/tunesformer?style=social' /></a>
<a style="display:inline-block" href="https://huggingface.co/datasets/sander-wood/massive_abcnotation_dataset"><img src="https://img.shields.io/badge/huggingface-dataset-ffcc66.svg"></a>
<a style="display:inline-block" href="https://arxiv.org/pdf/2301.02884.pdf"><img src="https://img.shields.io/badge/arXiv-2301.02884-b31b1b.svg"></a>
<a class="duplicate-button" style="display:inline-block" target="_blank" href="https://huggingface.co/spaces/sander-wood/tunesformer?duplicate=true"><img style="margin-top:0;margin-bottom:0" src="https://huggingface.co/datasets/huggingface/badges/raw/main/duplicate-this-space-md-dark.svg" alt="Duplicate Space"></a>
</div>

## ℹ️ How to use this demo?
1. Enter the control codes to set the musical form of the generated music. For details, please refer to the below "Control Codes" section. (optional)
2. Enter the prefix of the generated music. You can set the ABC header (i.e., note length, tempo, meter, and key) and the beginning of the music. (optional)
2. You can set the parameters (i.e., number of tunes, maximum length, top-p, temperature, and random seed) for the generation. (optional)
3. Click "Submit" and wait for the result.
4. The generated ABC notation can be converted to MIDI or PDF using [EasyABC](https://sourceforge.net/projects/easyabc/), you can also use this [online renderer](https://ldzhangyx.github.io/abc/) to render the ABC notation.

## 📝 Control Codes
There are 3 types of control codes: section bar `[BARS_NB]`, `[SECS_NS]`, and similarity `[SIM_EDS]`. The control codes can be used to set the musical form of the generated music:

- `[BARS_NB]`: controls the number of bars in a section of the melody. For example, users could specify that they want a section to contain 8 bars, and TunesFormer would generate a section that fits within that structure. It counts on the bar symbol `|`. The value of `NB` ranges from 1 to 32, where 1 means one bar and 32 means thirty-two bars.
- `[SECS_NS]`: controls the number of sections in the entire melody. This can be used to create a sense of structure and coherence within the melody, as different sections can be used to create musical themes or motifs. It counts on several symbols that are commonly used in ABC notation and can be used to represent section boundaries: `[|`,`||`,`|]`,`|:`,`::`, and `:|`. The value of `NS` ranges from 1 to 8, where 1 means one section and 8 means eight sections.
- `[SIM_EDS]`: controls the similarity between the generated music and the prefix. The value of `EDS` is the edit distance similarity level between the current section and a previous section in the melody. The value of EDS ranges from 0 to 10, where 0 means no similarity and 10 means the same section.

To provide a clearer understanding of the control codes, we provide an example below:

`[SECS_3][BARS_8][SIM_3][BARS_8][SIM_10][SIM_3][BARS_8]`

The first control code `[SECS_3]` specifies there are 3 sections in the tune, and the following control code `[BARS_8]` indicates the first section has 8 bars. The next two control codes `[SIM_3]` and `[BAR_8]` indicate that the EDS between section II and section I is approximately 0.3, and section II has 8 bars. The last three control codes `[SIM_10]`, `[SIM_3]` and `[BARS_8]` specify that section III is identical to section I while dissimilar to section II, and has 8 bars.

## ❕Notice
- If the prefix is given, you must also provide the control codes.
- The order of the ABC header cannot be changed (i.e., note length, tempo, meter, and key), and also do not support other headers (e.g., composer, title, etc.) as they are not used in the model. For details, please refer to the official [ABC notation](https://abcnotation.com/wiki/abc:standard:v2.1#description_of_information_fields).
- The demo is based on GPT2-small and trained from scratch on the [Massive ABC Notation Dataset](https://huggingface.co/datasets/sander-wood/massive_abcnotation_dataset).
- The demo is still in the early stage, and the generated music is not perfect. If you have any suggestions, please feel free to contact me via [email](mailto:shangda@mail.ccom.edu.cn).


"""

class Patchilizer:
    """
    A class for converting music bars to patches and vice versa. 
    """
    def __init__(self):
        self.delimiters = ["|:", "::", ":|", "[|", "||", "|]", "|"]
        self.regexPattern = '(' + '|'.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

def generate_abc(prompt, num_tunes, max_patch, top_p, top_k, temperature, seed, show_control_code):

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

    patchilizer = Patchilizer()

    patch_config = GPT2Config(num_hidden_layers=PATCH_NUM_LAYERS, 
                        max_length=PATCH_LENGTH, 
                        max_position_embeddings=PATCH_LENGTH,
                        vocab_size=1)
    char_config = GPT2Config(num_hidden_layers=CHAR_NUM_LAYERS, 
                        max_length=PATCH_SIZE, 
                        max_position_embeddings=PATCH_SIZE,
                        vocab_size=128)
    model = TunesFormer(patch_config, char_config, share_weights=SHARE_WEIGHTS)

    filename = "weights.pth"

    if os.path.exists(filename):
        print(f"Weights already exist at '{filename}'. Loading...")
    else:
        print(f"Downloading weights to '{filename}' from huggingface.co...")
        try:
            url = 'https://huggingface.co/sander-wood/tunesformer/resolve/main/weights.pth'
            response = requests.get(url, stream=True)

            total_size = int(response.headers.get('content-length', 0))
            chunk_size = 1024

            with open(filename, 'wb') as file, tqdm(
                desc=filename,
                total=total_size,
                unit='B',
                unit_scale=True,
                unit_divisor=1024,
            ) as bar:
                for data in response.iter_content(chunk_size=chunk_size):
                    size = file.write(data)
                    bar.update(size)
        except Exception as e:
            print(f"Error: {e}")
            exit()
            
    checkpoint = torch.load('weights.pth')
    model.load_state_dict(checkpoint['model'])
    model = model.to(device)
    model.eval()

    tunes = ""

    print("\n"+" OUTPUT TUNES ".center(60, "#"))

    start_time = time.time()

    for i in range(num_tunes):
        tune = "X:"+str(i+1) + "\n" + prompt
        lines = re.split(r'(\n)', tune)
        tune = ""
        skip = False
        for line in lines:
            if show_control_code or line[:2] not in ["S:", "B:", "E:"]:
                if not skip:
                    print(line, end="")
                    tune += line
                skip = False
            else:
                skip = True

        input_patches = torch.tensor([patchilizer.encode(prompt, add_special_patches=True)[:-1]], device=device)
        if tune=="":
            tokens = None
        else:
            prefix = patchilizer.decode(input_patches[0])
            remaining_tokens = prompt[len(prefix):]
            tokens  = torch.tensor([patchilizer.bos_token_id]+[ord(c) for c in remaining_tokens], device=device)
        
        while input_patches.shape[1]<max_patch:
            predicted_patch, seed = model.generate(input_patches,
                                                    tokens,
                                                    top_p=top_p,
                                                    top_k=top_k,
                                                    temperature=temperature,
                                                    seed=seed)
            tokens = None
            if predicted_patch[0]!=patchilizer.eos_token_id:
                next_bar = patchilizer.decode([predicted_patch])
                if show_control_code or next_bar[:2] not in ["S:", "B:", "E:"]:
                    print(next_bar, end="")
                    tune += next_bar
                if next_bar=="":
                    break
                next_bar = remaining_tokens+next_bar
                remaining_tokens = ""
                predicted_patch = torch.tensor(patchilizer.bar2patch(next_bar), device=device).unsqueeze(0)
                input_patches = torch.cat([input_patches, predicted_patch.unsqueeze(0)], dim=1)
            else:
                break

        tunes += tune+"\n\n"
    
    return tunes

default_prompt = """S:2
B:9
E:4
B:9
L:1/8
M:3/4
K:D
 de |"D" """

input_prompt = gr.inputs.Textbox(lines=5, label="Prefix", default=default_prompt)
input_num_tunes = gr.inputs.Slider(minimum=1, maximum=10, step=1, default=1, label="Number of Tunes")
input_max_patch = gr.inputs.Slider(minimum=10, maximum=128, step=1, default=128, label="Max Patch")
input_top_p = gr.inputs.Slider(minimum=0.0, maximum=1.0, step=0.05, default=0.8, label="Top P")
input_top_k = gr.inputs.Slider(minimum=1, maximum=20, step=1, default=8, label="Top K")
input_temperature = gr.inputs.Slider(minimum=0.0, maximum=2.0, step=0.05, default=1.2, label="Temperature")
input_seed = gr.inputs.Textbox(lines=1, label="Seed (int)", default="None")
output_abc = gr.outputs.Textbox(label="Generated Tunes")

gr.Interface(generate_abc,
                [input_prompt, input_num_tunes, input_max_patch, input_top_p, input_top_k, input_temperature, input_seed],
                output_abc,
            title="TunesFormer: Forming Irish Tunes with Control Codes by Bar Patching",
            description=description).launch()