# Copyright (c) 2023 Amphion.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.


import collections
import glob
import os
import random
import time
import argparse
from collections import OrderedDict

import json5
import numpy as np
import glob
from torch.nn import functional as F


try:
    from ruamel.yaml import YAML as yaml
except:
    from ruamel_yaml import YAML as yaml

import torch

from utils.hparam import HParams
import logging
from logging import handlers


def str2bool(v):
    """Used in argparse.ArgumentParser.add_argument to indicate
    that a type is a bool type and user can enter

        - yes, true, t, y, 1, to represent True
        - no, false, f, n, 0, to represent False

    See https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse  # noqa
    """
    if isinstance(v, bool):
        return v
    if v.lower() in ("yes", "true", "t", "y", "1"):
        return True
    elif v.lower() in ("no", "false", "f", "n", "0"):
        return False
    else:
        raise argparse.ArgumentTypeError("Boolean value expected.")


def find_checkpoint_of_mapper(mapper_ckpt_dir):
    mapper_ckpts = glob.glob(os.path.join(mapper_ckpt_dir, "ckpts/*.pt"))

    # Select the max steps
    mapper_ckpts.sort()
    mapper_weights_file = mapper_ckpts[-1]
    return mapper_weights_file


def pad_f0_to_tensors(f0s, batched=None):
    # Initialize
    tensors = []

    if batched == None:
        # Get the max frame for padding
        size = -1
        for f0 in f0s:
            size = max(size, f0.shape[-1])

        tensor = torch.zeros(len(f0s), size)

        for i, f0 in enumerate(f0s):
            tensor[i, : f0.shape[-1]] = f0[:]

        tensors.append(tensor)
    else:
        start = 0
        while start + batched - 1 < len(f0s):
            end = start + batched - 1

            # Get the max frame for padding
            size = -1
            for i in range(start, end + 1):
                size = max(size, f0s[i].shape[-1])

            tensor = torch.zeros(batched, size)

            for i in range(start, end + 1):
                tensor[i - start, : f0s[i].shape[-1]] = f0s[i][:]

            tensors.append(tensor)

            start = start + batched

        if start != len(f0s):
            end = len(f0s)

            # Get the max frame for padding
            size = -1
            for i in range(start, end):
                size = max(size, f0s[i].shape[-1])

            tensor = torch.zeros(len(f0s) - start, size)

            for i in range(start, end):
                tensor[i - start, : f0s[i].shape[-1]] = f0s[i][:]

            tensors.append(tensor)

    return tensors


def pad_mels_to_tensors(mels, batched=None):
    """
    Args:
        mels: A list of mel-specs
    Returns:
        tensors: A list of tensors containing the batched mel-specs
        mel_frames: A list of tensors containing the frames of the original mel-specs
    """
    # Initialize
    tensors = []
    mel_frames = []

    # Split mel-specs into batches to avoid cuda memory exceed
    if batched == None:
        # Get the max frame for padding
        size = -1
        for mel in mels:
            size = max(size, mel.shape[-1])

        tensor = torch.zeros(len(mels), mels[0].shape[0], size)
        mel_frame = torch.zeros(len(mels), dtype=torch.int32)

        for i, mel in enumerate(mels):
            tensor[i, :, : mel.shape[-1]] = mel[:]
            mel_frame[i] = mel.shape[-1]

        tensors.append(tensor)
        mel_frames.append(mel_frame)
    else:
        start = 0
        while start + batched - 1 < len(mels):
            end = start + batched - 1

            # Get the max frame for padding
            size = -1
            for i in range(start, end + 1):
                size = max(size, mels[i].shape[-1])

            tensor = torch.zeros(batched, mels[0].shape[0], size)
            mel_frame = torch.zeros(batched, dtype=torch.int32)

            for i in range(start, end + 1):
                tensor[i - start, :, : mels[i].shape[-1]] = mels[i][:]
                mel_frame[i - start] = mels[i].shape[-1]

            tensors.append(tensor)
            mel_frames.append(mel_frame)

            start = start + batched

        if start != len(mels):
            end = len(mels)

            # Get the max frame for padding
            size = -1
            for i in range(start, end):
                size = max(size, mels[i].shape[-1])

            tensor = torch.zeros(len(mels) - start, mels[0].shape[0], size)
            mel_frame = torch.zeros(len(mels) - start, dtype=torch.int32)

            for i in range(start, end):
                tensor[i - start, :, : mels[i].shape[-1]] = mels[i][:]
                mel_frame[i - start] = mels[i].shape[-1]

            tensors.append(tensor)
            mel_frames.append(mel_frame)

    return tensors, mel_frames


