import datetime
import glob
import json
import logging
import os
import pathlib
import shutil
import signal
import sys
from collections import deque
from random import randint, shuffle
from time import time as ttime

import numpy as np
from tqdm import tqdm

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

# Zluda hijack
import ultimate_rvc.rvc.lib.zluda
from ultimate_rvc.common import TRAINING_MODELS_DIR
from ultimate_rvc.rvc.common import RVC_TRAINING_MODELS_DIR
from ultimate_rvc.rvc.lib.algorithm import commons
from ultimate_rvc.rvc.train.losses import (
    discriminator_loss,
    feature_loss,
    generator_loss,
    kl_loss,
)
from ultimate_rvc.rvc.train.mel_processing import (
    MultiScaleMelSpectrogramLoss,
    mel_spectrogram_torch,
    spec_to_mel_torch,
)
from ultimate_rvc.rvc.train.process.extract_model import extract_model
from ultimate_rvc.rvc.train.utils import (
    HParams,
    latest_checkpoint_path,
    load_checkpoint,
    load_wav_to_torch,
    plot_spectrogram_to_numpy,
    remove_sox_libmso6_from_ld_preload,
    save_checkpoint,
    summarize,
)

logging.getLogger("torch").setLevel(logging.ERROR)
logger = logging.getLogger(__name__)

torch.backends.cudnn.deterministic = False
torch.backends.cudnn.benchmark = True
torch.multiprocessing.set_start_method("spawn", force=True)
os.environ["USE_LIBUV"] = "0" if sys.platform == "win32" else "1"

randomized = True
optimizer = "AdamW"  # "RAdam"
d_lr_coeff = 1.0
g_lr_coeff = 1.0
global_step = 0
lowest_g_value = {"value": float("inf"), "epoch": 0}
lowest_d_value = {"value": float("inf"), "epoch": 0}
consecutive_increases_gen = 0
consecutive_increases_disc = 0

avg_losses = {
    "grad_d_50": deque(maxlen=50),
    "grad_g_50": deque(maxlen=50),
    "disc_loss_50": deque(maxlen=50),
    "fm_loss_50": deque(maxlen=50),
    "kl_loss_50": deque(maxlen=50),
    "mel_loss_50": deque(maxlen=50),
    "gen_loss_50": deque(maxlen=50),
}


class EpochRecorder:
    """
    Records the time elapsed per epoch.
    """

    def __init__(self):
        self.last_time = ttime()

    def record(self):
        """
        Records the elapsed time and returns a formatted string.
        """
        now_time = ttime()
        elapsed_time = now_time - self.last_time
        self.last_time = now_time
        elapsed_time = round(elapsed_time, 1)
        elapsed_time_str = str(datetime.timedelta(seconds=int(elapsed_time)))
        current_time = datetime.datetime.now().strftime("%H:%M:%S")
        return f"time={current_time} | speed={elapsed_time_str}"


def verify_checkpoint_shapes(checkpoint_path: str, model: torch.nn.Module) -> None:
    """
    Verify that the checkpoint and model have the same
    architecture.
    """
    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
    checkpoint_state_dict = checkpoint["model"]
    try:
        if hasattr(model, "module"):
            model_state_dict = model.module.load_state_dict(checkpoint_state_dict)
        else:
            model_state_dict = model.load_state_dict(checkpoint_state_dict)
    except RuntimeError:
        logger.error(  # noqa: TRY400
            "The parameters of the pretrain model such as the sample rate or"
            " architecture do not match the selected model.",
        )
        sys.exit(1)
    else:
        del checkpoint
        del checkpoint_state_dict
        del model_state_dict


