import os

os.environ["CUDA_VISIBLE_DEVICES"] = ""
import random
import json
import numpy as np
import tqdm
import torch
import funcy
import time
import gc
from scipy.io import wavfile
import tempfile
import collections
from collections import defaultdict
from joblib import Parallel, delayed

from suno_utils.utils.s3 import _apply_mp
from suno_utils.audio import Audio
from suno_utils.tasks.data_loader import load_audio_mp
from suno_utils.utils.text import write_jsonl, read_jsonl, write_json, read_json
from suno_utils.utils.s3 import read_from_s3, check_s3_file_exists, open_from_s3
from suno_utils.audio.conversion import convert_audio_files

SAMPLE_RATE = 24_000
EMBEDDING_RATE = 25
N_CODEBOOKS = 8

IN_DATA_DIR = "/app/suno/data/mert_25hz_speech_short"
IN_AUDIO_DIR = os.path.join(IN_DATA_DIR, "audio")
IN_TSV_DIR = os.path.join(IN_DATA_DIR, "audio_tsv")
IN_LABEL_DIR = os.path.join(IN_DATA_DIR, "label")

OUT_DATA_DIR = "/app/suno/data/mert_25hz_short"
OUT_AUDIO_DIR = os.path.join(OUT_DATA_DIR, "audio")
OUT_TSV_DIR = os.path.join(OUT_DATA_DIR, "audio_tsv")
OUT_LABEL_DIR = os.path.join(OUT_DATA_DIR, "label")
OUT_TEMP_DIR = os.path.join(OUT_DATA_DIR, "temp")

os.makedirs(OUT_DATA_DIR, exist_ok=True)
os.makedirs(OUT_AUDIO_DIR, exist_ok=True)
os.makedirs(OUT_TSV_DIR, exist_ok=True)
os.makedirs(OUT_LABEL_DIR, exist_ok=True)
os.makedirs(OUT_TEMP_DIR, exist_ok=True)

tsv_info = []
with open(os.path.join(IN_TSV_DIR, "train.tsv"), "r") as f:
    for line in f.read().strip().split("\n"):
        if len(line.strip()) == 0:
            continue
        tsv_info.append(line.strip().split("\t"))

tsv_data_path = tsv_info[0]
tsv_data = tsv_info[1:]

from tqdm.contrib.concurrent import process_map, thread_map
import shutil

jobs = []
for data in tqdm.tqdm(tsv_data):
    jobs.append(( os.path.join(IN_AUDIO_DIR, data[0]), os.path.join(OUT_AUDIO_DIR, data[0])))
def my_copy(job):
    shutil.copyfile(job[0], job[1])

thread_map(my_copy, jobs, max_workers=10, chunksize=1)
print("DONE")