# -*- coding: utf-8 -*-
"""Intelligent-VIRTUOSO.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1hJXMgndwOv0kD8MD6ao20WI6MClA9mGI

# Intelligent VIRTUOSO (ver. 1.3)

## "Music never allows falsehoods for even the deaf hear flat notes!" ---IV

***

## Chordified GPT2-based Symbolic Music Artificial Intelligence Model Creator/Trainer.

### Multi-Instrumental, with special TMIDI Processors

***

Credit for char-based GPT2 implementation used in this colab goes out to Andrej Karpathy: https://github.com/karpathy/minGPT

***

WARNING: This complete implementation is a functioning model of the Artificial Intelligence. Please excercise great humility, care, and respect.

***

#### Project Los Angeles

#### Tegridy Code 2020

***

# Setup Environment, clone needed repos, and install all required dependencies
"""

#@title Install all dependencies (run only once per session)
!git clone https://github.com/asigalov61/minGPT
!git clone https://github.com/asigalov61/tegridy-tools
!apt install fluidsynth #Pip does not work for some reason. Only apt works
!pip install midi2audio

# Commented out IPython magic to ensure Python compatibility.
#@title Import all needed modules

print('Loading needed modules. Please wait...')
import os
if not os.path.exists('/content/Dataset'):
    os.makedirs('/content/Dataset')

os.chdir('/content/minGPT')
# make deterministic
from mingpt.utils import set_seed
set_seed(42)

import tqdm.auto
import pickle
import numpy as np
import torchvision
import torch
import torch.nn as nn
from torch import optim
import torch.nn.functional as F
from torch.utils.data import Dataset

import keras
from keras.utils import to_categorical

import time
import math
import datetime
from datetime import datetime

from mingpt.model import GPT, GPTConfig
from mingpt.trainer import Trainer, TrainerConfig
from mingpt.utils import sample

os.chdir('/content/tegridy-tools/tegridy-tools')
import TMIDI
import MIDI

import tqdm.auto

import matplotlib
import matplotlib.pyplot as plt
# %matplotlib inline

from midi2audio import FluidSynth
from IPython.display import display, Javascript, HTML, Audio

from google.colab import output, drive

dtype = torch.float
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# Assume that we are on a CUDA machine, then this should print a CUDA device:
print('Available Processing Device is:', device)
os.chdir('/content/')
print('Loading complete. Enjoy! :)')

"""# Upload/download and process MIDI dataset"""

# Commented out IPython magic to ensure Python compatibility.
#@title (OPTION 1) Download processed TXT dataset
#@markdown Based upon special Tegridy Piano MIDI dataset below

#@markdown Specifically, this TXT dataset is based the Deep Relaxation Melodies from the MIDI dataset

#@markdown Works best stand-alone/as-is for the optimal results
# %cd /content/
!wget 'https://github.com/asigalov61/Intelligent-VIRTUOSO/raw/main/Intelligent-Virtuoso-Music-TXT-Dataset.zip'
!unzip -j '/content/Intelligent-Virtuoso-Music-TXT-Dataset.zip'
!rm '/content/Intelligent-Virtuoso-Music-TXT-Dataset.zip'
# %cd /content/

# Commented out IPython magic to ensure Python compatibility.
#@title (OPTION 2) Start from scratch
#@markdown Download special Tegridy Piano MIDI dataset

#@markdown Works best stand-alone/as-is for the optimal results
# %cd /content/Dataset/
!rm *.mid 
!rm *.midi
!wget 'https://github.com/asigalov61/Tegridy-MIDI-Dataset/raw/master/Tegridy-Piano-CC-BY-NC-SA.zip'
!unzip -j '/content/Dataset/Tegridy-Piano-CC-BY-NC-SA.zip'
!rm '/content/Dataset/Tegridy-Piano-CC-BY-NC-SA.zip'
# %cd /content/

"""# If you are not sure where to start or what settings to select, please use original defaults"""

