Data_Generation2.py

"""Generate synthetic training data (v2).

Process:
1) Create a random non-collinear point set.
2) Apply one or two random constructions to enrich the configuration.
3) Run DDAR before/after the constructions.
4) Emit training pairs for newly derivable facts enabled by the constructions.

Outputs are written to training_data_v2.json by default.
"""

import copy
import json
import random
import logging
from typing import List, Dict, Tuple, Any, Set

import matplotlib
matplotlib.use("Agg")  # headless backend
import matplotlib.pyplot as plt

from Problem import GeometricProblem
from Constructions import GeometricConstructor
from ddar import DDARSystem
from relations import geometric, predicate


class DataGeneratorV2:
    def __init__(self, max_ddar_iterations: int = 5, log_file: str = "construction_log_v2.txt"):
        self.max_ddar_iterations = max_ddar_iterations
        self.ddar = DDARSystem(max_iterations=max_ddar_iterations, check_diagram=False)

        # logging
        self.logger = logging.getLogger("DataGeneratorV2")
        self.logger.setLevel(logging.INFO)
        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)

        # subset of constructions, favoring ones that create new points/relations
        self.constructions = [
            {'method': 'construct_midpoint', 'num_points': 2, 'name_param': 'name', 'mark_param': 'mark_point'},
            {'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_intersect_lines', 'num_points': 4, 'name_param': 'name', 'mark_param': 'mark_points'},
            {'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_incenter2', 'num_points': 3, 'name_param': 'name', 'mark_param': 'mark_points'},
            {'method': 'construct_excenter2', 'num_points': 3, 'name_param': 'name', 'mark_param': 'mark_points'},
            {'method': 'construct_orthocenter2', 'num_points': 3, 'name_param': 'name', 'mark_param': None},
            {'method': 'construct_tangents', 'num_points': 3, 'name_param': 'name1', 'mark_param': None},
        ]

    # ---------- helpers ----------
    def _random_points(self, num_points: int) -> Dict[str, geometric.Point]:
        """Create num_points random non-collinear points (retry until not all collinear)."""
        attempt = 0
        while True:
            attempt += 1
            pts = {}
            for i in range(num_points):
                name = chr(65 + i)
                pts[name] = geometric.Point(name, random.uniform(0, 10), random.uniform(0, 10))
            # quick collinearity check using area of triangle of first three
            names = list(pts.keys())
            if len(names) < 3:
                return pts
            a, b, c = (pts[names[0]], pts[names[1]], pts[names[2]])
            area2 = (b.x - a.x) * (c.y - a.y) - (b.y - a.y) * (c.x - a.x)
            if abs(area2) > 1e-6:
                return pts
            if attempt > 10:
                return pts

    def _predicate_involves_constructed_points(self, pred: Any, constructed: Set[geometric.Point]) -> bool:
        return hasattr(pred, 'get_points') and any(pt in constructed for pt in pred.get_points())

    def _is_redundant_goal(self, pred: Any) -> bool:
        if isinstance(pred, predicate.Cong):
            return pred.segments[0] == pred.segments[1]
        if isinstance(pred, predicate.Eqangle):
            return pred.angle1 == pred.angle2
        if isinstance(pred, predicate.Eqratio):
            r1_num, r1_den = pred.ratios[0]
            r2_num, r2_den = pred.ratios[1]
            if r1_num == r1_den and r2_num == r2_den:
                return True
            if r1_num == r2_num and r1_den == r2_den:
                return True
        return False

    def _apply_random_constructions(self, problem: GeometricProblem, constructor: GeometricConstructor, k: int) -> Tuple[str, List[Any], Set[geometric.Point]]:
        """Apply up to k random constructions that fit available points."""
        applied_code = []
        all_new_preds: List[Any] = []
        constructed_pts: Set[geometric.Point] = set()

        for _ in range(k):
            available_points = list(problem.points.values())
            applicable = [c for c in self.constructions if len(available_points) >= c['num_points']]
            if not applicable:
                break
            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, 9999)}"
                kwargs = {construction['name_param']: new_name}
                if construction.get('mark_param'):
                    kwargs[construction['mark_param']] = False
                elif construction.get('has_mark_points'):
                    kwargs['mark_points'] = False

                result = method(*selected_points, **kwargs)
                if isinstance(result, tuple) and len(result) >= 2:
                    new_preds = result[-1] if isinstance(result[-1], list) else []
                    new_pts = {item for item in result if isinstance(item, geometric.Point)}
                    for pred in new_preds:
                        if not any(pred == a for a in problem.assumptions):
                            problem.add_assumption(pred)
                            all_new_preds.append(pred)
                    for p in new_pts:
                        if p.name not in problem.points:
                            problem.add_point(p.name, p.x, p.y)
                        constructed_pts.add(p)
                    applied_code.append(f"constructor.{construction['method']}({', '.join(pt.name for pt in selected_points)})")
            except Exception as exc:
                self.logger.info(f"Construction {construction['method']} failed: {exc}")
                continue

        return '; '.join(applied_code), all_new_preds, constructed_pts

    # ---------- public API ----------
    def generate_training_pair(self, problem: GeometricProblem, constructor: GeometricConstructor) -> List[Dict]:
        setup_predicates = problem.assumptions.copy()

        # baseline
        try:
            before = copy.deepcopy(problem)
            _, _, facts_before = self.ddar.solve_problem(before)
            plt.close('all')
        except Exception as exc:
            self.logger.info(f"DDAR before failed: {exc}")
            plt.close('all')
            return []

        # constructions
        construction_code, new_predicates, constructed_pts = self._apply_random_constructions(
            problem, constructor, k=random.choice([1, 2])
        )
        if not construction_code or not new_predicates:
            return []

        # after
        try:
            after = copy.deepcopy(problem)
            _, _, facts_after = self.ddar.solve_problem(after)
            plt.close('all')
        except Exception as exc:
            self.logger.info(f"DDAR after failed: {exc}")
            plt.close('all')
            return []

        pairs = []
        for deduced in facts_after:
            if deduced in facts_before:
                continue
            if self._predicate_involves_constructed_points(deduced, constructed_pts):
                continue
            if self._is_redundant_goal(deduced):
                continue
            pairs.append({
                "input_text": {
                    "set_up": ", ".join(str(p) for p in setup_predicates),
                    "goal": str(deduced)
                },
                "output_text": construction_code
            })
        return pairs

    def generate_dataset(self,
                         num_iterations: int = 50,
                         constructions_per_problem: int = 2,
                         output_file: str = "training_data_v2.json",
                         num_points_range: Tuple[int, int] = (4, 6),
                         batch_size: int = 10) -> List[Dict]:
        all_pairs: List[Dict] = []
        print(f"Generating {num_iterations} problems (v2)...")
        for i in range(num_iterations):
            problem = GeometricProblem()
            num_points = random.randint(num_points_range[0], num_points_range[1])
            pts = self._random_points(num_points)
            for p in pts.values():
                problem.add_point(p.name, p.x, p.y)

            for _ in range(constructions_per_problem):
                constructor = GeometricConstructor(problem)
                pairs = self.generate_training_pair(problem, constructor)
                if pairs:
                    all_pairs.extend(pairs)

            if (i + 1) % 10 == 0:
                print(f"[{i+1}/{num_iterations}] pairs so far: {len(all_pairs)}")
            if output_file and (i + 1) % batch_size == 0:
                with open(output_file, 'w') as f:
                    json.dump(all_pairs, f, indent=2)
                print(f"  [batch saved] {len(all_pairs)} pairs to {output_file}")

        print(f"Total pairs: {len(all_pairs)}")
        if output_file and all_pairs:
            with open(output_file, 'w') as f:
                json.dump(all_pairs, f, indent=2)
            print(f"Saved to {output_file}")
        return all_pairs


def main():
    generator = DataGeneratorV2(max_ddar_iterations=5)
    training_data = generator.generate_dataset()
    if training_data:
        print("Sample pairs:")
        for pair in training_data[:3]:
            print(f"  Goal: {pair['input_text']['goal']} -> {pair['output_text']}")
    else:
        print("No training pairs generated.")


if __name__ == "__main__":
    main()