import matplotlib
import torch
import torch.distributions
import torch.optim
import torch.utils.data

import utils
import utils.infer_utils
from basics.base_dataset import BaseDataset
from basics.base_task import BaseTask
from basics.base_vocoder import BaseVocoder
from modules.aux_decoder import build_aux_loss
from modules.losses import DiffusionLoss, RectifiedFlowLoss
from modules.toplevel import DiffSingerAcoustic, ShallowDiffusionOutput
from modules.vocoders.registry import get_vocoder_cls
from utils.hparams import hparams
from utils.plot import spec_to_figure

matplotlib.use('Agg')


class AcousticDataset(BaseDataset):
    def __init__(self, prefix, preload=False):
        super(AcousticDataset, self).__init__(prefix, hparams['dataset_size_key'], preload)
        self.required_variances = {}  # key: variance name, value: padding value
        if hparams['use_energy_embed']:
            self.required_variances['energy'] = 0.0
        if hparams['use_breathiness_embed']:
            self.required_variances['breathiness'] = 0.0
        if hparams['use_voicing_embed']:
            self.required_variances['voicing'] = 0.0
        if hparams['use_tension_embed']:
            self.required_variances['tension'] = 0.0

        self.need_key_shift = hparams['use_key_shift_embed']
        self.need_speed = hparams['use_speed_embed']
        self.need_spk_id = hparams['use_spk_id']
        self.need_lang_id = hparams['use_lang_id']

    def collater(self, samples):
        batch = super().collater(samples)
        if batch['size'] == 0:
            return batch

        tokens = utils.collate_nd([s['tokens'] for s in samples], 0)
        f0 = utils.collate_nd([s['f0'] for s in samples], 0.0)
        mel2ph = utils.collate_nd([s['mel2ph'] for s in samples], 0)
        mel = utils.collate_nd([s['mel'] for s in samples], 0.0)
        batch.update({
            'tokens': tokens,
            'mel2ph': mel2ph,
            'mel': mel,
            'f0': f0,
        })
        for v_name, v_pad in self.required_variances.items():
            batch[v_name] = utils.collate_nd([s[v_name] for s in samples], v_pad)
        if self.need_key_shift:
            batch['key_shift'] = torch.FloatTensor([s['key_shift'] for s in samples])[:, None]
        if self.need_speed:
            batch['speed'] = torch.FloatTensor([s['speed'] for s in samples])[:, None]
        if self.need_spk_id:
            spk_ids = torch.LongTensor([s['spk_id'] for s in samples])
            batch['spk_ids'] = spk_ids
        if self.need_lang_id:
            languages = utils.collate_nd([s['languages'] for s in samples], 0)
            batch['languages'] = languages
        return batch