# Commented out IPython magic to ensure Python compatibility.
#@title Process MIDIs to special MIDI dataset with Tegridy MIDI Processor
#@markdown NOTES:

#@markdown 1) Dataset MIDI file names are used as song names. Feel free to change it to anything you like.

#@markdown 2) Best results are achieved with the single-track, single-channel, single-instrument MIDI 0 files with plain English names (avoid special or sys/foreign chars)

#@markdown 3) MIDI Channel = -1 means all MIDI channels. Otherwise, only single indicated MIDI channel will be processed.

full_path_to_output_dataset_to = "/content/Intelligent-Virtuoso-Music-MIDI-Dataset.npy" #@param {type:"string"}
desired_MIDI_channel_to_process = 0 #@param {type:"slider", min:-1, max:15, step:1}
MIDI_events_time_denominator = 10 #@param {type:"slider", min:1, max:100, step:1}
melody_notes_in_chords = True #@param {type:"boolean"}

debug = False #@param {type:"boolean"}

print('TMIDI Processor')
print('Starting up...')

chords_list = []

###########

average_note_pitch = 0
min_note = 127
max_note = 0

files_count = 0

ev = 0

chords_list_f = []
melody_list_f = []

chords_list = []
chords_count = 0

melody_chords = []
melody_count = 0

song_names = []

###########

def list_average(num):
  sum_num = 0
  for t in num:
      sum_num = sum_num + t           

  avg = sum_num / len(num)
  return avg

###########

print('Loading MIDI files...')
print('This may take a while on a large dataset in particular.')


# %cd /content/
dataset_addr = "/content/Dataset/"
os.chdir(dataset_addr)
filez = os.listdir(dataset_addr)

print('Creating list of songs from MIDI file names...')
for f in filez: # Based on input MIDI file names w/o the extension
  fn = os.path.basename(f)
  fno = fn.split('.')[0]
  song_names.append(fno)

print('Processing MIDI files. Please wait...')
for f in tqdm.auto.tqdm(filez):
  #print(f)
  files_count += 1
  chords_list, melody = TMIDI.Tegridy_MIDI_Processor(f, 
                                                     desired_MIDI_channel_to_process, 
                                                     MIDI_events_time_denominator,
                                                     )

  fn = os.path.basename(f)
  fno = fn.split('.')[0].replace(' ', '_')

  chords_l, melody_l = TMIDI.Tegridy_Chords_Converter(chords_list, 
                                                      melody, 
                                                      fno, 
                                                      melody_notes_in_chords)
  
  chords_list_f.extend(chords_l)
  melody_list_f.extend(melody_l)
  chords_count += len(chords_l)
  melody_count += len(melody_l)

average_note_pitch = int((min_note + max_note) / 2)

print('Task complete :)')
print('==================================================')
print('Number of processed dataset MIDI files:', files_count)
print('Average note pitch =', average_note_pitch)
#print('Min note pitch =', min_note)
#print('Max note pitch =', max_note)
#print('Number of MIDI events recorded:', len(events_matrix))
print('Number of MIDI chords recorded:', chords_count)
print('The longest chord:', len(max(chords_list_f, key=len)), 'notes') 
print(max(chords_list_f, key=len))
print('Number of recorded melody events:', len(melody_list_f))
print('First melody event:', melody_list_f[0], 'Last Melody event:', melody_list_f[-1])
print('Total number of MIDI events recorded:', len(chords_list_f))
# Dataset
MusicDataset = [chords_list_f, melody_list_f]

with open(full_path_to_output_dataset_to, 'wb') as filehandle:
    # store the data as binary data stream
    pickle.dump(MusicDataset, filehandle)
print('Dataset was saved at:', full_path_to_output_dataset_to)

# MIDI/files stats code from the previous version.

print('Task complete. Enjoy! :)')

