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

# # Spectrogram Channels U-net
#
# * Spectrogram-Channels U-Net: A Source Separation Model Viewing Each Channel as the Spectrogram of Each Source, [arXiv:1810.11520](https://arxiv.org/abs/1810.11520)
#   * <NAME>∗, <NAME>∗, <NAME>
#   
# * Implemented by `TensorFlow v2.0`

# ## Import modules

# +
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import sys
import time
import glob

import collections

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

import librosa
import librosa.display

from IPython import display

import tensorflow as tf
from tensorflow.keras import layers

tf.enable_eager_execution()

os.environ["CUDA_VISIBLE_DEVICES"]="0"

# +
# Training Flags (hyperparameter configuration)
model_name = 'spectrogram_channels_unet'
train_dir = 'train/' + model_name + '/exp1/'
max_epochs = 20
save_model_epochs = 20
print_steps = 1
batch_size = 8
learning_rate = 1e-3
# Total number of samples in train_dataset
N = 4525 # N_test = 2469 for test_dataset

BUFFER_SIZE = N
# -

# ## Set up dataset with `tf.data`

# +
data_path = './datasets/tfrecords/'
train_data_filenames = [os.path.join(data_path, 'train', name)
                        for name in os.listdir(os.path.join(data_path, 'train')) if 'tfrecord' in name]
for name in train_data_filenames:
  print(name)
print('--------')
  
test_data_filenames = [os.path.join(data_path, 'test', name)
                       for name in os.listdir(os.path.join(data_path, 'test')) if 'tfrecord' in name]
for name in test_data_filenames:
  print(name)


# -

# Transforms a scalar string `example_proto` into a pair of a scalar string and
# a scalar integer, representing an spectragram and corresponding informations, respectively.
def _parse_function(example_proto):
  features = {'wav_concat': tf.io.FixedLenFeature([], tf.string, default_value=""),
              'time_step': tf.io.FixedLenFeature([], tf.int64, default_value=0),
              'track_number': tf.io.FixedLenFeature([], tf.int64, default_value=0),
              'split_number': tf.io.FixedLenFeature([], tf.int64, default_value=0),}
  
  parsed_features = tf.io.parse_single_example(example_proto, features)

  wav_concat = tf.io.decode_raw(parsed_features["wav_concat"], out_type=tf.float32)
  time_step = tf.cast(parsed_features["time_step"], dtype=tf.int32)
  #track_number = tf.cast(parsed_features["track_number"], dtype=tf.int32)
  #split_number = tf.cast(parsed_features["split_number"], dtype=tf.int32)
  
  num_sources = 6 # for [mixtures, vocals, drums, basses, others, accompaniments]
  num_channels = 2 # for left, right
  wav_concat = tf.reshape(wav_concat, shape=[num_sources, num_channels, time_step])

  return wav_concat, time_step


def _augmentation_function(wav_concat, time_step):
  """Random cropping for data augmentation
  We set that the input (spectrogram) shape is (num_frames, frequency_bin) = (128, 1024).
  Number of samples to needed: 128 x 512 (hop_size) = 131072
  Time of wave to needed: 131072 / 44100 = 2.97 sec
  """
  target_time_step = 131072
  available_time_step = time_step - target_time_step
  
  crop_index = tf.random.uniform(shape=[]) * tf.cast(available_time_step, dtype=tf.float32)
  crop_index = tf.cast(crop_index, dtype=tf.int32)
  wav_concat_crop = wav_concat[:, :, crop_index:crop_index+target_time_step]
  
  return wav_concat_crop


def _stft(wav, sources_to_wanted):
  """Short time Fourier Transform
  
  Args:
    wav: Tens
    sources_to_wanted: list of name of sources to wanted
        - full source names: ['vocals', 'drums', 'basses', 'others', 'accompaniments']
        
  Returns:
    stft: The results of stft of input wav
  """
  num = collections.OrderedDict()
  num['vocals'] = 1
  num['drums'] = 2
  num['basses'] = 3
  num['others'] = 4
  num['accompaniments'] = 5
  
  wav_to_wanted = collections.OrderedDict()
  wav_to_wanted['mixtures'] = wav[0]
  for key, value in num.items():
    if key in sources_to_wanted:
      wav_to_wanted[key] = wav[value]  
      
  stfts_db = []
  for key, wav in wav_to_wanted.items():
    stft = tf.contrib.signal.stft(wav,
                          frame_length=2048,
                          frame_step=512,
                          fft_length=2048,
                          pad_end=True)[..., :1024]
    
    # [num_channels, num_frames, frequency_bin]
    # -> [num_frames, frequency_bin, num_channels]
    stft = tf.transpose(stft, perm=[1, 2, 0])

    def _amp_to_db(x):
      log_offset = 1e-6
      return 20 * tf.math.log(x + log_offset) / tf.math.log(10.0) # natural logarithm -> logarithm based on 10
    # for our suggestion
    min_level_db = -110
    ref_level_db = 40
    stft_db = _amp_to_db(tf.math.abs(stft)) - ref_level_db
    stfts_db.append(stft_db)

  return stfts_db[0], tf.concat(stfts_db[1:], axis=-1)


