import torch
import torch.nn as nn
from torch.nn import functional as F

from modules.commons.common_layers import (
    NormalInitEmbedding as Embedding,
    SinusoidalPosEmb,
    AdamWLinear,
)
from modules.fastspeech.tts_modules import FastSpeech2Encoder, mel2ph_to_dur, StretchRegulator
from utils.hparams import hparams
from utils.phoneme_utils import PAD_INDEX


class FastSpeech2Acoustic(nn.Module):
    def __init__(self, vocab_size):
        super().__init__()
        self.txt_embed = Embedding(vocab_size, hparams['hidden_size'], PAD_INDEX)
        self.use_lang_id = hparams.get('use_lang_id', False)
        if self.use_lang_id:
            self.lang_embed = Embedding(hparams['num_lang'] + 1, hparams['hidden_size'], padding_idx=0)

        self.use_stretch_embed = hparams.get('use_stretch_embed', False)
        if self.use_stretch_embed:
            self.sr = StretchRegulator()
            self.stretch_embed = nn.Sequential(
                SinusoidalPosEmb(hparams['hidden_size']),
                nn.Linear(hparams['hidden_size'], hparams['hidden_size'] * 4),
                nn.GELU(),
                nn.Linear(hparams['hidden_size'] * 4, hparams['hidden_size']),
            )
            self.stretch_embed_rnn = nn.GRU(hparams['hidden_size'], hparams['hidden_size'], 1, batch_first=True)
            self._stretch_embed_rnn_flattened = False

        self.dur_embed = AdamWLinear(1, hparams['hidden_size'])
        self.use_mix_ln = hparams.get('use_mix_ln', False)
        if self.use_mix_ln:
            self.mix_ln_layer = hparams['mix_ln_layer']
        else:
            self.mix_ln_layer = []
        self.encoder = FastSpeech2Encoder(
            hidden_size=hparams['hidden_size'], num_layers=hparams['enc_layers'],
            ffn_kernel_size=hparams['enc_ffn_kernel_size'], ffn_act=hparams['ffn_act'],
            dropout=hparams['dropout'], num_heads=hparams['num_heads'],
            use_pos_embed=hparams['use_pos_embed'], rel_pos=hparams.get('rel_pos', False),
            use_rope=hparams.get('use_rope', False), rope_interleaved=hparams.get('rope_interleaved', True),
            mix_ln_layer=self.mix_ln_layer
        )

        self.pitch_embed = AdamWLinear(1, hparams['hidden_size'])
        self.variance_embed_list = []
        self.use_energy_embed = hparams.get('use_energy_embed', False)
        self.use_breathiness_embed = hparams.get('use_breathiness_embed', False)
        self.use_voicing_embed = hparams.get('use_voicing_embed', False)
        self.use_tension_embed = hparams.get('use_tension_embed', False)
        if self.use_energy_embed:
            self.variance_embed_list.append('energy')
        if self.use_breathiness_embed:
            self.variance_embed_list.append('breathiness')
        if self.use_voicing_embed:
            self.variance_embed_list.append('voicing')
        if self.use_tension_embed:
            self.variance_embed_list.append('tension')

        self.use_variance_embeds = len(self.variance_embed_list) > 0
        if self.use_variance_embeds:
            self.variance_embeds = nn.ModuleDict({
                v_name: AdamWLinear(1, hparams['hidden_size'])
                for v_name in self.variance_embed_list
            })

        self.use_variance_scaling = hparams.get('use_variance_scaling', False)
        if self.use_variance_scaling:
            self.variance_scaling_factor = {
                'energy': 1. / 96,  # 96 dB — max dynamic range of 16-bit audio
                'breathiness': 1. / 96,
                'voicing': 1. / 96,
                'tension': 0.1,  # 1 / 10; tension logits are roughly [-10, 10]
                'key_shift': 1. / 12,  # one octave — max key shift in most editors
                'speed': 1.
            }
        else:
            self.variance_scaling_factor = {
                'energy': 1.,
                'breathiness': 1.,
                'voicing': 1.,
                'tension': 1.,
                'key_shift': 1.,
                'speed': 1.
            }

        self.use_key_shift_embed = hparams.get('use_key_shift_embed', False)
        if self.use_key_shift_embed:
            self.key_shift_embed = AdamWLinear(1, hparams['hidden_size'])

        self.use_speed_embed = hparams.get('use_speed_embed', False)
        if self.use_speed_embed:
            self.speed_embed = AdamWLinear(1, hparams['hidden_size'])

        self.use_spk_id = hparams['use_spk_id']
        if self.use_spk_id:
            self.spk_embed = Embedding(hparams['num_spk'], hparams['hidden_size'])

    def forward_variance_embedding(self, condition, key_shift=None, speed=None, **variances):
        if self.use_variance_embeds:
            variance_embeds = torch.stack([
                self.variance_embeds[v_name](variances[v_name][:, :, None] * self.variance_scaling_factor[v_name])
                for v_name in self.variance_embed_list
            ], dim=-1).sum(-1)
            condition += variance_embeds

        if self.use_key_shift_embed:
            key_shift_embed = self.key_shift_embed(key_shift[:, :, None] * self.variance_scaling_factor['key_shift'])
            condition += key_shift_embed

        if self.use_speed_embed:
            speed_embed = self.speed_embed(speed[:, :, None] * self.variance_scaling_factor['speed'])
            condition += speed_embed

        return condition

    def forward(
            self, txt_tokens, mel2ph, f0,
            key_shift=None, speed=None,
            spk_embed_id=None, languages=None,
            **kwargs
    ):
        spk_embed = None
        if self.use_spk_id:
            spk_mix_embed = kwargs.get('spk_mix_embed')
            if spk_mix_embed is not None:
                spk_embed = spk_mix_embed
            else:
                spk_embed = self.spk_embed(spk_embed_id)[:, None, :]
        txt_embed = self.txt_embed(txt_tokens)
        dur = mel2ph_to_dur(mel2ph, txt_tokens.shape[1])
        if self.use_variance_scaling:
            dur_embed = self.dur_embed(torch.log(1 + dur[:, :, None].float()))
        else:
            dur_embed = self.dur_embed(dur[:, :, None].float())
        if self.use_lang_id:
            lang_embed = self.lang_embed(languages)
            extra_embed = dur_embed + lang_embed
        else:
            extra_embed = dur_embed
        encoder_out = self.encoder(txt_embed, extra_embed, txt_tokens == 0, spk_embed)

        encoder_out = F.pad(encoder_out, [0, 0, 1, 0])
        mel2ph_ = mel2ph[..., None].repeat([1, 1, encoder_out.shape[-1]])
        condition = torch.gather(encoder_out, 1, mel2ph_)

        if self.use_stretch_embed:
            stretch = torch.round(1000 * self.sr(mel2ph, dur))
            if self.training and stretch.numel() > 1000:
                # construct a phoneme stretching index lookup table with a total of 1001 indexes (0~1000)
                table = self.stretch_embed(torch.arange(0, 1001, device=stretch.device))
                stretch_embed = torch.index_select(table, 0, stretch.view(-1).long()).view_as(condition)
            else:
                stretch_embed = self.stretch_embed(stretch)
            condition += stretch_embed
            # flatten_parameters fuses the GRU weights into a contiguous buffer for cuDNN.
            # It only needs to happen once after weight init, device change, or load_state_dict.
            # We guard with a flag to avoid the redundant call on every forward.
            # Limitation: the flag lives on this module and is invisible to PyTorch. After
            # load_state_dict() or model.to(device) replaces the GRU weights, the flag stays
            # True and flatten_parameters is skipped — cuDNN will fall back to the slower path.
            # To restore the fast path, reset the flag manually: model._stretch_embed_rnn_flattened = False
            if not self._stretch_embed_rnn_flattened:
                self.stretch_embed_rnn.flatten_parameters()
                self._stretch_embed_rnn_flattened = True
            stretch_embed_rnn_out, _ = self.stretch_embed_rnn(condition)
            condition = condition + stretch_embed_rnn_out

        if self.use_spk_id:
            condition += spk_embed

        f0_mel = (1 + f0 / 700).log()
        pitch_embed = self.pitch_embed(f0_mel[:, :, None])
        condition += pitch_embed

        condition = self.forward_variance_embedding(
            condition, key_shift=key_shift, speed=speed, **kwargs
        )

        return condition