def load_model_config(args):
    """Load model configurations (in args.json under checkpoint directory)

    Args:
        args (ArgumentParser): arguments to run bins/preprocess.py

    Returns:
        dict: dictionary that stores model configurations
    """
    if args.checkpoint_dir is None:
        assert args.checkpoint_file is not None
        checkpoint_dir = os.path.split(args.checkpoint_file)[0]
    else:
        checkpoint_dir = args.checkpoint_dir
    config_path = os.path.join(checkpoint_dir, "args.json")
    print("config_path: ", config_path)

    config = load_config(config_path)
    return config


def remove_and_create(dir):
    if os.path.exists(dir):
        os.system("rm -r {}".format(dir))
    os.makedirs(dir, exist_ok=True)


def has_existed(path, warning=False):
    if not warning:
        return os.path.exists(path)

    if os.path.exists(path):
        answer = input(
            "The path {} has existed. \nInput 'y' (or hit Enter) to skip it, and input 'n' to re-write it [y/n]\n".format(
                path
            )
        )
        if not answer == "n":
            return True

    return False


def remove_older_ckpt(saved_model_name, checkpoint_dir, max_to_keep=5):
    if os.path.exists(os.path.join(checkpoint_dir, "checkpoint")):
        with open(os.path.join(checkpoint_dir, "checkpoint"), "r") as f:
            ckpts = [x.strip() for x in f.readlines()]
    else:
        ckpts = []
    ckpts.append(saved_model_name)
    for item in ckpts[:-max_to_keep]:
        if os.path.exists(os.path.join(checkpoint_dir, item)):
            os.remove(os.path.join(checkpoint_dir, item))
    with open(os.path.join(checkpoint_dir, "checkpoint"), "w") as f:
        for item in ckpts[-max_to_keep:]:
            f.write("{}\n".format(item))


def set_all_random_seed(seed: int):
    random.seed(seed)
    np.random.seed(seed)
    torch.random.manual_seed(seed)


def save_checkpoint(
    args,
    generator,
    g_optimizer,
    step,
    discriminator=None,
    d_optimizer=None,
    max_to_keep=5,
):
    saved_model_name = "model.ckpt-{}.pt".format(step)
    checkpoint_path = os.path.join(args.checkpoint_dir, saved_model_name)

    if discriminator and d_optimizer:
        torch.save(
            {
                "generator": generator.state_dict(),
                "discriminator": discriminator.state_dict(),
                "g_optimizer": g_optimizer.state_dict(),
                "d_optimizer": d_optimizer.state_dict(),
                "global_step": step,
            },
            checkpoint_path,
        )
    else:
        torch.save(
            {
                "generator": generator.state_dict(),
                "g_optimizer": g_optimizer.state_dict(),
                "global_step": step,
            },
            checkpoint_path,
        )

    print("Saved checkpoint: {}".format(checkpoint_path))

    if os.path.exists(os.path.join(args.checkpoint_dir, "checkpoint")):
        with open(os.path.join(args.checkpoint_dir, "checkpoint"), "r") as f:
            ckpts = [x.strip() for x in f.readlines()]
    else:
        ckpts = []
    ckpts.append(saved_model_name)
    for item in ckpts[:-max_to_keep]:
        if os.path.exists(os.path.join(args.checkpoint_dir, item)):
            os.remove(os.path.join(args.checkpoint_dir, item))
    with open(os.path.join(args.checkpoint_dir, "checkpoint"), "w") as f:
        for item in ckpts[-max_to_keep:]:
            f.write("{}\n".format(item))