def _stft_for_test(wav, sources_to_wanted):
  """Short time Fourier Transform
  
  Args:
    wav: Tens
    sources_to_wanted: list of name of sources to wanted
        - full source names: ['vocals', 'drums', 'basses', 'others', 'accompaniments']
        
  Returns:
    mixture_stft: The real part of stft of input wav
    targets_stft: The real part of stft of target (each source) wav
    targets_stft_imag: The imaginary part of stft of target (each source) wav
  """
  num = collections.OrderedDict()
  num['vocals'] = 1
  num['drums'] = 2
  num['basses'] = 3
  num['others'] = 4
  num['accompaniments'] = 5
  
  wav_to_wanted = collections.OrderedDict()
  wav_to_wanted['mixtures'] = wav[0]
  for key, value in num.items():
    if key in sources_to_wanted:
      wav_to_wanted[key] = wav[value]  
      
  stfts_db = []
  stfts_angle = []
  for key, wav in wav_to_wanted.items():
    stft = tf.contrib.signal.stft(wav,
                          frame_length=2048,
                          frame_step=512,
                          fft_length=2048,
                          pad_end=True)[..., :1024]
    
    # [num_channels, num_frames, frequency_bin]
    # -> [num_frames, frequency_bin, num_channels]
    stft = tf.transpose(stft, perm=[1, 2, 0])
    
    # separate real part and imaginary part
    stft_angle = tf.math.angle(stft)
    stfts_angle.append(stft_angle)
    
    def _amp_to_db(x):
      log_offset = 1e-6
      return 20 * tf.math.log(x + log_offset) / tf.math.log(10.0) # natural logarithm -> logarithm based on 10
    # for our suggestion
    min_level_db = -110
    ref_level_db = 40
    stft_db = _amp_to_db(tf.math.abs(stft)) - ref_level_db
    stfts_db.append(stft_db)

  return stfts_db[0], tf.concat(stfts_db[1:], axis=-1), tf.concat(stfts_angle[1:], axis=-1)


train_dataset = tf.data.TFRecordDataset(train_data_filenames)
train_dataset = train_dataset.map(_parse_function)
train_dataset = train_dataset.map(_augmentation_function)
extracted_sources_list = ['vocals', 'accompaniments']
train_dataset = train_dataset.map(lambda x: _stft(x, extracted_sources_list))
train_dataset = train_dataset.shuffle(BUFFER_SIZE)
train_dataset = train_dataset.batch(batch_size)

# To test for batch
for mixtures, targets in train_dataset.take(3):
  start_time = time.time()
  print(mixtures.shape)
  print(targets.shape)
  print("duration: {} sec".format(time.time() - start_time))

# ### Plot for spectrogram of mixtures

plt.figure(figsize=(16, 4))
librosa.display.specshow(mixtures[0, :, :, 0].numpy().T)
plt.colorbar()
plt.show()

test_dataset = tf.data.TFRecordDataset(test_data_filenames)
test_dataset = test_dataset.map(_parse_function)
test_dataset = test_dataset.map(_augmentation_function)
extracted_sources_list = ['vocals', 'accompaniments']
test_dataset = test_dataset.map(lambda x: _stft_for_test(x, extracted_sources_list))
test_dataset = test_dataset.batch(1)

# To test for batch
for mixtures, targets, targets_imaginary in test_dataset.take(3):
  start_time = time.time()
  print(mixtures.shape)
  print(targets.shape)
  print(targets_imaginary.shape)
  print("duration: {} sec".format(time.time() - start_time))


# ## Build the model

