import os
import gradio as gr
from pytube import YouTube
from pydub import AudioSegment
import numpy as np
import faiss
from sklearn.cluster import MiniBatchKMeans
import traceback
from random import shuffle
import json
import pathlib
from subprocess import Popen, PIPE, STDOUT

# Define the function for training
def click_train(
    exp_dir1,
    sr2,
    if_f0_3,
    spk_id5,
    save_epoch10,
    total_epoch11,
    batch_size12,
    if_save_latest13,
    pretrained_G14,
    pretrained_D15,
    gpus16,
    if_cache_gpu17,
    if_save_every_weights18,
    version19,
):
    now_dir = os.getcwd()
    exp_dir = f"{now_dir}/logs/{exp_dir1}"
    os.makedirs(exp_dir, exist_ok=True)
    gt_wavs_dir = f"{exp_dir}/0_gt_wavs"
    feature_dir = (
        f"{exp_dir}/3_feature256" if version19 == "v1" else f"{exp_dir}/3_feature768"
    )
    
    if if_f0_3:
        f0_dir = f"{exp_dir}/2a_f0"
        f0nsf_dir = f"{exp_dir}/2b-f0nsf"
        names = (
            set([name.split(".")[0] for name in os.listdir(gt_wavs_dir)])
            & set([name.split(".")[0] for name in os.listdir(feature_dir)])
            & set([name.split(".")[0] for name in os.listdir(f0_dir)])
            & set([name.split(".")[0] for name in os.listdir(f0nsf_dir)])
        )
    else:
        names = set([name.split(".")[0] for name in os.listdir(gt_wavs_dir)]) & set(
            [name.split(".")[0] for name in os.listdir(feature_dir)]
        )
    
    opt = []
    for name in names:
        if if_f0_3:
            opt.append(
                f"{gt_wavs_dir.replace('\\', '\\\\')}/{name}.wav|{feature_dir.replace('\\', '\\\\')}/{name}.npy|{f0_dir.replace('\\', '\\\\')}/{name}.wav.npy|{f0nsf_dir.replace('\\', '\\\\')}/{name}.wav.npy|{spk_id5}"
            )
        else:
            opt.append(
                f"{gt_wavs_dir.replace('\\', '\\\\')}/{name}.wav|{feature_dir.replace('\\', '\\\\')}/{name}.npy|{spk_id5}"
            )
    
    fea_dim = 256 if version19 == "v1" else 768
    if if_f0_3:
        for _ in range(2):
            opt.append(
                f"{now_dir}/logs/mute/0_gt_wavs/mute{sr2}.wav|{now_dir}/logs/mute/3_feature{fea_dim}/mute.npy|{now_dir}/logs/mute/2a_f0/mute.wav.npy|{now_dir}/logs/mute/2b-f0nsf/mute.wav.npy|{spk_id5}"
            )
    else:
        for _ in range(2):
            opt.append(
                f"{now_dir}/logs/mute/0_gt_wavs/mute{sr2}.wav|{now_dir}/logs/mute/3_feature{fea_dim}/mute.npy|{spk_id5}"
            )
    
    shuffle(opt)
    with open(f"{exp_dir}/filelist.txt", "w") as f:
        f.write("\n".join(opt))

    print("Write filelist done")
    print("Use gpus:", str(gpus16))
    if pretrained_G14 == "":
        print("No pretrained Generator")
    if pretrained_D15 == "":
        print("No pretrained Discriminator")
    
    if version19 == "v1" or sr2 == "40k":
        config_path = f"configs/v1/{sr2}.json"
    else:
        config_path = f"configs/v2/{sr2}.json"
    
    config_save_path = os.path.join(exp_dir, "config.json")
    if not pathlib.Path(config_save_path).exists():
        with open(config_save_path, "w", encoding="utf-8") as f:
            with open(config_path, "r") as config_file:
                config_data = json.load(config_file)
                json.dump(
                    config_data,
                    f,
                    ensure_ascii=False,
                    indent=4,
                    sort_keys=True,
                )
            f.write("\n")

    cmd = (
        f'python infer/modules/train/train.py -e "{exp_dir1}" -sr {sr2} -f0 {1 if if_f0_3 else 0} -bs {batch_size12} -g {gpus16} -te {total_epoch11} -se {save_epoch10} {"-pg " + pretrained_G14 if pretrained_G14 != "" else ""} {"-pd " + pretrained_D15 if pretrained_D15 != "" else ""} -l {1 if if_save_latest13 else 0} -c {1 if if_cache_gpu17 else 0} -sw {1 if if_save_every_weights18 else 0} -v {version19}'
    )

    p = Popen(cmd, shell=True, cwd=now_dir, stdout=PIPE, stderr=STDOUT, bufsize=1, universal_newlines=True)
    
    for line in p.stdout:
        print(line.strip())
    
    p.wait()
    return "After the training is completed, you can view the console training log or train.log under the experiment folder"



