Data_Generation.py

"""Generate synthetic training data for language model.

Training pairs format:
- Input: setup predicates + goal predicate
- Output: construction code that helps prove the goal
"""

import random
import json
import copy
import logging
import matplotlib
matplotlib.use('Agg')  # Use non-interactive backend to avoid display issues
import matplotlib.pyplot as plt
from typing import List, Dict, Tuple, Any
from Problem import GeometricProblem
from Constructions import GeometricConstructor
from ddar import DDARSystem
from relations import geometric, predicate


class DataGenerator:
    """Generates synthetic training data by applying random constructions."""
    
    def __init__(self, max_ddar_iterations: int = 5, log_file: str = "construction_log.txt", max_ar_equations: int = None):
        self.max_ddar_iterations = max_ddar_iterations
        self.ddar = DDARSystem(max_iterations=max_ddar_iterations, check_diagram=False, max_ar_equations=max_ar_equations)
        
        # Setup logging
        self.logger = logging.getLogger('DataGenerator')
        self.logger.setLevel(logging.INFO)
        # Clear existing handlers to avoid duplicate logs if re-instantiated
        if self.logger.handlers:
            self.logger.handlers.clear()
            
        fh = logging.FileHandler(log_file, mode='w')
        formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
        fh.setFormatter(formatter)
        self.logger.addHandler(fh)
        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_intersect_lines', 'num_points': 4, 'name_param': 'name', 'mark_param': 'mark_points'},
            {'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_reflection', 'num_points': 3, 'name_param': 'name1', 'mark_param': None, 'has_mark_points': True},
            {'method': 'construct_on_dia', 'num_points': 2, 'name_param': 'name', 'mark_param': None},
            {'method': 'construct_tangents', 'num_points': 3, 'name_param': 'name1', 'mark_param': None},
        ]
    
    def _predicate_involves_point(self, pred: Any, point: geometric.Point) -> bool:
        """Check if predicate involves a specific point."""
        return hasattr(pred, 'get_points') and any(pt == point for pt in pred.get_points())
    
    def _predicate_involves_constructed_points(self, pred: Any, constructed_points: set) -> bool:
        """Check if predicate involves any constructed points."""
        return any(self._predicate_involves_point(pred, pt) for pt in constructed_points)
    
    def _is_redundant_goal(self, pred: Any) -> bool:
        """Check if a goal predicate is trivially true (redundant)."""
        if isinstance(pred, predicate.Cong):
            # cong A B C D is redundant if {A,B} == {C,D}
            return pred.segments[0] == pred.segments[1]
        
        elif isinstance(pred, predicate.Eqangle):
            # eqangle A B C D E F is redundant if (A,B,C) == (D,E,F)
            return pred.angle1 == pred.angle2
            
        elif isinstance(pred, predicate.Eqratio):
            # eqratio A B C D E F G H means AB/CD = EF/GH
            # Redundant if AB/CD == EF/GH is trivially true.
            # Case 1: AB=EF and CD=GH (ratios are identical)
            # Case 2: AB=CD and EF=GH (both ratios are 1)
            # The segments are stored as frozensets in pred.ratios
            r1_num, r1_den = pred.ratios[0]
            r2_num, r2_den = pred.ratios[1]
            
            # Check if ratio 1 is identical to ratio 2
            if r1_num == r2_num and r1_den == r2_den:
                return True
            
            # Check if both ratios are 1 (numerator == denominator)
            if r1_num == r1_den and r2_num == r2_den:
                return True
            
        return False

    def _apply_random_construction(self, problem: GeometricProblem, 
                                   constructor: GeometricConstructor, num_iters: int = 20) -> Tuple[str, List[Any], set]:
        """Apply a random construction and return construction code, new predicates, and constructed points."""
        available_points = list(problem.points.values())
        if len(available_points) < 2:
            return None, [], set()
        
        applicable = [c for c in self.constructions if len(available_points) >= c['num_points']]
        if not applicable:
            msg = f"No applicable constructions for {len(available_points)} points"
            print(f"    {msg}")
            self.logger.info(msg)
            return None, [], set()
        
        for _ in range(num_iters):
            construction = random.choice(applicable)
            selected_points = random.sample(available_points, construction['num_points'])
            
            try:
                method = getattr(constructor, construction['method'])
                point_names = ''.join(pt.name for pt in selected_points)
                new_name = f"{construction['method']}_{point_names}_{random.randint(0, 999)}"
                
                # Build kwargs based on construction requirements
                kwargs = {construction['name_param']: new_name}
                if construction.get('mark_param'):
                    kwargs[construction['mark_param']] = False
                elif 'has_mark_points' in construction and construction['has_mark_points']:
                     kwargs['mark_points'] = False
                
                msg = f"Attempting construction: {construction['method']} on {point_names}"
                print(f"    {msg}")
                self.logger.info(msg)
                result = method(*selected_points, **kwargs)
                
                if isinstance(result, tuple) and len(result) >= 2:
                    # Relations list is always the LAST element in the tuple
                    new_predicates = result[-1] if isinstance(result[-1], list) else []
                    # A small note here that we could have each method in constructions
                    # return a list of points, so accessing them is easier.
                    # It should not matter much for efficiency since the list is always short.
                    constructed_points = {item for item in result if isinstance(item, geometric.Point)}
                    point_args = ', '.join(pt.name for pt in selected_points)
                    construction_code = f"constructor.{construction['method']}({point_args})"
                    
                    msg = f"Construction successful: {construction_code}, added {len(new_predicates)} predicates"
                    print(f"    {msg}")
                    self.logger.info(msg)
                    return construction_code, new_predicates, constructed_points
                else:
                    msg = f"Construction returned invalid result format: {type(result)}"
                    print(f"    {msg}")
                    self.logger.warning(msg)
            except (ValueError, Exception) as e:
                # Log failed construction
                point_args = ', '.join(pt.name for pt in selected_points)
                msg = f"Construction failed: {construction['method']} on {point_args} - Error: {e}"
                print(f"    {msg}")
                self.logger.error(msg)
                continue
        
        msg = "Failed to apply any construction after 20 attempts"
        print(f"    {msg}")
        self.logger.info(msg)
        return None, [], set()
    
    def generate_training_pair(self, problem: GeometricProblem, 
                               constructor: GeometricConstructor) -> List[Dict]:
        """Generate training pairs from a single construction application."""
        setup_predicates = problem.assumptions.copy()
        
        # STEP 1: Run DDAR BEFORE construction to get baseline facts
        try:
            problem_before = copy.deepcopy(problem)
            _, _, facts_before_construction = self.ddar.solve_problem(problem_before)
            print(f"  DDAR before: {len(facts_before_construction)} facts derived")
            plt.close('all')  # Clean up any matplotlib figures
        except Exception as e:
            # If DDAR fails on the base problem, skip this problem entirely
            print(f"  WARNING: DDAR failed on base problem: {e}")
            plt.close('all')
            return []
        
        # STEP 2: Apply construction
        construction_code, new_predicates, constructed_points = self._apply_random_construction(problem, constructor)
        if not construction_code:
            print("  No construction code generated")
            return []
        
        # If construction added no new predicates, skip it
        if not new_predicates:
            print("  Construction added no new predicates")
            return []
        
        constructed_pred_strs = set()
        constructed_predicates = []
        for pred in new_predicates:
            if not any(pred == a for a in problem.assumptions):
                problem.add_assumption(pred)
                constructed_pred_strs.add(str(pred))
                constructed_predicates.append(pred)
        # Add the constructed points
        for point in constructed_points:
            if point.name not in problem.points:
                problem.add_point(point.name, point.x, point.y)
        
        # If no new predicates were actually added (all were duplicates), skip
        if not constructed_pred_strs:
            print("  All new predicates were duplicates")
            return []
        
        print(f"  Added {len(constructed_pred_strs)} unique predicates from construction")

        # STEP 3: Run DDAR AFTER construction
        try:
            problem_after = copy.deepcopy(problem)
            _, _, facts_after_construction = self.ddar.solve_problem(problem_after)
            print(f"  DDAR after: {len(facts_after_construction)} facts derived")
            plt.close('all')  # Clean up any matplotlib figures, helps with speed a lot
        except Exception as e:
            print(f"  WARNING: DDAR failed after construction: {e}")
            plt.close('all')
            return []
        
        # STEP 4: Find facts that are NEWLY derivable (not derivable before construction)
        # If a fact wasn't derivable before but is now, the construction enabled it
        training_pairs = []
        new_facts_count = 0
        redundant_count = 0
        max_pairs_per_problem = 3
        
        for deduced_pred in facts_after_construction:
            # Limit to 3 training pairs per problem
            if len(training_pairs) >= max_pairs_per_problem:
                break
                
            deduced_str = str(deduced_pred)
            
            # Skip if this fact was already derivable before construction
            if deduced_pred in facts_before_construction:
                continue
            
            new_facts_count += 1
            
            # Skip redundant goals
            if self._is_redundant_goal(deduced_pred):
                redundant_count += 1
                # print(f"    Skipping redundant goal: {deduced_str}")
                continue
            
            # This is a new fact enabled by the construction (includes facts about constructed points)
            # Include both setup predicates and constructed predicates in facts
            all_facts = setup_predicates + constructed_predicates
            print(f"  Found valid training pair! Goal: {deduced_str}")
            training_pairs.append({
                "input_text": {
                    "set_up": ", ".join(str(p) for p in all_facts),
                    "goal": deduced_str
                },
                "output_text": construction_code
            })
        
        print(f"  New facts after construction: {new_facts_count}, Redundant: {redundant_count}, Valid pairs: {len(training_pairs)}")
        if not training_pairs:
            print("  No valid training pairs created")

        return training_pairs
    
    def generate_dataset(self, num_iterations: int = 100, 
                        constructions_per_problem: int = 3,
                        output_file: str = "training_data.json",
                        num_points_range: Tuple[int, int] = (3, 4),
                        num_predicates_range: Tuple[int, int] = (1, 2),
                        predicate_types: List[str] = None,
                        batch_size: int = 10) -> List[Dict]:
        """Generate dataset by creating random problems and applying constructions."""
        if predicate_types is None:
            predicate_types = ['cong', 'perp']
        
        all_training_pairs = []
        print(f"Generating {num_iterations} random problems...")
        print(f"  Points: {num_points_range[0]}-{num_points_range[1]}, "
              f"Predicates: {num_predicates_range[0]}-{num_predicates_range[1]}")
        print(f"  Types: {predicate_types}")
        print(f"  Constructions per problem: {constructions_per_problem}")
        print("=" * 60)
        
        for i in range(num_iterations):
            base_problem = self._create_random_problem(
                num_points_range, num_predicates_range, predicate_types
            )
            
            for _ in range(constructions_per_problem):
                problem = copy.deepcopy(base_problem)
                constructor = GeometricConstructor(problem)
                pairs = self.generate_training_pair(problem, constructor)
                if pairs:
                    all_training_pairs.extend(pairs)
            
            if (i + 1) % 10 == 0:
                print(f"[{i+1}/{num_iterations}] Generated {len(all_training_pairs)} pairs")

            # Save in batches
            if output_file and (i + 1) % batch_size == 0:
                with open(output_file, 'w') as f:
                    json.dump(all_training_pairs, f, indent=2)
                print(f"  [Batch Save] Saved {len(all_training_pairs)} pairs to {output_file}")
        
        print("=" * 60)
        print(f"Total training pairs: {len(all_training_pairs)}")
        
        if output_file and all_training_pairs:
            with open(output_file, 'w') as f:
                json.dump(all_training_pairs, f, indent=2)
            print(f"Saved to {output_file}")
        
        return all_training_pairs
    
    def _is_consistent(self, problem: GeometricProblem, new_pred: Any) -> bool:
        """Check if adding new_pred creates a contradiction with existing assumptions."""
        
        # Check if predicate already exists
        if any(new_pred == p for p in problem.assumptions):
            return False

        # Check 1: Cyclic vs Collinear AND Cyclic vs Circle
        if isinstance(new_pred, predicate.Cyclic):
            cyclic_points = set(new_pred.points)
            for assum in problem.assumptions:
                if isinstance(assum, predicate.Col):
                    col_points = set(assum.points)
                    if col_points.issubset(cyclic_points):
                        return False
                
                elif isinstance(assum, predicate.Circle):
                    center = assum.center
                    on_circle = set(assum.triangle)
                    if center in cyclic_points and on_circle.issubset(cyclic_points):
                        return False
        
        elif isinstance(new_pred, predicate.Col):
            col_points = set(new_pred.points)
            for assum in problem.assumptions:
                if isinstance(assum, predicate.Cyclic):
                    cyclic_points = set(assum.points)
                    if col_points.issubset(cyclic_points):
                        return False
                
                # Check 2: Parallel vs Collinear (when adding Col)
                if isinstance(assum, predicate.Para):
                    lines_list = list(assum.lines)
                    l1_pts = set(lines_list[0])
                    l2_pts = set(lines_list[1])
                    
                    if l1_pts.issubset(col_points) and not l2_pts.isdisjoint(col_points):
                        return False
                    if l2_pts.issubset(col_points) and not l1_pts.isdisjoint(col_pts):
                        return False

        # Check 2: Parallel vs Collinear (when adding Para)
        elif isinstance(new_pred, predicate.Para):
            lines_list = list(new_pred.lines)
            l1_pts = set(lines_list[0])
            l2_pts = set(lines_list[1])
            
            for assum in problem.assumptions:
                # Check Para vs Perp
                if isinstance(assum, predicate.Perp):
                    if assum.lines == new_pred.lines:
                        return False

                if isinstance(assum, (predicate.Col, predicate.Midp)):
                    col_pts = set(assum.points)
                    if l1_pts.issubset(col_pts) and not l2_pts.isdisjoint(col_pts):
                        return False
                    if l2_pts.issubset(col_points) and not l1_pts.isdisjoint(col_pts):
                        return False

        # Check 6: Perpendicular vs Parallel (when adding Perp)
        elif isinstance(new_pred, predicate.Perp):
            for assum in problem.assumptions:
                if isinstance(assum, predicate.Para):
                    if assum.lines == new_pred.lines:
                        return False

        # Check 5: Midpoint (implies Collinear)
        elif isinstance(new_pred, predicate.Midp):
            col_points = set(new_pred.points)
            for assum in problem.assumptions:
                if isinstance(assum, predicate.Cyclic):
                    cyclic_points = set(assum.points)
                    if col_points.issubset(cyclic_points):
                        return False
                
                # Check Parallel vs Midpoint (treated as Collinear)
                if isinstance(assum, predicate.Para):
                    lines_list = list(assum.lines)
                    l1_pts = set(lines_list[0])
                    l2_pts = set(lines_list[1])
                    
                    if l1_pts.issubset(col_points) and not l2_pts.isdisjoint(col_points):
                        return False
                    if l2_pts.issubset(col_points) and not l1_pts.isdisjoint(col_pts):
                        return False

        # Check 3: Circle vs Cyclic (when adding Circle)
        elif isinstance(new_pred, predicate.Circle):
            center = new_pred.center
            on_circle = set(new_pred.triangle)
            new_points = set(new_pred.points)
            
            for assum in problem.assumptions:
                # Check for Circle vs Circle (same points)
                if isinstance(assum, predicate.Circle):
                    if set(assum.points) == new_points:
                        return False

                if isinstance(assum, predicate.Cyclic):
                    cyclic_points = set(assum.points)
                    # If center is in cyclic points AND the points on circle are also in cyclic points
                    if center in cyclic_points and on_circle.issubset(cyclic_points):
                        return False

        # Check 4: Simtri1/Contri1 vs Simtri2/Contri2
        # We cannot have both direct and opposite similarity/congruence for the same pair of triangles (sets of points)
        elif isinstance(new_pred, (predicate.Simtri1, predicate.Contri1, predicate.Simtri2, predicate.Contri2)):
            t1_points = set(new_pred.points[0:3])
            t2_points = set(new_pred.points[3:6])
            
            # Determine if new_pred is Type 1 (direct) or Type 2 (opposite)
            is_type1 = isinstance(new_pred, (predicate.Simtri1, predicate.Contri1))
            
            for assum in problem.assumptions:
                # Check against conflicting type
                is_assum_type1 = isinstance(assum, (predicate.Simtri1, predicate.Contri1))
                is_assum_type2 = isinstance(assum, (predicate.Simtri2, predicate.Contri2))
                
                if (is_type1 and is_assum_type2) or (not is_type1 and is_assum_type1):
                    a_t1_points = set(assum.points[0:3])
                    a_t2_points = set(assum.points[3:6])
                    
                    # Check if the sets of points match (either order)
                    if (t1_points == a_t1_points and t2_points == a_t2_points) or \
                       (t1_points == a_t2_points and t2_points == a_t1_points):
                        return False
                            
        return True

    def _create_random_problem(self, 
                              num_points_range: Tuple[int, int] = (3, 4),
                              num_predicates_range: Tuple[int, int] = (1, 2),
                              predicate_types: List[str] = None) -> GeometricProblem:
        """Create a random geometric problem."""
        if predicate_types is None:
            # Why did we choose these as defaults?
            predicate_types = ['cong', 'perp']
        
        problem = GeometricProblem()
        
        # Create random points
        num_points = random.randint(num_points_range[0], num_points_range[1])
        point_names = [chr(i) for i in range(65, 65 + num_points + 1)]

        for name in point_names:
            problem.add_point(name, random.uniform(0, 10), random.uniform(0, 10))
        
        # Add random predicates
        num_predicates = random.randint(num_predicates_range[0], num_predicates_range[1])
        points_list = list(problem.points.values())
        
        for _ in range(num_predicates):
            pred_type = random.choice(predicate_types)
            try:
                new_pred = None
                if pred_type == 'cong' and len(points_list) >= 4:
                    new_pred = predicate.Cong(*random.sample(points_list, 4))
                elif pred_type == 'perp' and len(points_list) >= 4:
                    new_pred = predicate.Perp(*random.sample(points_list, 4))
                elif pred_type == 'para' and len(points_list) >= 4:
                    new_pred = predicate.Para(*random.sample(points_list, 4))
                elif pred_type == 'col' and len(points_list) >= 3:
                    new_pred = predicate.Col(*random.sample(points_list, 3))
                elif pred_type == 'cyclic' and len(points_list) >= 4:
                    new_pred = predicate.Cyclic(*random.sample(points_list, 4))
                elif pred_type == 'eqangle' and len(points_list) >= 6:
                    new_pred = predicate.Eqangle(*random.sample(points_list, 6))
                elif pred_type == 'midp' and len(points_list) >= 3:
                    new_pred = predicate.Midp(*random.sample(points_list, 3))
                elif pred_type == 'circle' and len(points_list) >= 4:
                    new_pred = predicate.Circle(*random.sample(points_list, 4))
                elif pred_type == 'eqratio' and len(points_list) >= 8:
                    new_pred = predicate.Eqratio(*random.sample(points_list, 8))
                elif pred_type == 'simtri1' and len(points_list) >= 6:
                    new_pred = predicate.Simtri1(*random.sample(points_list, 6))
                elif pred_type == 'simtri2' and len(points_list) >= 6:
                    new_pred = predicate.Simtri2(*random.sample(points_list, 6))
                elif pred_type == 'contri1' and len(points_list) >= 6:
                    new_pred = predicate.Contri1(*random.sample(points_list, 6))
                elif pred_type == 'contri2' and len(points_list) >= 6:
                    new_pred = predicate.Contri2(*random.sample(points_list, 6))
                elif pred_type == 'sameclock' and len(points_list) >= 6:
                    new_pred = predicate.Sameclock(*random.sample(points_list, 6))
                
                if new_pred and self._is_consistent(problem, new_pred):
                    problem.add_assumption(new_pred)
            except (ValueError, Exception):
                pass
        
        print(f"Created problem with {len(problem.points)} points and {len(problem.assumptions)} predicates")
        return problem