class ConvBlock(tf.keras.Model):
  def __init__(self, num_filters, size):
    super(ConvBlock, self).__init__()
    self.conv1 = layers.Conv2D(filters=num_filters,
                               kernel_size=(size, size),
                               padding='same',
                               use_bias=False)
    self.batchnorm1 = layers.BatchNormalization()
    self.conv2 = layers.Conv2D(filters=num_filters,
                               kernel_size=(size, size),
                               padding='same',
                               use_bias=False)
    self.batchnorm2 = layers.BatchNormalization()
  
  def call(self, x, training=True):
    x = self.conv1(x)
    x = self.batchnorm1(x, training=training)
    x = tf.nn.relu(x)
    x = self.conv2(x)
    x = self.batchnorm2(x, training=training)
    x = tf.nn.relu(x)
    
    return x


class EncoderBlock(tf.keras.Model):
  def __init__(self, num_filters, size):
    super(EncoderBlock, self).__init__()
    self.conv_block = ConvBlock(num_filters, 3)
    
  def call(self, x, training=True):
    encoder = self.conv_block(x, training=training)
    encoder_pool = layers.MaxPooling2D(pool_size=(2, 2),
                                       strides=(2, 2))(encoder)
  
    return encoder_pool, encoder


class ConvTransposeBlock(tf.keras.Model):
  def __init__(self, num_filters, size):
    super(ConvTransposeBlock, self).__init__()
    self.convT1 = layers.Conv2DTranspose(filters=num_filters,
                                         kernel_size=(size, size),
                                         padding='same',
                                         use_bias=False)
    self.batchnorm1 = layers.BatchNormalization()
    self.convT2 = layers.Conv2DTranspose(filters=num_filters,
                                         kernel_size=(size, size),
                                         padding='same',
                                         use_bias=False)
    self.batchnorm2 = layers.BatchNormalization()
  
  def call(self, x, training=True):
    x = self.convT1(x)
    x = self.batchnorm1(x, training=training)
    x = tf.nn.relu(x)    
    x = self.convT2(x)
    x = self.batchnorm2(x, training=training)
    x = tf.nn.relu(x)
    
    return x


class DecoderBlock(tf.keras.Model):
  def __init__(self, num_filters, size):
    super(DecoderBlock, self).__init__()
    self.convT = layers.Conv2DTranspose(filters=num_filters,
                                        kernel_size=(size+2, size+2),
                                        strides=(2, 2),
                                        padding='same',
                                        use_bias=False)
    self.batchnorm = layers.BatchNormalization()
    self.dropout = layers.Dropout(0.4)
    self.convT_block = ConvTransposeBlock(num_filters, size)
  
  def call(self, input_tensor, concat_tensor, training=True):
    # Upsampling
    x = self.convT(input_tensor)
    x = self.batchnorm(x, training=training)
    x = tf.nn.relu(x)
    x = self.dropout(x, training=training)
    
    # concatenate
    x = tf.concat([x, concat_tensor], axis=-1)
    
    # just two consecutive conv_transpose
    x = self.convT_block(x, training=training)
    
    return x


class SpectrogramChannelsUNet(tf.keras.Model):
  def __init__(self):
    super(SpectrogramChannelsUNet, self).__init__()
    self.down1 = EncoderBlock(32, 3)
    self.down2 = EncoderBlock(64, 3)
    self.down3 = EncoderBlock(128, 3)
    self.down4 = EncoderBlock(256, 3)
    self.center = ConvBlock(512, 3)

    self.up1 = DecoderBlock(256, 3)
    self.up2 = DecoderBlock(128, 3)
    self.up3 = DecoderBlock(64, 3)
    self.up4 = DecoderBlock(32, 3)

    self.last = layers.Conv2D(filters=4,
                              kernel_size=(1, 1),
                              padding='same')
  
  def call(self, x, training):
    # x shape == (bs, 1024, 256, 1)
    x1_pool, x1 = self.down1(x, training=training) # (bs, 512, 128, 32)
    x2_pool, x2 = self.down2(x1_pool, training=training) # (bs, 256, 64, 64)
    x3_pool, x3 = self.down3(x2_pool, training=training) # (bs, 128, 32, 128)
    x4_pool, x4 = self.down4(x3_pool, training=training) # (bs, 64, 16, 256)
    x_center = self.center(x4_pool, training=training) # (bs, 64, 16, 512)

    x5 = self.up1(x_center, x4, training=training) # (bs, 128, 32, 256)
    x6 = self.up2(x5, x3, training=training) # (bs, 256, 64, 128)
    x7 = self.up3(x6, x2, training=training) # (bs, 512, 128, 64)
    x8 = self.up4(x7, x1, training=training) # (bs, 1024, 256, 32)

    x_last = self.last(x8) # (bs, 1024, 256, 5)
    x_last = tf.math.sigmoid(x_last)

    return x_last


