from chainer import serializers
import chainer
import chainer.functions as F
import chainer.links as L
import librosa
import numpy as np
import pandas as pd
from tqdm import tqdm
import os
import sys
import glob
import shutil
import json

import network
import config as C

# labelからbpmを検索
def SearchBpm(label, labelf=C.label):
    bpm_label = pd.read_csv(labelf, encoding="utf_8", sep=",")
    return bpm_label.iat[bpm_label.query("label==@label").index[0], 1]

# networkのロード
def LoadModel(model_path="{}/simple_tempo.model".format(C.log_path)):
    path = os.path.join(os.path.dirname(model_path), "param.txt")
    with open(path) as f:
        net = f.read()
    path = os.path.join(os.path.dirname(C.label), "tmp.py")
    with open(path, "w") as f:
        f.write("import chainer" + "\n" + "import chainer.functions as F" + "\n" + "import chainer.links as L" + "\n" + "import config as C" + "\n" + net)
    import tmp
    model = L.Classifier(tmp.myNN())
    serializers.load_npz(model_path, model)
    os.remove(path)

    return model

# 予測値をsoftmaxで返す
def ReturnPred(model_path, spec):
    model = LoadModel(model_path)
    with chainer.using_config('train', False), chainer.using_config('enable_backprop', False):
        pred = model.predictor(spec)
    pred = F.softmax(pred)
    
    return pred

# 0秒-C.time秒のスペクトログラムを返す
def ReturnSpec(fname):
    y, sr = librosa.load(fname, sr=None)
    y = y[np.nonzero(y)[0][0]:]
    min_length = min([y.size])
    y = librosa.resample(y[:min_length], sr, C.sr)
    y = y[:C.time*C.sr]
    # 2分に満たないものの処理
    if len(y) < C.time*C.sr:
        y = np.append(y, [.0]*(C.time*C.sr-len(y)))

    spec = librosa.stft(y, n_fft=C.fft_size, hop_length=C.hop).astype(np.float32)
    norm = spec.max()
    spec /= norm
    
    # ここで(channel, height, width)に
    spec = spec[np.newaxis,:]
    spec = spec[np.newaxis,:]
    
    return spec

def UseGiantStep(num=-1 ,meta_dir=C.meta_path, wav_dir=C.test_path, model_path="{}/simple_tempo.model".format(C.log_path)):
    print("use model:{}".format(model_path))
    files = glob.glob("{}/*.wav".format(wav_dir))

    if not type(num) is int:
        print("num is int")
        sys.exit()
    if num != -1:
        files = files[:num]
    
    cnt = 0
    raw_list = []
    pred_list = []
    for fname in tqdm(files):
        raw_list.append(os.path.basename(os.path.splitext(fname)[0]))

        # 真値
        meta_data = os.path.join(meta_dir, "{}.json".format(os.path.basename(os.path.splitext(fname)[0])))
        try:
            with open(meta_data) as f:
                meta_list = json.load(f)
            g_truth = meta_list["bpm"]
        except json.decoder.JSONDecodeError:
            g_truth = None

        # networkの予測
        spec = ReturnSpec(fname)
        net_pred = ReturnPred(model_path ,spec)
        net_pred = SearchBpm(np.argmax(net_pred.data.data))
        
    pred_list = pd.DataFrame(pred_list,index=[raw_list], columns=["truth", "network"])
    pred_list.to_csv(os.path.join(os.path.dirname(model_path), "predictions.csv"))
    print("Save Predictions:" + os.path.join(os.path.dirname(model_path), "predictions.csv"))