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

# ==============================================================================
# 1. Basic Layers & Activations
# ==============================================================================

class Swish(nn.Module):
    """Swish activation function: x * sigmoid(x)."""
    def forward(self, x):
        return x * torch.sigmoid(x)

class GLU(nn.Module):
    """Gated Linear Unit for feature fusion (Eq. 5)."""
    def __init__(self, input_dim):
        super(GLU, self).__init__()
        self.dim = input_dim
        # W_f and V_f in Equation 5
        self.linear_out = nn.Linear(input_dim, input_dim)
        self.linear_gate = nn.Linear(input_dim, input_dim)

    def forward(self, x):
        # x shape: [Batch, Length, Dim]
        out = torch.tanh(self.linear_out(x))
        gate = torch.sigmoid(self.linear_gate(x))
        return out * gate

# ==============================================================================
# 2. Input Processing: Beat Synchronization & Fusion (Section 3.1)
# ==============================================================================

class BeatSynchronousPooling(nn.Module):
    """
    Aligns frame-level features (MERT/Acoustic) to a Beat Grid (Eq. 3).
    """
    def forward(self, frame_features, beat_indices):
        """
        Args:
            frame_features: [Batch, Time_Frames, Feat_Dim]
            beat_indices: List of tensors containing frame indices for beats.
                          Example: [[0, 20, 45, ...], [0, 22, 50...]]
        Returns:
            beat_features: [Batch, Max_Beats, Feat_Dim]
        """
        batch_size, _, feat_dim = frame_features.shape
        # Find max beats in the batch for padding
        max_beats = max([len(b) for b in beat_indices]) - 1
        
        pooled_batch = []
        
        for b in range(batch_size):
            feats = frame_features[b]
            beats = beat_indices[b]
            beat_segments = []
            
            for k in range(len(beats) - 1):
                start = beats[k]
                end = beats[k+1]
                # Adaptive Average Pooling within the beat interval
                if end > start:
                    segment = feats[start:end].mean(dim=0)
                else:
                    segment = feats[start] # Fallback for single frame
                beat_segments.append(segment)
            
            # Stack and Pad
            beat_tensor = torch.stack(beat_segments)
            pad_len = max_beats - beat_tensor.size(0)
            if pad_len > 0:
                pad = torch.zeros(pad_len, feat_dim, device=feats.device)
                beat_tensor = torch.cat([beat_tensor, pad], dim=0)
            
            pooled_batch.append(beat_tensor)
            
        return torch.stack(pooled_batch)

class GatedFeatureFusion(nn.Module):
    """
    Fuses Semantic (MERT) and Acoustic features using GLU (Eq. 4 & 5).
    """
    def __init__(self, sem_dim, ac_dim, embed_dim):
        super().__init__()
        self.concat_proj = nn.Linear(sem_dim + ac_dim, embed_dim)
        self.glu = GLU(embed_dim)
        self.layer_norm = nn.LayerNorm(embed_dim)

    def forward(self, sem_beats, ac_beats):
        # Concatenate: [Batch, Beats, Sem_Dim + Ac_Dim]
        cat_feat = torch.cat([sem_beats, ac_beats], dim=-1)
        # Project and Apply Gate
        x = self.concat_proj(cat_feat)
        x = self.glu(x)
        return self.layer_norm(x)

# ==============================================================================
# 3. Backbone: Conformer Block (Section 3.2)
# ==============================================================================

class FeedForward(nn.Module):
    """Half-Step Feed Forward Network (Eq. 6)."""
    def __init__(self, dim, hidden_dim, dropout=0.1):
        super().__init__()
        self.net = nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, hidden_dim),
            Swish(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim),
            nn.Dropout(dropout)
        )
    def forward(self, x):
        return self.net(x)

