import os
import math
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import torchaudio
import torchaudio.transforms as T
from huggingface_hub import PyTorchModelHubMixin

# =====================================================================
# 1. ARCHITECTURE: Multi-Masking U-Net
# =====================================================================

class UNetBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.GroupNorm(num_groups=4, num_channels=out_channels),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.GroupNorm(num_groups=4, num_channels=out_channels),
            nn.LeakyReLU(0.2, inplace=True)
        )

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


class MultiStemSeparationModel(nn.Module, PyTorchModelHubMixin):
    """
    An advanced 2D U-Net that predicts multiple independent masks.
    By setting num_stems (e.g., 6), the model outputs 6 channels.
    We apply a Softmax activation across the stem channels so that the
    predicted masks dynamically sum up to 1.0, preserving audio energy.
    """
    def __init__(self, in_channels: int = 1, num_stems: int = 6, base_features: int = 16):
        super().__init__()
        
        self.in_channels = in_channels
        self.num_stems = num_stems
        self.base_features = base_features

        # Downsampling Encoder
        self.enc1 = UNetBlock(in_channels, base_features)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        
        self.enc2 = UNetBlock(base_features, base_features * 2)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        
        self.enc3 = UNetBlock(base_features * 2, base_features * 4)
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)

        # Bottleneck
        self.bottleneck = UNetBlock(base_features * 4, base_features * 8)

        # Upsampling Decoder
        self.up3 = nn.ConvTranspose2d(base_features * 8, base_features * 4, kernel_size=2, stride=2)
        self.dec3 = UNetBlock(base_features * 8, base_features * 4)
        
        self.up2 = nn.ConvTranspose2d(base_features * 4, base_features * 2, kernel_size=2, stride=2)
        self.dec2 = UNetBlock(base_features * 4, base_features * 2)
        
        self.up1 = nn.ConvTranspose2d(base_features * 2, base_features, kernel_size=2, stride=2)
        self.dec1 = UNetBlock(base_features * 2, base_features)

        # Output projection to predict N masks
        self.final_conv = nn.Conv2d(base_features, num_stems, kernel_size=1)
        
        # Softmax ensures that at any given frequency/time bin, the sum of all 
        # mask values equals 1.0 (prevents instruments bleeding over each other).
        self.activation = nn.Softmax(dim=1)

    def forward(self, x):
        """
        Input x: [Batch, 1, Freqs, TimeSteps]
        Output:
            separated_stems: [Batch, num_stems, Freqs, TimeSteps]
            masks: [Batch, num_stems, Freqs, TimeSteps]
        """
        # Encoder
        x1 = self.enc1(x)
        p1 = self.pool1(x1)
        
        x2 = self.enc2(p1)
        p2 = self.pool2(x2)
        
        x3 = self.enc3(p2)
        p3 = self.pool3(x3)

        # Bottleneck
        b = self.bottleneck(p3)

        # Decoder
        u3 = self.up3(b)
        if u3.shape != x3.shape:
            u3 = nn.functional.interpolate(u3, size=x3.shape[2:])
        m3 = torch.cat([u3, x3], dim=1)
        d3 = self.dec3(m3)
        
        u2 = self.up2(d3)
        if u2.shape != x2.shape:
            u2 = nn.functional.interpolate(u2, size=x2.shape[2:])
        m2 = torch.cat([u2, x2], dim=1)
        d2 = self.dec2(m2)
        
        u1 = self.up1(d2)
        if u1.shape != x1.shape:
            u1 = nn.functional.interpolate(u1, size=x1.shape[2:])
        m1 = torch.cat([u1, x1], dim=1)
        d1 = self.dec1(m1)

        # Predict multi-channel masks
        masks = self.activation(self.final_conv(d1))
        
        # Multiply each mask channel with the original input mixture spectrogram
        # x is [B, 1, F, T], masks is [B, num_stems, F, T]
        separated_stems = x * masks
        return separated_stems, masks


# =====================================================================
# 2. DATASET: Dynamic Multi-Stem Loader
# =====================================================================