def main(
    model_name: str,
    sample_rate: int,
    vocoder: str,
    total_epoch: int,
    batch_size: int,
    save_every_epoch: int,
    save_only_latest: bool,
    save_every_weights: bool,
    pretrain_g: str,
    pretrain_d: str,
    overtraining_detector: bool,
    overtraining_threshold: int,
    cleanup: bool,
    cache_data_in_gpu: bool,
    checkpointing: bool,
    device_type: str,
    gpus: set[int] | None,
) -> None:
    """
    Start the training process.

    Raises:
        RuntimeError: If the sample rate of the pretrained model does not match the dataset audio sample rate.

    """
    remove_sox_libmso6_from_ld_preload()
    experiment_dir = os.path.join(TRAINING_MODELS_DIR, model_name)
    config_save_path = os.path.join(experiment_dir, "config.json")

    # Use a Manager to create a shared list
    manager = mp.Manager()
    global_gen_loss = manager.list([0] * total_epoch)
    global_disc_loss = manager.list([0] * total_epoch)

    with pathlib.Path(config_save_path).open() as f:
        config = json.load(f)
    config = HParams(**config)
    config.data.training_files = os.path.join(experiment_dir, "filelist.txt")

    # Set up distributed training environment for master node.
    # master node is localhost because we are running on a single local
    # machine. master port is randomly selected
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = str(randint(20000, 55555))
    logger.info("MASTER_PORT: %s", os.environ["MASTER_PORT"])
    # Check sample rate
    wavs = glob.glob(
        os.path.join(experiment_dir, "sliced_audios", "*.wav"),
    )
    if wavs:
        _, sr = load_wav_to_torch(wavs[0])
        if sr != sample_rate:
            error_message = (
                f"Error: Pretrained model sample rate ({sample_rate} Hz) does not match"
                f" dataset audio sample rate ({sr} Hz)."
            )
            raise RuntimeError(error_message)
    else:
        logger.warning("No wav file found.")

    device = torch.device(device_type)
    gpus = gpus or {0}
    n_gpus = len(gpus)

    if device.type == "cpu":
        logger.warning("Training with CPU, this will take a long time.")

    def start() -> None:
        """Start the training process with multi-GPU support or CPU."""
        children = []
        pid_data = {"process_pids": []}
        with pathlib.Path(config_save_path).open() as pid_file:
            try:
                existing_data = json.load(pid_file)
                pid_data.update(existing_data)
            except json.JSONDecodeError:
                pass
        with pathlib.Path(config_save_path).open("w") as pid_file:
            for rank, device_id in enumerate(gpus):
                subproc = mp.Process(
                    target=run,
                    args=(
                        rank,
                        n_gpus,
                        experiment_dir,
                        pretrain_g,
                        pretrain_d,
                        total_epoch,
                        save_every_weights,
                        config,
                        device,
                        device_id,
                        model_name,
                        sample_rate,
                        vocoder,
                        batch_size,
                        save_every_epoch,
                        save_only_latest,
                        overtraining_detector,
                        overtraining_threshold,
                        checkpointing,
                        cache_data_in_gpu,
                        global_gen_loss,
                        global_disc_loss,
                    ),
                )
                children.append(subproc)
                subproc.start()
                pid_data["process_pids"].append(subproc.pid)
            json.dump(pid_data, pid_file, indent=4)
        cancel_signal = signal.SIGTERM if os.name == "nt" else -signal.SIGTERM
        error_codes = []
        for i in range(n_gpus):
            children[i].join()
            exit_code = children[i].exitcode
            if exit_code != 0:
                logger.warning(
                    "Process running on device %s exited with code %s.",
                    device_id,
                    exit_code,
                )
                if exit_code != cancel_signal:
                    error_codes.append(exit_code)
        if error_codes:
            err_msg = (
                "One or more training processes failed. See the logs or console for"
                " details."
            )
            raise RuntimeError(err_msg)

    if cleanup:
        logger.info("Removing files from the prior training attempt...")

        # Clean up unnecessary files
        for entry in os.scandir(os.path.join(TRAINING_MODELS_DIR, model_name)):
            if entry.is_file():
                _, file_extension = os.path.splitext(entry.name)
                if file_extension in {".0", ".pth", ".index"}:
                    pathlib.Path(entry.path).unlink()
            elif entry.is_dir() and entry.name == "eval":
                shutil.rmtree(entry.path)

        logger.info("Cleanup done!")
    start()


