ar.py

"""
ar_best.py
Optimized Algebraic Reasoner
- Diagram Kernel FIXED: Uses arctan2 for correct angle signs.
- Deterministic & Modulo-180 Robust (Virtual PI).
"""

import numpy as np
from typing import Dict, List, Tuple, Set, Optional
from relations import predicate, Variable, LinearEquation
from Problem import GeometricProblem as GP
from augmented_rref import IncrementalAugmentedRREF
from fractions import Fraction

try:
    from numba import njit  # type: ignore
except ImportError:
    njit = None

def _compute_angles_kernel(xs: np.ndarray, ys: np.ndarray, angle_triplets: np.ndarray) -> np.ndarray:
    """Batched angle computation kernel using ARCTAN2 (Corrects quadrant issues)."""
    def _kernel(xs_loc: np.ndarray, ys_loc: np.ndarray, triplets_loc: np.ndarray) -> np.ndarray:
        n = triplets_loc.shape[0]
        out = np.empty(n, dtype=np.float64)
        for i in range(n):
            A_idx = triplets_loc[i, 0]
            O_idx = triplets_loc[i, 1]
            B_idx = triplets_loc[i, 2]
            ax, ay = xs_loc[A_idx], ys_loc[A_idx]
            ox, oy = xs_loc[O_idx], ys_loc[O_idx]
            bx, by = xs_loc[B_idx], ys_loc[B_idx]
            
            # Vectors OA and OB
            oax, oay = ax - ox, ay - oy
            obx, oby = bx - ox, by - oy
            
            # Cross product (Determinant) and Dot product
            det = oax * oby - oay * obx
            dot = oax * obx + oay * oby
            
            if det == 0.0 and dot == 0.0:
                 out[i] = 0.0 # Undefined, default to 0
                 continue

            # atan2(y, x) -> returns angle in (-pi, pi]
            theta_rad = np.arctan2(det, dot)
            out[i] = np.degrees(theta_rad)
        return out

    if njit is not None:
        return njit(_kernel)(xs, ys, angle_triplets)
    return _kernel(xs, ys, angle_triplets)

