import numpy as np
import torch
import torch.distributions
import torch.nn.functional as F
import torch.optim
import torch.utils.data

from tasks.speech_editing.dataset_utils import StutterSpeechDataset
from tasks.tts.speech_base import SpeechBaseTask
from utils.audio.align import mel2token_to_dur
from utils.audio.pitch.utils import denorm_f0
from utils.commons.hparams import hparams
from utils.feature_extractor.gst import GST
from utils.eval.mcd import get_metrics_mels


class SpeechEditingBaseTask(SpeechBaseTask):
    def __init__(self):
        super().__init__()
        self.dataset_cls = StutterSpeechDataset
        self.sil_ph = self.token_encoder.sil_phonemes()
        self.gst = GST

        self.mcd_dict = {'mcd_total': 0, 'num': 0}

    def build_tts_model(self):
        dict_size = len(self.token_encoder)
        self.model = SpeechEditingBaseTask(dict_size, hparams)

    def run_model(self, sample, infer=False, *args, **kwargs):
        txt_tokens = sample['txt_tokens']  # [B, T_t]
        spk_embed = sample.get('spk_embed')
        spk_id = sample.get('spk_ids')
        if not infer:
            target = sample['mels']  # [B, T_s, 80]
            mel2ph = sample['mel2ph']  # [B, T_s]
            f0 = sample.get('f0')
            uv = sample.get('uv')
            output = self.model(txt_tokens, mel2ph=mel2ph, spk_embed=spk_embed, spk_id=spk_id,
                                f0=f0, uv=uv, infer=False)
            losses = {}
            self.add_mel_loss(output['mel_out'], target, losses)
            self.add_dur_loss(output['dur'], mel2ph, txt_tokens, losses=losses)
            if hparams['use_pitch_embed']:
                self.add_pitch_loss(output, sample, losses)
            return losses, output
        else:
            use_gt_dur = kwargs.get('infer_use_gt_dur', hparams['use_gt_dur'])
            use_gt_f0 = kwargs.get('infer_use_gt_f0', hparams['use_gt_f0'])
            mel2ph, uv, f0 = None, None, None
            if use_gt_dur:
                mel2ph = sample['mel2ph']
            if use_gt_f0:
                f0 = sample['f0']
                uv = sample['uv']
            output = self.model(txt_tokens, mel2ph=mel2ph, spk_embed=spk_embed, spk_id=spk_id,
                                f0=f0, uv=uv, infer=True)
            return output

    def add_dur_loss(self, dur_pred, mel2ph, txt_tokens, losses=None):
        """

        :param dur_pred: [B, T], float, log scale
        :param mel2ph: [B, T]
        :param txt_tokens: [B, T]
        :param losses:
        :return:
        """
        B, T = txt_tokens.shape
        nonpadding = (txt_tokens != 0).float()
        dur_gt = mel2token_to_dur(mel2ph, T).float() * nonpadding
        is_sil = torch.zeros_like(txt_tokens).bool()
        for p in self.sil_ph:
            is_sil = is_sil | (txt_tokens == self.token_encoder.encode(p)[0])
        is_sil = is_sil.float()  # [B, T_txt]
        losses['pdur'] = F.mse_loss((dur_pred + 1).log(), (dur_gt + 1).log(), reduction='none')
        losses['pdur'] = (losses['pdur'] * nonpadding).sum() / nonpadding.sum()
        losses['pdur'] = losses['pdur'] * hparams['lambda_ph_dur']
        # use linear scale for sentence and word duration
        if hparams['lambda_word_dur'] > 0:
            word_id = (is_sil.cumsum(-1) * (1 - is_sil)).long()
            word_dur_p = dur_pred.new_zeros([B, word_id.max() + 1]).scatter_add(1, word_id, dur_pred)[:, 1:]
            word_dur_g = dur_gt.new_zeros([B, word_id.max() + 1]).scatter_add(1, word_id, dur_gt)[:, 1:]
            wdur_loss = F.mse_loss((word_dur_p + 1).log(), (word_dur_g + 1).log(), reduction='none')
            word_nonpadding = (word_dur_g > 0).float()
            wdur_loss = (wdur_loss * word_nonpadding).sum() / word_nonpadding.sum()
            losses['wdur'] = wdur_loss * hparams['lambda_word_dur']
        if hparams['lambda_sent_dur'] > 0:
            sent_dur_p = dur_pred.sum(-1)
            sent_dur_g = dur_gt.sum(-1)
            sdur_loss = F.mse_loss((sent_dur_p + 1).log(), (sent_dur_g + 1).log(), reduction='mean')
            losses['sdur'] = sdur_loss.mean() * hparams['lambda_sent_dur']

    def add_pitch_loss(self, output, sample, losses):
        mel2ph = sample['mel2ph']  # [B, T_s]
        f0 = sample['f0']
        uv = sample['uv']
        nonpadding = (mel2ph != 0).float() if hparams['pitch_type'] == 'frame' \
            else (sample['txt_tokens'] != 0).float()
        p_pred = output['pitch_pred']
        assert p_pred[..., 0].shape == f0.shape
        if hparams['use_uv'] and hparams['pitch_type'] == 'frame':
            assert p_pred[..., 1].shape == uv.shape, (p_pred.shape, uv.shape)
            losses['uv'] = (F.binary_cross_entropy_with_logits(
                p_pred[:, :, 1], uv, reduction='none') * nonpadding).sum() \
                           / nonpadding.sum() * hparams['lambda_uv']
            nonpadding = nonpadding * (uv == 0).float()
        f0_pred = p_pred[:, :, 0]
        losses['f0'] = (F.l1_loss(f0_pred, f0, reduction='none') * nonpadding).sum() \
                       / nonpadding.sum() * hparams['lambda_f0']

    def add_conn_loss(self, output, time_mel_masks, target, sample, losses, postfix=''):
        losses[f'conn'] = getattr(self, f'conn_loss')(output, time_mel_masks, target)

    def add_perceptual_loss(self, output, time_mel_masks, target, losses, postfix=''):
        losses[f'perceptual'] = getattr(self, f'perceptual_loss')(output, time_mel_masks, target)

    def conn_loss(self, output, time_mel_masks, target, postfix=''):
        intervals = []
        conn_loss = torch.tensor(0.005).cuda()
        for i in range(output.shape[0]):
            nonzero_indices = torch.nonzero(time_mel_masks[i])
            modi_left = nonzero_indices[0, :][0]
            modi_right = nonzero_indices[-1, :][0]
            # 计算第左帧的方差
            variance_output_left = torch.var(output[i, modi_left, :], axis=0)
            # 计算左-1帧的方差
            variance_output_left1 = torch.var(output[i, modi_left - 1, :], axis=0)
            # 计算第右帧的方差
            variance_output_right = torch.var(output[i, modi_right, :], axis=0)
            # 计算第左帧的方差
            variance_target_left = torch.var(target[i, modi_left, :], axis=0)
            # 计算左-1帧的方差
            variance_target_left1 = torch.var(target[i, modi_left - 1, :], axis=0)
            # 计算第右帧的方差
            variance_target_right = torch.var(target[i, modi_right, :], axis=0)
            output_conn_left = variance_output_left1 - variance_output_left
            target_conn_left = variance_target_left1 - variance_target_left
            conn_left_loss = F.mse_loss(output_conn_left, target_conn_left, reduction='none')
            # weights = weights_nonzero_speech(target_conn_left)
            # conn_left_loss = (mse_loss_left * weights).sum() / weights.sum()
            if modi_right != output.size(1) - 1:
                # 计算右+1帧的方差
                variance_output_right1 = torch.var(output[i, modi_right + 1, :], axis=0)
                # 计算右+1帧的方差
                variance_target_right1 = torch.var(target[i, modi_right + 1, :], axis=0)
                output_conn_right = variance_output_right1 - variance_output_right
                target_conn_right = variance_target_right1 - variance_target_right
                conn_loss_right = F.mse_loss(output_conn_right, target_conn_right, reduction='none')
                # weights = weights_nonzero_speech(target_conn_right)
                # conn_loss_right = (mse_loss_right * weights).sum() / weights.sum()
                conn_loss += (conn_left_loss + conn_loss_right).cuda()

            else:
                conn_loss += conn_left_loss.cuda()
        conn_loss = conn_loss / output.shape[0]
        return conn_loss

    def cconn_loss(self, output, time_mel_masks, target):
        conn_loss = torch.tensor(0.005).cuda()
        for i in range(output.shape[0]):
            nonzero_indices = torch.nonzero(time_mel_masks[i])
            modi_left = nonzero_indices[0, :][0]
            modi_right = nonzero_indices[-1, :][0]
            conn_left_loss = F.mse_loss(output[i, modi_left, :] - output[i, modi_left - 1, :],
                                        target[i, modi_left, :] - target[i, modi_left - 1, :], reduction='none').mean()
            if modi_right != output.size(1) - 1:
                conn_loss_right = F.mse_loss(output[i, modi_right + 1, :] - output[i, modi_right, :],
                                             target[i, modi_right + 1, :] - target[i, modi_right, :],
                                             reduction='none').mean()
                conn_loss += (conn_left_loss + conn_loss_right).cuda()
            else:
                conn_loss += conn_left_loss.cuda()
        conn_loss = conn_loss / output.shape[0]
        return conn_loss

    def perceptual_loss(self, output, time_mel_masks, target):
        preceptual_loss = torch.tensor(0.000001).cuda()
        feature = self.gst(256, 7)
        for i in range(output.shape[0]):
            nonzero_indices = torch.nonzero(time_mel_masks[i, :, :])
            start = nonzero_indices[0, :][0]
            end = nonzero_indices[-1, :][0]
            masked_predict = feature(output[i, start:end, :].unsqueeze(0).cpu())[0].squeeze(0)
            out_target = feature(target[i].unsqueeze(0).cpu())[0].squeeze(0)
            preceptual_loss += torch.mean(F.mse_loss(masked_predict, out_target, reduction='none').cuda())
        preceptual_loss = preceptual_loss / output.size(0)
        return preceptual_loss

    def save_valid_result(self, sample, batch_idx, model_out):
        sr = hparams['audio_sample_rate']
        f0_gt = None
        mel_out = model_out['mel_out']
        if sample.get('f0') is not None:
            f0_gt = denorm_f0(sample['f0'][0].cpu(), sample['uv'][0].cpu())
        self.plot_mel(batch_idx, sample['mels'], mel_out, f0s=f0_gt)
        if self.global_step > 0:
            wav_pred = self.vocoder.spec2wav(mel_out[0].cpu(), f0=f0_gt)
            self.logger.add_audio(f'wav_val_{batch_idx}', wav_pred, self.global_step, sr)
            # with gt duration
            model_out = self.run_model(sample, infer=True, infer_use_gt_dur=True)
            dur_info = self.get_plot_dur_info(sample, model_out)
            del dur_info['dur_pred']
            wav_pred = self.vocoder.spec2wav(model_out['mel_out'][0].cpu(), f0=f0_gt)
            self.logger.add_audio(f'wav_gdur_{batch_idx}', wav_pred, self.global_step, sr)
            self.plot_mel(batch_idx, sample['mels'], model_out['mel_out'][0], f'mel_gdur_{batch_idx}',
                          dur_info=dur_info, f0s=f0_gt)

            # with pred duration
            if not hparams['use_gt_dur']:
                model_out = self.run_model(sample, infer=True, infer_use_gt_dur=False)
                dur_info = self.get_plot_dur_info(sample, model_out)
                self.plot_mel(batch_idx, sample['mels'], model_out['mel_out'][0], f'mel_pdur_{batch_idx}',
                              dur_info=dur_info, f0s=f0_gt)
                wav_pred = self.vocoder.spec2wav(model_out['mel_out'][0].cpu(), f0=f0_gt)
                self.logger.add_audio(f'wav_pdur_{batch_idx}', wav_pred, self.global_step, sr)
        # gt wav
        if self.global_step <= hparams['valid_infer_interval']:
            mel_gt = sample['mels'][0].cpu()
            wav_gt = self.vocoder.spec2wav(mel_gt, f0=f0_gt)
            self.logger.add_audio(f'wav_gt_{batch_idx}', wav_gt, self.global_step, sr)

    def get_plot_dur_info(self, sample, model_out):
        T_txt = sample['txt_tokens'].shape[1]
        dur_gt = mel2token_to_dur(sample['mel2ph'], T_txt)[0]
        dur_pred = model_out['dur'] if 'dur' in model_out else dur_gt
        txt = self.token_encoder.decode(sample['txt_tokens'][0].cpu().numpy())
        txt = txt.split(" ")
        return {'dur_gt': dur_gt, 'dur_pred': dur_pred, 'txt': txt}

    def test_step(self, sample, batch_idx):
        assert sample['txt_tokens'].shape[0] == 1, 'only support batch_size=1 in inference'
        outputs = self.run_model(sample, infer=True)
        text = sample['text'][0]
        item_name = sample['item_name'][0]
        tokens = sample['txt_tokens'][0].cpu().numpy()
        mel_gt = sample['mels'][0].cpu().numpy()
        mel_pred = outputs['mel_out'][0].cpu().numpy()
        mel2ph = sample['mel2ph'][0].cpu().numpy()
        time_mel_masks = sample['time_mel_masks'][0].cpu().numpy()
        mel2ph_pred = None
        str_phs = self.token_encoder.decode(tokens, strip_padding=True)
        base_fn = f'[{batch_idx:06d}][{item_name.replace("%", "_")}][%s]'
        if text is not None:
            base_fn += text.replace(":", "$3A")[:80]
        base_fn = base_fn.replace(' ', '_')
        gen_dir = self.gen_dir
        wav_pred = self.vocoder.spec2wav(mel_pred)
        self.saving_result_pool.add_job(self.save_result, args=[
            wav_pred, mel_pred, base_fn % 'P', gen_dir, str_phs, mel2ph_pred])
        mel_pred_seg = mel_pred[time_mel_masks == 1]
        wav_pred_seg = self.vocoder.spec2wav(mel_pred_seg)
        self.saving_result_pool.add_job(self.save_result, args=[
            wav_pred_seg, mel_pred_seg, base_fn % 'P_SEG', gen_dir, None, None])
        if hparams['save_gt']:
            wav_gt = self.vocoder.spec2wav(mel_gt)
            self.saving_result_pool.add_job(self.save_result, args=[
                wav_gt, mel_gt, base_fn % 'G', gen_dir, str_phs, mel2ph])
            mel_gt_seg = mel_gt[time_mel_masks == 1]
            wav_gt_seg = self.vocoder.spec2wav(mel_gt_seg)
            self.saving_result_pool.add_job(self.save_result, args=[
                wav_gt_seg, mel_gt_seg, base_fn % 'G_SEG', gen_dir, None, None])

        # print(f"Pred_shape: {mel_pred.shape}, gt_shape: {mel_gt.shape}")
        return {
            'item_name': item_name,
            'text': text,
            'ph_tokens': self.token_encoder.decode(tokens.tolist()),
            'wav_fn_pred': base_fn % 'P',
            'wav_fn_gt': base_fn % 'G',
            'wav_fn_orig': sample['wav_fn'][0],
        }
