# import pytorch_lightning as pl

import torch.utils
from voicefixer.tools.mel_scale import MelScale
import torch.utils.data
import matplotlib.pyplot as plt
import librosa.display
from voicefixer.vocoder.base import Vocoder
from voicefixer.tools.pytorch_util import *
from voicefixer.restorer.model_kqq_bn import UNetResComplex_100Mb
from voicefixer.tools.random_ import *
from voicefixer.tools.wav import *
from voicefixer.tools.modules.fDomainHelper import FDomainHelper

from voicefixer.tools.io import load_json, write_json
from matplotlib import cm

os.environ["KMP_DUPLICATE_LIB_OK"] = "True"
EPS = 1e-8


class BN_GRU(torch.nn.Module):
    def __init__(
        self,
        input_dim,
        hidden_dim,
        layer=1,
        bidirectional=False,
        batchnorm=True,
        dropout=0.0,
    ):
        super(BN_GRU, self).__init__()
        self.batchnorm = batchnorm
        if batchnorm:
            self.bn = nn.BatchNorm2d(1)
        self.gru = torch.nn.GRU(
            input_size=input_dim,
            hidden_size=hidden_dim,
            num_layers=layer,
            bidirectional=bidirectional,
            dropout=dropout,
            batch_first=True,
        )
        self.init_weights()

    def init_weights(self):
        for m in self.modules():
            if type(m) in [nn.GRU, nn.LSTM, nn.RNN]:
                for name, param in m.named_parameters():
                    if "weight_ih" in name:
                        torch.nn.init.xavier_uniform_(param.data)
                    elif "weight_hh" in name:
                        torch.nn.init.orthogonal_(param.data)
                    elif "bias" in name:
                        param.data.fill_(0)

    def forward(self, inputs):
        # (batch, 1, seq, feature)
        if self.batchnorm:
            inputs = self.bn(inputs)
        out, _ = self.gru(inputs.squeeze(1))
        return out.unsqueeze(1)


class Generator(nn.Module):
    def __init__(self, n_mel, hidden, channels):
        super(Generator, self).__init__()
        # todo the currently running trail don't have dropout
        self.denoiser = nn.Sequential(
            nn.BatchNorm2d(1),
            nn.Linear(n_mel, n_mel * 2),
            nn.ReLU(inplace=True),
            nn.BatchNorm2d(1),
            nn.Linear(n_mel * 2, n_mel * 4),
            nn.Dropout(0.5),
            nn.ReLU(inplace=True),
            BN_GRU(
                input_dim=n_mel * 4,
                hidden_dim=n_mel * 2,
                bidirectional=True,
                layer=2,
                batchnorm=True,
            ),
            BN_GRU(
                input_dim=n_mel * 4,
                hidden_dim=n_mel * 2,
                bidirectional=True,
                layer=2,
                batchnorm=True,
            ),
            nn.BatchNorm2d(1),
            nn.ReLU(inplace=True),
            nn.Linear(n_mel * 4, n_mel * 4),
            nn.Dropout(0.5),
            nn.BatchNorm2d(1),
            nn.ReLU(inplace=True),
            nn.Linear(n_mel * 4, n_mel),
            nn.Sigmoid(),
        )

        self.unet = UNetResComplex_100Mb(channels=channels)

    def forward(self, sp, mel_orig):
        # Denoising
        noisy = mel_orig.clone()
        clean = self.denoiser(noisy) * noisy
        x = to_log(clean.detach())
        unet_in = torch.cat([to_log(mel_orig), x], dim=1)
        # unet_in = lstm_out
        unet_out = self.unet(unet_in)["mel"]
        # masks
        mel = unet_out + x
        # todo mel and addition here are in log scales
        return {
            "mel": mel,
            "lstm_out": unet_out,
            "unet_out": unet_out,
            "noisy": noisy,
            "clean": clean,
        }