def run(
    rank,
    n_gpus,
    experiment_dir,
    pretrain_g,
    pretrain_d,
    custom_total_epoch,
    custom_save_every_weights,
    config,
    device,
    device_id,
    model_name,
    sample_rate,
    vocoder,
    batch_size,
    save_every_epoch,
    save_only_latest,
    overtraining_detector,
    overtraining_threshold,
    checkpointing,
    cache_data_in_gpu,
    global_gen_loss,
    global_disc_loss,
):
    """
    Runs the training loop on a specific GPU or CPU.

    Args:
        rank (int): The rank of the current process within the distributed training setup.
        n_gpus (int): The total number of GPUs available for training.
        experiment_dir (str): The directory where experiment logs and checkpoints will be saved.
        pretrain_g (str): Path to the pre-trained generator model.
        pretrain_d (str): Path to the pre-trained discriminator model.
        custom_total_epoch (int): The total number of epochs for training.
        custom_save_every_weights (int): The interval (in epochs) at which to save model weights.
        config (object): Configuration object containing training parameters.
        device (torch.device): The device to use for training (CPU or GPU).

    """
    global global_step, optimizer, lowest_d_value, lowest_g_value, consecutive_increases_gen, consecutive_increases_disc

    if rank == 0:
        writer_eval = SummaryWriter(log_dir=os.path.join(experiment_dir, "eval"))
    else:
        writer_eval = None

    # Initialize distributed training environment for child node.
    dist.init_process_group(
        backend="gloo" if sys.platform == "win32" or device.type != "cuda" else "nccl",
        init_method="env://",
        world_size=n_gpus if device.type == "cuda" else 1,
        rank=rank if device.type == "cuda" else 0,
    )

    torch.manual_seed(config.train.seed)

    if device.type == "cuda":
        torch.cuda.set_device(device_id)

    # Create datasets and dataloaders
    from ultimate_rvc.rvc.train.data_utils import (
        DistributedBucketSampler,
        TextAudioCollateMultiNSFsid,
        TextAudioLoaderMultiNSFsid,
    )

    train_dataset = TextAudioLoaderMultiNSFsid(config.data)
    collate_fn = TextAudioCollateMultiNSFsid()
    train_sampler = DistributedBucketSampler(
        train_dataset,
        batch_size * n_gpus,
        [50, 100, 200, 300, 400, 500, 600, 700, 800, 900],
        num_replicas=n_gpus,
        rank=rank,
        shuffle=True,
    )

    train_loader = DataLoader(
        train_dataset,
        num_workers=4,
        shuffle=False,
        pin_memory=True,
        collate_fn=collate_fn,
        batch_sampler=train_sampler,
        persistent_workers=True,
        prefetch_factor=8,
    )
    if len(train_loader) < 3:
        logger.error(
            "Not enough data in the training set. Perhaps you forgot to slice the"
            " audio files in preprocess?",
        )
        sys.exit(1)
    else:
        g_file = latest_checkpoint_path(experiment_dir, "G_*.pth")
        if g_file != None:
            logger.info("Checking saved weights...")
            g = torch.load(g_file, map_location="cpu", weights_only=False)
            if (
                optimizer == "RAdam"
                and "amsgrad" in g["optimizer"]["param_groups"][0].keys()
            ):
                optimizer = "AdamW"
                logger.info(
                    "Optimizer choice has been reverted to %s to match the saved D/G"
                    " weights.",
                    optimizer,
                )
            elif (
                optimizer == "AdamW"
                and "decoupled_weight_decay" in g["optimizer"]["param_groups"][0].keys()
            ):
                optimizer = "RAdam"
                logger.info(
                    "Optimizer choice has been reverted to %s to match the saved D/G"
                    " weights.",
                    optimizer,
                )
            del g

    # Initialize models and optimizers
    from ultimate_rvc.rvc.lib.algorithm.discriminators import MultiPeriodDiscriminator
    from ultimate_rvc.rvc.lib.algorithm.synthesizers import Synthesizer

    # NOTE checkingpointing here means whether or not activations are
    # saved during forward pass for backpropagation during backward pass

    net_g = Synthesizer(
        config.data.filter_length // 2 + 1,
        config.train.segment_size // config.data.hop_length,
        **config.model,
        use_f0=True,
        sr=sample_rate,
        vocoder=vocoder,
        checkpointing=checkpointing,
        randomized=randomized,
    )

    net_d = MultiPeriodDiscriminator(
        config.model.use_spectral_norm,
        checkpointing=checkpointing,
    )

    if device.type == "cuda":
        net_g = net_g.cuda(device_id)
        net_d = net_d.cuda(device_id)
    else:
        net_g = net_g.to(device)
        net_d = net_d.to(device)

    if optimizer == "AdamW":
        optimizer = torch.optim.AdamW
    elif optimizer == "RAdam":
        optimizer = torch.optim.RAdam

    optim_g = optimizer(
        net_g.parameters(),
        config.train.learning_rate * g_lr_coeff,
        betas=config.train.betas,
        eps=config.train.eps,
    )
    optim_d = optimizer(
        net_d.parameters(),
        config.train.learning_rate * d_lr_coeff,
        betas=config.train.betas,
        eps=config.train.eps,
    )

    fn_mel_loss = MultiScaleMelSpectrogramLoss(sample_rate=sample_rate)

    # Wrap models with DDP for multi-gpu processing
    if n_gpus > 1 and device.type == "cuda":
        net_g = DDP(net_g, device_ids=[device_id])
        net_d = DDP(net_d, device_ids=[device_id])

    # Load checkpoint if available
    try:
        logger.info("Starting training...")
        _, _, _, epoch_str, lowest_d_value, consecutive_increases_disc = (
            load_checkpoint(
                latest_checkpoint_path(experiment_dir, "D_*.pth"),
                net_d,
                optim_d,
            )
        )
        _, _, _, epoch_str, lowest_g_value, consecutive_increases_gen = load_checkpoint(
            latest_checkpoint_path(experiment_dir, "G_*.pth"),
            net_g,
            optim_g,
        )
        epoch_str += 1
        global_step = (epoch_str - 1) * len(train_loader)

    except Exception:
        epoch_str = 1
        global_step = 0
        if pretrain_g not in {"", "None"}:
            if rank == 0:
                verify_checkpoint_shapes(pretrain_g, net_g)
                logger.info("Loaded pretrained (G) '%s'", pretrain_g)
            if hasattr(net_g, "module"):
                net_g.module.load_state_dict(
                    torch.load(pretrain_g, map_location="cpu", weights_only=False)[
                        "model"
                    ],
                )
            else:
                net_g.load_state_dict(
                    torch.load(pretrain_g, map_location="cpu", weights_only=False)[
                        "model"
                    ],
                )

        if pretrain_d not in {"", "None"}:
            if rank == 0:
                verify_checkpoint_shapes(pretrain_d, net_d)
                logger.info("Loaded pretrained (D) '%s'", pretrain_d)
            if hasattr(net_d, "module"):
                net_d.module.load_state_dict(
                    torch.load(pretrain_d, map_location="cpu", weights_only=False)[
                        "model"
                    ],
                )
            else:
                net_d.load_state_dict(
                    torch.load(pretrain_d, map_location="cpu", weights_only=False)[
                        "model"
                    ],
                )

    # Initialize schedulers
    scheduler_g = torch.optim.lr_scheduler.ExponentialLR(
        optim_g,
        gamma=config.train.lr_decay,
        last_epoch=epoch_str - 2,
    )
    scheduler_d = torch.optim.lr_scheduler.ExponentialLR(
        optim_d,
        gamma=config.train.lr_decay,
        last_epoch=epoch_str - 2,
    )

    cache = []
    # get the first sample as reference for tensorboard evaluation
    # custom reference temporarily disabled
    if (
        True == False
        and pathlib.Path(
            os.path.join(RVC_TRAINING_MODELS_DIR, "reference", f"ref{sample_rate}.wav")
        ).is_file()
    ):
        phone = np.load(
            os.path.join(
                RVC_TRAINING_MODELS_DIR,
                "reference",
                f"ref{sample_rate}_feats.npy",
            ),
        )
        # expanding x2 to match pitch size
        phone = np.repeat(phone, 2, axis=0)
        phone = torch.FloatTensor(phone).unsqueeze(0).to(device)
        phone_lengths = torch.LongTensor(phone.size(0)).to(device)
        pitch = np.load(
            os.path.join(
                RVC_TRAINING_MODELS_DIR,
                "reference",
                f"ref{sample_rate}_f0c.npy",
            ),
        )
        # removed last frame to match features
        pitch = torch.LongTensor(pitch[:-1]).unsqueeze(0).to(device)
        pitchf = np.load(
            os.path.join(
                RVC_TRAINING_MODELS_DIR,
                "reference",
                f"ref{sample_rate}_f0f.npy",
            ),
        )
        # removed last frame to match features
        pitchf = torch.FloatTensor(pitchf[:-1]).unsqueeze(0).to(device)
        sid = torch.LongTensor([0]).to(device)
        reference = (
            phone,
            phone_lengths,
            pitch,
            pitchf,
            sid,
        )
    else:
        for info in train_loader:
            phone, phone_lengths, pitch, pitchf, _, _, _, _, sid = info
            if device.type == "cuda":
                reference = (
                    phone.cuda(device_id, non_blocking=True),
                    phone_lengths.cuda(device_id, non_blocking=True),
                    pitch.cuda(device_id, non_blocking=True),
                    pitchf.cuda(device_id, non_blocking=True),
                    sid.cuda(device_id, non_blocking=True),
                )
            else:
                reference = (
                    phone.to(device),
                    phone_lengths.to(device),
                    pitch.to(device),
                    pitchf.to(device),
                    sid.to(device),
                )
            break

    if epoch_str > custom_total_epoch:
        cleanup_training_processes(experiment_dir)
    for epoch in range(epoch_str, custom_total_epoch + 1):
        train_and_evaluate(
            rank,
            epoch,
            config,
            [net_g, net_d],
            [optim_g, optim_d],
            [scheduler_g, scheduler_d],
            [train_loader, None],
            [writer_eval],
            cache,
            custom_save_every_weights,
            custom_total_epoch,
            device,
            device_id,
            reference,
            fn_mel_loss,
            model_name,
            experiment_dir,
            sample_rate,
            vocoder,
            save_every_epoch,
            save_only_latest,
            overtraining_detector,
            overtraining_threshold,
            cache_data_in_gpu,
            global_gen_loss,
            global_disc_loss,
        )


