#!/usr/bin/env python
"""Script to modify per-scene kick/snare loading for Fases 11-15"""

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()
    
    # The old kick/snare loading code to replace
    old_code = '''            # Kick — only in drum sections
            if flags.get("drums"):
                sample = _pick_for_scene(all_kicks, si, total_scenes)
                if sample and _load_audio(track_map["kick"], sample, si):
                    samples_loaded += 1

            # Snare — only in drum sections
            if flags.get("drums"):
                sample = _pick_for_scene(all_snares, si, total_scenes)
                if sample and _load_audio(track_map["snare"], sample, si):
                    samples_loaded += 1'''
    
    # New code with Fases 11-15 implementation
    new_code = '''            # ================================================================
            # FASES 11-15: VARIACION MASIVA DE KICKS Y SNARES
            # ================================================================
            
            # Scene 0 (Intro): NO kicks/snares loaded
            if si == 0:
                # Intro scene - skip all drum samples
                pass
            elif flags.get("drums"):
                # Get velocity ranges based on energy
                kick_vel_min, kick_vel_max = _get_velocity_for_energy(energy, "kick")
                snare_vel_min, snare_vel_max = _get_velocity_for_energy(energy, "snare")
                
                # Determine how many kicks/snares to load based on energy
                if energy > 0.8:
                    num_kicks = 3   # High energy: 3 kicks
                    num_snares = 2  # High energy: 2 snares
                elif energy > 0.5:
                    num_kicks = 2   # Medium energy: 2 kicks
                    num_snares = 2  # Medium energy: 2 snares
                else:
                    num_kicks = 2   # Low energy: 2 kicks
                    num_snares = 1  # Low energy: 1 snare
                
                # Get previous scene samples to avoid repetition
                prev_kicks = _prev_scene_samples.get("kicks", [])
                prev_snares = _prev_scene_samples.get("snares", [])
                
                current_scene_kicks = []
                current_scene_snares = []
                
                # Load multiple kicks per scene with advanced picker
                for kick_idx in range(num_kicks):
                    sample = _pick_for_scene_advanced(
                        all_kicks, si, total_scenes, energy, prev_kicks if kick_idx == 0 else current_scene_kicks, 
                        sample_type="kick"
                    )
                    if sample:
                        # Determine which track to load into
                        # Use multiple kick tracks if available, otherwise use main kick track
                        kick_track_key = "kick" if kick_idx == 0 else "kick_%d" % (kick_idx + 1)
                        if kick_track_key in track_map:
                            tidx = track_map[kick_track_key]
                        else:
                            tidx = track_map.get("kick", 0)
                        
                        if _load_audio(tidx, sample, si):
                            samples_loaded += 1
                            current_scene_kicks.append(sample)
                            # Apply velocity based on energy
                            try:
                                t = self._song.tracks[tidx]
                                if slot.has_clip and hasattr(slot.clip, 'velocity'):
                                    import random
                                    slot.clip.velocity = random.randint(kick_vel_min, kick_vel_max)
                            except:
                                pass
                
                # Load multiple snares per scene with advanced picker
                for snare_idx in range(num_snares):
                    sample = _pick_for_scene_advanced(
                        all_snares, si, total_scenes, energy, prev_snares if snare_idx == 0 else current_scene_snares,
                        sample_type="snare"
                    )
                    if sample:
                        # Determine which track to load into
                        snare_track_key = "snare" if snare_idx == 0 else "snare_%d" % (snare_idx + 1)
                        if snare_track_key in track_map:
                            tidx = track_map[snare_track_key]
                        else:
                            tidx = track_map.get("snare", 0)
                        
                        if _load_audio(tidx, sample, si):
                            samples_loaded += 1
                            current_scene_snares.append(sample)
                            # Apply velocity based on energy
                            try:
                                t = self._song.tracks[tidx]
                                if slot.has_clip and hasattr(slot.clip, 'velocity'):
                                    import random
                                    slot.clip.velocity = random.randint(snare_vel_min, snare_vel_max)
                            except:
                                pass
                
                # Update previous scene samples for next iteration
                _prev_scene_samples["kicks"] = current_scene_kicks[:]
                _prev_scene_samples["snares"] = current_scene_snares[:]
                
                # Log scene details
                log.append("scene %d (%s): kicks=%d, snares=%d, energy=%.2f, kick_vel=%d-%d, snare_vel=%d-%d" % (
                    si, scene_name, len(current_scene_kicks), len(current_scene_snares),
                    energy, kick_vel_min, kick_vel_max, snare_vel_min, snare_vel_max
                ))'''
    
    if old_code not in content:
        print("ERROR: Could not find the old kick/snare loading code!")
        # Try to find approximate location
        idx = content.find('# Kick')
        if idx >= 0:
            print(f"Found '# Kick' at position {idx}")
            print("Context:", repr(content[idx:idx+500]))
        return 1
    
    # Replace
    new_content = content.replace(old_code, new_code)
    
    # Write back
    with open(filepath, 'w', encoding='utf-8') as f:
        f.write(new_content)
    
    print("SUCCESS: Replaced kick/snare loading with Fases 11-15 implementation")
    print(f"File size changed from {len(content)} to {len(new_content)}")
    return 0

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