Adaptive Stem Alignment

def align_stem_duration(input_filepath, output_filepath, target_duration, sample_rate, 
                        subtype='FLOAT', strategy='auto', stem_type=None):
    """
    Align stem duration to target using appropriate strategy.
    
    Args:
        input_filepath: Path to input stem
        output_filepath: Path for aligned output
        target_duration: Target duration in seconds
        sample_rate: Sample rate
        subtype: Audio subtype (FLOAT, PCM_24, etc.)
        strategy: 'zero_pad', 'loop', 'crossfade', 'auto'
        stem_type: Type of stem ('rhythmic', 'melodic', 'ambient', 'percussion')
    
    Returns:
        dict with success status and alignment info
    """
    if not os.path.exists(input_filepath):
        return {'success': False, 'error': f'Input file not found: {input_filepath}'}
    
    # Load audio
    data, sr = sf.read(input_filepath)
    current_duration = len(data) / sr
    
    # Check if alignment needed (allow 0.5s tolerance)
    if abs(current_duration - target_duration) < 0.5:
        print(f"  Duration already aligned: {current_duration:.2f}s ≈ {target_duration:.2f}s")
        # Just copy file
        sf.write(output_filepath, data, sample_rate, subtype=subtype, format='WAV')
        return {'success': True, 'strategy': 'none', 'original_duration': current_duration}
    
    # Auto-detect strategy if not specified
    if strategy == 'auto':
        if stem_type in ['rhythmic', 'percussion', 'drums']:
            strategy = 'loop'
        elif stem_type in ['ambient', 'pad', 'drone', 'fx']:
            strategy = 'zero_pad'
        else:  # melodic, vocals, guitars, etc.
            strategy = 'crossfade'
    
    print(f"  Aligning duration: {current_duration:.2f}s → {target_duration:.2f}s using '{strategy}'")
    
    # Calculate frames needed
    target_frames = int(target_duration * sample_rate)
    current_frames = len(data)
    
    if strategy == 'zero_pad':
        # Simple zero-padding
        if current_frames < target_frames:
            aligned_data = np.zeros(target_frames, dtype=data.dtype)
            aligned_data[:current_frames] = data
        else:
            # Truncate with fade-out
            fade_frames = int(0.5 * sample_rate)  # 500ms fade
            aligned_data = data[:target_frames].copy()
            if target_frames < current_frames:
                fade_start = max(0, target_frames - fade_frames)
                fade_curve = np.linspace(1, 0, target_frames - fade_start)
                aligned_data[fade_start:] *= fade_curve
    
    elif strategy == 'loop':
        # Loop to fill duration
        aligned_data = np.zeros(target_frames, dtype=data.dtype)
        loop_count = (target_frames // current_frames) + 1
        
        for i in range(loop_count):
            start = i * current_frames
            end = min(start + current_frames, target_frames)
            copy_len = end - start
            
            if copy_len > 0:
                aligned_data[start:end] = data[:copy_len]
        
        # Apply crossfade at loop points to avoid clicks
        crossfade_frames = int(0.05 * sample_rate)  # 50ms crossfade
        if current_frames > crossfade_frames * 2:
            for i in range(1, loop_count):
                loop_start = i * current_frames
                if loop_start < target_frames:
                    # Crossfade region
                    cf_end = min(loop_start + crossfade_frames, target_frames)
                    cf_start = max(loop_start - crossfade_frames, 0)
                    if cf_end > cf_start:
                        fade_in = np.linspace(0, 1, cf_end - cf_start)
                        fade_out = np.linspace(1, 0, cf_end - cf_start)
                        aligned_data[cf_start:cf_end] = (
                            aligned_data[cf_start:cf_end] * fade_out +
                            np.roll(aligned_data[cf_start:cf_end], -current_frames) * fade_in
                        )
    
    elif strategy == 'crossfade':
        # Crossfade-based extension with smooth transition
        if current_frames < target_frames:
            # Extend with crossfaded loop
            extension_frames = target_frames - current_frames
            fade_frames = min(int(2.0 * sample_rate), extension_frames // 2)  # Max 2s fade
            
            # Create extension from looped content with fade
            extension_data = np.zeros(extension_frames, dtype=data.dtype)
            
            if extension_frames <= current_frames:
                # Just take from beginning with fade-in
                extension_data[:extension_frames] = data[:extension_frames]
                if fade_frames > 0:
                    fade_in = np.linspace(0, 1, min(fade_frames, extension_frames))
                    extension_data[:len(fade_in)] *= fade_in
            else:
                # Loop multiple times with crossfades
                loop_data = np.tile(data, (extension_frames // current_frames) + 2)[:extension_frames]
                
                # Apply fade-in to extension
                if fade_frames > 0:
                    fade_in = np.linspace(0, 1, fade_frames)
                    loop_data[:fade_frames] *= fade_in
                
                extension_data = loop_data
            
            # Concatenate with crossfade
            aligned_data = np.zeros(target_frames, dtype=data.dtype)
            aligned_data[:current_frames] = data
            
            # Crossfade region at junction
            if fade_frames > 0:
                junction_start = current_frames - fade_frames
                junction_end = min(current_frames + fade_frames, target_frames)
                
                if junction_end > junction_start:
                    crossfade_len = junction_end - junction_start
                    fade_out = np.linspace(1, 0, crossfade_len)
                    fade_in = np.linspace(0, 1, crossfade_len)
                    
                    aligned_data[junction_start:junction_end] = (
                        aligned_data[junction_start:junction_end] * fade_out +
                        extension_data[:crossfade_len] * fade_in
                    )
                else:
                    aligned_data[current_frames:current_frames + extension_frames] = extension_data
            else:
                aligned_data[current_frames:] = extension_data
        else:
            # Truncate with fade-out
            fade_frames = int(2.0 * sample_rate)
            aligned_data = data[:target_frames].copy()
            fade_start = max(0, target_frames - fade_frames)
            fade_curve = np.linspace(1, 0, target_frames - fade_start)
            aligned_data[fade_start:] *= fade_curve
    
    else:
        return {'success': False, 'error': f'Unknown strategy: {strategy}'}
    
    # Clip to prevent overload
    aligned_data = np.clip(aligned_data, -1, 1)
    
    # Export
    sf.write(output_filepath, aligned_data, sample_rate, subtype=subtype, format='WAV')
    
    # Verify
    result = verify_stem(output_filepath, sample_rate, subtype, target_duration)
    if result['success']:
        return {
            'success': True,
            'strategy': strategy,
            'original_duration': current_duration,
            'aligned_duration': result['info'].duration
        }
    else:
        return result

# Apply duration alignment to all stems
print("\n=== Aligning stem durations ===")
aligned_stems = []

for stem_name in STEM_NAMES:
    input_file = f'{stem_name}_stem.wav'
    output_file = f'{stem_name}_aligned.wav'
    
    # Determine stem type for strategy selection
    stem_type_map = {
        'bass': 'rhythmic',
        'guitars': 'melodic',
        'synths': 'ambient',
        'bridge': 'melodic'
    }
    stem_type = stem_type_map.get(stem_name, 'melodic')
    
    print(f"Aligning {stem_name} (type: {stem_type})...")
    result = align_stem_duration(
        input_file, output_file, TARGET_DURATION, SAMPLE_RATE,
        subtype=SUBTYPE, strategy='auto', stem_type=stem_type
    )
    
    if result['success']:
        if result['strategy'] != 'none':
            print(f"✓ {stem_name} aligned: {result['original_duration']:.2f}s → {result['aligned_duration']:.2f}s ({result['strategy']})")
        else:
            print(f"✓ {stem_name} already aligned")
        aligned_stems.append(output_file)
    else:
        print(f"✗ {stem_name} alignment FAILED: {result['error']}")
        raise RuntimeError(f"Stem alignment failed for {stem_name}: {result['error']}")

# Align drums separately
drums_aligned = 'drums_aligned.wav'
print(f"Aligning drums (type: percussion)...")
drums_result = align_stem_duration(
    'drums_stem.wav', drums_aligned, TARGET_DURATION, SAMPLE_RATE,
    subtype=SUBTYPE, strategy='auto', stem_type='percussion'
)
if drums_result['success']:
    if drums_result['strategy'] != 'none':
        print(f"✓ Drums aligned: {drums_result['original_duration']:.2f}s → {drums_result['aligned_duration']:.2f}s ({drums_result['strategy']})")
    else:
        print(f"✓ Drums already aligned")
    aligned_stems.append(drums_aligned)
else:
    raise RuntimeError(f"Drums alignment failed: {drums_result['error']}")

# Final duration verification - all stems must match
print("\n=== Verifying duration alignment ===")
final_durations = {}
for stem_file in aligned_stems:
    info = sf.info(stem_file)
    stem_name = os.path.basename(stem_file).replace('_aligned.wav', '')
    final_durations[stem_name] = info.duration
    duration_diff = abs(info.duration - TARGET_DURATION)
    
    if duration_diff > 0.5:
        print(f"✗ WARNING: {stem_name} duration mismatch: {info.duration:.2f}s vs target {TARGET_DURATION:.2f}s")
    else:
        print(f"✓ {stem_name}: {info.duration:.2f}s (Δ{duration_diff:.2f}s)")

max_duration_diff = max(abs(d - TARGET_DURATION) for d in final_durations.values())
if max_duration_diff > 0.5:
    raise RuntimeError(f"Duration alignment incomplete: max deviation {max_duration_diff:.2f}s exceeds tolerance")
print(f"\nAll stems aligned within tolerance (max deviation: {max_duration_diff:.2f}s)")

from scipy import signal

def apply_lowpass_filter(audio_data, sample_rate, cutoff_freq=8000):
    """Apply lowpass filter using scipy.signal."""
    nyquist = sample_rate / 2
    normalized_cutoff = cutoff_freq / nyquist
    b, a = signal.butter(4, normalized_cutoff, btype='low')
    return signal.filtfilt(b, a, audio_data)

def apply_effects_and_verify(input_filepath, output_filepath, sample_rate, subtype):
    """Apply effects to stem and verify output."""
    data, sr = sf.read(input_filepath)
    
    # Apply effects
    processed = apply_lowpass_filter(data, sr, cutoff_freq=8000)
    processed = np.clip(processed, -1, 1)
    
    # Export
    sf.write(output_filepath, processed, sample_rate, subtype=subtype, format='WAV')
    
    # Verify
    result = verify_stem(output_filepath, sample_rate, subtype, TARGET_DURATION)
    return result, processed

print("\n=== Applying effects to all stems ===")
processed_stems = []

for stem_name in STEM_NAMES:
    input_file = f'{stem_name}_aligned.wav'
    output_file = f'{stem_name}_processed.wav'
    
    print(f"Processing {stem_name}...")
    result, _ = apply_effects_and_verify(input_file, output_file, SAMPLE_RATE, SUBTYPE)
    
    if result['success']:
        print(f"✓ {stem_name} processed and verified")
        processed_stems.append(output_file)
    else:
        print(f"✗ {stem_name} processing FAILED: {result['error']}")
        raise RuntimeError(f"Effects processing failed for {stem_name}")

# Process drums separately
drums_output = 'drums_processed.wav'
drums_result, _ = apply_effects_and_verify('drums_aligned.wav', drums_output, SAMPLE_RATE, SUBTYPE)
if drums_result['success']:
    print(f"✓ Drums processed and verified")
    processed_stems.append(drums_output)
else:
    raise RuntimeError(f"Drums processing failed: {drums_result['error']}")

def create_master_track(stem_files, output_filepath, sample_rate, subtype):
    """Create master track from verified stems with gain staging."""
    # Load first stem to get dimensions
    first_data, sr = sf.read(stem_files[0])
    master_audio = np.zeros(len(first_data), dtype=np.float32)
    
    # Verify all stems have matching length
    for stem_file in stem_files:
        data, _ = sf.read(stem_file)
        if len(data) != len(first_data):
            raise ValueError(f"Stem length mismatch: {stem_file} has {len(data)} frames, expected {len(first_data)}")
    
    # Mix all stems with gain staging
    gain_per_stem = 0.4  # Prevent clipping with 5 stems
   
    for i, stem_file in enumerate(stem_files):
        data, sr = sf.read(stem_file)
        master_audio += data * gain_per_stem
        print(f"  Mixed {os.path.basename(stem_file)} (gain: {gain_per_stem})")
    
    # Apply master bus limiting
    master_audio = np.clip(master_audio, -1, 1)
    
    # Soft clip for warmth
    master_audio = np.tanh(master_audio * 1.2) / 1.2
    
    # Export
    sf.write(output_filepath, master_audio, sample_rate, subtype=subtype, format='WAV')
    
    # Verify
    info = sf.info(output_filepath)
    print(f"Master exported: {info.duration:.2f}s @ {info.samplerate}Hz, {info.channels}ch")
    
    return output_filepath, master_audio

print("\n=== Creating master track ===")
master_filepath, master_data = create_master_track(processed_stems, 'master.wav', SAMPLE_RATE, SUBTYPE)

import json
from datetime import datetime

def create_archive_manifest(stem_files, master_file, output_dir='deliverables'):
    """Create archive manifest with comprehensive verification."""
    os.makedirs(output_dir, exist_ok=True)
    
    manifest = {
        'created': datetime.now().isoformat(),
        'target_duration': TARGET_DURATION,
        'sample_rate': SAMPLE_RATE,
        'subtype': SUBTYPE,
        'stems': [],
        'master': None,
        'verification': {
            'all_stems_aligned': True,
            'all_stems_verified': True,
            'master_verified': True
        }
    }
    
    # Verify each stem
    for stem_file in stem_files:
        if not os.path.exists(stem_file):
            manifest['verification']['all_stems_verified'] = False
            continue
        
        info = sf.info(stem_file)
        stem_name = os.path.basename(stem_file)
        duration_diff = abs(info.duration - TARGET_DURATION)
        
        stem_info = {
            'file': stem_name,
            'duration': info.duration,
            'sample_rate': info.samplerate,
            'channels': info.channels,
            'duration_aligned': duration_diff < 0.5
        }
        manifest['stems'].append(stem_info)
        
        if duration_diff >= 0.5:
            manifest['verification']['all_stems_aligned'] = False
            print(f"WARNING: {stem_name} duration misaligned by {duration_diff:.2f}s")
    
    # Verify master
    if os.path.exists(master_file):
        info = sf.info(master_file)
        manifest['master'] = {
            'file': os.path.basename(master_file),
            'duration': info.duration,
            'sample_rate': info.samplerate,
            'channels': info.channels,
            'subtype': info.subtype
        }
        
        # Check master duration matches target
        if abs(info.duration - TARGET_DURATION) > 1.0:
            manifest['verification']['master_verified'] = False
            print(f"WARNING: Master duration {info.duration:.2f}s differs from target {TARGET_DURATION:.2f}s")
    else:
        manifest['verification']['master_verified'] = False
    
    # Save manifest
    manifest_path = os.path.join(output_dir, 'manifest.json')
    with open(manifest_path, 'w') as f:
        json.dump(manifest, f, indent=2)
    
    # Copy files to deliverables
    import shutil
    for stem_file in stem_files:
        shutil.copy(stem_file, output_dir)
    shutil.copy(master_file, output_dir)
    
    return manifest_path, manifest

print("\n=== Creating archive ===")
manifest_path, manifest = create_archive_manifest(processed_stems, master_filepath)
print(f"Archive manifest created: {manifest_path}")

# Final summary
print("\n" + "="*60)
print("PRODUCTION COMPLETE")
print("="*60)
print(f"Target duration: {TARGET_DURATION:.2f}s")
print(f"Sample rate: {SAMPLE_RATE}Hz")
print(f"Stems processed: {len(processed_stems)}")
print(f"All stems aligned: {manifest['verification']['all_stems_aligned']}")
print(f"Master verified: {manifest['verification']['master_verified']}")
print(f"Deliverables: ./deliverables/")
print("="*60)