#!/usr/bin/env python3
import sys
import os
import argparse
import re

def modify_masked_vision_encoder_index(file_path, mask_indices):
    """
    Modify the masked_vision_encoder_index in the deepseek_vl.py file
    
    Args:
        file_path: Path to the deepseek_vl.py file
        mask_indices: List of indices to mask (e.g., [0, 1, 2] or [] for no masking)
    """
    try:
        # Read the file
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
        
        # Pattern to match the masked_vision_encoder_index parameter
        pattern = r'masked_vision_encoder_index=\[\]'
        replacement = f'masked_vision_encoder_index={mask_indices}'
        
        # Replace the pattern
        modified_content = re.sub(pattern, replacement, content)
        
        # Check if replacement was made
        if content == modified_content:
            print(f"Warning: No changes made. Pattern not found in {file_path}")
            return False
        
        # Write back to file
        with open(file_path, 'w', encoding='utf-8') as f:
            f.write(modified_content)
        
        print(f"Successfully modified masked_vision_encoder_index to {mask_indices}")
        return True
        
    except Exception as e:
        print(f"Error modifying file: {e}")
        return False

def main():
    parser = argparse.ArgumentParser(description='Modify masked_vision_encoder_index in DeepSeek VL model')
    parser.add_argument('--file', '-f', required=True, 
                       help='Path to deepseek_vl.py file')
    parser.add_argument('--indices', '-i', nargs='*', type=int, default=[],
                       help='Vision encoder indices to mask (space-separated integers)')
    parser.add_argument('--restore', '-r', action='store_true',
                       help='Restore to empty list (no masking)')
    
    args = parser.parse_args()
    
    # Determine mask indices
    if args.restore:
        mask_indices = []
    else:
        mask_indices = args.indices
    
    # Modify the file
    success = modify_masked_vision_encoder_index(args.file, mask_indices)
    
    if success:
        print(f"File modified successfully with mask indices: {mask_indices}")
        sys.exit(0)
    else:
        print("Failed to modify file")
        sys.exit(1)

if __name__ == "__main__":
    main()