class ConformerConvModule(nn.Module):
    """Convolution Module for local detail extraction (Eq. 9 & 10)."""
    def __init__(self, dim, kernel_size=31, dropout=0.1):
        super().__init__()
        self.layer_norm = nn.LayerNorm(dim)
        # Pointwise 1
        self.pointwise1 = nn.Conv1d(dim, dim * 2, kernel_size=1)
        self.glu = nn.GLU(dim=1)
        # Depthwise
        self.depthwise = nn.Conv1d(dim, dim, kernel_size, padding=(kernel_size-1)//2, groups=dim)
        self.batch_norm = nn.BatchNorm1d(dim)
        self.swish = Swish()
        # Pointwise 2
        self.pointwise2 = nn.Conv1d(dim, dim, kernel_size=1)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # x: [Batch, Time, Dim] -> Permute to [Batch, Dim, Time]
        inputs = x
        x = self.layer_norm(x)
        x = x.transpose(1, 2)
        
        x = self.pointwise1(x)
        x = self.glu(x)
        x = self.depthwise(x)
        x = self.batch_norm(x)
        x = self.swish(x)
        x = self.pointwise2(x)
        x = self.dropout(x)
        
        x = x.transpose(1, 2)
        return inputs + x # Residual connection

class ConformerBlock(nn.Module):
    """Macaron-style Conformer Block (Eq. 11)."""
    def __init__(self, dim, heads, ffn_dim, conv_kernel=31, dropout=0.1):
        super().__init__()
        self.ffn1 = FeedForward(dim, ffn_dim, dropout)
        self.attn_norm = nn.LayerNorm(dim)
        self.attn = nn.MultiheadAttention(embed_dim=dim, num_heads=heads, dropout=dropout, batch_first=True)
        self.conv = ConformerConvModule(dim, conv_kernel, dropout)
        self.ffn2 = FeedForward(dim, ffn_dim, dropout)
        self.final_norm = nn.LayerNorm(dim)

    def forward(self, x):
        # 1. Half-Step FFN
        x = x + 0.5 * self.ffn1(x)
        
        # 2. Self-Attention
        res = x
        x = self.attn_norm(x)
        x, _ = self.attn(x, x, x) # Standard SA for demo (Paper implies relative pos, simplified here)
        x = res + x
        
        # 3. Convolution
        x = self.conv(x)
        
        # 4. Half-Step FFN
        x = x + 0.5 * self.ffn2(x)
        
        return self.final_norm(x)

# ==============================================================================
# 4. Hierarchical Attention Mechanism (HAM) (Section 3.3)
# ==============================================================================

class HierarchicalAttention(nn.Module):
    """
    HAM: Motif Level (Local) -> Section Level (Global) (Eq. 12-15).
    """
    def __init__(self, dim, window_size=5, num_heads=4):
        super().__init__()
        self.dim = dim
        self.window_size = window_size
        
        # Motif Level: Local Attention (Implemented via masked MultiheadAttention)
        self.motif_attn = nn.MultiheadAttention(dim, num_heads, batch_first=True)
        self.motif_norm = nn.LayerNorm(dim)
        
        # Section Level: Global Attention
        self.section_attn = nn.MultiheadAttention(dim, num_heads, batch_first=True)
        self.section_norm = nn.LayerNorm(dim)

    def create_local_mask(self, length, device):
        """Creates a diagonal band mask for local attention."""
        mask = torch.ones(length, length, device=device)
        mask = torch.triu(mask, diagonal=-self.window_size) * \
               torch.tril(mask, diagonal=self.window_size)
        # Convert to float mask: 0 for allow, -inf for block
        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
        return mask

    def forward(self, x):
        batch, length, _ = x.shape
        
        # --- 1. Motif Level (Local) ---
        mask = self.create_local_mask(length, x.device)
        # Apply local attention with mask (Eq. 14)
        m_out, _ = self.motif_attn(x, x, x, attn_mask=mask)
        m_out = self.motif_norm(x + m_out) # Residual
        
        # --- 2. Section Level (Global) ---
        # Apply global attention on motif features (Eq. 15)
        s_out, _ = self.section_attn(m_out, m_out, m_out)
        s_out = self.section_norm(m_out + s_out) # Residual
        
        return s_out

# ==============================================================================
# 5. Full MCH-BFNet Model
# ==============================================================================

class MCH_BFNet(nn.Module):
    def __init__(self, 
                 mert_dim=768, 
                 ac_dim=142, 
                 embed_dim=256, 
                 num_layers=4, 
                 num_classes=8):
        super().__init__()
        
        print(f"Initializing MCH-BFNet with Embed Dim: {embed_dim}")
        
        # 1. Feature Handling
        self.beat_pool = BeatSynchronousPooling()
        self.fusion = GatedFeatureFusion(mert_dim, ac_dim, embed_dim)
        
        # 2. Conformer Backbone
        self.layers = nn.ModuleList([
            ConformerBlock(dim=embed_dim, heads=4, ffn_dim=embed_dim*4)
            for _ in range(num_layers)
        ])
        
        # 3. Hierarchical Attention
        self.ham = HierarchicalAttention(embed_dim, window_size=3)
        
        # 4. Decoding Heads
        # Boundary Detection Head (Binary Regression) (Eq. 16)
        self.boundary_head = nn.Sequential(
            nn.Linear(embed_dim, embed_dim // 2),
            nn.ReLU(),
            nn.Linear(embed_dim // 2, 1),
            nn.Sigmoid()
        )
        
        # Function Classification Head (Multi-class) (Eq. 17)
        self.classification_head = nn.Sequential(
            nn.Linear(embed_dim, embed_dim // 2),
            nn.ReLU(),
            nn.Linear(embed_dim // 2, num_classes)
            # Softmax is applied in CrossEntropyLoss during training
        )

    def forward(self, mert_frames, ac_frames, beat_indices):
        """
        Args:
            mert_frames: [B, T_frames, 768] (Simulated MERT output)
            ac_frames: [B, T_frames, 142] (Simulated MFCC/Chroma)
            beat_indices: List of lists containing beat frame indices
        """
        # 1. Beat Sync Pooling (Section 3.1)
        sem_beats = self.beat_pool(mert_frames, beat_indices)
        ac_beats = self.beat_pool(ac_frames, beat_indices)
        
        # 2. Fusion (Section 3.1)
        x = self.fusion(sem_beats, ac_beats) # [B, T_beats, Embed_Dim]
        
        # 3. Backbone Encoding (Section 3.2)
        for layer in self.layers:
            x = layer(x)
            
        # 4. Hierarchical Mapping (Section 3.3)
        struct_feats = self.ham(x)
        
        # 5. Output Heads
        boundary_prob = self.boundary_head(struct_feats) # [B, T_beats, 1]
        class_logits = self.classification_head(struct_feats) # [B, T_beats, Num_Classes]
        
        return boundary_prob, class_logits

# ==============================================================================
# 6. Runnable Verification Instance
# ==============================================================================

if __name__ == "__main__":
    # --- Experiment Settings ---
    BATCH_SIZE = 2
    FRAME_LEN = 1000   # e.g., 10 seconds of audio frames
    MERT_DIM = 768     # Standard MERT output
    AC_DIM = 142       # Standard Acoustic dim (e.g., MFCC+Chroma+Spectral)
    NUM_CLASSES = 6    # e.g., Intro, Verse, Chorus, Bridge, Outro, Silence
    
    # --- 1. Simulate Input Data ---
    print(">>> Generating Simulated Data...")
    # Random tensor simulating pre-trained MERT output
    sim_mert_frames = torch.randn(BATCH_SIZE, FRAME_LEN, MERT_DIM)
    # Random tensor simulating acoustic features
    sim_ac_frames = torch.randn(BATCH_SIZE, FRAME_LEN, AC_DIM)
    
    # Simulate Beat Indices (Beat Tracking Output)
    # Creating fake beats roughly every 20 frames
    sim_beat_indices = []
    for _ in range(BATCH_SIZE):
        beats = list(range(0, FRAME_LEN, 20)) 
        # Ensure the last frame is included for boundaries
        if beats[-1] != FRAME_LEN - 1:
            beats.append(FRAME_LEN - 1)
        sim_beat_indices.append(beats)
        
    print(f"Input Frame Shape: {sim_mert_frames.shape}")
    print(f"Num Beats (Batch 0): {len(sim_beat_indices[0])}")
    
    # --- 2. Initialize Model ---
    model = MCH_BFNet(
        mert_dim=MERT_DIM, 
        ac_dim=AC_DIM, 
        embed_dim=128,  # Reduced for demo speed
        num_layers=2,
        num_classes=NUM_CLASSES
    )
    
    # --- 3. Forward Pass ---
    print("\n>>> Running Forward Pass...")
    model.eval() # Set to evaluation mode
    with torch.no_grad():
        boundaries, classes = model(sim_mert_frames, sim_ac_frames, sim_beat_indices)
        
    # --- 4. Output Analysis ---
    print("\n>>> Output Shapes:")
    print(f"Boundary Probabilities: {boundaries.shape} (Expected: [B, Beats, 1])")
    print(f"Classification Logits:  {classes.shape}   (Expected: [B, Beats, {NUM_CLASSES}])")
    
    print("\n>>> Sample Output (First 3 beats of Batch 0):")
    print(f"Boundary Prob: \n{boundaries[0, :3, 0]}")
    print(f"Class Preds:   \n{torch.argmax(classes[0, :3, :], dim=-1)}")
    
    print("\n>>> Success! MCH-BFNet architecture is functional.")