import subprocess
import faiss
import os
import numpy as np
from random import shuffle
from sklearn.cluster import MiniBatchKMeans
import sys

os.setpgrp()

print(f"************ Trainer starts ************")

trainset_dir = sys.argv[1]
sr = 48000
sr_k = "48k"
exp_root = sys.argv[3]
exp_name = sys.argv[2]
exp_dir = f"{exp_root}/logs/{exp_name}"
is_half = False

world_size = 1
rank = 0
gpu_id = 0
batch_size = int(os.getenv("RVC_BATCH_SIZE", "12"))

pretrained_G = "assets/pretrained_v2/f0G48k.pth"
pretrained_D = "assets/pretrained_v2/f0D48k.pth"

os.makedirs(exp_dir, exist_ok=True)
os.makedirs(f"{exp_root}/assets/weights", exist_ok=True)

gt_wavs_dir = f"{exp_dir}/0_gt_wavs"
f0_dir = f"{exp_dir}/2a_f0"
f0nsf_dir = f"{exp_dir}/2b-f0nsf"
feature_dir = f"{exp_dir}/3_feature768"

subprocess.run(
    [
        "python",
        "infer/modules/train/preprocess.py",
        trainset_dir,
        f"{sr}",
        "1",
        exp_dir,
        "True",
        f"{3.0 if is_half else 3.7}",
    ],
    check=True,
)

subprocess.run(
    [
        "python",
        "infer/modules/train/extract/extract_f0_rmvpe.py",
        f"{world_size}",
        f"{rank}",
        f"{gpu_id}",
        exp_dir,
        f"{is_half}",
    ],
    check=True,
)

subprocess.run(
    [
        "python",
        "infer/modules/train/extract_feature_print.py",
        "cuda",
        f"{world_size}",
        f"{rank}",
        f"{gpu_id}",
        exp_dir,
        "v2",
    ],
    check=True,
)

train_names = set.intersection(
    *(
        set(name.split(".")[0] for name in os.listdir(d))
        for d in [gt_wavs_dir, feature_dir, f0_dir, f0nsf_dir]
    )
)

train_set = [
    f"{gt_wavs_dir}/{name}.wav|{feature_dir}/{name}.npy|{f0_dir}/{name}.wav.npy|{f0nsf_dir}/{name}.wav.npy|0"
    for name in train_names
]

for _ in range(2):
    train_set.append(
        f"logs/mute/0_gt_wavs/mute{sr_k}.wav|logs/mute/3_feature768/mute.npy|logs/mute/2a_f0/mute.wav.npy|logs/mute/2b-f0nsf/mute.wav.npy|0"
    )

# train for about 10,000 clip-epochs (max of 200 epochs)
total_epoch = min(200, 10000 // len(train_set) + 1)
save_epoch = total_epoch

print(f"************ Training for {total_epoch} epochs ************")

shuffle(train_set)

with open(f"{exp_dir}/filelist.txt", "w") as f:
    f.write("\n".join(train_set))

with open(f"configs/v2/{sr_k}.json", "r") as config_in, open(
    f"{exp_dir}/config.json", "w"
) as config_out:
    config_out.write(config_in.read())

subprocess.run(
    [
        "python",
        "infer/modules/train/train.py",
        "-r",
        exp_root,
        "-e",
        exp_name,
        "-sr",
        f"{sr_k}",
        "-f0",
        "1",
        "-bs",
        f"{batch_size}",
        "-g",
        "1",
        "-te",
        f"{total_epoch}",
        "-se",
        f"{save_epoch}",
        "-pg",
        pretrained_G,
        "-pd",
        pretrained_D,
        "-l",
        "1",
        "-c",
        "1",  # cache entire dataset in gpu memory
        "-sw",
        "0",
        "-v",
        "v2",
    ],
    check=True,
)

big_npy = np.concatenate(
    [np.load(f"{feature_dir}/{f}") for f in os.listdir(feature_dir)], 0
)
big_npy_idx = np.arange(big_npy.shape[0])
np.random.shuffle(big_npy_idx)
big_npy = big_npy[big_npy_idx]

if big_npy.shape[0] > 2e5:
    big_npy = (
        MiniBatchKMeans(
            n_clusters=10000,
            verbose=True,
            batch_size=256,
            compute_labels=False,
            init="random",
        )
        .fit(big_npy)
        .cluster_centers_
    )

n_ivf = min(int(16 * np.sqrt(big_npy.shape[0])), big_npy.shape[0] // 39)
index = faiss.index_factory(768, f"IVF{n_ivf},Flat")

index_ivf = faiss.extract_index_ivf(index)
index_ivf.nprobe = 1
index.train(big_npy)
faiss.write_index(
    index,
    f"{exp_dir}/trained_IVF{n_ivf}_Flat_nprobe_{index_ivf.nprobe}_{exp_name}_v2.index",
)

batch_size_add = 8192
for i in range(0, big_npy.shape[0], batch_size_add):
    index.add(big_npy[i : i + batch_size_add])
faiss.write_index(
    index,
    f"{exp_dir}/added_IVF{n_ivf}_Flat_nprobe_{index_ivf.nprobe}_{exp_name}_v2.index",
)