import os
import torch
import numpy as np
from torchvision.transforms import Compose, Resize, Normalize
from huggingface_hub import snapshot_download

MODEL_DIR = snapshot_download(
    "ccmusic-database/Guzheng_Tech99",
    cache_dir="./__pycache__",
)


def toCUDA(x):
    if hasattr(x, "cuda"):
        if torch.cuda.is_available():
            return x.cuda()

    return x


def find_files(folder_path=f"{MODEL_DIR}/examples", ext=".flac"):
    audio_files = []
    for root, _, files in os.walk(folder_path):
        for file in files:
            if file.endswith(ext):
                file_path = os.path.join(root, file)
                audio_files.append(file_path)

    return audio_files


def get_modelist(model_dir=MODEL_DIR, assign_model=""):
    pt_files = []
    for _, _, files in os.walk(model_dir):
        for file in files:
            if file.endswith(".pt"):
                model = os.path.basename(file)[:-3]
                if assign_model and assign_model.lower() in model:
                    pt_files.insert(0, model)
                else:
                    pt_files.append(model)

    return pt_files


def embed(input: list, img_size: int):
    compose = Compose(
        [
            Resize([img_size, img_size]),
            Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
        ]
    )
    inputs = []
    for x in input:
        x = np.array(x).transpose(2, 0, 1)
        x = torch.from_numpy(x).repeat(3, 1, 1)
        inputs.append(compose(x).float())

    return toCUDA(torch.tensor(np.array(inputs)))