def attempt_to_restore(
    generator, g_optimizer, checkpoint_dir, discriminator=None, d_optimizer=None
):
    checkpoint_list = os.path.join(checkpoint_dir, "checkpoint")
    if os.path.exists(checkpoint_list):
        checkpoint_filename = open(checkpoint_list).readlines()[-1].strip()
        checkpoint_path = os.path.join(checkpoint_dir, "{}".format(checkpoint_filename))
        print("Restore from {}".format(checkpoint_path))
        checkpoint = torch.load(checkpoint_path, map_location="cpu")
        if generator:
            if not list(generator.state_dict().keys())[0].startswith("module."):
                raw_dict = checkpoint["generator"]
                clean_dict = OrderedDict()
                for k, v in raw_dict.items():
                    if k.startswith("module."):
                        clean_dict[k[7:]] = v
                    else:
                        clean_dict[k] = v
                generator.load_state_dict(clean_dict)
            else:
                generator.load_state_dict(checkpoint["generator"])
        if g_optimizer:
            g_optimizer.load_state_dict(checkpoint["g_optimizer"])
        global_step = 100000
        if discriminator and "discriminator" in checkpoint.keys():
            discriminator.load_state_dict(checkpoint["discriminator"])
            global_step = checkpoint["global_step"]
            print("restore discriminator")
        if d_optimizer and "d_optimizer" in checkpoint.keys():
            d_optimizer.load_state_dict(checkpoint["d_optimizer"])
            print("restore d_optimizer...")
    else:
        global_step = 0
    return global_step


class ExponentialMovingAverage(object):
    def __init__(self, decay):
        self.decay = decay
        self.shadow = {}

    def register(self, name, val):
        self.shadow[name] = val.clone()

    def update(self, name, x):
        assert name in self.shadow
        update_delta = self.shadow[name] - x
        self.shadow[name] -= (1.0 - self.decay) * update_delta


def apply_moving_average(model, ema):
    for name, param in model.named_parameters():
        if name in ema.shadow:
            ema.update(name, param.data)


def register_model_to_ema(model, ema):
    for name, param in model.named_parameters():
        if param.requires_grad:
            ema.register(name, param.data)


class YParams(HParams):
    def __init__(self, yaml_file):
        if not os.path.exists(yaml_file):
            raise IOError("yaml file: {} is not existed".format(yaml_file))
        super().__init__()
        self.d = collections.OrderedDict()
        with open(yaml_file) as fp:
            for _, v in yaml().load(fp).items():
                for k1, v1 in v.items():
                    try:
                        if self.get(k1):
                            self.set_hparam(k1, v1)
                        else:
                            self.add_hparam(k1, v1)
                        self.d[k1] = v1
                    except Exception:
                        import traceback

                        print(traceback.format_exc())

    # @property
    def get_elements(self):
        return self.d.items()


def override_config(base_config, new_config):
    """Update new configurations in the original dict with the new dict

    Args:
        base_config (dict): original dict to be overridden
        new_config (dict): dict with new configurations

    Returns:
        dict: updated configuration dict
    """
    for k, v in new_config.items():
        if type(v) == dict:
            if k not in base_config.keys():
                base_config[k] = {}
            base_config[k] = override_config(base_config[k], v)
        else:
            base_config[k] = v
    return base_config