model = SpectrogramChannelsUNet()

# To fetch inputs to model in order to set to input shape
for mixtures, targets in train_dataset.take(1):
  predictions = model(mixtures, training=True)

model.summary()


# ## Defining loss functions and optimizers

def singing_voice_loss(y_true, y_pred, alpha=1.0):
  # Use the mean_absolute_error
  #mae = tf.keras.losses.MeanAbsoluteError()
  #loss_vocal = mae(y_true[..., :2], y_pred[..., :2])
  #loss_accompaniments = mae(y_true[..., 2:], y_pred[..., 2:])
  
  loss_vocal = tf.losses.absolute_difference(y_true[..., :2], y_pred[..., :2])
  loss_accompaniments = tf.losses.absolute_difference(y_true[..., 2:], y_pred[..., 2:])
  
  return alpha * loss_vocal + (1. - alpha) * loss_accompaniments


#optimizer = tf.keras.optimizers.Adam(learning_rate)
optimizer = tf.train.AdamOptimizer(learning_rate)

checkpoint_dir = train_dir
#if not tf.io.gfile.exists(checkpoint_dir):
#  tf.io.gfile.makedirs(checkpoint_dir)
if not tf.gfile.Exists(checkpoint_dir):
  tf.gfile.MakeDirs(checkpoint_dir)
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
checkpoint = tf.train.Checkpoint(optimizer=optimizer, model=model)


def inverse_stft(test_predictions, test_targets_angle):
  target_shape = test_predictions.shape
  zero_padding = tf.zeros(shape=[target_shape[0], target_shape[1], 1, target_shape[3]])
  pred = tf.concat([test_predictions, zero_padding], axis=2)
  angle = tf.concat([test_targets_angle, zero_padding], axis=2)
  
  def _db_to_amp(x):
    return tf.math.pow(10.0, x * 0.05)

  ref_level_db = 40
  magnitude = _db_to_amp(pred + ref_level_db)
  
  magnitude_complex = tf.dtypes.complex(magnitude, 0.)
  angle_complex = tf.dtypes.complex(0., angle)
  
  stfts = magnitude_complex * angle_complex
  
  # [batch_size, num_frames, frequency_bin, num_channels]
  # -> [batch_size, num_channels, num_frames, frequency_bin]
  stfts = tf.transpose(stfts, perm=[0, 3, 1, 2])
  
  inv_stfts = tf.contrib.signal.inverse_stft(stfts=stfts,
                                     frame_length=2048,
                                     frame_step=512)#,
#                                     fft_length=2048)

  #inv_stfts = librosa.core.istft(stft[0,0, :,:].numpy().T, hop_length=512, win_length=2048)
  
  return inv_stfts


# +
print('Start Training.')
global_step = 1

#for epoch in range(max_epochs):
for epoch in range(1):
  
  for mixtures, targets in train_dataset.take(1):
    start_time = time.time()
    
    with tf.GradientTape() as tape:
      predictions = model(mixtures, training=True)
      loss = singing_voice_loss(targets, predictions)

    gradients = tape.gradient(loss, model.variables)
    optimizer.apply_gradients(zip(gradients, model.variables))
        
    epochs = global_step * batch_size / float(N)
    duration = time.time() - start_time

    if global_step % print_steps == 0:
      display.clear_output(wait=True)
      examples_per_sec = batch_size / float(duration)
      print("Epochs: {:.2f} global_step: {} loss: {:.3f} ({:.2f} examples/sec; {:.3f} sec/batch)".format(
                epochs, global_step, loss, examples_per_sec, duration))
      for test_mixtures, test_targets, test_targets_angle in test_dataset.take(1):
        test_predictions = model(test_mixtures, training=False)
        test_wav = inverse_stft(test_predictions, test_targets_angle)
    global_step = global_step + 1

  # saving (checkpoint) the model every save_epochs
  if (epoch + 1) % save_model_epochs == 0:
    checkpoint.save(file_prefix=checkpoint_prefix)
# -

for test_mixtures, test_targets, test_targets_angle in test_dataset.take(1):
  test_predictions = model(test_mixtures, training=False)
  test_wav = inverse_stft(test_predictions, test_targets_angle)

display.Audio(test_wav[0, :2, :], rate=44100)

# ## Restore the latest checkpoint and test

# restoring the latest checkpoint in checkpoint_dir
checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))

# ## Testing on the entire test dataset



StarcoderdataJupyterScriptsDedupFiltered