from typing import Any, Dict, Optional, Tuple

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

from speechtokenizer.modules.seanet import SEANetEncoder
from util import chunks_to_bits


class SkipGatedBlock(nn.Module):
    def __init__(self, c_in, c_out, kernel_size, stride, padding):
        super().__init__()
        self.conv = nn.Conv2d(c_in, c_out, kernel_size=kernel_size, stride=stride, padding=padding, bias=True)
        self.gate = nn.Conv2d(c_in, c_out, kernel_size=kernel_size, stride=stride, padding=padding, bias=True)
        self.skip_connection = c_in == c_out and (stride == 1 if isinstance(stride, int) else all(s == 1 for s in stride))
        self.has_skip_proj = (c_in == c_out) and not self.skip_connection

    def forward(self, x):
        conv_output = self.conv(x)
        gated_output = torch.sigmoid(self.gate(x))
        output = conv_output * gated_output
        if self.skip_connection:
            output.add_(x)
        elif self.has_skip_proj:
            skip_output = F.adaptive_avg_pool2d(x, output.shape[2:])
            output.add_(skip_output)
        return output


class WatermarkDecoder(nn.Module):
    def __init__(self, input_channels, nbits, nchunk_size, hidden_dim, d_model):
        super().__init__()
        self.nchunk_size = nchunk_size
        assert nbits % nchunk_size == 0
        self.nbits = nbits
        self.d_model = d_model
        self.nchunks = nbits // nchunk_size

        self.proj = nn.Linear(input_channels, d_model)
        self.detect_encoder = nn.Sequential(
            nn.Conv1d(d_model, d_model, kernel_size=3, padding=1), nn.GELU(),
            nn.Conv1d(d_model, d_model, kernel_size=3, padding=1), nn.GELU(),
            nn.Conv1d(d_model, d_model, kernel_size=3, padding=1), nn.GELU(),
            nn.Conv1d(d_model, d_model, kernel_size=3, padding=1), nn.GELU(),
            nn.Conv1d(d_model, 1, kernel_size=3, padding=1),
        )

        self.message_encoder = nn.ModuleList([
            nn.Conv2d(1, 16, kernel_size=(5, 3), stride=(2, 1), padding=(0, 1)), nn.GELU(),
            nn.Conv2d(16, 32, kernel_size=(5, 3), stride=(2, 1), padding=(0, 1)), nn.GELU(),
            nn.Conv2d(32, 64, kernel_size=(5, 3), stride=(2, 1), padding=(0, 1)), nn.GELU(),
            nn.Conv2d(64, 128, kernel_size=(5, 3), stride=(2, 1), padding=(0, 1)), nn.GELU(),
            nn.Conv2d(128, 256, kernel_size=(5, 3), stride=(2, 1), padding=(0, 1)), nn.GELU(),
            nn.Conv2d(256, 512, kernel_size=(5, 3), stride=(1, 1), padding=(0, 1)),
        ])

        self.d_proj_16 = nn.Linear(126, 1)
        self.d_proj_32 = nn.Linear(61, 1)
        self.d_proj_64 = nn.Linear(29, 1)
        self.d_proj_128 = nn.Linear(13, 1)
        self.d_proj_256 = nn.Linear(5, 1)
        self.d_projs = [self.d_proj_16, self.d_proj_32, self.d_proj_64, self.d_proj_128, self.d_proj_256]
        self.message_head = nn.Sequential(
            nn.Linear(512 + 256 + 128 + 64 + 32 + 16, 1024), nn.GELU(), nn.Linear(1024, 2**nchunk_size * self.nchunks),
        )

    def forward(self, x):
        batch_size = x.shape[0]
        x = self.proj(x.transpose(-1, -2)).transpose(-1, -2)
        frame_logits = self.detect_encoder(x).squeeze(1)
        temperature = 2.0
        frame_weights = torch.sigmoid(frame_logits / temperature).clamp(min=0.01, max=1.0)

        current_feat = x.unsqueeze(1)
        multi_scale_features = []
        d_proj_idx = 0
        for layer in self.message_encoder:
            current_feat = layer(current_feat)
            if isinstance(layer, nn.Conv2d):
                B, C, D, T = current_feat.shape
                if D != 1:
                    feat_trans = current_feat.permute(0, 1, 3, 2)
                    feat_flat = feat_trans.reshape(-1, D)
                    d_proj = self.d_projs[d_proj_idx]
                    feat_proj = d_proj(feat_flat).reshape(B, C, T, 1).squeeze(-1)
                else:
                    feat_proj = current_feat.squeeze(2)
                multi_scale_features.append(feat_proj)
                d_proj_idx += 1

        combined_features = torch.cat(multi_scale_features, dim=1).permute(0, 2, 1)
        time_logits = self.message_head(combined_features)
        weighted_logits = time_logits * frame_weights.unsqueeze(-1)
        weight_sum = frame_weights.sum(dim=1, keepdim=True).clamp(min=1e-6)
        avg_logits = weighted_logits.sum(dim=1) / weight_sum
        chunk_logits = avg_logits.reshape(batch_size, self.nchunks, 2 ** self.nchunk_size)
        return frame_logits, chunk_logits

    def detect_watermark(self, x):
        frame_logits, chunk_logits = self.forward(x)
        chunk_probs = F.softmax(chunk_logits, dim=-1)
        chunk_indices = torch.argmax(chunk_probs, dim=-1)
        chunk_values = [chunk_indices[:, i] for i in range(self.nchunks)]
        binary_message = chunks_to_bits(chunk_values, self.nchunk_size)
        return frame_logits, chunk_logits, binary_message


