"""Dual-Stream Spectrogram and Physics Neural Network Architecture for Audio Deepfake Detection."""

from typing import Dict, Optional, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F

from src.features import HybridFeatureExtractor, SpectrogramExtractor
from src.physics_features import StackedPhysicsFeatureExtractor


class SpectrogramBranchEncoder(nn.Module):
    """Stream A: 2D Convolutional encoder for Log-Mel spectrogram features."""

    def __init__(self, in_channels: int = 1, embed_dim: int = 128):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, 32, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(32),
            nn.GELU(),
            nn.MaxPool2d((2, 2)),
            nn.Dropout(0.1),

            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.GELU(),
            nn.MaxPool2d((2, 2)),
            nn.Dropout(0.1),

            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(128),
            nn.GELU(),
        )

        # Pool over frequency bins (height), retain temporal variation
        self.freq_pool = nn.AdaptiveAvgPool2d((1, None))

        self.proj = nn.Sequential(
            nn.Linear(256, embed_dim),
            nn.LayerNorm(embed_dim),
            nn.GELU(),
            nn.Dropout(0.2),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Args:
            x: (B, 1, n_mels, T)
        Returns:
            embedding: (B, embed_dim)
        """
        feats = self.conv(x)                      # (B, 128, F', T')
        feats = self.freq_pool(feats).squeeze(2)  # (B, 128, T')
        mean_pool = torch.mean(feats, dim=-1)
        max_pool = torch.max(feats, dim=-1)[0]
        pooled = torch.cat([mean_pool, max_pool], dim=-1)  # (B, 256)
        return self.proj(pooled)


class PhysicsStreamEncoder(nn.Module):
    """Stream B: Multi-modal physics encoder unifying 1D dynamics, 2D phase, and 2D bicoherence."""

    def __init__(self, embed_dim: int = 128):
        super().__init__()
        # 1. 1D Temporal Dynamics (TKEO stats + Glottal residual stats: 6 channels)
        self.enc_1d = nn.Sequential(
            nn.Conv1d(6, 32, kernel_size=5, stride=1, padding=2, bias=False),
            nn.BatchNorm1d(32),
            nn.GELU(),
            nn.MaxPool1d(2),
            nn.Dropout(0.1),

            nn.Conv1d(32, 64, kernel_size=5, stride=1, padding=2, bias=False),
            nn.BatchNorm1d(64),
            nn.GELU(),
        )
        self.proj_1d = nn.Sequential(
            nn.Linear(128, 64),
            nn.LayerNorm(64),
            nn.GELU()
        )

        # 2. 2D Modified Group Delay (Phase Continuity)
        self.enc_mgd = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(32),
            nn.GELU(),
            nn.MaxPool2d((2, 2)),
            nn.Dropout(0.1),

            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.GELU(),
            nn.AdaptiveAvgPool2d((1, None)),
        )
        self.proj_mgd = nn.Sequential(
            nn.Linear(128, 64),
            nn.LayerNorm(64),
            nn.GELU()
        )

        # 3. 2D Bicoherence (Quadratic Phase Coupling)
        self.enc_bico = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(32),
            nn.GELU(),
            nn.MaxPool2d(2),
            nn.Dropout(0.1),

            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.GELU(),
            nn.AdaptiveAvgPool2d((2, 2)),
        )
        self.proj_bico = nn.Sequential(
            nn.Linear(256, 64),
            nn.LayerNorm(64),
            nn.GELU()
        )

        # Unified Physics Latent Projection (64 * 3 = 192 -> embed_dim)
        self.out_proj = nn.Sequential(
            nn.Linear(192, embed_dim),
            nn.LayerNorm(embed_dim),
            nn.GELU(),
            nn.Dropout(0.2)
        )

    def forward(self, feats_1d: torch.Tensor, mgd: torch.Tensor, bico: torch.Tensor) -> torch.Tensor:
        """Args:
            feats_1d: (B, 6, T)
            mgd:      (B, 1, 64, T)
            bico:     (B, 1, 32, 32)
        Returns:
            h_phys: (B, embed_dim)
        """
        # 1D Temporal Pooling
        f1d = self.enc_1d(feats_1d)
        p1d = torch.cat([torch.mean(f1d, dim=-1), torch.max(f1d, dim=-1)[0]], dim=-1)
        h_1d = self.proj_1d(p1d)

        # 2D MGD Phase Pooling
        fmgd = self.enc_mgd(mgd).squeeze(2)
        pmgd = torch.cat([torch.mean(fmgd, dim=-1), torch.max(fmgd, dim=-1)[0]], dim=-1)
        h_mgd = self.proj_mgd(pmgd)

        # 2D Bicoherence Pooling
        fbico = self.enc_bico(bico).flatten(1)
        h_bico = self.proj_bico(fbico)

        # Assemble and project
        h_stacked = torch.cat([h_1d, h_mgd, h_bico], dim=-1)  # (B, 192)
        return self.out_proj(h_stacked)


class CrossModalGatedFusion(nn.Module):
    """Attention-based Gating Fusion layer uniting Spectrogram and Physics streams."""

    def __init__(self, embed_dim: int = 128):
        super().__init__()
        self.gate_fc = nn.Sequential(
            nn.Linear(embed_dim * 2, embed_dim),
            nn.Sigmoid()
        )
        self.fuse_proj = nn.Sequential(
            nn.Linear(embed_dim * 4, embed_dim * 2),
            nn.LayerNorm(embed_dim * 2),
            nn.GELU(),
            nn.Dropout(0.3)
        )

    def forward(self, h_spec: torch.Tensor, h_phys: torch.Tensor) -> torch.Tensor:
        """Args:
            h_spec: (B, embed_dim)
            h_phys: (B, embed_dim)
        Returns:
            fused: (B, embed_dim * 2)
        """
        concat = torch.cat([h_spec, h_phys], dim=-1)
        gate = self.gate_fc(concat)  # (B, embed_dim)
        gated_blend = gate * h_spec + (1.0 - gate) * h_phys

        # Multi-scale fusion: gated blend + individual streams
        full_repr = torch.cat([gated_blend, h_spec, h_phys, gate], dim=-1)
        return self.fuse_proj(full_repr)


class TandemPhysicsDualStreamDetector(nn.Module):
    """End-to-end Dual-Stream Deepfake Detector under Tandem Telephony Compression.
    
    Combines:
    - Stream A: Log-Mel Spectrogram representation (2D CNN).
    - Stream B: Physics & Math representation (TKEO, GIF jitter/shimmer, MGD phase, Bicoherence QPC).
    - Fusion: Cross-modal gated attention.
    - Classification: Multi-layer perceptron.
    """

    def __init__(self, embed_dim: int = 128, n_mels: int = 64, **kwargs):
        super().__init__()
        # Feature Extractors
        self.spec_extractor = SpectrogramExtractor(sample_rate=8000, n_mels=n_mels)
        self.physics_extractor = StackedPhysicsFeatureExtractor(sample_rate=8000)

        # Encoders
        self.spec_encoder = SpectrogramBranchEncoder(in_channels=1, embed_dim=embed_dim)
        self.physics_encoder = PhysicsStreamEncoder(embed_dim=embed_dim)

        # Fusion & Classifier
        self.fusion = CrossModalGatedFusion(embed_dim=embed_dim)
        self.classifier = nn.Sequential(
            nn.Linear(embed_dim * 2, 64),
            nn.GELU(),
            nn.Dropout(0.3),
            nn.Linear(64, 1)
        )

    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
        """Forward pass directly from 8 kHz lossy compressed audio (B, N).
        
        Returns:
            logits: (B,) log-odds for binary deepfake classification.
        """
        # 1. Feature Extraction
        spec_feats = self.spec_extractor(waveform)      # (B, 1, 64, T)
        phys_dict = self.physics_extractor(waveform)   # feats_1d, mgd, bico

        # 2. Dual-Stream Encodings
        h_spec = self.spec_encoder(spec_feats)
        h_phys = self.physics_encoder(
            feats_1d=phys_dict["feats_1d"],
            mgd=phys_dict["mgd"],
            bico=phys_dict["bico"]
        )

        # 3. Cross-Modal Fusion
        fused = self.fusion(h_spec, h_phys)

        # 4. Classification Logit
        logits = self.classifier(fused).squeeze(-1)
        return logits


# Backward compatibility alias
DualBranchTelephonyDetector = TandemPhysicsDualStreamDetector