batch_processor.py

"""
Batch Processing Script for Image Prompt Variations
Automates generation of multiple prompt variations from input templates
"""

import json
from pathlib import Path
from typing import List, Dict, Optional
from core.orchestrator import PromptOrchestrator
from core.modifiers import ModifierCategory


class BatchProcessor:
    """Handles batch processing of image prompts with various automation features"""
    
    def __init__(self, output_dir: str = "output"):
        """
        Initialize batch processor.
        
        Args:
            output_dir: Directory to save output files
        """
        self.orchestrator = PromptOrchestrator()
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(exist_ok=True)
    
    def process_single_prompt(self, base_prompt: str, config: Dict) -> List[Dict]:
        """
        Process a single prompt with specified configuration.
        
        Args:
            base_prompt: The base prompt to vary
            config: Configuration dict with keys:
                - lighting_ids: List of modifier IDs (or "all", "dry", "wet")
                - styles: List of style names to apply
                - export_format: Output format (txt, json, csv)
        
        Returns:
            List of variation dictionaries
        """
        # Determine which lighting modifiers to use
        if config.get("lighting_ids") == "all":
            modifier_ids = list(range(1, 21))
        elif config.get("lighting_ids") == "dry":
            dry_mods = self.orchestrator.get_modifiers_by_category(ModifierCategory.DRY)
            modifier_ids = [m.id for m in dry_mods]
        elif config.get("lighting_ids") == "wet":
            wet_mods = self.orchestrator.get_modifiers_by_category(ModifierCategory.WET)
            modifier_ids = [m.id for m in wet_mods]
        else:
            modifier_ids = config.get("lighting_ids", [])
        
        # Generate variations
        variations = self.orchestrator.generate_variation_batch(
            base_prompt=base_prompt,
            modifier_ids=modifier_ids if modifier_ids else None,
            include_styles=config.get("styles", None)
        )
        
        return variations
    
    def process_multiple_prompts(self, prompts: List[str], config: Dict,
                                output_filename: str = "batch_output") -> str:
        """
        Process multiple base prompts and save results.
        
        Args:
            prompts: List of base prompts
            config: Configuration dictionary
            output_filename: Name for output file (without extension)
        
        Returns:
            Path to output file
        """
        all_results = []
        
        for idx, prompt in enumerate(prompts, 1):
            variations = self.process_single_prompt(prompt, config)
            all_results.append({
                "base_prompt_id": idx,
                "base_prompt": prompt,
                "variations": variations
            })
        
        # Export results
        export_format = config.get("export_format", "json")
        output_path = self.output_dir / f"{output_filename}.{export_format}"
        
        if export_format == "json":
            with open(output_path, "w", encoding="utf-8") as f:
                json.dump(all_results, f, indent=2)
        elif export_format == "txt":
            with open(output_path, "w", encoding="utf-8") as f:
                for result in all_results:
                    f.write(f"\n{'=' * 80}\n")
                    f.write(f"BASE PROMPT #{result['base_prompt_id']}: {result['base_prompt']}\n")
                    f.write(f"{'=' * 80}\n\n")
                    for var in result['variations']:
                        f.write(f"[{var['id']}] {var['type'].upper()} - {var['category']}\n")
                        f.write(f"  Description: {var['description']}\n")
                        f.write(f"  Prompt: {var['prompt']}\n\n")
        elif export_format == "csv":
            with open(output_path, "w", encoding="utf-8") as f:
                f.write("base_prompt_id,base_prompt,variation_id,type,category,description,prompt\n")
                for result in all_results:
                    for var in result['variations']:
                        f.write(f"{result['base_prompt_id']},\"{result['base_prompt']}\",")
                        f.write(f"{var['id']},{var['type']},{var['category']},")
                        f.write(f"\"{var['description']}\",\"{var['prompt']}\"\n")
        
        return str(output_path)
    
    def create_preset_collection(self, base_prompt: str, preset: str) -> List[Dict]:
        """
        Apply a preset collection of variations.
        
        Presets:
            - "full_day_cycle": All DRY modifiers (complete day cycle)
            - "wet_night": All WET modifiers (dramatic night scenes)
            - "key_moments": Selected dramatic lighting moments
            - "commercial": Best modifiers for commercial photography
        
        Args:
            base_prompt: The base prompt
            preset: Preset name
        
        Returns:
            List of variations
        """
        preset_configs = {
            "full_day_cycle": {
                "lighting_ids": "dry",
                "styles": None
            },
            "wet_night": {
                "lighting_ids": "wet",
                "styles": None
            },
            "key_moments": {
                "lighting_ids": [1, 7, 8, 10, 18, 20],  # Sunrise, golden hour, sunset, blue hours, wet night
                "styles": None
            },
            "commercial": {
                "lighting_ids": [3, 4, 7, 11],  # Crisp lighting for product shots
                "styles": None
            },
            "artistic": {
                "lighting_ids": [7, 8, 9, 10],
                "styles": ["impressionist", "minimalist", "art_deco"]
            }
        }
        
        if preset not in preset_configs:
            raise ValueError(f"Unknown preset: {preset}. Available: {list(preset_configs.keys())}")
        
        return self.process_single_prompt(base_prompt, preset_configs[preset])
    
    def generate_comparison_set(self, base_prompt: str, 
                               modifier_pairs: List[tuple]) -> List[Dict]:
        """
        Generate side-by-side comparison prompts.
        
        Args:
            base_prompt: Base prompt
            modifier_pairs: List of (modifier_id_1, modifier_id_2) tuples
        
        Returns:
            List of comparison variation sets
        """
        comparisons = []
        
        for mod_a, mod_b in modifier_pairs:
            prompt_a = self.orchestrator.apply_lighting_modifier(base_prompt, mod_a)
            prompt_b = self.orchestrator.apply_lighting_modifier(base_prompt, mod_b)
            
            comparisons.append({
                "comparison_type": f"modifier_{mod_a}_vs_{mod_b}",
                "variant_a": {
                    "modifier_id": mod_a,
                    "description": self.orchestrator.modifiers[mod_a].description,
                    "prompt": prompt_a
                },
                "variant_b": {
                    "modifier_id": mod_b,
                    "description": self.orchestrator.modifiers[mod_b].description,
                    "prompt": prompt_b
                }
            })
        
        return comparisons