class WatermarkEmbedder(nn.Module):
    def __init__(self, nbits, input_dim, hidden_dim, d_model):
        super().__init__()
        self.nbits = nbits
        self.hidden_dim = hidden_dim
        self.input_dim = input_dim
        self.proj = nn.Linear(input_dim, d_model)
        self.msg_embedding = nn.Embedding(2 * nbits, d_model)
        self.conv_layers = nn.Sequential(
            SkipGatedBlock(1 + 1, hidden_dim, kernel_size=3, stride=1, padding=1),
            SkipGatedBlock(hidden_dim, hidden_dim, kernel_size=3, stride=1, padding=1),
            SkipGatedBlock(hidden_dim, hidden_dim, kernel_size=3, stride=1, padding=1),
            SkipGatedBlock(hidden_dim, hidden_dim, kernel_size=3, stride=1, padding=1),
            SkipGatedBlock(hidden_dim, hidden_dim, kernel_size=3, stride=1, padding=1),
            SkipGatedBlock(hidden_dim, 1, kernel_size=3, stride=1, padding=1),
        )
        self.out_proj = nn.Conv1d(d_model, input_dim, kernel_size=1)

    def embed_bits(self, message):
        idx = message + torch.arange(self.nbits, device=message.device) * 2
        emb = self.msg_embedding(idx)
        return emb.sum(dim=1).unsqueeze(1)

    def forward(self, hidden, msg):
        seq_len = hidden.shape[-1]
        hidden_orig = hidden
        hidden = self.proj(hidden.transpose(-1, -2)).transpose(-1, -2)
        msg_emb = self.embed_bits(msg)
        if hidden.dim() == 3:
            hidden = hidden.unsqueeze(1)
        combined_input = torch.cat([hidden, msg_emb.unsqueeze(-1).expand(-1, -1, -1, seq_len)], dim=1)
        x = self.conv_layers(combined_input)
        output = x.squeeze(1)
        output = self.out_proj(output)
        output = hidden_orig + output
        return output


class WatermarkModel(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.nbits = config.nbits
        self.nfft = config.wm_mb.nfft
        self.sr = config.wm_mb.sr
        self.latent_dim = 1024

        self.encoder = SEANetEncoder(
            n_filters=64,
            dimension=self.latent_dim,
            ratios=[8, 5, 4, 2],
            lstm=2,
            dilation_base=2,
            residual_kernel_size=3,
            n_residual_layers=1,
            activation="ELU",
            bidirectional=True,
        )

        self.embedder = WatermarkEmbedder(
            nbits=config.nbits, input_dim=self.latent_dim, hidden_dim=32, d_model=256
        )
        self.detector = WatermarkDecoder(
            self.latent_dim, config.nbits, nchunk_size=config.wm_mb.nchunk_size, hidden_dim=32, d_model=256
        )

    def decode_watermark(self, x: torch.Tensor) -> Tuple[Any, ...]:
        embedding = self.encoder(x)
        frame_logits, chunk_logits, binary_message = self.detector.detect_watermark(embedding)
        return embedding, (frame_logits, chunk_logits), binary_message

    def forward(self, feat: torch.Tensor, message: Optional[torch.Tensor]) -> Dict[str, torch.Tensor]:
        feat_wm = self.embedder(feat, message)
        return feat_wm


class AudioFusionModel(nn.Module):
    def __init__(self, n_fft=256, hop_length=64, win_length=256, hidden_dim=64, nbits=16):
        super().__init__()
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.win_length = win_length
        self.hidden_dim = hidden_dim
        self.weight_net = nn.Sequential(
            nn.Conv2d(4, hidden_dim, kernel_size=(3, 3), padding=(1, 1)), nn.LeakyReLU(0.1),
            nn.Conv2d(hidden_dim, hidden_dim, kernel_size=(3, 3), padding=(1, 1)), nn.LeakyReLU(0.1),
            nn.Conv2d(hidden_dim, hidden_dim, kernel_size=(3, 3), padding=(1, 1)), nn.LeakyReLU(0.1),
            nn.Conv2d(hidden_dim, 2, kernel_size=(3, 3), padding=(1, 1)), nn.Sigmoid(),
        )

    def forward(self, wav_orig, wav_wm):
        wav_orig = wav_orig.detach()
        stft_orig = torch.stft(
            wav_orig.squeeze(1), self.n_fft, hop_length=self.hop_length, win_length=self.win_length,
            window=torch.hann_window(self.win_length).to(wav_orig.device), return_complex=True
        )
        real_orig, imag_orig = stft_orig.real, stft_orig.imag
        stft_wm = torch.stft(
            wav_wm.squeeze(1), self.n_fft, hop_length=self.hop_length, win_length=self.win_length,
            window=torch.hann_window(self.win_length).to(wav_wm.device), return_complex=True
        )
        real_wm, imag_wm = stft_wm.real, stft_wm.imag
        complex_input = torch.stack([real_orig, imag_orig, real_wm, imag_wm], dim=1)
        alpha_weights = self.weight_net(complex_input)
        alpha_real = alpha_weights[:, 0, :, :]
        alpha_imag = alpha_weights[:, 1, :, :]
        real_fused = real_orig * alpha_real + real_wm * (1 - alpha_real)
        imag_fused = imag_orig * alpha_imag + imag_wm * (1 - alpha_imag)
        final_stft = torch.complex(real_fused, imag_fused)
        wav_fused = torch.istft(
            final_stft.to(wav_orig.device), n_fft=self.n_fft, hop_length=self.hop_length, win_length=self.win_length,
            window=torch.hann_window(self.win_length).to(wav_orig.device), length=wav_orig.shape[-1]
        ).unsqueeze(1)
        return wav_fused