rabbit_generation.py

"""
Rabbit Generation - Construction generation when DDAR reaches a fixed point

This module provides functionality to generate constructions when DDAR fails to
derive new facts, either through user input or LLM generation.
"""

from typing import Tuple, List, Any
from Problem import GeometricProblem
from Constructions import GeometricConstructor
import matplotlib.pyplot as plt
import random


class RabbitGenerator:
    """Generates constructions when DDAR reaches a fixed point"""
    
    def __init__(self, mode: str = 'user', max_rabbit_calls: int = 1, num_random_constructions: int = 3, sanitize_on_retry: bool = False):
        """
        Initialize the RabbitGenerator.
        
        Args:
            mode: Generation mode - 'user' (prompt user), 'llm' (use LLM), or 'random' (random constructions)
            max_rabbit_calls: Maximum number of times rabbit generation can be called
            num_random_constructions: Number of random constructions to try when mode='random'
            sanitize_on_retry: If True, remove constructed points when retrying after a failed round
        """
        if mode not in ['user', 'llm', 'random']:
            raise ValueError(f"Invalid mode '{mode}'. Must be 'user', 'llm', or 'random'")
        
        self.mode = mode
        self.max_rabbit_calls = max_rabbit_calls
        self.rabbit_call_count = 0
        self.num_random_constructions = num_random_constructions
        self.sanitize_on_retry = sanitize_on_retry
        
        # Track all constructed points across all rabbit calls
        self.constructed_points = set()
        self.constructed_point_names = set()
        
        # Construction definitions (same as in Data_Generation.py)
        self.constructions = [
            {'method': 'construct_midpoint', 'num_points': 2, 'name_param': 'name', 'mark_param': 'mark_point'},
            {'method': 'construct_angle_bisector', 'num_points': 3, 'name_param': 'pname', 'mark_param': 'mark_points'},
            {'method': 'construct_external_angle_bisector', 'num_points': 3, 'name_param': 'pname', 'mark_param': 'mark_points'},
            {'method': 'construct_foot', 'num_points': 3, 'name_param': 'name', 'mark_param': 'mark_points'},
            {'method': 'construct_circle', 'num_points': 3, 'name_param': 'cname', 'mark_param': None},
            {'method': 'construct_incenter', 'num_points': 3, 'name_param': 'name', 'mark_param': None},
            {'method': 'construct_incenter2', 'num_points': 3, 'name_param': 'name', 'mark_param': 'mark_points'},
            {'method': 'construct_excenter', 'num_points': 3, 'name_param': 'name', 'mark_param': None},
            {'method': 'construct_excenter2', 'num_points': 3, 'name_param': 'name', 'mark_param': 'mark_points'},
            {'method': 'construct_centroid', 'num_points': 3, 'name_param': 'name', 'mark_param': None},
            {'method': 'construct_orthocenter', 'num_points': 3, 'name_param': 'name', 'mark_param': None},
            {'method': 'construct_orthocenter2', 'num_points': 3, 'name_param': 'name', 'mark_param': None},
            {'method': 'construct_mirror', 'num_points': 2, 'name_param': 'name', 'mark_param': 'mark_points'},
            {'method': 'construct_on_dia', 'num_points': 2, 'name_param': 'name', 'mark_param': None},
            {'method': 'construct_intersect_lines', 'num_points': 4, 'name_param': 'name', 'mark_param': 'mark_points'},
            {'method': 'construct_reflection', 'num_points': 3, 'name_param': 'name1', 'mark_param': None},
            {'method': 'construct_tangents', 'num_points': 3, 'name_param': 'name1', 'mark_param': None},
        ]
        
    def reset_count(self):
        """Reset the rabbit call counter and constructed points tracking"""
        self.rabbit_call_count = 0
        self.constructed_points = set()
        self.constructed_point_names = set()
        
    def can_generate(self) -> bool:
        """Check if we can still call rabbit generation"""
        return self.rabbit_call_count < self.max_rabbit_calls
    
    def generate_construction(self, problem: GeometricProblem, 
                            derived_facts: List[Any], sanitize_before: bool = False) -> Tuple[bool, List[Any], List[dict]]:
        """
        Generate a construction to help DDAR progress.
        
        Args:
            problem: The geometric problem being solved
            derived_facts: List of facts derived so far
            sanitize_before: If True, remove previously constructed points before generating new construction
            
        Returns:
            Tuple of (success, new_predicates, reasoning_steps) where:
                - success: True if construction was successfully applied
                - new_predicates: List of new predicates from the construction
                - reasoning_steps: List of reasoning dictionaries for the new predicates
        """
        if not self.can_generate():
            print("\n✗ Maximum rabbit generation calls reached")
            return False, [], []
        
        # Sanitize problem if requested (remove previously constructed points)
        if sanitize_before and self.sanitize_on_retry and self.constructed_point_names:
            print("\n🧹 Sanitizing problem: removing previously constructed points...")
            num_removed = self._sanitize_problem(problem)
            print(f"  Removed {num_removed} facts involving constructed points")
        
        self.rabbit_call_count += 1
        print(f"\n{'='*60}")
        print(f"RABBIT GENERATION (Call {self.rabbit_call_count}/{self.max_rabbit_calls}) - Mode: {self.mode.upper()}")
        print(f"{'='*60}")
        
        # Get construction code based on mode
        if self.mode == 'llm':
            construction_code = self._llm_generate_construction(problem, derived_facts)
        elif self.mode == 'random':
            construction_code = self._random_generate_construction(problem, derived_facts)
        else:  # mode == 'user'
            construction_code = self._user_input_construction(problem, derived_facts)
        
        if not construction_code:
            print("✗ No construction code provided")
            return False, [], []
        
        # Apply the construction
        success, new_predicates, reasoning_steps = self._apply_construction(
            problem, construction_code
        )
        
        if success:
            print(f"✓ Construction applied successfully: {len(new_predicates)} new predicates")
        else:
            print("✗ Construction failed to apply")
            
        return success, new_predicates, reasoning_steps
    
    def _user_input_construction(self, problem: GeometricProblem, 
                                 derived_facts: List[Any]) -> str:
        """
        Prompt user to provide construction code.
        
        Args:
            problem: The geometric problem
            derived_facts: Current derived facts
            
        Returns:
            Construction code string
        """
        print("\nCurrent problem state:")
        print(f"  Derived facts: {len(derived_facts)}")
        print(f"  Goals: {len(problem.goals)}")
        
        # Display all available points prominently
        point_names = sorted(problem.points.keys())
        print(f"\n{'='*60}")
        print(f"AVAILABLE POINTS: {', '.join(point_names)}")
        print(f"{'='*60}")
        
        print("\nSee the documentation for construction methods.")
        print("\nEnter construction code (or press Enter to skip):")
        construction_code = input("> ").strip()
        
        return construction_code
    
    def _random_generate_construction(self, problem: GeometricProblem,
                                      derived_facts: List[Any]) -> str:
        """
        Generate random construction code.
        
        Args:
            problem: The geometric problem
            derived_facts: Current derived facts
            
        Returns:
            Construction code string, or None if no construction could be generated
        """
        print(f"\nAttempting {self.num_random_constructions} random construction(s)...")
        
        available_points = list(problem.points.values())
        if len(available_points) < 2:
            print("  Not enough points for constructions")
            return None
        
        applicable = [c for c in self.constructions if len(available_points) >= c['num_points']]
        if not applicable:
            print(f"  No applicable constructions for {len(available_points)} points")
            return None
        
        # Try multiple random constructions
        for attempt in range(self.num_random_constructions):
            construction = random.choice(applicable)
            selected_points = random.sample(available_points, construction['num_points'])
            
            point_names = ''.join(pt.name for pt in selected_points)
            new_name = f"{construction['method']}_{point_names}_{random.randint(0, 999)}"
            
            # Build construction code string
            point_args = ', '.join(pt.name for pt in selected_points)
            kwargs = [f"{construction['name_param']}='{new_name}'"]
            
            if construction.get('mark_param'):
                kwargs.append(f"{construction['mark_param']}=False")
            
            construction_code = f"constructor.{construction['method']}({point_args}, {', '.join(kwargs)})"
            
            print(f"  Attempt {attempt + 1}: {construction['method']} on {point_names}")
            print(f"    Code: {construction_code}")
            
            # Test if this construction would work
            try:
                # Create a temporary constructor to validate
                temp_constructor = GeometricConstructor(problem)
                method = getattr(temp_constructor, construction['method'])
                
                # Build kwargs
                kwargs_dict = {construction['name_param']: new_name}
                if construction.get('mark_param'):
                    kwargs_dict[construction['mark_param']] = False
                
                # Try to execute (with add_results=False to not modify the problem yet)
                result = method(*selected_points, **kwargs_dict, add_results=False)
                
                # Clean up
                plt.close('all')
                
                # If we got here, construction worked
                if result and len(result) >= 2 and result[1]:
                    print(f"  ✓ Valid construction found")
                    return construction_code
                    
            except Exception as e:
                print(f"    ✗ Failed: {e}")
                plt.close('all')
                continue
        
        print(f"  No valid construction found after {self.num_random_constructions} attempts")
        return None
    
    def _llm_generate_construction(self, problem: GeometricProblem, 
                                   derived_facts: List[Any]) -> str:
        """
        Use LLM to generate construction code.
        
        Args:
            problem: The geometric problem
            derived_facts: Current derived facts
            
        Returns:
            Construction code string
        """
        # TODO: Implement LLM-based construction generation
        # This is a placeholder for future implementation
        print("\n⚠ LLM construction generation not yet implemented")
        print("Falling back to user input...")
        return self._user_input_construction(problem, derived_facts)
    
    def _apply_construction(self, problem: GeometricProblem, 
                           construction_code: str) -> Tuple[bool, List[Any], List[dict]]:
        """
        Apply the construction code to the problem.
        
        Args:
            problem: The geometric problem
            construction_code: Python code for the construction
            
        Returns:
            Tuple of (success, new_predicates, reasoning_steps)
        """
        if not construction_code:
            return False, [], []
        
        try:
            # Create constructor
            constructor = GeometricConstructor(problem)
            
            # Prepare namespace for executing construction code
            namespace = {
                'constructor': constructor,
                **problem.points  # Make all points available by name
            }
            
            # Execute the construction code
            result = eval(construction_code, namespace)
            
            # Clean up matplotlib figures
            plt.close('all')
            
            # Extract new predicates and points from the result
            if isinstance(result, tuple) and len(result) >= 2:
                constructed_objects, new_predicates = result[0], result[1]
                
                # Track constructed points
                if constructed_objects:
                    # Handle single point or tuple of points
                    if isinstance(constructed_objects, tuple):
                        for obj in constructed_objects:
                            if hasattr(obj, 'name'):
                                self.constructed_points.add(obj)
                                self.constructed_point_names.add(obj.name)
                    elif hasattr(constructed_objects, 'name'):
                        self.constructed_points.add(constructed_objects)
                        self.constructed_point_names.add(constructed_objects.name)
                
                # Handle different return types
                if not isinstance(new_predicates, list):
                    new_predicates = [new_predicates] if new_predicates else []
                
                # Create reasoning steps for the new predicates
                reasoning_steps = []
                for pred in new_predicates:
                    reasoning_steps.append({
                        'predicate': pred,
                        'dependencies': {
                            'type': 'RABBIT_CONSTRUCTION',
                            'construction_code': construction_code,
                            'premises': []
                        }
                    })
                
                print(f"\n  Construction result: {len(new_predicates)} new predicates")
                for pred in new_predicates:
                    print(f"    → {pred}")
                
                if self.constructed_point_names:
                    print(f"  Tracked constructed points: {', '.join(sorted(self.constructed_point_names))}")
                
                return True, new_predicates, reasoning_steps
            else:
                print(f"\n  Unexpected construction result type: {type(result)}")
                return False, [], []
                
        except Exception as e:
            print(f"\n✗ Error applying construction: {e}")
            import traceback
            traceback.print_exc()
            return False, [], []
    
    def _sanitize_problem(self, problem: GeometricProblem) -> int:
        """
        Remove all facts involving constructed points from the problem.
        
        Args:
            problem: The geometric problem to sanitize
            
        Returns:
            Number of facts removed
        """
        if not self.constructed_point_names:
            return 0
        
        def predicate_involves_constructed(pred: Any) -> bool:
            """Check if a predicate involves any constructed points"""
            if not hasattr(pred, 'get_points'):
                return False
            
            pred_points = pred.get_points()
            return any(pt.name in self.constructed_point_names for pt in pred_points)
        
        # Remove facts involving constructed points from derived_facts
        original_count = len(problem.derived_facts)
        problem.derived_facts = [f for f in problem.derived_facts 
                                if not predicate_involves_constructed(f)]
        
        # Remove from dd_derived
        if hasattr(problem, 'dd_derived'):
            problem.dd_derived = [f for f in problem.dd_derived 
                                 if not predicate_involves_constructed(f)]
        
        # Remove from assumptions (constructed predicates)
        problem.assumptions = [a for a in problem.assumptions 
                              if not predicate_involves_constructed(a)]
        
        # Remove constructed points from problem.points
        for point_name in list(problem.points.keys()):
            if point_name in self.constructed_point_names:
                del problem.points[point_name]
        
        # Clean up reasoning steps if they exist
        if hasattr(problem, 'reasoning_steps'):
            for level in problem.reasoning_steps:
                # Filter out reasoning involving constructed points
                filtered_level = []
                for reasoning in level:
                    pred = reasoning.get('predicate')
                    if not predicate_involves_constructed(pred):
                        filtered_level.append(reasoning)
                level[:] = filtered_level
        
        removed_count = original_count - len(problem.derived_facts)
        return removed_count