from flask import Flask, jsonify, request
from librosa.feature import mfcc
import numpy as np
from scipy.io.wavfile import read
from tsai.all import *
from werkzeug.utils import secure_filename

app = Flask(__name__)
UPLOAD_FOLDER = "./"
app.config["UPLOAD_FOLDER"] = UPLOAD_FOLDER

def preprocess(audio):
    sr, a = audio
    m = mfcc(y=a.astype(float), sr=sr)
    l = []
    for a in m:
        l.append([])
        flag = True
        for b in a:
            if abs(b) < 1e-3 and flag:
                continue
            else:
                flag = False
                l[len(l) - 1].append(b)
    listy = []
    for a in l.reverse():
        listy.insert(0, [])
        for b in a.reverse():
            if abs(b) < 1e-3 and flag:
                continue
            else:
                flag = False
                listy[0].insert(0, b)
    temp = []
    for a in listy:
        if len(a) < 236:
            a.extend([0 for i in range(236 - len(a))])
            temp.append(a)
        elif len(a) > 236:
            c = (len(a) - 236) / 2
            temp.append(a[c : c + 236])
        else:
            temp.append(a)
    m = np.array(temp)
    return m
    
@app.route("/tone-vowel", methods=["POST"])
def tone_vowel_pred():
    if "blob" in request.files:
        file = request.files["blob"]
        filename = secure_filename(file.filename)
        filepath = os.path.join(app.config["UPLOAD_FOLDER"], filename)
        file.save(filepath)
        audio = read(filepath)
        m = preprocess(audio)
        t = tone.predict(m).item(0)
        v = vowel.predict(m).item(0)
        return jsonify({"vowel": v, "tone": t})

if __name__ == "__main__":
    tone = load_minirocket("minirocket-tone-8981")
    vowel = load_minirocket("minirocket-vowel-9814")
    app.run(debug=True)