class AcousticTask(BaseTask):
    def __init__(self):
        super().__init__()
        self.dataset_cls = AcousticDataset
        self.diffusion_type = hparams['diffusion_type']
        assert self.diffusion_type in ['ddpm', 'reflow'], f"Unknown diffusion type: {self.diffusion_type}"
        self.use_shallow_diffusion = hparams['use_shallow_diffusion']
        if self.use_shallow_diffusion:
            self.shallow_args = hparams['shallow_diffusion_args']
            self.train_aux_decoder = self.shallow_args['train_aux_decoder']
            self.train_diffusion = self.shallow_args['train_diffusion']

        self.use_vocoder = hparams['infer'] or hparams['val_with_vocoder']
        if self.use_vocoder:
            self.vocoder: BaseVocoder = get_vocoder_cls(hparams)()
        self.logged_gt_wav = set()
        self.required_variances = []
        if hparams['use_energy_embed']:
            self.required_variances.append('energy')
        if hparams['use_breathiness_embed']:
            self.required_variances.append('breathiness')
        if hparams['use_voicing_embed']:
            self.required_variances.append('voicing')
        if hparams['use_tension_embed']:
            self.required_variances.append('tension')
        super()._finish_init()

    def _build_model(self):
        return DiffSingerAcoustic(
            vocab_size=len(self.phoneme_dictionary),
            out_dims=hparams['audio_num_mel_bins']
        )

    # noinspection PyAttributeOutsideInit
    def build_losses_and_metrics(self):
        if self.use_shallow_diffusion:
            self.aux_mel_loss = build_aux_loss(self.shallow_args['aux_decoder_arch'])
            self.lambda_aux_mel_loss = hparams['lambda_aux_mel_loss']
            self.register_validation_loss('aux_mel_loss')
        if self.diffusion_type == 'ddpm':
            self.mel_loss = DiffusionLoss(loss_type=hparams['main_loss_type'])
        elif self.diffusion_type == 'reflow':
            self.mel_loss = RectifiedFlowLoss(
                loss_type=hparams['main_loss_type'], log_norm=hparams['main_loss_log_norm']
            )
        else:
            raise ValueError(f"Unknown diffusion type: {self.diffusion_type}")
        self.register_validation_loss('mel_loss')

    def run_model(self, sample, infer=False):
        txt_tokens = sample['tokens']  # [B, T_ph]
        target = sample['mel']  # [B, T_s, M]
        mel2ph = sample['mel2ph']  # [B, T_s]
        f0 = sample['f0']
        variances = {
            v_name: sample[v_name]
            for v_name in self.required_variances
        }
        key_shift = sample.get('key_shift')
        speed = sample.get('speed')

        if hparams['use_spk_id']:
            spk_embed_id = sample['spk_ids']
        else:
            spk_embed_id = None
        if hparams['use_lang_id']:
            languages = sample['languages']
        else:
            languages = None
        output: ShallowDiffusionOutput = self.model(
            txt_tokens, mel2ph=mel2ph, f0=f0, **variances,
            key_shift=key_shift, speed=speed,
            spk_embed_id=spk_embed_id, languages=languages,
            gt_mel=target, infer=infer
        )

        if infer:
            return output
        else:
            losses = {}

            if output.aux_out is not None:
                aux_out = output.aux_out
                norm_gt = self.model.aux_decoder.norm_spec(target)
                aux_mel_loss = self.lambda_aux_mel_loss * self.aux_mel_loss(aux_out, norm_gt)
                losses['aux_mel_loss'] = aux_mel_loss

            non_padding = (mel2ph > 0).unsqueeze(-1).float()
            if output.diff_out is not None:
                if self.diffusion_type == 'ddpm':
                    x_recon, x_noise = output.diff_out
                    mel_loss = self.mel_loss(x_recon, x_noise, non_padding=non_padding)
                elif self.diffusion_type == 'reflow':
                    v_pred, v_gt, t = output.diff_out
                    mel_loss = self.mel_loss(v_pred, v_gt, t=t, non_padding=non_padding)
                else:
                    raise ValueError(f"Unknown diffusion type: {self.diffusion_type}")
                losses['mel_loss'] = mel_loss

            return losses

    def on_train_start(self):
        if self.use_vocoder and self.vocoder.get_device() != self.device:
            self.vocoder.to_device(self.device)

    def _on_validation_start(self):
        if self.use_vocoder and self.vocoder.get_device() != self.device:
            self.vocoder.to_device(self.device)

    def _validation_step(self, sample, batch_idx):
        losses = self.run_model(sample, infer=False)
        if sample['size'] > 0 and min(sample['indices']) < hparams['num_valid_plots']:
            mel_out: ShallowDiffusionOutput = self.run_model(sample, infer=True)
            for i in range(len(sample['indices'])):
                data_idx = sample['indices'][i].item()
                if data_idx < hparams['num_valid_plots']:
                    if self.use_vocoder:
                        self.plot_wav(
                            data_idx, sample['mel'][i],
                            mel_out.aux_out[i] if mel_out.aux_out is not None else None,
                            mel_out.diff_out[i],
                            sample['f0'][i]
                        )
                    if mel_out.aux_out is not None:
                        self.plot_mel(data_idx, sample['mel'][i], mel_out.aux_out[i], 'auxmel')
                    if mel_out.diff_out is not None:
                        self.plot_mel(data_idx, sample['mel'][i], mel_out.diff_out[i], 'diffmel')
        return losses, sample['size']

    ############
    # validation plots
    ############
    def plot_wav(self, data_idx, gt_mel, aux_mel, diff_mel, f0):
        f0_len = self.valid_dataset.metadata['f0'][data_idx]
        mel_len = self.valid_dataset.metadata['mel'][data_idx]
        gt_mel = gt_mel[:mel_len].unsqueeze(0)
        if aux_mel is not None:
            aux_mel = aux_mel[:mel_len].unsqueeze(0)
        if diff_mel is not None:
            diff_mel = diff_mel[:mel_len].unsqueeze(0)
        f0 = f0[:f0_len].unsqueeze(0)
        if data_idx not in self.logged_gt_wav:
            gt_wav = self.vocoder.spec2wav_torch(gt_mel, f0=f0)
            self.logger.all_rank_experiment.add_audio(
                f'gt_{data_idx}', gt_wav,
                sample_rate=hparams['audio_sample_rate'],
                global_step=self.global_step
            )
            self.logged_gt_wav.add(data_idx)
        if aux_mel is not None:
            aux_wav = self.vocoder.spec2wav_torch(aux_mel, f0=f0)
            self.logger.all_rank_experiment.add_audio(
                f'aux_{data_idx}', aux_wav,
                sample_rate=hparams['audio_sample_rate'],
                global_step=self.global_step
            )
        if diff_mel is not None:
            diff_wav = self.vocoder.spec2wav_torch(diff_mel, f0=f0)
            self.logger.all_rank_experiment.add_audio(
                f'diff_{data_idx}', diff_wav,
                sample_rate=hparams['audio_sample_rate'],
                global_step=self.global_step
            )

    def plot_mel(self, data_idx, gt_spec, out_spec, name_prefix='mel'):
        vmin = hparams['mel_vmin']
        vmax = hparams['mel_vmax']
        mel_len = self.valid_dataset.metadata['mel'][data_idx]
        spec_cat = torch.cat([(out_spec - gt_spec).abs() + vmin, gt_spec, out_spec], -1)
        title_text = f"{self.valid_dataset.metadata['spk_names'][data_idx]} - {self.valid_dataset.metadata['names'][data_idx]}"
        self.logger.all_rank_experiment.add_figure(f'{name_prefix}_{data_idx}', spec_to_figure(
            spec_cat[:mel_len], vmin, vmax, title_text
        ), global_step=self.global_step)
