#!/usr/bin/env python
"""Script to add Fases 11-15 advanced sample picker to __init__.py"""

import sys

def main():
    filepath = r'C:\ProgramData\Ableton\Live 12 Suite\Resources\MIDI Remote Scripts\AbletonMCP_AI\__init__.py'
    
    # Read the file
    with open(filepath, 'r', encoding='utf-8') as f:
        content = f.read()
    
    # Check if already exists
    if '_pick_for_scene_advanced' in content:
        print("ERROR: _pick_for_scene_advanced already exists!")
        return 1
    
    # The old function to find
    old_function = '''        def _pick_for_scene(all_samples, scene_idx, total_scenes):
            """Distribute samples across scenes so each gets a different one."""
            if not all_samples:
                return None
            if len(all_samples) <= total_scenes:
                return all_samples[scene_idx % len(all_samples)]
            step = len(all_samples) / total_scenes
            idx = int(scene_idx * step) % len(all_samples)
            return all_samples[idx]

        # Sort drum loops by BPM proximity to tempo'''
    
    # The new function to add
    new_function = '''        def _pick_for_scene(all_samples, scene_idx, total_scenes):
            """Distribute samples across scenes so each gets a different one."""
            if not all_samples:
                return None
            if len(all_samples) <= total_scenes:
                return all_samples[scene_idx % len(all_samples)]
            step = len(all_samples) / total_scenes
            idx = int(scene_idx * step) % len(all_samples)
            return all_samples[idx]

        # ================================================================
        # FASES 11-15: SISTEMA AVANZADO DE VARIACION MASIVA DE KICKS Y SNARES
        # ================================================================
        # Track samples used in previous scene to avoid repetition
        _prev_scene_samples = {"kicks": [], "snares": []}
        _scene_sample_usage = {"kicks": {}, "snares": {}}  # Track usage count per sample
        _all_kicks_used = []  # Track order of all kicks used
        _all_snares_used = []  # Track order of all snares used

        def _pick_for_scene_advanced(all_samples, scene_idx, total_scenes, energy, prev_samples, sample_type="kick"):
            """
            Advanced sample selection with energy-based filtering and no-repetition policy.
            
            Args:
                all_samples: List of all available sample paths
                scene_idx: Current scene index
                total_scenes: Total number of scenes
                energy: Energy level (0.0-1.0)
                prev_samples: List of samples used in previous scene
                sample_type: "kick" or "snare" for logging
            
            Returns:
                Selected sample path or None
            """
            if not all_samples:
                return None
            
            # Energy-based keyword filtering
            soft_keywords = ["soft", "light", "minimal", "gentle", "quiet", "smooth"]
            hard_keywords = ["hard", "heavy", "punch", "kick", "strong", "aggressive", "tight", "solid"]
            
            # Filter samples based on energy level
            if energy < 0.3:
                # Low energy: prefer soft/light samples
                filtered = [s for s in all_samples if any(kw in s.lower() for kw in soft_keywords)]
                selection_pool = filtered if filtered else all_samples
            elif energy > 0.8:
                # High energy: prefer hard/heavy/punch samples
                filtered = [s for s in all_samples if any(kw in s.lower() for kw in hard_keywords)]
                selection_pool = filtered if filtered else all_samples
            else:
                # Medium energy: use all samples
                selection_pool = all_samples
            
            # Remove samples used in previous scene (no repetition policy)
            available = [s for s in selection_pool if s not in prev_samples]
            
            # If not enough samples after filtering, fall back to all samples (excluding prev)
            if len(available) < 1:
                available = [s for s in all_samples if s not in prev_samples]
            
            # If still no samples (all were used in prev), use full pool
            if not available:
                available = all_samples
            
            # Select sample with least usage count for even rotation
            min_usage = float('inf')
            best_candidates = []
            
            usage_dict = _scene_sample_usage.get(sample_type, {})
            for sample in available:
                usage_count = usage_dict.get(sample, 0)
                if usage_count < min_usage:
                    min_usage = usage_count
                    best_candidates = [sample]
                elif usage_count == min_usage:
                    best_candidates.append(sample)
            
            # Pick first from best candidates (they have equal lowest usage)
            selected = best_candidates[0] if best_candidates else available[0] if available else None
            
            if selected:
                # Update usage tracking
                usage_dict[selected] = usage_dict.get(selected, 0) + 1
                _scene_sample_usage[sample_type] = usage_dict
                
                # Track global usage order
                if sample_type == "kick" and selected not in _all_kicks_used:
                    _all_kicks_used.append(selected)
                elif sample_type == "snare" and selected not in _all_snares_used:
                    _all_snares_used.append(selected)
            
            return selected

        def _get_velocity_for_energy(energy, drum_type="kick"):
            """
            Get velocity range based on energy level and drum type.
            
            Args:
                energy: Energy level (0.0-1.0)
                drum_type: "kick" or "snare"
            
            Returns:
                Tuple of (min_velocity, max_velocity)
            """
            if energy < 0.4:
                # Low energy: softer velocities
                if drum_type == "kick":
                    return (70, 80)
                else:  # snare
                    return (65, 75)
            elif energy <= 0.7:
                # Medium energy
                if drum_type == "kick":
                    return (85, 85)  # Fixed at 85
                else:  # snare
                    return (80, 80)  # Fixed at 80
            else:
                # High energy: loud velocities
                if drum_type == "kick":
                    return (95, 110)
                else:  # snare
                    return (90, 100)

        # Sort drum loops by BPM proximity to tempo'''
    
    if old_function not in content:
        print("ERROR: Could not find the old function!")
        # Try to find it
        idx = content.find('def _pick_for_scene')
        if idx >= 0:
            print(f"Found at position {idx}")
            print("Context:", repr(content[idx:idx+300]))
        return 1
    
    # Replace
    new_content = content.replace(old_function, new_function)
    
    # Write back
    with open(filepath, 'w', encoding='utf-8') as f:
        f.write(new_content)
    
    print("SUCCESS: Added _pick_for_scene_advanced and _get_velocity_for_energy")
    print(f"File size changed from {len(content)} to {len(new_content)}")
    return 0

if __name__ == '__main__':
    sys.exit(main())