class AlgebraicReasoner:
    def __init__(self, check_diagram: bool = True):
        self.problem: GP = None
        self.variables: Set[Variable] = set()
        self.new_angles: List[Variable] = []
        self.equations: Set[LinearEquation] = set()
        
        self.PI_VAR_INDEX = 0
        self.n_angles: int = 1 
        
        self.angle_variables: Dict[Variable, int] = {}
        self.segment_variables: Dict[Variable, int] = {}
        self.angle_equations: Dict[int, LinearEquation] = {}
        self.segment_equations: Dict[int, LinearEquation] = {} 
        self.n_angle_eq: int = 0
        self.n_segments: int = 0
        self.n_segment_eq: int = 0
        
        self.angle_matrix: IncrementalAugmentedRREF = IncrementalAugmentedRREF(initial_num_vars=1)
        self.segment_matrix: IncrementalAugmentedRREF = IncrementalAugmentedRREF(initial_num_vars=0)
        
        self.check_diagram = check_diagram
        self._angle_values = {}

    def set_problem(self, problem: GP):
        self.problem = problem

    def clear(self):
        self.variables.clear()
        self.equations.clear()
        self.angle_variables.clear()
        self.segment_variables.clear()
        self.angle_equations.clear()
        self.segment_equations.clear()
        self.n_angles = 1
        self.n_segments = 0
        self.n_angle_eq = 0
        self.n_segment_eq = 0
        self.angle_matrix = IncrementalAugmentedRREF(initial_num_vars=1)
        self.segment_matrix = IncrementalAugmentedRREF(initial_num_vars=0)
        self.check_diagram = True
        self._angle_values = {}

    def add_equation(self, equation: LinearEquation):
        if equation in self.equations: return
        self.equations.add(equation)
        
        is_angle_eq = False
        coef_eq = {}
        sorted_coeffs = sorted(equation.coefficients.items(), key=lambda item: str(item[0]))
        
        for v, c in sorted_coeffs:
            if v.var_type == "angle":
                is_angle_eq = True
                if v not in self.angle_variables:
                    self.angle_variables[v] = self.n_angles
                    self.n_angles += 1
                coef_eq[self.angle_variables[v]] = c
                if v not in self.variables:
                    self.new_angles.append(v)
            else:
                if v not in self.segment_variables:
                    self.segment_variables[v] = self.n_segments
                    self.n_segments += 1
                coef_eq[self.segment_variables[v]] = c
            self.variables.add(v)
            
        if is_angle_eq:
            if equation.constant != 0:
                pi_coeff = Fraction(-equation.constant, 180)
                coef_eq[self.PI_VAR_INDEX] = coef_eq.get(self.PI_VAR_INDEX, 0) + pi_coeff
            self.angle_matrix.add_equation(coef_eq)
            self.angle_equations[self.n_angle_eq] = equation
            self.n_angle_eq += 1
        else:
            self.segment_matrix.add_equation(coef_eq)
            self.segment_equations[self.n_segment_eq] = equation
            self.n_segment_eq += 1

    def convert_predicate_to_equations(self, pred) -> List[LinearEquation]:
        equations = []
        try:
            if hasattr(pred, "equations"):
                equations.extend(pred.equations)
        except Exception: pass
        return equations

    def add_assumptions(self, assumptions: List):
        sorted_assumptions = sorted(assumptions, key=str)
        for assumption in sorted_assumptions:
            equations = self.convert_predicate_to_equations(assumption)
            for eq in equations:
                self.add_equation(eq)

        angles = sorted([var for var in self.variables if var.var_type == "angle"], key=str)
        for a in angles:
            a_c_points = a.points[::-1]
            a_c = next((b for b in angles if b.points == a_c_points), None)
            if not a_c:
                a_c = Variable(hash(tuple(a_c_points)), a_c_points, "angle")
            self.add_equation(LinearEquation(None, [(a, 1), (a_c, 1)], 180, "angle_sum_pi"))
    
    def create_goal_coefs(self, goal) -> List[tuple[bool, Dict[int, object], object, LinearEquation]]:
        goal_equations = self.convert_predicate_to_equations(goal)
        if not goal_equations: return []
        goal_coef_eqs = []
        for eq in goal_equations:
            is_angle_eq = False
            coef_eq = {}
            for v, c in eq.coefficients.items():
                if v.var_type == "angle":
                    is_angle_eq = True
                    if v not in self.angle_variables: return []
                    coef_eq[self.angle_variables[v]] = c
                else:
                    if v not in self.segment_variables: return []
                    coef_eq[self.segment_variables[v]] = c
            goal_coef_eqs.append((is_angle_eq, coef_eq, eq.constant, eq))
        return goal_coef_eqs
    
    def _build_angle_value_table(self) -> None:
        if not self.new_angles: return
        all_points = set()
        for v in self.new_angles:
            for p in v.points: all_points.add(p)
        sorted_points = sorted(list(all_points), key=lambda p: p.name if hasattr(p, 'name') else id(p))
        point_index = {p: i for i, p in enumerate(sorted_points)}

        n_points = len(point_index)
        xs = np.empty(n_points, dtype=np.float64)
        ys = np.empty(n_points, dtype=np.float64)
        for p, idx in point_index.items():
            xs[idx], ys[idx] = p.x, p.y

        triplets = np.empty((len(self.new_angles), 3), dtype=np.int64)
        for i, v in enumerate(self.new_angles):
            triplets[i, 0] = point_index[v.points[0]]
            triplets[i, 1] = point_index[v.points[1]]
            triplets[i, 2] = point_index[v.points[2]]

        values = _compute_angles_kernel(xs, ys, triplets)
        for v, val in zip(self.new_angles, values):
            self._angle_values[v] = float(val)

    def check_angle_equation_diagram(self, eq: LinearEquation) -> bool:
        lhs = 0.0
        for v, c in eq.coefficients.items():
            if c == 0: continue
            val = self._angle_values.get(v)
            if val is None: return False
            lhs += val * c
        rhs = eq.constant 
        diff = abs(lhs - rhs)
        remainder = diff % 180.0
        # Relaxed tolerance for noisy diagram checks
        return min(remainder, 180.0 - remainder) < 1e-2

    def can_derive(self, goal) -> Tuple[bool, str, predicate, Dict]:
        goal_triplets = self.create_goal_coefs(goal)
        if not goal_triplets: return False, "Cannot convert goal"
        try:
            all_derivable = True
            derivable_count = 0
            premises = []
            sol_vecs = []
            if self.check_diagram:
                self._build_angle_value_table()
                self.new_angles = []
            
            for is_angle, coef_eq, rhs, eq in goal_triplets:
                sol_vec = None
                if is_angle:
                    if rhs != 0:
                        coef_eq[self.PI_VAR_INDEX] = coef_eq.get(self.PI_VAR_INDEX, 0) - Fraction(rhs, 180)
                    reduced_row, prov = self.angle_matrix.reduce(coef_eq)
                    all_zero_except_pi = True
                    for i, val in enumerate(reduced_row):
                        if i == self.PI_VAR_INDEX: continue
                        if val != 0: 
                            all_zero_except_pi = False; break
                    if all_zero_except_pi:
                        if reduced_row[self.PI_VAR_INDEX].denominator == 1:
                            sol_vec = self.angle_matrix._format_provenance(prov)
                else:
                    reduced_row, prov = self.segment_matrix.reduce(coef_eq)
                    if all(v == 0 for v in reduced_row):
                        sol_vec = self.segment_matrix._format_provenance(prov)

                if sol_vec:
                    derivable_count += 1
                    curr_matrix = self.angle_equations if is_angle else self.segment_equations
                    for eq_idx, weight in enumerate(sol_vec):
                         if weight != 0 and eq_idx in curr_matrix:
                             premises.append(curr_matrix[eq_idx].predicate)
                    sol_vecs.append(sol_vec)
                else:
                    all_derivable = False
            
            if all_derivable and derivable_count == len(goal_triplets):
                return True, {'predicate': goal, 'dependencies': {'type': 'AR', 'premises': premises, 'sol_vecs': sol_vecs}}
            return False, {}
        except Exception: return False, {}

    def _get_eqangle(self, v1, v2) -> Tuple[predicate, Dict]:
        coeffs = {self.angle_variables[v1]: 1, self.angle_variables[v2]: -1}
        reduced_row, prov = self.angle_matrix.reduce(coeffs)
        for i, val in enumerate(reduced_row):
            if i == self.PI_VAR_INDEX: continue
            if val != 0: return None, None
        if reduced_row[self.PI_VAR_INDEX].denominator != 1: return None, None
        
        if self.check_diagram:
            val1 = self._angle_values.get(v1)
            val2 = self._angle_values.get(v2)
            if val1 is not None and val2 is not None:
                diff = abs(val1 - val2) % 180.0
                if min(diff, 180.0 - diff) > 1e-2: return None, None 

        sol_vec = self.angle_matrix._format_provenance(prov)
        premises = [self.angle_equations[i].predicate for i, c in enumerate(sol_vec) if c != 0]
        try:
            pred = predicate.Eqangle(*(v1.points + v2.points))
            return pred, {"predicate": pred, "dependencies": {"type": "AR", "premises": premises, "sol_vec": sol_vec}}
        except ValueError: return None, None

    def _get_cong(self, v1, v2) -> Tuple[predicate, Dict]:
        coeffs = {self.segment_variables[v1]: 1, self.segment_variables[v2]: -1}
        reduced_row, prov = self.segment_matrix.reduce(coeffs)
        if not all(v == 0 for v in reduced_row): return None, None
        sol_vec = self.segment_matrix._format_provenance(prov)
        premises = [self.segment_equations[i].predicate for i, c in enumerate(sol_vec) if c != 0]
        try:
            pred = predicate.Cong(*(v1.points + v2.points))
            return pred, {"predicate": pred, "dependencies": {"type": "AR", "premises": premises, "sol_vec": sol_vec}}
        except ValueError: return None, None
        
    def _get_perp(self, v) -> Tuple[predicate, Dict]:
        coeffs = {self.angle_variables[v]: 1}
        reduced_row, prov = self.angle_matrix.reduce(coeffs)
        for i, val in enumerate(reduced_row):
            if i == self.PI_VAR_INDEX: continue
            if val != 0: return None, None
        
        pi_val = reduced_row[self.PI_VAR_INDEX]
        if pi_val.denominator == 2 and (pi_val.numerator % 2 != 0):
             pass
        else:
             return None, None
             
        if self.check_diagram:
            val = self._angle_values.get(v)
            if val is not None:
                diff = abs(val - 90.0) % 180.0
                if min(diff, 180.0 - diff) > 1e-2: return None, None

        sol_vec = self.angle_matrix._format_provenance(prov)
        premises = [self.angle_equations[i].predicate for i, c in enumerate(sol_vec) if c != 0]
        try:
            pred = predicate.Perp(*[v.points[0], v.points[1], v.points[1], v.points[2]])
            return pred, {"predicate": pred, "dependencies": {"type": "AR", "premises": premises, "sol_vec": sol_vec}}
        except ValueError: return None, None

    def _get_para(self, v1, v2, deps) -> Tuple[bool, List]:
        points = []
        if v1.points[0] == v2.points[1] and v1.points[1] == v2.points[0]:
            points = [v1.points[1], v1.points[2], v2.points[1], v2.points[2]]
        elif v1.points[1] == v2.points[2] and v1.points[2] == v2.points[1]:
            points = [v1.points[0], v1.points[1], v2.points[0], v2.points[1]]
        else: return None, None
        try:
            pred = predicate.Para(*points)
            return pred, {"predicate": pred, "dependencies": deps}
        except ValueError: return None, None

    def check_variable_pairs(self) -> List[Tuple[predicate, Dict]]:
        facts: List[Tuple[predicate, Dict]] = []
        angle_vars = sorted(list(self.angle_variables.keys()), key=str)
        seg_vars = sorted(list(self.segment_variables.keys()), key=str)
        if self.check_diagram: self._build_angle_value_table()
        
        for i, var1 in enumerate(angle_vars):
            val1 = self._angle_values.get(var1)
            for j in range(i + 1, len(angle_vars)):
                var2 = angle_vars[j]
                val2 = self._angle_values.get(var2)
                if val1 is not None and val2 is not None:
                    diff = abs(val1 - val2) % 180.0
                    if min(diff, 180.0 - diff) > 1e-2: continue 
                pred, info = self._get_eqangle(var1, var2)
                if pred:
                    facts.append((pred, info))
                    ppred, pinfo = self._get_para(var1, var2, info["dependencies"])
                    if ppred: facts.append((ppred, pinfo))

        for i, var1 in enumerate(seg_vars):
            for j in range(i + 1, len(seg_vars)):
                var2 = seg_vars[j]
                pred, info = self._get_cong(var1, var2)
                if pred: facts.append((pred, info))
        return facts
    
    def check_perp_lines(self) -> List[Tuple[predicate, Dict]]:
        facts = []
        angle_vars = sorted(list(self.angle_variables.keys()), key=str)
        checked_point_sets = set()
        for v in angle_vars:
            val = self._angle_values.get(v)
            if val is not None:
                diff = abs(val - 90.0) % 180.0
                if min(diff, 180.0 - diff) > 1e-2: continue
            if any(set(v.points) == set(c_p) for c_p in checked_point_sets): continue
            pred, info = self._get_perp(v)
            if pred:
                facts.append((pred, info))
                checked_point_sets.add(tuple(v.points))
        return facts
    
    def find_all_derivable_facts(self) -> List[Tuple]:
        if self.check_diagram:
            self._build_angle_value_table()
            self.new_angles = []
        try:
            return self.check_perp_lines() + self.check_variable_pairs()
        except Exception as e:
            print(f"AR Error: {e}")
            return []