def get_lowercase_keys_config(cfg):
    """Change all keys in cfg to lower case

    Args:
        cfg (dict): dictionary that stores configurations

    Returns:
        dict: dictionary that stores configurations
    """
    updated_cfg = dict()
    for k, v in cfg.items():
        if type(v) == dict:
            v = get_lowercase_keys_config(v)
        updated_cfg[k.lower()] = v
    return updated_cfg


def _load_config(config_fn, lowercase=False):
    """Load configurations into a dictionary

    Args:
        config_fn (str): path to configuration file
        lowercase (bool, optional): whether changing keys to lower case. Defaults to False.

    Returns:
        dict: dictionary that stores configurations
    """
    with open(config_fn, "r") as f:
        data = f.read()
    config_ = json5.loads(data)
    if "base_config" in config_:
        # load configurations from new path
        p_config_path = os.path.join(os.getenv("WORK_DIR"), config_["base_config"])
        p_config_ = _load_config(p_config_path)
        config_ = override_config(p_config_, config_)
    if lowercase:
        # change keys in config_ to lower case
        config_ = get_lowercase_keys_config(config_)
    return config_


def load_config(config_fn, lowercase=False):
    """Load configurations into a dictionary

    Args:
        config_fn (str): path to configuration file
        lowercase (bool, optional): _description_. Defaults to False.

    Returns:
        JsonHParams: an object that stores configurations
    """
    config_ = _load_config(config_fn, lowercase=lowercase)
    # create an JsonHParams object with configuration dict
    cfg = JsonHParams(**config_)
    return cfg


def save_config(save_path, cfg):
    """Save configurations into a json file

    Args:
        save_path (str): path to save configurations
        cfg (dict): dictionary that stores configurations
    """
    with open(save_path, "w") as f:
        json5.dump(
            cfg, f, ensure_ascii=False, indent=4, quote_keys=True, sort_keys=True
        )


class JsonHParams:
    def __init__(self, **kwargs):
        for k, v in kwargs.items():
            if type(v) == dict:
                v = JsonHParams(**v)
            self[k] = v

    def keys(self):
        return self.__dict__.keys()

    def items(self):
        return self.__dict__.items()

    def values(self):
        return self.__dict__.values()

    def __len__(self):
        return len(self.__dict__)

    def __getitem__(self, key):
        return getattr(self, key)

    def __setitem__(self, key, value):
        return setattr(self, key, value)

    def __contains__(self, key):
        return key in self.__dict__

    def __repr__(self):
        return self.__dict__.__repr__()


class ValueWindow:
    def __init__(self, window_size=100):
        self._window_size = window_size
        self._values = []

    def append(self, x):
        self._values = self._values[-(self._window_size - 1) :] + [x]

    @property
    def sum(self):
        return sum(self._values)

    @property
    def count(self):
        return len(self._values)

    @property
    def average(self):
        return self.sum / max(1, self.count)

    def reset(self):
        self._values = []


class Logger(object):
    def __init__(
        self,
        filename,
        level="info",
        when="D",
        backCount=10,
        fmt="%(asctime)s : %(message)s",
    ):
        self.level_relations = {
            "debug": logging.DEBUG,
            "info": logging.INFO,
            "warning": logging.WARNING,
            "error": logging.ERROR,
            "crit": logging.CRITICAL,
        }
        if level == "debug":
            fmt = "%(asctime)s - %(pathname)s[line:%(lineno)d] - %(levelname)s: %(message)s"
        self.logger = logging.getLogger(filename)
        format_str = logging.Formatter(fmt)
        self.logger.setLevel(self.level_relations.get(level))
        sh = logging.StreamHandler()
        sh.setFormatter(format_str)
        th = handlers.TimedRotatingFileHandler(
            filename=filename, when=when, backupCount=backCount, encoding="utf-8"
        )
        th.setFormatter(format_str)
        self.logger.addHandler(sh)
        self.logger.addHandler(th)
        self.logger.info(
            "==========================New Starting Here=============================="
        )


