Problem.py

"""
Problem.py - Data structure for storing geometric problems

This module defines how we store the initial assumptions and goals of a geometric problem.
"""

from typing import List, Dict, Tuple, Any
from relations import predicate, geometric


class GeometricProblem:
    """
    Represents a geometric problem with assumptions and goals.
    """

    def __init__(self):
        """Initialize an empty geometric problem."""
        self.points: Dict[str, geometric.Point] = {}
        self.assumptions: List[Any] = []  # List of predicate instances
        self.derived_facts: List[Any] = [] # List of facts derived through the solution of the problem
        self.dd_derived: List[Any] = []
        self.reasoning_steps: List[List[dict]] = [] # List of explanations of derived facts in a dict form
        self.proved_goals: List[dict] = [] # List of goals that have been already proved
        self.goals: List[Any] = []  # List of predicate instances to prove
        self.points: Dict[str, geometric.Point] = {}
        self.lines: Dict[str, Tuple[float, float, float]] = {} # we should add line and circle objects for this
        self.circles: Dict[str, Tuple[float, float, float]] = {}
        
    def add_point(self, name: str, x: float, y: float):
        """Add a point to the problem."""
        self.points[name] = geometric.Point(name, x, y)
        
    def add_assumption(self, predicate_instance):
        """Add an assumption (given predicate) to the problem."""
        self.assumptions.append(predicate_instance)
        
    def add_goal(self, predicate_instance):
        """Add a goal (predicate to prove) to the problem."""
        self.goals.append(predicate_instance)

    def remove_goal(self, predicate_instance):
        """Remove a goal from the problem."""
        self.goals.remove(predicate_instance)
        
    def set_geometric_data(self, points_dict, lines_dict, circles_dict):
        """Set the geometric data from parsed GeoGebra file."""
        # Add points
        for name, (x, y) in points_dict.items():
            self.add_point(name, x, y)
            
        self.lines = lines_dict
        self.circles = circles_dict
        
    def get_point(self, name: str) -> geometric.Point:
        """Get a point by name."""
        if name not in self.points:
            raise ValueError(f"Point {name} not found in problem")
        return self.points[name]
        
    def __str__(self):
        """String representation of the problem."""
        result = f"Points: {list(self.points.keys())}\n"
        result += f"Assumptions: {len(self.assumptions)}\n"
        result += f"Goals: {len(self.goals)}\n"
        result += f"Lines: {len(self.lines)}\n"
        result += f"Circles: {len(self.circles)}\n"
        return result

    def pretty_print_reasoning_steps(self):
        """Assumes the problem has been solved.
           Prints the reasoning steps leading to each goal
           using problem.reasoning_steps.
           We assume that reasoning_steps is a list of lists of dictionaries."""
        if not self.reasoning_steps:
            print("No reasoning steps available.")
            return

        # Flatten indexes: by object identity; keep list for __eq__ comparison; string as last resort
        pred_index_obj = {}
        pred_index_str = {}
        pred_entries = []  # list of tuples (pred_obj, level, deps)
        def _iter_reasoning_dicts(container):
            """Yield all dict entries, flattening any nested lists/tuples."""
            stack = [container]
            while stack:
                item = stack.pop()
                if item is None:
                    continue
                if isinstance(item, dict):
                    yield item
                elif isinstance(item, (list, tuple)):
                    # push children to handle multi-level nesting such as [[{...}]]
                    stack.extend(item)
                # ignore other types (e.g., stray predicates) silently

        for level, level_list in enumerate(self.reasoning_steps):
            for entry in level_list:
                for sub in _iter_reasoning_dicts(entry):
                    if 'predicate' in sub:
                        pred_obj = sub.get('predicate')
                        deps = sub.get('dependencies', {}) or {}
                        if pred_obj is None:
                            continue
                        pred_index_obj[pred_obj] = (level, pred_obj, deps)
                        pred_index_str.setdefault(str(pred_obj), []).append((level, pred_obj, deps))
                        pred_entries.append((pred_obj, level, deps))
                    else:
                        for pred_obj, deps in sub.items():
                            deps = deps or {}
                            if pred_obj is None:
                                continue
                            pred_index_obj[pred_obj] = (level, pred_obj, deps)
                            pred_index_str.setdefault(str(pred_obj), []).append((level, pred_obj, deps))
                            pred_entries.append((pred_obj, level, deps))

        # Add proved goals and derived facts as fallback nodes (unknown deps) if missing
        fallback_nodes = []
        for pg in getattr(self, 'proved_goals', []) or []:
            if isinstance(pg, dict) and 'predicate' in pg:
                fallback_nodes.append(pg['predicate'])
        for df in getattr(self, 'derived_facts', []) or []:
            fallback_nodes.append(df)

        for pred_obj in fallback_nodes:
            if pred_obj is None:
                continue
            if pred_obj not in pred_index_obj:
                pred_index_obj[pred_obj] = (None, pred_obj, {'type': 'UNKNOWN', 'premises': []})
                pred_index_str.setdefault(str(pred_obj), []).append((None, pred_obj, {'type': 'UNKNOWN', 'premises': []}))
                pred_entries.append((pred_obj, None, {'type': 'UNKNOWN', 'premises': []}))
        def find_entry(pred):
            # Exact object match
            entry = pred_index_obj.get(pred)
            if entry:
                return entry
            # Equality match using __eq__
            for p_obj, lvl, deps in pred_entries:
                try:
                    if pred == p_obj:
                        return (lvl, p_obj, deps)
                except Exception:
                    continue
            # String match as last resort
            candidates = pred_index_str.get(str(pred))
            return candidates[0] if candidates else None

        if not (pred_index_obj or pred_index_str):
            print("No reasoning steps available.")
            return

        for goal in self.goals:
            queue = [(goal, True)]  # (predicate, should_recurse)
            visited = set()
            collected = []

            while queue:
                current, allow_recurse = queue.pop(0)
                if current is None:
                    continue
                key = str(current)
                if key in visited:
                    continue
                visited.add(key)

                entry = find_entry(current)
                if entry:
                    level, pred_obj, deps = entry
                else:
                    level, pred_obj, deps = (None, current, {'type': 'UNKNOWN', 'premises': []})

                if pred_obj is None:
                    continue

                dep_type = deps.get('type', 'unknown') if isinstance(deps, dict) else 'unknown'
                premises = deps.get('premises', []) if isinstance(deps, dict) else []
                premises = [p for p in premises if p is not None]
                collected.append((level, pred_obj, dep_type, premises))

                # Recurse through DD and rabbit constructions so their premises show in the trace
                dep_type_lower = dep_type.lower() if isinstance(dep_type, str) else ''
                should_expand = allow_recurse and (dep_type_lower == 'dd' or dep_type_lower.startswith('rabbit'))
                for prem in premises or []:
                    queue.append((prem, should_expand))

            # Ensure rabbit-derived steps are visible even if they are not on the goal's dependency chain
            rabbit_extras = []
            for p_obj, lvl, deps in pred_entries:
                dep_type = deps.get('type', 'unknown') if isinstance(deps, dict) else 'unknown'
                dep_type_lower = dep_type.lower() if isinstance(dep_type, str) else ''
                if not dep_type_lower.startswith('rabbit'):
                    continue
                prem_list = deps.get('premises', []) if isinstance(deps, dict) else []
                prem_list = [p for p in prem_list if p is not None]
                # avoid duplicates by predicate identity/eq
                already = any(p_obj == existing for _, existing, _, _ in collected)
                if not already:
                    rabbit_extras.append((lvl, p_obj, dep_type, prem_list))

            collected.extend(rabbit_extras)

            # Sort by level (None last) then by predicate string for stability
            collected.sort(key=lambda x: (float('inf') if x[0] is None else x[0], str(x[1])))

            # Store or print
            print(f"Reasoning for goal: {goal}")
            for lvl, pred_obj, dep_type, premises in collected:
                print(f"  [level {lvl}] {pred_obj} via {dep_type}")
                if premises:
                    prem_str = ', '.join(str(p) for p in premises)
                    print(f"    premises: {prem_str}")
            print()