# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %load_ext autoreload
# %autoreload 2

# +
from collections import Counter
import pickle as pkl
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, SequentialSampler, RandomSampler, Dataset, SubsetRandomSampler

from sklearn.model_selection import train_test_split

from tqdm.notebook import tqdm

from torchtext.vocab import Vocab

from chord_rec.models.seq2seq.Seq2Seq import BaseSeq2Seq
from chord_rec.models.seq2seq.Encoder import BaseEncoder
from chord_rec.models.seq2seq.Decoder import BaseDecoder
# -

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

class Vec21Dataset(Dataset):
    def __init__(self, note_vec_seq, chord_seq, vocab):
        'Initialization'
        self.note_vec_seq = note_vec_seq
        self.chord_seq = chord_seq
        self.vocab = vocab

    def __len__(self):
        'Get the total length of the dataset'
        return len(self.note_vec_seq)
    
    def __getitem__(self, index):
        'Generates one sample of data'
        # Select sample
        
        
        return self.note_vec_seq[index], self.vec_encode(self.chord_seq[index])
    
    def encode(self, x):
        return self.vocab.stoi[x]


    def vec_encode(self, x):
        return np.vectorize(self.encode)(x)
    

    def decode(self, x):
        return self.vocab.itos[x]


    def vec_decode(self, x):
        return np.vectorize(self.decode)(x)
    

# +
data = pkl.load(open('bach_reduced_seq2seq_4mm_new.pkl',"rb"))

note_seq, chord_seq = [],[]
max_seq_len = 0
data_num = 0
for file in data:
    data_num += len(file)
    for window in file:
        note_seq.append(window[0])
        chord_seq.append(window[1])
        max_seq_len = max(max_seq_len, len(window[1]))

# +
note_padding_vec = np.full(len(note_seq[0][0]), -1).reshape(1,-1) # should be 45; not sure if -1 is good
note_ending_vec = np.ones(len(note_seq[0][0])).reshape(1,-1) # should be 45
note_starting_vec = np.zeros(len(note_seq[0][0])).reshape(1,-1) # should be 45

chord_start = "<sos>"
chord_padding = "<pad>"
chord_end = "<eos>"

padded_note_seq = []
padded_chord_seq = []
for i in range(len(note_seq)):
    len_diff = max_seq_len - len(note_seq[i])

    temp_note_vec = np.vstack((note_starting_vec, np.array(note_seq[i]), note_ending_vec, np.repeat(note_padding_vec, len_diff , axis = 0)))
    padded_note_seq.append(temp_note_vec)
    
    temp_chord_vec = np.hstack((chord_start, np.array(chord_seq[i]), chord_end, np.repeat(chord_padding, len_diff , axis = 0)))
    padded_chord_seq.append(temp_chord_vec)
# -

stacked_note_seq = np.stack(padded_note_seq, axis = 0)
stacked_chord_seq = np.vstack(padded_chord_seq)

SEED = 0
VAL_SIZE = 0.2
TEST_SIZE = 0.2

note_vec = np.asarray(stacked_note_seq, dtype = np.float32)

chord_vocab = Vocab(Counter(list(stacked_chord_seq.flatten())))

vec_size = note_vec.shape[-1]
vocab_size = len(chord_vocab.stoi)

# +
note_train, note_test, chord_train, chord_test \
    = train_test_split(note_vec, stacked_chord_seq, test_size=TEST_SIZE, random_state=SEED)

note_train, note_val, chord_train, chord_val \
    = train_test_split(note_vec, stacked_chord_seq, test_size= VAL_SIZE/ (1-TEST_SIZE), random_state=SEED)
# -

train_dataset = Vec21Dataset(note_train, chord_train, chord_vocab)
val_dataset = Vec21Dataset(note_val, chord_val, chord_vocab)
test_dataset = Vec21Dataset(note_test, chord_test, chord_vocab)

batch_size = 32
train_loader = DataLoader(train_dataset, batch_size = batch_size, drop_last = True)
val_loader = DataLoader(val_dataset, batch_size =  batch_size, drop_last = True)
test_loader =  DataLoader(test_dataset, batch_size =  batch_size, drop_last = True)