def calculate_audio_duration(file_path):
    duration_seconds = len(AudioSegment.from_file(file_path)) / 1000.0
    return duration_seconds

def youtube_to_wav(url, dataset_folder):
    try:
        yt = YouTube(url).streams.get_audio_only().download(output_path=dataset_folder)
        mp4_path = os.path.join(dataset_folder, 'audio.mp4')
        wav_path = os.path.join(dataset_folder, 'audio.wav')
        os.rename(yt, mp4_path)
        os.system(f'ffmpeg -i {mp4_path} -acodec pcm_s16le -ar 44100 {wav_path}')
        os.remove(mp4_path)
        return f'Audio downloaded and converted to WAV: {wav_path}'
    except Exception as e:
        return f"Error: {e}"

def create_training_files(model_name, dataset_folder, youtube_link):
    if youtube_link:
        youtube_to_wav(youtube_link, dataset_folder)

    if not os.listdir(dataset_folder):
        return "Your dataset folder is empty."

    os.makedirs(f'./logs/{model_name}', exist_ok=True)
    
    os.system(f'python infer/modules/train/preprocess.py {dataset_folder} 32000 2 ./logs/{model_name} False 3.0 > /dev/null 2>&1')
    
    with open(f'./logs/{model_name}/preprocess.log', 'r') as f:
        if 'end preprocess' in f.read():
            return "Preprocessing Success"
        else:
            return "Error preprocessing data... Make sure your dataset folder is correct."

def extract_features(model_name, f0method):
    os.system(f'python infer/modules/train/extract/extract_f0_rmvpe.py 1 0 0 ./logs/{model_name} True' if f0method == "rmvpe_gpu" else 
              f'python infer/modules/train/extract/extract_f0_print.py ./logs/{model_name} 2 {f0method}')
    os.system(f'python infer/modules/train/extract_feature_print.py cuda:0 1 0 ./logs/{model_name} v2 True')
    
    with open(f'./logs/{model_name}/extract_f0_feature.log', 'r') as f:
        if 'all-feature-done' in f.read():
            return "Feature Extraction Success"
        else:
            return "Error in feature extraction... Make sure your data was preprocessed."

