import pretty_midi
import soundfile as sf
from tqdm import tqdm
import librosa
import numpy as np
import os
import sys
import shutil

import config as C


# 親ディレクトリの設定
def MakeHome(home_dir=C.fft_path):
    if os.path.exists(home_dir):
        shutil.rmtree(home_dir)
    os.mkdir(home_dir)

# 使用するピッチのリストを作成
def NoteList():
    note_names = []
    single_note = ['C', 'D', 'E', 'F', 'G', 'A', 'B']
    for i in range(C.octave):
        for note in single_note:
            note_names.append(note + str(i+2)) # "note" + "i"
    
    return note_names

# 楽譜を作成してwavを生成
# "note数"x"楽器数"の出力
def MakeNpy(fft_dir, tempo):
    
    # 楽譜を作成
    note_names = NoteList()
    for i in tqdm(range(C.inst_num)):
        for j, note_name in enumerate(note_names):
            pm = pretty_midi.PrettyMIDI(initial_tempo=tempo)
            instrument = pretty_midi.Instrument(i)
            
            # tempoの設定
            start_time = 0
            end_time = 60/tempo
            while(1):
                if start_time+end_time >= 30:
                    end_time = 30
                note_number = pretty_midi.note_name_to_number(note_name)
                note = pretty_midi.Note(velocity=100, pitch=note_number, start=start_time, end=start_time+end_time)
                start_time = start_time + end_time
                instrument.notes.append(note)
                if start_time >= 30:
                    break
            pm.instruments.append(instrument)
            
            # fluidsynthでwavに変換
            y = pm.fluidsynth(sf2_path='/usr/share/sounds/sf2/FluidR3_GM.sf2')
            
            # npyを生成してwavを削除
            # fluidsynthのサンプリングレートが44100Hzなので合わせる
            time = C.time
            rate = C.orig_sr
            y = y[:time*rate]
            fname = "{}/{}_{}.wav".format(fft_dir, i, j)
            sf.write(fname, y, rate)
            y, sr = librosa.load(fname, sr=None)
            os.remove(fname)
            
            # STFT情報ををnpyで保存
            y = y[np.nonzero(y)[0][0]:]
            min_length = min([y.size])
            y = librosa.resample(y[:min_length], C.orig_sr, C.sr)
            fname = fname.split("/")[-1].split(".")[0]
            spec = librosa.stft(y, n_fft=C.fft_size, hop_length=C.hop).astype(np.float32)
            norm = spec.max()
            spec /= norm
            np.save(os.path.join("{}/{}".format(C.fft_path, tempo), fname+".npy"), spec)
            

# tempoを30-285で設定
def Generate(home, num):
    
    # 親ディレクトリの作成
    MakeHome()
    
    # 諸情報の表示
    print("Save npy")
    print("tempo:{}-{}".format(C.def_tempo, C.def_tempo+num-1))
    print("instruments num:", C.inst_num)
    print("pitch num:", C.octave*7)
    print("output one dir wavs:", C.octave*7*C.inst_num)
    print("output all wavs:", C.octave*7*C.inst_num*(C.def_tempo+num-C.def_tempo))
    
    # 音楽ファイルの生成
    def_tempo = C.def_tempo
    for i in tqdm(range(num)):
        tempo_dir = '{}/{}'.format(home, str(i+def_tempo))
        if not os.path.exists('{}/{}'.format(home, str(i+def_tempo))):
            os.mkdir('{}/{}'.format(home, str(i+def_tempo)))
        MakeNpy(tempo_dir, i+def_tempo)