def main():
    """Generate training data from random problems."""
    generator = DataGenerator(max_ddar_iterations=2, max_ar_equations=110)
    
    print("=" * 60)
    print("SYNTHETIC TRAINING DATA GENERATION")
    print("=" * 60 + "\n")
    
    training_data = generator.generate_dataset(
        num_iterations=1000,
        constructions_per_problem=1,
        output_file="training_data.json",
        num_points_range=(3, 7),
        num_predicates_range=(2, 6),
        predicate_types=['cong', 'perp', 'para', 'col', 'cyclic', 'eqangle', 'midp', 'circle', 'eqratio', 'simtri1', 'simtri2', 'contri1', 'contri2']
    )
    
    # Show sample pairs
    if training_data:
        print(f"\n{'=' * 60}")
        print("SAMPLE TRAINING PAIRS:")
        print(f"{'=' * 60}")
        for i, pair in enumerate(training_data[:5], 1):
            setup = pair['input_text']['set_up']
            setup = setup[:70] + "..." if len(setup) > 70 else setup
            print(f"\nPair {i}:")
            print(f"  Setup: {setup}")
            print(f"  Goal: {pair['input_text']['goal']}")
            print(f"  Construction: {pair['output_text']}")
    else:
        print("\nNo training pairs generated.")


if __name__ == "__main__":
    main()