#@title Process MIDI Dataset to TXT Dataset (w/Tegridy MIDI-TXT Processor)
#@markdown Output file is a plain UTF-8 TXT file, readable by any TXT application. On Windows you can open it even with the Notepad.
#@markdown line_by_line_dataset option is for creating new types of IV datasets. Useful for AI BPE tokenizers.
full_path_to_TXT_dataset = "/content/Intelligent-Virtuoso-Music-TXT-Dataset.txt" #@param {type:"string"}
line_by_line_dataset = True #@param {type:"boolean"}
represent_events_every_nth_chord = 0 #@param {type:"slider", min:0, max:100, step:1}
chords_durations_multiplier = 1 #@param {type:"slider", min:0.1, max:2, step:0.1}
chords_beat_divider = 10 #@param {type:"slider", min:1, max:100, step:1}
pad_chords_with_rests = False #@param {type:"boolean"}
simulate_velocity = True #@param {type:"boolean"}

# MIDI Dataset to txt dataset converter 
print('TMIDI-TXT Processor')
print('Starting up...')
if os.path.exists(full_path_to_TXT_dataset):
  os.remove(full_path_to_TXT_dataset)
  print('Removing old Dataset...')
else:
  print("Creating new Dataset file...")
if simulate_velocity:
  print('Simulated velocity mode is enabled.')

TXT = ''
number_of_chords = 0
number_of_bad_chords = 0
dataset_name = 'DATASET=Intelligent_VIRTUOSO_TXT_Music_Dataset'

file = open(full_path_to_TXT_dataset, 'a')
TXT, number_of_chords, number_of_bad_chords = TMIDI.Tegridy_MIDI_TXT_Processor(dataset_name, 
                                                                               chords_list_f, 
                                                                               melody_list_f, 
                                                                               simulate_velocity, 
                                                                               line_by_line_dataset,
                                                                               represent_events_every_nth_chord,
                                                                               chords_durations_multiplier,
                                                                               pad_chords_with_rests,
                                                                               chords_beat_divider)
file.write(TXT)
file.close()
print('Number of chords recorded: ', number_of_chords)
print('Number of bad/skipped chords: ', number_of_bad_chords)
print('Done!')

"""# Setup and Intialize the Model

## YOU MUST RUN ALL CELLS/CODE IN THIS SECTION to init the model. Does not matter if the model is empty or pre-trained.

## DO NOT EXECUTE TRAIN CELL/CODE UNLESS YOU INTEND TO TRAIN FROM SCRATCH
"""

#@title Setup functions and procedures
model_attention_span_in_tokens = 768 #@param {type:"slider", min:0, max:1024, step:16}

class CharDataset(Dataset):

    def __init__(self, data, block_size):
        chars = sorted(list(set(data)))
        data_size, vocab_size = len(data), len(chars)
        print('data has %d characters, %d unique.' % (data_size, vocab_size))
        
        self.stoi = { ch:i for i,ch in enumerate(chars) }
        self.itos = { i:ch for i,ch in enumerate(chars) }
        self.block_size = block_size
        self.vocab_size = vocab_size
        self.data = data
    
    def __len__(self):
        return len(self.data) - self.block_size

    def __getitem__(self, idx):
        # grab a chunk of (block_size + 1) characters from the data
        chunk = self.data[idx:idx + self.block_size + 1]
        # encode every character to an integer
        dix = [self.stoi[s] for s in chunk]
        
        x = torch.tensor(dix[:-1], dtype=torch.long)
        y = torch.tensor(dix[1:], dtype=torch.long)
        return x, y

        
block_size = model_attention_span_in_tokens # spatial extent of the model for its context

#@title Specify full path to the processed TMIDI-TXT dataset file
full_path_to_training_text_file = "/content/Intelligent-Virtuoso-Music-TXT-Dataset.txt" #@param {type:"string"}
text = open(full_path_to_training_text_file, 'r').read() # don't worry we won't run out of file handles
train_dataset = CharDataset(text, block_size) # one line of poem is roughly 50 characters

