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

class ConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=(3, 3), stride=(2, 1), padding=(1, 1)):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
        self.bn = nn.BatchNorm2d(out_channels)
        self.act = nn.LeakyReLU(0.2)

    def forward(self, x):
        return self.act(self.bn(self.conv(x)))

class DeconvBlock(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=(3, 3), stride=(2, 1), padding=(1, 1), output_padding=(1, 0)):
        super().__init__()
        self.deconv = nn.ConvTranspose2d(
            in_channels, out_channels, kernel_size,
            stride=stride, padding=padding, output_padding=output_padding
        )
        self.bn = nn.BatchNorm2d(out_channels)
        self.act = nn.LeakyReLU(0.2)

    def forward(self, x):
        return self.act(self.bn(self.deconv(x)))

class ComplexCRNDenoiser(nn.Module):
    """
    Complex Convolutional Recurrent Network (CRN) for Audio Noise Reduction.
    Operates on complex STFT representations (Real + Imaginary) to preserve both
    magnitude and phase, completely eliminating phase distortion.
    """
    def __init__(
        self,
        n_fft=512,
        hop_length=128,
        win_length=512,
        channels=[16, 32, 64, 128, 256],
        rnn_hidden=256,
        rnn_layers=2
    ):
        super().__init__()
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.win_length = win_length
        self.register_buffer("window", torch.hann_window(win_length))

        # Encoder (Frequency downsampling by 2 at each step)
        # Input shape: (B, 2, 256, T) -> channels[0] -> channels[1] -> ... -> channels[4]
        # F: 256 -> 128 -> 64 -> 32 -> 16 -> 8
        self.enc1 = ConvBlock(2, channels[0])
        self.enc2 = ConvBlock(channels[0], channels[1])
        self.enc3 = ConvBlock(channels[1], channels[2])
        self.enc4 = ConvBlock(channels[2], channels[3])
        self.enc5 = ConvBlock(channels[3], channels[4])

        bottleneck_f = 8
        self.rnn_in_dim = channels[4] * bottleneck_f
        self.rnn = nn.GRU(
            input_size=self.rnn_in_dim,
            hidden_size=rnn_hidden,
            num_layers=rnn_layers,
            batch_first=True,
            bidirectional=True
        )
        self.rnn_fc = nn.Linear(rnn_hidden * 2, self.rnn_in_dim)

        # Decoder with Skip-Connections
        self.dec5 = DeconvBlock(channels[4] * 2, channels[3])
        self.dec4 = DeconvBlock(channels[3] * 2, channels[2])
        self.dec3 = DeconvBlock(channels[2] * 2, channels[1])
        self.dec2 = DeconvBlock(channels[1] * 2, channels[0])
        self.dec1 = nn.ConvTranspose2d(
            channels[0] * 2, 2, kernel_size=(3, 3), stride=(2, 1),
            padding=(1, 1), output_padding=(1, 0)
        )

    def forward_stft(self, x):
        # x is (B, 2, 256, T)
        e1 = self.enc1(x)  # (B, C0, 128, T)
        e2 = self.enc2(e1) # (B, C1, 64, T)
        e3 = self.enc3(e2) # (B, C2, 32, T)
        e4 = self.enc4(e3) # (B, C3, 16, T)
        e5 = self.enc5(e4) # (B, C4, 8, T)

        B, C, F_dim, T_dim = e5.shape
        # Reshape for GRU: (B, T, C * F)
        rnn_in = e5.permute(0, 3, 1, 2).reshape(B, T_dim, C * F_dim)
        rnn_out, _ = self.rnn(rnn_in)
        rnn_proj = self.rnn_fc(rnn_out)
        # Reshape back to (B, C, F, T)
        d_in = rnn_proj.reshape(B, T_dim, C, F_dim).permute(0, 2, 3, 1)

        # Skip connection decoding
        d5 = self.dec5(torch.cat([d_in, e5], dim=1))
        d4 = self.dec4(torch.cat([d5, e4], dim=1))
        d3 = self.dec3(torch.cat([d4, e3], dim=1))
        d2 = self.dec2(torch.cat([d3, e2], dim=1))
        mask = self.dec1(torch.cat([d2, e1], dim=1)) # (B, 2, 256, T)

        # Complex mask bounded via tanh
        mask = torch.tanh(mask)
        return mask

    def forward(self, noisy_wav):
        """
        Takes raw 1D audio waveform: (Batch, AudioSamples)
        Returns cleaned 1D audio waveform: (Batch, AudioSamples)
        """
        orig_len = noisy_wav.shape[-1]
        if self.window.device != noisy_wav.device:
            self.window = self.window.to(noisy_wav.device)

        # 1. STFT
        stft_res = torch.stft(
            noisy_wav,
            n_fft=self.n_fft,
            hop_length=self.hop_length,
            win_length=self.win_length,
            window=self.window,
            return_complex=True
        ) # (B, 257, T)

        # Separate Real and Imaginary components
        real = stft_res.real
        imag = stft_res.imag

        # Drop highest frequency bin (257 -> 256) for power-of-2 conv pooling
        real_256 = real[:, :256, :]
        imag_256 = imag[:, :256, :]
        stft_in = torch.stack([real_256, imag_256], dim=1) # (B, 2, 256, T)

        # 2. Predict Complex Mask (Mr, Mi)
        mask = self.forward_stft(stft_in)
        mask_r = mask[:, 0, :, :]
        mask_i = mask[:, 1, :, :]

        # 3. Complex multiplication: (Yr + j Yi) = (Xr + j Xi) * (Mr + j Mi)
        # Yr = Xr * Mr - Xi * Mi
        # Yi = Xr * Mi + Xi * Mr
        clean_real_256 = real_256 * mask_r - imag_256 * mask_i
        clean_imag_256 = real_256 * mask_i + imag_256 * mask_r

        # Pad 257th frequency back with zeros
        clean_real = F.pad(clean_real_256, (0, 0, 0, 1))
        clean_imag = F.pad(clean_imag_256, (0, 0, 0, 1))

        clean_stft = torch.complex(clean_real, clean_imag)

        # 4. Inverse STFT to time domain
        clean_wav = torch.istft(
            clean_stft,
            n_fft=self.n_fft,
            hop_length=self.hop_length,
            win_length=self.win_length,
            window=self.window,
            length=orig_len
        )
        return clean_wav