import torch, numpy as np, librosa
from transformers import ASTForAudioClassification, ASTFeatureExtractor

SR=16000; SEG_LEN=SR*14; HOP_LEN=SR*7

def rms_normalize(x):
    return x*(0.15/(np.sqrt((x**2).mean())+1e-8))

def get_segments(y):
    segs=[]
    for s in range(0,len(y)-SEG_LEN+1,HOP_LEN):
        segs.append(y[s:s+SEG_LEN])
    if not segs: segs=[y[:SEG_LEN]]
    return segs

class Model:
    def __init__(self, path):
        self.fe = ASTFeatureExtractor.from_pretrained(path)
        self.model = ASTForAudioClassification.from_pretrained(path)
        self.model.eval()

    def __call__(self, fp):
        y,_=librosa.load(fp,sr=SR)
        if len(y)<SEG_LEN:
            y=np.tile(y,int(np.ceil(SEG_LEN/len(y))))
        probs=[]
        for seg in get_segments(y):
            seg=rms_normalize(seg)
            x=self.fe(seg,sampling_rate=SR,return_tensors="pt")["input_values"]
            p=torch.softmax(self.model(x).logits,dim=1).detach().numpy()[0]
            probs.append(p)
        return np.median(probs,0)