#@title Create GPT2 model
model_embed_size = 256 #@param {type:"slider", min:0, max:1024, step:64}
number_of_heads = 16 #@param {type:"slider", min:1, max:16, step:1}
number_of_layers = 4 #@param {type:"slider", min:1, max:16, step:1}


mconf = GPTConfig(train_dataset.vocab_size, 
                  train_dataset.block_size,
                  n_layer=number_of_layers, 
                  n_head=number_of_heads, 
                  n_embd=model_embed_size)

model = GPT(mconf)

#@title Setup all training parameters
number_of_training_epochs = 3 #@param {type:"slider", min:1, max:5, step:1}
training_batch_size = 24 #@param {type:"slider", min:0, max:160, step:4}
model_learning_rate = 6e-4 #@param {type:"number"}
# initialize a trainer instance and kick off training

tconf = TrainerConfig(max_epochs=number_of_training_epochs, 
                      batch_size=training_batch_size, 
                      learning_rate=model_learning_rate,
                      num_workers=4)
trainer = Trainer(model, train_dataset, None, tconf)

"""# Train the model or Load/Re-load the existing pre-trained model checkpoint"""

# Commented out IPython magic to ensure Python compatibility.
#@title (OPTION 1) Train the model
# %cd /content/
trainer.train()

#@title Plot Positional Embeddings

# visualize some of the learned positional embeddings, maybe they contain structure
plt.figure(figsize=(18, 1))  
ci = model.pos_emb.data[0, :, 0].cpu()
zci = torch.cat((torch.tensor([0.0]), ci)) # pre-cat a zero
plt.imshow(zci.view(1, block_size+1).numpy())
plt.axis('off')

# Commented out IPython magic to ensure Python compatibility.
#@title Save/Re-Save the model from memory
#@markdown Standard PyTorch AI models file extension is PTH
full_path_to_save_model_to = "/content/Intelligent-Virtuoso-Trained-Model.pth" #@param {type:"string"}
# %cd /content/
torch.save(model, full_path_to_save_model_to)

#@title (OPTION 2) Load existing model/checkpoint
full_path_to_model_checkpoint = "/content/Intelligent-Virtuoso-Trained-Model.pth" #@param {type:"string"}
model = torch.load(full_path_to_model_checkpoint)
model.eval()

"""# Generate, download, plot, and listen to the output"""

#@title Generate and download the composition as TXT file.
#@markdown PLEASE NOTE IMPORTANT POINTS: 

#@markdown 0) If you are not sure where to start/what settings to set, please use original defaults.

#@markdown 1) Model primes from the dataset !!!

#@markdown 2) Model's first output may be empty or garbled so please try several times before discarting the model

#@markdown 3) You can now communicate to the model desired length of the output composition by suffixing input_prompt with number of chords.

#@markdown I.e. SONG=Relax_with_900_Chords

#@markdown 3) Coherence of GPT2 Models is inversly proportional to the length of the generated composition, so the best resutls are achieved with shorter compositions and/or continuation routines use (which be implemented in the future version of Intelligent VIRTUOSO)

print('Intelligent VIRTUOSO Model Generator')
print('Starting up...')
number_of_tokens_to_generate = 16768 #@param {type:"slider", min:0, max:32768, step:128}
creativity_temperature = 0.8 #@param {type:"slider", min:0.05, max:4, step:0.05}
top_k_prob = 24 #@param {type:"slider", min:0, max:50, step:1}
input_prompt = "SONG=Deep" #@param {type:"string"}
debug = False #@param {type:"boolean"}

os.chdir('/content/')

model.to(device)

context = input_prompt
x = torch.tensor([train_dataset.stoi[s] for s in context], dtype=torch.long)[None,...].to(trainer.device)
y = sample(model, x, number_of_tokens_to_generate, temperature=creativity_temperature, sample=True, top_k=top_k_prob)[0]
completion = ''.join([train_dataset.itos[int(i)] for i in y])