def init_weights(m, mean=0.0, std=0.01):
    classname = m.__class__.__name__
    if classname.find("Conv") != -1:
        m.weight.data.normal_(mean, std)


def get_padding(kernel_size, dilation=1):
    return int((kernel_size * dilation - dilation) / 2)


def slice_segments(x, ids_str, segment_size=4):
    ret = torch.zeros_like(x[:, :, :segment_size])
    for i in range(x.size(0)):
        idx_str = ids_str[i]
        idx_end = idx_str + segment_size
        ret[i] = x[i, :, idx_str:idx_end]
    return ret


def rand_slice_segments(x, x_lengths=None, segment_size=4):
    b, d, t = x.size()
    if x_lengths is None:
        x_lengths = t
    ids_str_max = x_lengths - segment_size + 1
    ids_str = (torch.rand([b]).to(device=x.device) * ids_str_max).to(dtype=torch.long)
    ret = slice_segments(x, ids_str, segment_size)
    return ret, ids_str


def subsequent_mask(length):
    mask = torch.tril(torch.ones(length, length)).unsqueeze(0).unsqueeze(0)
    return mask


@torch.jit.script
def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
    n_channels_int = n_channels[0]
    in_act = input_a + input_b
    t_act = torch.tanh(in_act[:, :n_channels_int, :])
    s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
    acts = t_act * s_act
    return acts


def convert_pad_shape(pad_shape):
    l = pad_shape[::-1]
    pad_shape = [item for sublist in l for item in sublist]
    return pad_shape


def sequence_mask(length, max_length=None):
    if max_length is None:
        max_length = length.max()
    x = torch.arange(max_length, dtype=length.dtype, device=length.device)
    return x.unsqueeze(0) < length.unsqueeze(1)


def generate_path(duration, mask):
    """
    duration: [b, 1, t_x]
    mask: [b, 1, t_y, t_x]
    """
    device = duration.device

    b, _, t_y, t_x = mask.shape
    cum_duration = torch.cumsum(duration, -1)

    cum_duration_flat = cum_duration.view(b * t_x)
    path = sequence_mask(cum_duration_flat, t_y).to(mask.dtype)
    path = path.view(b, t_x, t_y)
    path = path - F.pad(path, convert_pad_shape([[0, 0], [1, 0], [0, 0]]))[:, :-1]
    path = path.unsqueeze(1).transpose(2, 3) * mask
    return path


def clip_grad_value_(parameters, clip_value, norm_type=2):
    if isinstance(parameters, torch.Tensor):
        parameters = [parameters]
    parameters = list(filter(lambda p: p.grad is not None, parameters))
    norm_type = float(norm_type)
    if clip_value is not None:
        clip_value = float(clip_value)

    total_norm = 0
    for p in parameters:
        param_norm = p.grad.data.norm(norm_type)
        total_norm += param_norm.item() ** norm_type
        if clip_value is not None:
            p.grad.data.clamp_(min=-clip_value, max=clip_value)
    total_norm = total_norm ** (1.0 / norm_type)
    return total_norm


def get_current_time():
    pass


def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
    """
    Args:
      lengths:
        A 1-D tensor containing sentence lengths.
      max_len:
        The length of masks.
    Returns:
      Return a 2-D bool tensor, where masked positions
      are filled with `True` and non-masked positions are
      filled with `False`.

    >>> lengths = torch.tensor([1, 3, 2, 5])
    >>> make_pad_mask(lengths)
    tensor([[False,  True,  True,  True,  True],
            [False, False, False,  True,  True],
            [False, False,  True,  True,  True],
            [False, False, False, False, False]])
    """
    assert lengths.ndim == 1, lengths.ndim
    max_len = max(max_len, lengths.max())
    n = lengths.size(0)
    seq_range = torch.arange(0, max_len, device=lengths.device)
    expaned_lengths = seq_range.unsqueeze(0).expand(n, max_len)

    return expaned_lengths >= lengths.unsqueeze(-1)