# +
input_size = vec_size
emb_size = vec_size
encoder_hidden_size = 128
decoder_hidden_size = 128
encoder_dropout = 0.5
decoder_dropout = 0.5
n_layers = 2
output_size = vocab_size
model_type = "LSTM"

model = nn.Transformer()
# -

vocab_size

criterion = nn.CrossEntropyLoss(ignore_index = chord_vocab.stoi["<pad>"])
optimizer = optim.AdamW(model.parameters(), lr = 1e-3)

# +
# model.state_dict = torch.load("sears_reduce_20warm80epoch_s2s.sdict")
# -

train_pgbar = tqdm(train_loader)
for idx, (note, chord) in enumerate(train_pgbar):
    chord = chord.long().to(device)

    target = chord[:,:-1].permute(1,0)
    gt = chord[:, 1:].permute(1,0)
    note = note.permute(1,0,2)
#     print(note.shape)
#     print(target.shape)
    
    output = model(note, target)
    break

# +
epochs = 100
train_losses = []
val_losses = []
clip = 1

for epoch in range(1,epochs+1):
    total_loss = 0
    
    model.train()
    train_pgbar = tqdm(train_loader)
    for idx, (note, chord) in enumerate(train_pgbar):
        chord = chord.long().to(device)
        
        target = chord[:,:-1]
        gt = chord[:, 1:]
        
        pred = model()
        pred = pred.permute(0,2,1)
        loss = criterion(pred, chord.long().to(device))

        optimizer.zero_grad()
        if loss.item() != 0:
            loss.backward()
#             torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
            optimizer.step()
        
#         s = "Epoch: {} Train Loss: {:.5f}".format(epoch, loss.item())
#         tqdm.write(s)
        train_pgbar.set_postfix({'Epoch': epoch, 'Train Loss': "{:.5f}".format(loss.item())})
        total_loss += loss.item()
    print("Total Training Loss: %.2f" % total_loss)
    train_losses.append(total_loss/len(train_loader))
    total_loss = 0
    
    model.eval()
    val_pgbar = tqdm(val_loader)
    for idx, (note, chord) in enumerate(val_pgbar):

        pred = model(note.to(device), chord.long().to(device), teacher_forcing = False, start_idx = chord_vocab.stoi["<sos>"])
        loss = criterion(pred.permute(0,2,1), chord.long().to(device))
        
#         s = "Epoch: {} Val Loss: {:.5f}".format(epoch, loss.item())
#         tqdm.write(s)
        val_pgbar.set_postfix({'Epoch': epoch, 'Val Loss': "{:.5f}".format(loss.item())})
        total_loss += loss.item()
    print("Total Validation Loss: %.2f" % total_loss)
    val_losses.append(total_loss/len(val_loader))
# -

import matplotlib.pyplot as plt

# +
fig = plt.figure()
ax = fig.add_subplot(111)
plt.plot(train_losses, color = "b", label = "train")
plt.plot(val_losses, color = "r", label = "val")
plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left', borderaxespad=0.)
ax.set_xlabel('epochs')
ax.set_ylabel('cross entropy loss') 
xint = range(len(train_losses))

plt.xticks(xint)
plt.show()
# -

model.eval()
all_pred2 = []
all_label2 = []
for idx, (note, chord) in enumerate(tqdm(val_loader)):
        pred = model(note.to(device), chord.long().to(device), teacher_forcing = False, start_idx = chord_vocab.stoi["<sos>"])
        pred = pred.detach().cpu().numpy().argmax(axis = -1)
        
        label = chord.detach().cpu().numpy()
        pred[:,0] = np.full(len(pred), chord_vocab.stoi["<sos>"])
        all_pred2.append(val_dataset.vec_decode(pred))
        all_label2.append(val_dataset.vec_decode(label))

all_pred2 = np.vstack(all_pred2)
all_label2 = np.vstack(all_label2)

all_pred2[0]

all_label2[0]