def example_workflows():
    """Demonstrate various workflow scenarios"""
    processor = BatchProcessor(output_dir="prompt_variations")
    
    print("BATCH PROCESSING WORKFLOWS\n")
    print("=" * 80)
    
    # Workflow 1: Single prompt with all lighting variations
    print("\n1. FULL LIGHTING CYCLE GENERATION")
    print("-" * 80)
    base_prompt_1 = "Luxury sedan in urban environment, architectural background"
    
    config_full = {
        "lighting_ids": "all",
        "export_format": "json"
    }
    
    output_1 = processor.process_multiple_prompts(
        prompts=[base_prompt_1],
        config=config_full,
        output_filename="full_lighting_cycle"
    )
    print(f"✓ Generated full lighting cycle → {output_1}")
    
    # Workflow 2: Multiple prompts with preset
    print("\n2. KEY MOMENTS PRESET (Multiple Prompts)")
    print("-" * 80)
    base_prompts = [
        "Sports car on mountain road, scenic vista",
        "Classic convertible on beach, sunset backdrop",
        "Electric SUV in modern city, glass buildings"
    ]
    
    variations = []
    for prompt in base_prompts:
        vars = processor.create_preset_collection(prompt, "key_moments")
        variations.extend(vars)
    
    print(f"✓ Generated {len(variations)} variations across {len(base_prompts)} base prompts")
    
    # Workflow 3: Style transformations
    print("\n3. STYLE TRANSFORMATION SET")
    print("-" * 80)
    base_prompt_3 = "Vintage motorcycle in garage, mechanical details"
    
    config_styles = {
        "lighting_ids": [7],  # Golden hour only
        "styles": ["cyberpunk", "art_deco", "minimalist", "noir"],
        "export_format": "txt"
    }
    
    output_3 = processor.process_multiple_prompts(
        prompts=[base_prompt_3],
        config=config_styles,
        output_filename="style_variations"
    )
    print(f"✓ Generated style variations → {output_3}")
    
    # Workflow 4: Comparison set
    print("\n4. COMPARISON SET GENERATION")
    print("-" * 80)
    base_prompt_4 = "Modern architecture, geometric forms"
    
    comparison_pairs = [
        (4, 7),   # Midday vs Golden Hour
        (11, 18), # Night dry vs Night wet
        (1, 8)    # Sunrise vs Sunset
    ]
    
    comparisons = processor.generate_comparison_set(base_prompt_4, comparison_pairs)
    
    # Save comparisons
    output_4 = processor.output_dir / "comparisons.json"
    with open(output_4, "w", encoding="utf-8") as f:
        json.dump(comparisons, f, indent=2)
    
    print(f"✓ Generated {len(comparisons)} comparison sets → {output_4}")
    
    # Workflow 5: Commercial photography set
    print("\n5. COMMERCIAL PHOTOGRAPHY PRESET")
    print("-" * 80)
    base_prompt_5 = "Product shot of smartwatch, clean background"
    
    commercial_vars = processor.create_preset_collection(base_prompt_5, "commercial")
    print(f"✓ Generated {len(commercial_vars)} commercial-grade variations")
    
    print("\n" + "=" * 80)
    print("All workflows completed successfully!")
    print(f"Output directory: {processor.output_dir}")


def quick_generation_example():
    """Quick example for immediate use"""
    orchestrator = PromptOrchestrator()
    
    # Your base prompt
    base = "Ferrari on coastal highway at sunset"
    
    print("\nQUICK GENERATION EXAMPLES")
    print("=" * 80)
    
    # Generate 5 quick variations
    variations = [
        orchestrator.apply_lighting_modifier(base, 1),  # Sunrise
        orchestrator.apply_lighting_modifier(base, 7),  # Golden hour
        orchestrator.apply_lighting_modifier(base, 10), # Blue hour
        orchestrator.apply_lighting_modifier(base, 18), # Blue hour wet
        orchestrator.apply_style_transformation(base, "cyberpunk")
    ]
    
    for i, var in enumerate(variations, 1):
        print(f"\n{i}. {var}")


if __name__ == "__main__":
    # Run example workflows
    example_workflows()
    
    # Quick generation
    quick_generation_example()