import os
import librosa
import gdown
import zipfile
import numpy as np
from panns_inference import SoundEventDetection, labels, AudioTagging

data_dir = './example_audio'
at = AudioTagging(checkpoint_path=None, device='cpu')

# 下载音频文件
def download_audio_data():
    url = 'https://drive.google.com/uc?id=1bKu21JWBfcZBuEuzFEvPoAX6PmRrgnUp'
    gdown.download(url)
    with zipfile.ZipFile('example_audio.zip', 'r') as zip_ref:
        zip_ref.extractall(data_dir)

# 构建音频特征向量并存储
def embed_and_save(path, at):
    
    audio, _ = librosa.core.load(path, sr=32000, mono=True)
    audio = audio[None, :]
    try:
        _, embedding = at.inference(audio)
        embedding = embedding/np.linalg.norm(embedding)
        embedding = embedding.tolist()[0]
        mr = collection.insert([[embedding]])
        ids = mr.primary_keys
        collection.load()
        red.set(str(ids[0]), path)
    except Exception as e:
        print("failed: " + path + "; error {}".format(e))

# 遍历与操作
print("Starting Insert")
download_audio_data()
for subdir, dirs, files in os.walk(data_dir):
    for file in files:
        path = os.path.join(subdir, file)
        embed_and_save(path, at)
print("Insert Done")

# 对待检索音频批量抽取特征，返回embedding
def get_embed(paths, at):
    embedding_list = []
    for x in paths:
        audio, _ = librosa.core.load(x, sr=32000, mono=True)
        audio = audio[None, :]
        try:
            _, embedding = at.inference(audio)
            embedding = embedding/np.linalg.norm(embedding)
            embedding_list.append(embedding)
        except:
            print("Embedding Failed: " + x)
    return np.array(embedding_list, dtype=np.float32).squeeze()

random_ids = [int(red.randomkey()) for x in range(2)]
search_clips = [x.decode("utf-8") for x in red.mget(random_ids)]
embeddings = get_embed(search_clips, at)
print(embeddings.shape)



import IPython.display as ipd
# 输出结果
def show_results(query, results, distances):
    print("Query: ")
    ipd.display(ipd.Audio(query))
    print("Results: ")
    for x in range(len(results)):
        print("Distance: " + str(distances[x]))
        ipd.display(ipd.Audio(results[x]))
    print("-"*50)

embeddings_list = embeddings.tolist()

# 检索参数配置
search_params = {"metric_type": "L2", "params": {"nprobe": 16}}

# 使用milvus进行向量检索并返回结果
try:
    start = time.time()
    results = collection.search(embeddings_list, anns_field="embedding", param=search_params, limit=3)
    end = time.time() - start
    print("Search took a total of: ", end)
    for x in range(len(results)):
        query_file = search_clips[x]
        result_files = [red.get(y.id).decode('utf-8') for y in results[x]]
        distances = [y.distance for y in results[x]]
        show_results(query_file, result_files, distances)
except Exception as e:
    print("Failed to search vectors in Milvus: {}".format(e))