class MultiStemDataset(Dataset):
    """
    Loads parallel directories representing multiple stems for training.
    Stems must be named exactly:
    - vocals.wav, backing_vocals.wav, drums.wav, bass.wav, guitar.wav, synth.wav
    """
    def __init__(self, dataset_dir, stems_list, sample_rate=16000, duration_seconds=4.0):
        super().__init__()
        self.dataset_dir = dataset_dir
        self.stems_list = stems_list
        self.sample_rate = sample_rate
        self.num_samples = int(sample_rate * duration_seconds)
        
        # Identify song folders
        self.song_folders = sorted([
            os.path.join(dataset_dir, d) for d in os.listdir(dataset_dir)
            if os.path.isdir(os.path.join(dataset_dir, d))
        ])
        
        self.spectrogram_extractor = T.Spectrogram(
            n_fft=512, win_length=512, hop_length=128, power=1.0, normalized=True
        )

    def __len__(self):
        return len(self.song_folders)

    def _load_and_pad_crop(self, filepath):
        if not os.path.exists(filepath):
            # Return silent tensor if a specific instrument doesn't exist in a song
            return torch.zeros((1, self.num_samples))
            
        waveform, sr = torchaudio.load(filepath)
        if sr != self.sample_rate:
            resampler = T.Resample(orig_freq=sr, new_freq=self.sample_rate)
            waveform = resampler(waveform)
        if waveform.shape[0] > 1:
            waveform = torch.mean(waveform, dim=0, keepdim=True)
            
        length = waveform.shape[1]
        if length > self.num_samples:
            waveform = waveform[:, :self.num_samples]
        elif length < self.num_samples:
            padding = self.num_samples - length
            waveform = torch.nn.functional.pad(waveform, (0, padding))
            
        return waveform

    def __getitem__(self, idx):
        song_path = self.song_folders[idx]
        stem_tensors = []
        
        # Load each individual stem
        for stem_name in self.stems_list:
            stem_file = os.path.join(song_path, f"{stem_name}.wav")
            stem_wav = self._load_and_pad_crop(stem_file)
            stem_tensors.append(stem_wav)
            
        # Combine stems to form the dynamic mixture (Input)
        mix_wav = torch.stack(stem_tensors, dim=0).sum(dim=0)
        
        # Normalization
        max_val = torch.max(torch.abs(mix_wav))
        if max_val > 0:
            mix_wav = mix_wav / max_val
            for i in range(len(stem_tensors)):
                stem_tensors[i] = stem_tensors[i] / max_val

        # Convert mixture to spectrogram magnitude -> shape [1, Freqs, Time]
        mix_spec = self.spectrogram_extractor(mix_wav)
        
        # Convert all target stems to spectrograms
        target_specs = []
        for stem_wav in stem_tensors:
            stem_spec = self.spectrogram_extractor(stem_wav)
            target_specs.append(stem_spec)
            
        # Stack targets into a single tensor: [num_stems, Freqs, Time]
        target_specs_tensor = torch.cat(target_specs, dim=0)
        
        return mix_spec, target_specs_tensor


# =====================================================================
# 3. TRAINING PIPELINE
# =====================================================================

def train_multi_stem_model(model, dataloader, epochs=5, lr=1e-3, device="cuda"):
    model = model.to(device)
    optimizer = optim.Adam(model.parameters(), lr=lr)
    criterion = nn.MSELoss()
    
    print(f"Starting Multi-Stem training on device: {device}")
    model.train()
    
    for epoch in range(1, epochs + 1):
        running_loss = 0.0
        for batch_idx, (mix_spec, targets_spec) in enumerate(dataloader):
            mix_spec = mix_spec.to(device)         # [B, 1, Freq, Time]
            targets_spec = targets_spec.to(device) # [B, num_stems, Freq, Time]
            
            optimizer.zero_grad()
            
            # Forward pass (Outputs estimated stems of shape [B, num_stems, Freq, Time])
            pred_stems, _ = model(mix_spec)
            
            loss = criterion(pred_stems, targets_spec)
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item()
            
        epoch_loss = running_loss / len(dataloader)
        print(f"Epoch [{epoch}/{epochs}] - Loss: {epoch_loss:.6f}")


if __name__ == "__main__":
    # Define our custom targets
    MY_STEMS = ["vocals", "backing_vocals", "drums", "bass", "guitar", "synth"]
    device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
    
    model = MultiStemSeparationModel(in_channels=1, num_stems=len(MY_STEMS), base_features=16)
    print(f"Initialized U-Net for {len(MY_STEMS)} stems separation.")