def train_index(exp_dir1, version19):
    exp_dir = f"logs/{exp_dir1}"
    os.makedirs(exp_dir, exist_ok=True)
    feature_dir = f"{exp_dir}/3_feature256" if version19 == "v1" else f"{exp_dir}/3_feature768"
    if not os.path.exists(feature_dir):
        return "Please perform feature extraction first!"

    listdir_res = list(os.listdir(feature_dir))
    if len(listdir_res) == 0:
        return "Please perform feature extraction first!"

    infos = []
    npys = []
    for name in sorted(listdir_res):
        phone = np.load(f"{feature_dir}/{name}")
        npys.append(phone)
    big_npy = np.concatenate(npys, 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:
        infos.append(f"Trying k-means with {big_npy.shape[0]} to 10k centers.")
        try:
            big_npy = MiniBatchKMeans(
                n_clusters=10000,
                verbose=True,
                batch_size=256,
                compute_labels=False,
                init="random",
            ).fit(big_npy).cluster_centers_
        except:
            info = traceback.format_exc()
            infos.append(info)
            return "\n".join(infos)

    np.save(f"{exp_dir}/total_fea.npy", big_npy)
    n_ivf = min(int(16 * np.sqrt(big_npy.shape[0])), big_npy.shape[0] // 39)
    infos.append(f"{big_npy.shape},{n_ivf}")

    index = faiss.index_factory(256 if version19 == "v1" else 768, f"IVF{n_ivf},Flat")
    infos.append("Training index")
    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_dir1}_{version19}.index")

    infos.append("Adding to 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_dir1}_{version19}.index")

    infos.append(f"Successfully built index: added_IVF{n_ivf}_Flat_nprobe_{index_ivf.nprobe}_{exp_dir1}_{version19}.index")
    return "\n".join(infos)

with gr.Blocks() as demo:
    with gr.Tab("Training"):
        with gr.Tab("CREATE TRANING FILES - This will process the data, extract the features and create your index file for you!"):
            with gr.Row():
                model_name = gr.Textbox(label="Model Name", value="My-Voice")
                dataset_folder = gr.Textbox(label="Dataset Folder", value="/content/dataset")
            youtube_link = gr.Textbox(label="YouTube Link (optional)")
            with gr.Row():
                start_button = gr.Button("Create Training Files")
                f0method = gr.Dropdown(["pm", "harvest", "rmvpe", "rmvpe_gpu"], label="F0 Method", value="rmvpe_gpu")
            extract_button = gr.Button("Extract Features")
            train_button = gr.Button("Train Index")

        output = gr.Textbox(label="Output")

        start_button.click(create_training_files, inputs=[model_name, dataset_folder, youtube_link], outputs=output)
        extract_button.click(extract_features, inputs=[model_name, f0method], outputs=output)
        train_button.click(train_index, inputs=[model_name, "v2"], outputs=output)
        with gr.Tab("train"):
            exp_dir1 = gr.Textbox(label="Experiment Directory", value="mymodel")
            sr2 = gr.Dropdown(choices=["32k", "40k", "48k"], label="Sample Rate", value="32k")
            if_f0_3 = gr.Checkbox(label="Use F0", value=True)
            spk_id5 = gr.Number(label="Speaker ID", value=0)
            save_epoch10 = gr.Slider(label="Save Frequency", minimum=5, maximum=50, step=5, value=25)
            total_epoch11 = gr.Slider(label="Total Epochs", minimum=10, maximum=2000, step=10, value=500)
            batch_size12 = gr.Slider(label="Batch Size", minimum=1, maximum=20, step=1, value=8)
            if_save_latest13 = gr.Checkbox(label="Save Latest", value=True)
            pretrained_G14 = gr.Textbox(label="Pretrained Generator File", value="/content/pre/assets/pretrained_v2/f0Ov2Super32kG.pth")
            pretrained_D15 = gr.Textbox(label="Pretrained Discriminator File", value="/content/pre/assets/pretrained_v2/f0Ov2Super32kD.pth")
            gpus16 = gr.Number(label="GPUs", value=0)
            if_cache_gpu17 = gr.Checkbox(label="Cache GPU", value=False)
            if_save_every_weights18 = gr.Checkbox(label="Save Every Weights", value=True)
            version19 = gr.Textbox(label="Version", value="v2")
            training_log = gr.Textbox(label="Training Log", interactive=False)
            train_button = gr.Button("Start Training")
            
            train_button.click(
                fn=click_train,
                inputs=[
                    exp_dir1, sr2, if_f0_3, spk_id5, save_epoch10, total_epoch11, batch_size12,
                    if_save_latest13, pretrained_G14, pretrained_D15, gpus16, if_cache_gpu17,
                    if_save_every_weights18, version19
                ],
                outputs=training_log
            )



demo.launch()






# beta state ......