"""
Advanced source separation using neural networks.
Requires: pip install asteroid-bin pyannote.audio
"""

import soundfile as sf
import os


def separate_with_neural(audio_path: str, output_dir: str, n_sources: int = 6):
    """
    Use asteroid or pyannote for better separation.
    """

    try:
        from asteroid.models import ConvTasNet

        print("Using Conv-TasNet for separation...")

        model = ConvTasNet.from_pretrained("mpariente/ConvTasNet_WHAM!_sepclean")

        audio, sr = sf.read(audio_path)
        if audio.ndim > 1:
            audio = audio.mean(axis=1)

        import torch

        audio_tensor = torch.from_numpy(audio).float().unsqueeze(0)

        with torch.no_grad():
            estimates = model(audio_tensor)

        os.makedirs(output_dir, exist_ok=True)

        for i in range(min(n_sources, estimates.shape[1])):
            source = estimates[0, i].numpy()
            sf.write(os.path.join(output_dir, f"neural_source_{i + 1}.wav"), source, sr)

        print(f"Saved {min(n_sources, estimates.shape[1])} neural-separated sources")

    except ImportError:
        print("Neural libraries not available. Install with:")
        print("  pip install asteroid-bin pyannote.audio")

        audio, sr = sf.read(audio_path)

        print(f"\nFallback: Using current ICA separation")
        print(f"  Audio: {audio_path}")
        print(f"  Duration: {len(audio) / sr:.1f}s")

        return False

    return True


if __name__ == "__main__":
    import sys

    audio_file = sys.argv[1] if len(sys.argv) > 1 else "../data/mixture.wav"
    output = sys.argv[2] if len(sys.argv) > 2 else "output_neural"

    separate_with_neural(audio_file, output, n_sources=6)