class VoiceFixer(nn.Module):
    def __init__(
        self,
        channels,
        type_target="vocals",
        nsrc=1,
        loss="l1",
        lr=0.002,
        gamma=0.9,
        batchsize=None,
        frame_length=None,
        sample_rate=None,
        warm_up_steps=1000,
        reduce_lr_steps=15000,
        # datas
        check_val_every_n_epoch=5,
    ):
        super(VoiceFixer, self).__init__()

        if sample_rate == 44100:
            window_size = 2048
            hop_size = 441
            n_mel = 128
        elif sample_rate == 24000:
            window_size = 768
            hop_size = 240
            n_mel = 80
        elif sample_rate == 16000:
            window_size = 512
            hop_size = 160
            n_mel = 80
        else:
            raise ValueError(
                "Error: Sample rate " + str(sample_rate) + " not supported"
            )

        center = (True,)
        pad_mode = "reflect"
        window = "hann"
        freeze_parameters = True

        # self.save_hyperparameters()
        self.nsrc = nsrc
        self.type_target = type_target
        self.channels = channels
        self.lr = lr
        self.generated = None
        self.gamma = gamma
        self.sample_rate = sample_rate
        self.sample_rate = sample_rate
        self.batchsize = batchsize
        self.frame_length = frame_length
        # self.hparams['channels'] = 2

        # self.am = AudioMetrics()
        # self.im = ImgMetrics()

        self.vocoder = Vocoder(sample_rate=44100)

        self.valid = None
        self.fake = None

        self.train_step = 0
        self.val_step = 0
        self.val_result_save_dir = None
        self.val_result_save_dir_step = None
        self.downsample_ratio = 2**6  # This number equals 2^{#encoder_blcoks}
        self.check_val_every_n_epoch = check_val_every_n_epoch

        self.f_helper = FDomainHelper(
            window_size=window_size,
            hop_size=hop_size,
            center=center,
            pad_mode=pad_mode,
            window=window,
            freeze_parameters=freeze_parameters,
        )

        hidden = window_size // 2 + 1

        self.mel = MelScale(n_mels=n_mel, sample_rate=sample_rate, n_stft=hidden)

        # masking
        self.generator = Generator(n_mel, hidden, channels)

        self.lr_lambda = lambda step: self.get_lr_lambda(
            step,
            gamma=self.gamma,
            warm_up_steps=warm_up_steps,
            reduce_lr_steps=reduce_lr_steps,
        )

        self.lr_lambda_2 = lambda step: self.get_lr_lambda(
            step, gamma=self.gamma, warm_up_steps=10, reduce_lr_steps=reduce_lr_steps
        )

        self.mel_weight_44k_128 = (
            torch.tensor(
                [
                    19.40951426,
                    19.94047336,
                    20.4859038,
                    21.04629067,
                    21.62194148,
                    22.21335214,
                    22.8210215,
                    23.44529231,
                    24.08660962,
                    24.74541882,
                    25.42234287,
                    26.11770576,
                    26.83212784,
                    27.56615283,
                    28.32007747,
                    29.0947679,
                    29.89060111,
                    30.70832636,
                    31.54828121,
                    32.41121487,
                    33.29780773,
                    34.20865341,
                    35.14437675,
                    36.1056621,
                    37.09332763,
                    38.10795802,
                    39.15039691,
                    40.22119881,
                    41.32154931,
                    42.45172373,
                    43.61293329,
                    44.80609379,
                    46.031602,
                    47.29070223,
                    48.58427549,
                    49.91327905,
                    51.27863232,
                    52.68119708,
                    54.1222372,
                    55.60274206,
                    57.12364703,
                    58.68617876,
                    60.29148652,
                    61.94081306,
                    63.63501986,
                    65.37562658,
                    67.16408954,
                    69.00109084,
                    70.88850318,
                    72.82736101,
                    74.81985537,
                    76.86654792,
                    78.96885475,
                    81.12900906,
                    83.34840929,
                    85.62810662,
                    87.97005418,
                    90.37689804,
                    92.84887686,
                    95.38872881,
                    97.99777002,
                    100.67862715,
                    103.43232942,
                    106.26140638,
                    109.16827015,
                    112.15470471,
                    115.22184756,
                    118.37439245,
                    121.6122689,
                    124.93877158,
                    128.35661454,
                    131.86761321,
                    135.47417938,
                    139.18059494,
                    142.98713744,
                    146.89771854,
                    150.91684347,
                    155.0446638,
                    159.28614648,
                    163.64270198,
                    168.12035831,
                    172.71749158,
                    177.44220154,
                    182.29556933,
                    187.28286676,
                    192.40502126,
                    197.6682721,
                    203.07516896,
                    208.63088733,
                    214.33770931,
                    220.19910108,
                    226.22363072,
                    232.41087124,
                    238.76803591,
                    245.30079083,
                    252.01064464,
                    258.90261676,
                    265.98474,
                    273.26010248,
                    280.73496362,
                    288.41440094,
                    296.30489752,
                    304.41180337,
                    312.7377183,
                    321.28877878,
                    330.07870237,
                    339.10812951,
                    348.38276173,
                    357.91393924,
                    367.70513992,
                    377.76413924,
                    388.09467408,
                    398.70920178,
                    409.61813793,
                    420.81980127,
                    432.33215467,
                    444.16083117,
                    456.30919947,
                    468.78589276,
                    481.61325588,
                    494.78824596,
                    508.31969844,
                    522.2238331,
                    536.51163441,
                    551.18859414,
                    566.26142988,
                    581.75006061,
                    597.66210737,
                ]
            )
            / 19.40951426
        )
        self.mel_weight_44k_128 = self.mel_weight_44k_128[None, None, None, ...]

        self.g_loss_weight = 0.01
        self.d_loss_weight = 1

    def get_vocoder(self):
        return self.vocoder

    def get_f_helper(self):
        return self.f_helper

    def get_lr_lambda(self, step, gamma, warm_up_steps, reduce_lr_steps):
        r"""Get lr_lambda for LambdaLR. E.g.,

        .. code-block: python
            lr_lambda = lambda step: get_lr_lambda(step, warm_up_steps=1000, reduce_lr_steps=10000)

            from torch.optim.lr_scheduler import LambdaLR
            LambdaLR(optimizer, lr_lambda)
        """
        if step <= warm_up_steps:
            return step / warm_up_steps
        else:
            return gamma ** (step // reduce_lr_steps)

    def init_weights(self, module: nn.Module):
        for m in module.modules():
            if type(m) in [nn.GRU, nn.LSTM, nn.RNN]:
                for name, param in m.named_parameters():
                    if "weight_ih" in name:
                        torch.nn.init.xavier_uniform_(param.data)
                    elif "weight_hh" in name:
                        torch.nn.init.orthogonal_(param.data)
                    elif "bias" in name:
                        param.data.fill_(0)

    def pre(self, input):
        sp, _, _ = self.f_helper.wav_to_spectrogram_phase(input)
        mel_orig = self.mel(sp.permute(0, 1, 3, 2)).permute(0, 1, 3, 2)
        return sp, mel_orig

    def forward(self, sp, mel_orig):
        """
        Args:
          input: (batch_size, channels_num, segment_samples)

        Outputs:
          output_dict: {
            'wav': (batch_size, channels_num, segment_samples),
            'sp': (batch_size, channels_num, time_steps, freq_bins)}
        """
        return self.generator(sp, mel_orig)

    def configure_optimizers(self):
        optimizer_g = torch.optim.Adam(
            [{"params": self.generator.parameters()}],
            lr=self.lr,
            amsgrad=True,
            betas=(0.5, 0.999),
        )
        optimizer_d = torch.optim.Adam(
            [{"params": self.discriminator.parameters()}],
            lr=self.lr,
            amsgrad=True,
            betas=(0.5, 0.999),
        )

        scheduler_g = {
            "scheduler": torch.optim.lr_scheduler.LambdaLR(optimizer_g, self.lr_lambda),
            "interval": "step",
            "frequency": 1,
        }
        scheduler_d = {
            "scheduler": torch.optim.lr_scheduler.LambdaLR(optimizer_d, self.lr_lambda),
            "interval": "step",
            "frequency": 1,
        }
        return [optimizer_g, optimizer_d], [scheduler_g, scheduler_d]

    def preprocess(self, batch, train=False, cutoff=None):
        if train:
            vocal = batch[self.type_target]  # final target
            noise = batch["noise_LR"]  # augmented low resolution audio with noise
            augLR = batch[
                self.type_target + "_aug_LR"
            ]  # # augment low resolution audio
            LR = batch[self.type_target + "_LR"]
            # embed()
            vocal, LR, augLR, noise = (
                vocal.float().permute(0, 2, 1),
                LR.float().permute(0, 2, 1),
                augLR.float().permute(0, 2, 1),
                noise.float().permute(0, 2, 1),
            )
            # LR, noise = self.add_random_noise(LR, noise)
            snr, scale = [], []
            for i in range(vocal.size()[0]):
                (
                    vocal[i, ...],
                    LR[i, ...],
                    augLR[i, ...],
                    noise[i, ...],
                    _snr,
                    _scale,
                ) = add_noise_and_scale_with_HQ_with_Aug(
                    vocal[i, ...],
                    LR[i, ...],
                    augLR[i, ...],
                    noise[i, ...],
                    snr_l=-5,
                    snr_h=45,
                    scale_lower=0.6,
                    scale_upper=1.0,
                )
                snr.append(_snr), scale.append(_scale)
            # vocal, LR = self.amp_to_original_f(vocal, LR)
            # noise = (noise * 0.0) + 1e-8 # todo
            return vocal, augLR, LR, noise + augLR
        else:
            if cutoff is None:
                LR_noisy = batch["noisy"]
                LR = batch["vocals"]
                vocals = batch["vocals"]
                vocals, LR, LR_noisy = (
                    vocals.float().permute(0, 2, 1),
                    LR.float().permute(0, 2, 1),
                    LR_noisy.float().permute(0, 2, 1),
                )
                return vocals, LR, LR_noisy, batch["fname"][0]
            else:
                LR_noisy = batch["noisy" + "LR" + "_" + str(cutoff)]
                LR = batch["vocals" + "LR" + "_" + str(cutoff)]
                vocals = batch["vocals"]
                vocals, LR, LR_noisy = (
                    vocals.float().permute(0, 2, 1),
                    LR.float().permute(0, 2, 1),
                    LR_noisy.float().permute(0, 2, 1),
                )
                return vocals, LR, LR_noisy, batch["fname"][0]

    def training_step(self, batch, batch_nb, optimizer_idx):
        # dict_keys(['vocals', 'vocals_aug', 'vocals_augLR', 'noise'])
        config = load_json("temp_path.json")
        if "g_loss_weight" not in config.keys():
            config["g_loss_weight"] = self.g_loss_weight
            config["d_loss_weight"] = self.d_loss_weight
            write_json(config, "temp_path.json")
        elif (
            config["g_loss_weight"] != self.g_loss_weight
            or config["d_loss_weight"] != self.d_loss_weight
        ):
            print(
                "Update d_loss weight, from",
                self.d_loss_weight,
                "to",
                config["d_loss_weight"],
            )
            print(
                "Update g_loss weight, from",
                self.g_loss_weight,
                "to",
                config["g_loss_weight"],
            )
            self.g_loss_weight = config["g_loss_weight"]
            self.d_loss_weight = config["d_loss_weight"]

        if optimizer_idx == 0:
            self.vocal, self.augLR, _, self.LR_noisy = self.preprocess(
                batch, train=True
            )

            for i in range(self.vocal.size()[0]):
                save_wave(
                    tensor2numpy(self.vocal[i, ...]),
                    str(i) + "vocal" + ".wav",
                    sample_rate=44100,
                )
                save_wave(
                    tensor2numpy(self.LR_noisy[i, ...]),
                    str(i) + "LR_noisy" + ".wav",
                    sample_rate=44100,
                )

            # all_mel_e2e in non-log scale
            _, self.mel_target = self.pre(self.vocal)
            self.sp_LR_target, self.mel_LR_target = self.pre(self.augLR)
            self.sp_LR_target_noisy, self.mel_LR_target_noisy = self.pre(self.LR_noisy)

            if self.valid is None or self.valid.size()[0] != self.mel_target.size()[0]:
                self.valid = torch.ones(
                    self.mel_target.size()[0], 1, self.mel_target.size()[2], 1
                )
                self.valid = self.valid.type_as(self.mel_target)
            if self.fake is None or self.fake.size()[0] != self.mel_target.size()[0]:
                self.fake = torch.zeros(
                    self.mel_target.size()[0], 1, self.mel_target.size()[2], 1
                )
                self.fake = self.fake.type_as(self.mel_target)

            self.generated = self(self.sp_LR_target_noisy, self.mel_LR_target_noisy)

            denoise_loss = self.l1loss(self.generated["clean"], self.mel_LR_target)
            targ_loss = self.l1loss(self.generated["mel"], to_log(self.mel_target))

            self.log(
                "targ-l",
                targ_loss,
                on_step=True,
                on_epoch=False,
                logger=True,
                sync_dist=True,
                prog_bar=True,
            )
            self.log(
                "noise-l",
                denoise_loss,
                on_step=True,
                on_epoch=False,
                logger=True,
                sync_dist=True,
                prog_bar=True,
            )

            loss = targ_loss + denoise_loss

            if self.train_step >= 18000:
                g_loss = self.bce_loss(
                    self.discriminator(self.generated["mel"]), self.valid
                )
                self.log(
                    "g_l",
                    g_loss,
                    on_step=True,
                    on_epoch=False,
                    logger=True,
                    sync_dist=True,
                    prog_bar=True,
                )
                # print("g_loss", g_loss)
                all_loss = loss + self.g_loss_weight * g_loss
                self.log(
                    "all_loss",
                    all_loss,
                    on_step=True,
                    on_epoch=True,
                    logger=True,
                    sync_dist=True,
                )
            else:
                all_loss = loss
            self.train_step += 0.5
            return {"loss": all_loss}

        elif optimizer_idx == 1:
            if self.train_step >= 16000:
                self.generated = self(self.sp_LR_target_noisy, self.mel_LR_target_noisy)
                self.train_step += 0.5
                real_loss = self.bce_loss(
                    self.discriminator(to_log(self.mel_target)), self.valid
                )
                self.log(
                    "r_l",
                    real_loss,
                    on_step=True,
                    on_epoch=False,
                    logger=True,
                    sync_dist=True,
                    prog_bar=True,
                )
                fake_loss = self.bce_loss(
                    self.discriminator(self.generated["mel"].detach()), self.fake
                )
                self.log(
                    "d_l",
                    fake_loss,
                    on_step=True,
                    on_epoch=False,
                    logger=True,
                    sync_dist=True,
                    prog_bar=True,
                )
                d_loss = self.d_loss_weight * (real_loss + fake_loss) / 2
                self.log(
                    "discriminator_loss",
                    d_loss,
                    on_step=True,
                    on_epoch=True,
                    logger=True,
                    sync_dist=True,
                )
                return {"loss": d_loss}

    def draw_and_save(
        self, mel: torch.Tensor, path, clip_max=None, clip_min=None, needlog=True
    ):
        plt.figure(figsize=(15, 5))
        if clip_min is None:
            clip_max, clip_min = self.clip(mel)
        mel = np.transpose(tensor2numpy(mel)[0, 0, ...], (1, 0))
        # assert np.sum(mel < 0) == 0, str(np.sum(mel < 0)) + str(np.sum(mel < 0))

        if needlog:
            assert np.sum(mel < 0) == 0, str(np.sum(mel < 0)) + "-" + path
            mel_log = np.log10(mel + EPS)
        else:
            mel_log = mel

        # plt.imshow(mel)
        librosa.display.specshow(
            mel_log,
            sr=44100,
            x_axis="frames",
            y_axis="mel",
            cmap=cm.jet,
            vmax=clip_max,
            vmin=clip_min,
        )
        plt.colorbar()
        plt.savefig(path)
        plt.close()

    def clip(self, *args):
        val_max, val_min = [], []
        for each in args:
            val_max.append(torch.max(each))
            val_min.append(torch.min(each))
        return max(val_max), min(val_min)