''' #You can also try auto-regressive sampling on the output 
#but it does not show any improvements cuz its essentially one model canibalizing itself. :(

context = input_prompt
x = torch.tensor([train_dataset.stoi[s] for s in context], dtype=torch.long)[None,...].to(trainer.device)

for i in range(10):
  context = completion
  x = torch.tensor([train_dataset.stoi[s] for s in context], dtype=torch.long)[None,...].to(trainer.device)
  y = sample(model, x, number_of_tokens_to_generate, temperature=creativity_temperature, sample=True, top_k=top_k_prob)[0]
  completion = ''.join([train_dataset.itos[int(i)] for i in y])'''

# Stuff for datetime stamp
now = ''
now_n = str(datetime.now())
now_n = now_n.replace(' ', '_')
now_n = now_n.replace(':', '_')
now = now_n.replace('.', '_')
    
fname = '/content/Intelligent-VIRTUOSO-Composition-' + 'generated-on-' + str(now) 


print('Done!')
print('Saving to', str(fname + '.txt'))
with open(fname + '.txt', "w") as text_file:
    print(completion, file=text_file)

print('Downloading TXT file...')
from google.colab import files
files.download(fname + '.txt')

#@title Convert to MIDI from TXT (w/Tegridy MIDI-TXT Processor)

#@markdown NOTE: For now, only the first generated performance is recorded as MIDI.
#@markdown The rest of generated output sequence is only recorded in the generated TXT file.
#@markdown Therefore, please check TXT file for other/better sequences that might've been generated.

#@markdown Standard MIDI timings are 400/120(80)

#@markdown dataset_MIDI_events_time_denominator MUST MATCH with the dataset's denominator!!!

#@markdown Silence offset is model specific and usually the same so most likely it will need to be set up only once. Otherwise, please play with the settings utill you will find the right ones for your particular model.

#@markdown line_by_line_dataset option is for support of legacy datasets (IV 1.0-1.1). From now on IV will be using new line-by-line datasets for compatibility with AI BPE tokenizers.

line_by_line_dataset = True #@param {type:"boolean"}
dataset_MIDI_events_time_denominator = 10 #@param {type:"slider", min:1, max:200, step:1}
number_of_ticks_per_quarter = 420 #@param {type:"slider", min:10, max:1000, step:10}
start_from_this_generated_event = 0 #@param {type:"slider", min:0, max:100, step:1}
remove_generated_silence_if_needed = False #@param {type:"boolean"}
silence_offset_from_start = 75000 #@param {type:"integer"}
simulate_velocity = False #@param {type:"boolean"}

song_score = []
n = 0
z = 0
detailed_stats = 0

#fname = '/content/untitled'

with open(str(fname + '.txt'), 'r') as file:
    input=file.read()

output_signature = 'Intelligent VIRTUOSO'

midi_data, n, z, detailed_stats = TMIDI.Tegridy_TXT_MIDI_Processor(input,
                                             line_by_line_dataset,
                                             dataset_MIDI_events_time_denominator,
                                             number_of_ticks_per_quarter,
                                             start_from_this_generated_event,
                                             remove_generated_silence_if_needed,
                                             silence_offset_from_start,
                                             simulate_velocity,
                                             output_signature)                         

with open(fname + '.mid', 'wb') as midi_file:
    midi_file.write(midi_data)
    midi_file.close()
print('Done!')

from google.colab import files
files.download(fname + '.mid')
print('Detailed MIDI stats:')
detailed_stats

#@title Listen to the last generated composition
#@markdown NOTE: May be very slow with the long compositions
print('Synthesizing the last output MIDI. Please stand-by... ')
FluidSynth("/usr/share/sounds/sf2/FluidR3_GM.sf2", 16000).midi_to_audio(str(fname + '.mid'), str(fname + '.wav'))
Audio(str(fname + '.wav'), rate=16000)

"""## Congrats! :) You did it :)"""