decoded_preds = all_pred2
decoded_chords = all_label2

decoded_preds

# +
mask = (decoded_preds != "<sos>") & (decoded_preds != "<eos>") & (decoded_preds != "<pad>")
masked_preds = decoded_preds[mask]
masked_chords = decoded_chords[mask]

print(np.sum(masked_preds == masked_chords) / len(masked_chords))
# -

print(masked_chords)

masked_preds == masked_chords

# +
# SEPERATE EVALUATION OF ROOT AND QUALITY AFTER DECODING
# seperate all pred 
root_preds = decoded_preds.copy()
quality_preds = decoded_preds.copy()
for r_id in range(decoded_preds.shape[0]):
    for c_id in range(decoded_preds.shape[1]):
        sp = decoded_preds[r_id, c_id].split(' ')
        root_preds[r_id, c_id] = sp[0]
        quality_preds[r_id, c_id] = ' '.join(sp[1:])
    
root_labels = decoded_chords.copy()
quality_labels = decoded_chords.copy()
for r_id in range(decoded_chords.shape[0]):
    for c_id in range(decoded_chords.shape[1]):
        sp = decoded_chords[r_id, c_id].split(' ')
        root_labels[r_id, c_id] = sp[0]
        quality_labels[r_id, c_id] = ' '.join(sp[1:])
# # seperate all lable 
# root_labels = []
# quality_labels = []
# for c in decoded_chords:
#     sp = c.split(' ')
#     root_labels.append(sp[0])
#     quality_labels.append(' '.join(sp[1:]))
    
# root_labels = np.asarray(root_labels)
# quality_labels = np.asarray(quality_labels)
# -

mask = (root_preds != "<sos>") & (root_preds != "<eos>") & (root_preds != "<pad>")
root_preds = root_preds[mask]
quality_preds = quality_preds[mask]
root_label = root_labels[mask]
quality_labels = quality_labels[mask]

np.sum(root_preds == root_label) / len(root_preds)

np.sum(quality_preds == quality_labels) / len(quality_preds)

torch.save(model.state_dict(), "bach_reduced_10warm80epoch50post_s2s.sdict")

# EVALUATE root AFTER DECODING
sum(1 for x,y in zip(root_preds, root_labels) if x == y) / len(root_labels)

# EVALUATE quality AFTER DECODING
sum(1 for x,y in zip(quality_preds, quality_labels) if x == y) / len(quality_labels)

co = Counter(chords)

most_common_chord = list(list(zip(*co.most_common(20)))[0])

for i in range(len(decoded_preds)):
    if decoded_preds[i] not in most_common_chord:
        decoded_preds[i] = "others"
    if decoded_chords[i] not in most_common_chord:
        decoded_chords[i] = "others"

from sklearn.metrics import confusion_matrix, plot_confusion_matrix
import seaborn as sn

cm = confusion_matrix(decoded_chords, decoded_preds, normalize = "true", labels = most_common_chord + ["others"])

# +
fig, ax = plt.subplots(figsize=(13,10)) 

sn.heatmap(cm, annot=False)
ax.set_xticklabels(most_common_chord + ["others"])
ax.set_yticklabels(most_common_chord + ["others"])
plt.yticks(rotation=0) 
plt.xticks(rotation="vertical") 
plt.show()
fig.savefig("confusion_bachhaydn_baseline.pdf", format = "pdf")
# -

torch.save(model.state_dict(), "baseline_bach_and_haydn.pt")

symbol, num = list(zip(*co.most_common(50)))
symbol = list(symbol)
num = list(num)

symbol += ['others']
num += [np.sum(list(co.values())) - np.sum(num)]

num/=np.sum(num)

# +
plt.subplots(figsize=(13,10)) 
x_pos = [i for i, _ in enumerate(symbol)]

plt.bar(x_pos, num)
plt.xlabel("Chord Symbol")
plt.ylabel("Occurance")
plt.title("bach chorales Chord Distribution")

plt.xticks(x_pos, symbol, rotation = "vertical")

plt.show()
# -



StarcoderdataJupyterScriptsDedupFiltered