def train_and_evaluate(
    rank,
    epoch,
    config,
    nets,
    optims,
    schedulers,
    loaders,
    writers,
    cache,
    custom_save_every_weights,
    custom_total_epoch,
    device,
    device_id,
    reference,
    fn_mel_loss,
    model_name,
    experiment_dir,
    sample_rate,
    vocoder,
    save_every_epoch,
    save_only_latest,
    overtraining_detector,
    overtraining_threshold,
    cache_data_in_gpu,
    global_gen_loss,
    global_disc_loss,
) -> None:
    """Train and evaluates the model for one epoch."""
    global global_step, lowest_g_value, lowest_d_value, consecutive_increases_gen, consecutive_increases_disc

    model_add = []
    checkpoint_idxs = []
    done = False

    net_g, net_d = nets
    optim_g, optim_d = optims
    scheduler_g, scheduler_d = schedulers
    skip_g_scheduler, skip_d_scheduler = False, False
    train_loader = loaders[0] if loaders is not None else None
    if writers is not None:
        writer = writers[0]

    train_loader.batch_sampler.set_epoch(epoch)

    net_g.train()
    net_d.train()

    # Data caching
    if device.type == "cuda" and cache_data_in_gpu:
        if cache == []:
            for batch_idx, info in enumerate(train_loader):
                # phone, phone_lengths, pitch, pitchf, spec, spec_lengths, wave, wave_lengths, sid
                info = [tensor.cuda(device_id, non_blocking=True) for tensor in info]
                cache.append((batch_idx, info))
        shuffle(cache)
        data_iterator = cache
    else:
        data_iterator = enumerate(train_loader)

    epoch_recorder = EpochRecorder()
    with tqdm(total=len(train_loader), leave=False) as pbar:
        for batch_idx, info in data_iterator:
            if device.type == "cuda" and not cache_data_in_gpu:
                info = [tensor.cuda(device_id, non_blocking=True) for tensor in info]
            elif device.type != "cuda":
                info = [tensor.to(device) for tensor in info]
            # else iterator is going thru a cached list with a device already assigned

            (
                phone,
                phone_lengths,
                pitch,
                pitchf,
                spec,
                spec_lengths,
                wave,
                wave_lengths,
                sid,
            ) = info

            # Forward pass
            model_output = net_g(
                phone,
                phone_lengths,
                pitch,
                pitchf,
                spec,
                spec_lengths,
                sid,
            )
            y_hat, ids_slice, x_mask, z_mask, (z, z_p, m_p, logs_p, m_q, logs_q) = (
                model_output
            )
            # slice of the original waveform to match a generate slice
            if randomized:
                wave = commons.slice_segments(
                    wave,
                    ids_slice * config.data.hop_length,
                    config.train.segment_size,
                    dim=3,
                )
            y_d_hat_r, y_d_hat_g, _, _ = net_d(wave, y_hat.detach())
            loss_disc, _, _ = discriminator_loss(y_d_hat_r, y_d_hat_g)
            # Discriminator backward and update
            global_disc_loss[epoch - 1] += loss_disc.item()
            optim_d.zero_grad()
            loss_disc.backward()
            grad_norm_d = commons.grad_norm(net_d.parameters())
            optim_d.step()

            # Generator backward and update
            _, y_d_hat_g, fmap_r, fmap_g = net_d(wave, y_hat)
            loss_mel = fn_mel_loss(wave, y_hat) * config.train.c_mel / 3.0
            loss_kl = kl_loss(z_p, logs_q, m_p, logs_p, z_mask) * config.train.c_kl
            loss_fm = feature_loss(fmap_r, fmap_g)
            loss_gen, _ = generator_loss(y_d_hat_g)
            loss_gen_all = loss_gen + loss_fm + loss_mel + loss_kl
            global_gen_loss[epoch - 1] += loss_gen_all.item()
            optim_g.zero_grad()
            loss_gen_all.backward()
            grad_norm_g = commons.grad_norm(net_g.parameters())
            optim_g.step()

            global_step += 1

            # queue for rolling losses over 50 steps
            avg_losses["grad_d_50"].append(grad_norm_d)
            avg_losses["grad_g_50"].append(grad_norm_g)
            avg_losses["disc_loss_50"].append(loss_disc.detach())
            avg_losses["fm_loss_50"].append(loss_fm.detach())
            avg_losses["kl_loss_50"].append(loss_kl.detach())
            avg_losses["mel_loss_50"].append(loss_mel.detach())
            avg_losses["gen_loss_50"].append(loss_gen_all.detach())

            if rank == 0 and global_step % 50 == 0:
                # logging rolling averages
                scalar_dict = {
                    "grad_avg_50/norm_d": (
                        sum(avg_losses["grad_d_50"]) / len(avg_losses["grad_d_50"])
                    ),
                    "grad_avg_50/norm_g": (
                        sum(avg_losses["grad_g_50"]) / len(avg_losses["grad_g_50"])
                    ),
                    "loss_avg_50/d/total": torch.mean(
                        torch.stack(list(avg_losses["disc_loss_50"])),
                    ),
                    "loss_avg_50/g/fm": torch.mean(
                        torch.stack(list(avg_losses["fm_loss_50"])),
                    ),
                    "loss_avg_50/g/kl": torch.mean(
                        torch.stack(list(avg_losses["kl_loss_50"])),
                    ),
                    "loss_avg_50/g/mel": torch.mean(
                        torch.stack(list(avg_losses["mel_loss_50"])),
                    ),
                    "loss_avg_50/g/total": torch.mean(
                        torch.stack(list(avg_losses["gen_loss_50"])),
                    ),
                }
                summarize(
                    writer=writer,
                    global_step=global_step,
                    scalars=scalar_dict,
                )

            pbar.update(1)
        # end of batch train
    # end of tqdm
    scheduler_d.step()
    scheduler_g.step()

    with torch.no_grad():
        torch.cuda.empty_cache()
    # Logging and checkpointing
    if rank == 0:
        avg_global_disc_loss = global_disc_loss[epoch - 1] / len(train_loader.dataset)
        avg_global_gen_loss = global_gen_loss[epoch - 1] / len(train_loader.dataset)

        min_delta = 0.004

        if avg_global_disc_loss < lowest_d_value["value"] - min_delta:
            lowest_d_value = {"value": avg_global_disc_loss, "epoch": epoch}
            consecutive_increases_disc = 0
        else:
            consecutive_increases_disc += 1

        if avg_global_gen_loss < lowest_g_value["value"] - min_delta:
            logger.info(
                "New best epoch %d with average generator loss %.3f and discriminator"
                " loss %.3f",
                epoch,
                avg_global_gen_loss,
                avg_global_disc_loss,
            )
            lowest_g_value = {"value": avg_global_gen_loss, "epoch": epoch}
            consecutive_increases_gen = 0
            model_add.append(
                os.path.join(experiment_dir, f"{model_name}_best.pth"),
            )
        else:
            consecutive_increases_gen += 1

        # used for tensorboard chart - all/mel
        mel = spec_to_mel_torch(
            spec,
            config.data.filter_length,
            config.data.n_mel_channels,
            config.data.sample_rate,
            config.data.mel_fmin,
            config.data.mel_fmax,
        )
        # used for tensorboard chart - slice/mel_org
        if randomized:
            y_mel = commons.slice_segments(
                mel,
                ids_slice,
                config.train.segment_size // config.data.hop_length,
                dim=3,
            )
        else:
            y_mel = mel
        # used for tensorboard chart - slice/mel_gen
        y_hat_mel = mel_spectrogram_torch(
            y_hat.float().squeeze(1),
            config.data.filter_length,
            config.data.n_mel_channels,
            config.data.sample_rate,
            config.data.hop_length,
            config.data.win_length,
            config.data.mel_fmin,
            config.data.mel_fmax,
        )

        lr = optim_g.param_groups[0]["lr"]

        scalar_dict = {
            "loss/g/total": loss_gen_all,
            "loss/d/total": loss_disc,
            "learning_rate": lr,
            "grad/norm_d": grad_norm_d,
            "grad/norm_g": grad_norm_g,
            "loss/g/fm": loss_fm,
            "loss/g/mel": loss_mel,
            "loss/g/kl": loss_kl,
        }

        image_dict = {
            "slice/mel_org": plot_spectrogram_to_numpy(y_mel[0].data.cpu().numpy()),
            "slice/mel_gen": plot_spectrogram_to_numpy(y_hat_mel[0].data.cpu().numpy()),
            "all/mel": plot_spectrogram_to_numpy(mel[0].data.cpu().numpy()),
        }
        overtrain_info = ""
        # Print training progress
        lowest_g_value_rounded = float(lowest_g_value["value"])
        lowest_g_value_rounded = round(lowest_g_value_rounded, 3)

        record = f"{model_name} | epoch={epoch} | {epoch_recorder.record()}"
        record += (
            f" | best avg-gen-loss={lowest_g_value_rounded:.3f} (epoch"
            f" {lowest_g_value['epoch']})"
        )
        # Check overtraining
        if overtraining_detector:
            overtrain_info = (
                f"Average epoch generator loss {avg_global_gen_loss:.3f} and"
                f" discriminator loss {avg_global_disc_loss:.3f}"
            )

            remaining_epochs_gen = max(
                overtraining_threshold - consecutive_increases_gen,
                0,
            )
            remaining_epochs_disc = max(
                overtraining_threshold * 2 - consecutive_increases_disc,
                0,
            )
            record += (
                " | overtrain countdown: g="
                f"{remaining_epochs_gen},d={remaining_epochs_disc} |"
                f" avg-gen-loss={avg_global_gen_loss:.3f} | avg-"
                f"disc-loss={avg_global_disc_loss:.3f}"
            )

            if remaining_epochs_disc == 0 or remaining_epochs_gen == 0:
                record += (
                    f"\nOvertraining detected at epoch {epoch} with average"
                    f" generator loss {avg_global_gen_loss:.3f} and discriminator loss"
                    f" {avg_global_disc_loss:.3f}"
                )
                done = True
        print(record)

        # Save weights, checkpoints and reference inference results
        # every N epochs
        if epoch % save_every_epoch == 0:
            with torch.no_grad():
                if hasattr(net_g, "module"):
                    o, *_ = net_g.module.infer(*reference)
                else:
                    o, *_ = net_g.infer(*reference)
            audio_dict = {f"gen/audio_{global_step:07d}": o[0, :, :]}
            summarize(
                writer=writer,
                global_step=global_step,
                images=image_dict,
                scalars=scalar_dict,
                audios=audio_dict,
                audio_sample_rate=config.data.sample_rate,
            )
            checkpoint_idxs.append(2333333)
            if not save_only_latest:
                checkpoint_idxs.append(epoch)

            if custom_save_every_weights:
                model_add.append(
                    os.path.join(experiment_dir, f"{model_name}_{epoch}.pth"),
                )
        else:
            summarize(
                writer=writer,
                global_step=global_step,
                images=image_dict,
                scalars=scalar_dict,
            )
        for idx in checkpoint_idxs:
            save_checkpoint(
                net_g,
                optim_g,
                config.train.learning_rate,
                epoch,
                lowest_g_value,
                consecutive_increases_gen,
                os.path.join(experiment_dir, f"G_{idx}.pth"),
            )
            save_checkpoint(
                net_d,
                optim_d,
                config.train.learning_rate,
                epoch,
                lowest_d_value,
                consecutive_increases_disc,
                os.path.join(experiment_dir, f"D_{idx}.pth"),
            )
        if model_add:
            ckpt = (
                net_g.module.state_dict()
                if hasattr(net_g, "module")
                else net_g.state_dict()
            )
            for m in model_add:
                extract_model(
                    ckpt=ckpt,
                    sr=sample_rate,
                    name=model_name,
                    model_dir=m,
                    epoch=epoch,
                    step=global_step,
                    hps=config,
                    overtrain_info=overtrain_info,
                    vocoder=vocoder,
                )
        # Check completion
        if epoch >= custom_total_epoch:
            lowest_g_value_rounded = float(lowest_g_value["value"])
            lowest_g_value_rounded = round(lowest_g_value_rounded, 3)
            print(
                f"Training has been successfully completed with {epoch} epoch(s),"
                f" {global_step} step(s) and {round(avg_global_gen_loss, 3)} average"
                " generator loss.",
            )
            print(
                f"Lowest average generator loss: {lowest_g_value_rounded} at epoch"
                f" {lowest_g_value['epoch']}",
            )

            done = True
        with torch.no_grad():
            torch.cuda.empty_cache()
        if done:
            cleanup_training_processes(experiment_dir)
            os._exit(0)


def cleanup_training_processes(experiment_dir) -> None:
    dist.destroy_process_group()
    pid_file_path = os.path.join(experiment_dir, "config.json")
    with pathlib.Path(pid_file_path).open() as pid_file:
        pid_data = json.load(pid_file)
    with pathlib.Path(pid_file_path).open("w") as pid_file:
        pid_data.pop("process_pids", None)
        json.dump(pid_data, pid_file, indent=4)
