#!/usr/bin/env python

# Base Package
import os
import time
from tqdm import tqdm
import numpy as np
from argparse import ArgumentParser
from pprint import pformat
import pandas as pd
import logging
import librosa 

# Pytorch
import torch
from torch.utils import data
from torch.autograd import Variable
from torch.nn.modules.loss import _WeightedLoss
from torch.utils.data import Dataset, DataLoader
from efficientnet_pytorch import EfficientNet

# Our Package
# from dataset import SoundDataset, get_melspec

# API Package
import flask
import json 


logger = logging.getLogger(__file__)

app = flask.Flask("predict-server")
def run_predict_server():
    app.run(host='0.0.0.0', port=8080, debug=True)

# 製作頻譜圖
def get_melspec(file_name):
    y, sr = librosa.load(file_name, sr = 8000)
    if len(y)<sr*5:
        print(prefix,file_name)
        l = sr*5-len(y)
        y = np.hstack([y,np.zeros(l)])
    mel_spec = librosa.power_to_db(librosa.feature.melspectrogram(y, sr=sr, n_mels=128, n_fft=1024, hop_length=256))
    return mel_spec
    
# def load_model(): 
#     prefix = '/opt/ml/'
#     model_name = "model_auc_test.pth"
#     model_path = os.path.join(prefix, 'model', model_name)  
#     model = EfficientNet.from_pretrained('efficientnet-b0', num_classes=10,in_channels=3).to(device)
#     model.load_state_dict(torch.load(model_path))
#     model.eval()
#     model = model.to(device)
    
#     return model 
    
def load_model(): 
    prefix = '/opt/ml/'
    model_list = []
    # for metric in ['loss','auc','acc']:
    for metric in ['auc']:
        for fold in [0,1,2,3,4]:
            model_name = f'model_{metric}_fold_{fold}.ckpt'
            model_path = os.path.join(prefix, 'model/eff_class10_0714_3', model_name)  
            model = EfficientNet.from_pretrained('efficientnet-b0', num_classes=10,in_channels=3).to(device)
            model.load_state_dict(torch.load(model_path))
            model.eval()
            model = model.to(device)
            model_list.append(model)
    return model_list 
    

@app.route('/ping', methods=['GET'])
def ping():
    """Determine if the container is working and healthy. In this sample container, we declare
    it healthy if we can load the model successfully."""
    status = 200 
    return flask.Response(response='\n', status=status, mimetype='application/json')

import uuid 
@app.route('/invocations', methods=['POST'])
def predict():
    data = flask.request.data
        
    path = '/{}.wav'.format(uuid.uuid4())
    tmpFile = open(path, 'wb')
    tmpFile.write(data)
    tmpFile.close()
    
    # Loading data
    mel_spec = get_melspec(path)
    mel_concat = np.stack((mel_spec, mel_spec, mel_spec))/255

    # Reshape data to tensor
    img_test = np.reshape(mel_concat, (1, mel_concat.shape[0], mel_concat.shape[1], mel_concat.shape[2]))
    wav = torch.tensor(img_test) 
    
    # model input format
    model_input = wav.to(device, dtype=torch.float)
    outputs_list = []
    model_num = 5
    for i in range(model_num):
        outputs = model_list[i](model_input)
        outputs_list.append(torch.nn.functional.softmax(outputs, dim=1))
    outputs = torch.stack(outputs_list, dim=0).sum(0)/model_num
    _, preds = torch.max(outputs, 1)
    

    # Request
    pred_label = preds.cpu().detach().numpy()
    outputs = outputs.cpu().detach().numpy()
    print(pred_label, outputs)
    results = {}
    results['label']=int(pred_label[0])
    results['probability']=outputs[0].tolist()
    return flask.Response(response=json.dumps(results), status=200, mimetype='text/json')

# params = ParameterSetting(sr=8000,nfft=200, hop=80, mel=64, normalize=None, preload=False)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_list = load_model()
run_predict_server() 
