query_generator.py

#!/usr/bin/env python3
"""
================================================================================
STANDALONE QUERY GENERATOR (v1.0)
================================================================================

DESCRIPTION:
   Takes an EXISTING graph.json file and generates queries1.json, queries2.json, 
   and queries3.json in the same directory.

FEATURES:
   - HIGHLY CUSTOMIZABLE Phase 1 counts (SP Time vs Dist, KNN types, Updates).
   - SHUFFLE: Interleaves updates with queries to test dynamic graph correctness.
   - Phase 2 & 3 custom counts.
   - Lightweight: No OSMNX/NetworkX required.

USAGE:
   python generate_queries_only.py path/to/graph.json [OPTIONS]
================================================================================
"""
import os
import json
import random
import argparse
import sys
from pathlib import Path

# ==============================================================================
# CONFIGURATION
# ==============================================================================
ALLOWED_POIS = ["restaurant", "petrol station", "hospital", "pharmacy", "hotel", "atm"]
ALLOWED_ROAD_TYPES = ["primary", "secondary", "tertiary", "local", "expressway"]

# ==============================================================================
# GENERATOR CLASS
# ==============================================================================
class QueryGenerator:
    def __init__(self, graph_path):
        self.graph_path = Path(graph_path)
        if not self.graph_path.exists():
            print(f"Error: File {self.graph_path} not found.")
            sys.exit(1)
            
        print(f"Loading graph from {self.graph_path}...")
        with open(self.graph_path, 'r') as f:
            self.graph_data = json.load(f)
            
        self.nodes = self.graph_data.get("nodes", [])
        self.edges = self.graph_data.get("edges", [])
        self.num_nodes = len(self.nodes)
        self.num_edges = len(self.edges)
        
        if self.num_nodes < 2:
            print("Error: Graph has fewer than 2 nodes.")
            sys.exit(1)
            
        print(f"   -> Loaded {self.num_nodes} nodes and {self.num_edges} edges.")

    def _get_rand_nodes(self, k=2):
        """Returns k random Node IDs"""
        # We assume node IDs are continuous 0 to N-1, but let's be safe and pick from actual list
        # Optimizing for speed: if IDs are 0..N-1, use random.sample(range)
        # Otherwise sample from loaded nodes
        return [n["id"] for n in random.sample(self.nodes, k)]

    def _get_rand_edge_id(self):
        if not self.edges: return None
        return random.choice(self.edges)["id"]

    def _get_rand_edge_data(self):
        if not self.edges: return None
        return random.choice(self.edges)

    # --------------------------------------------------------------------------
    # PHASE 1 GENERATION
    # --------------------------------------------------------------------------
    def generate_phase1(self, counts, shuffle=True):
        print("Generating Phase 1 Queries...")
        events = []
        eid = 1

        # 1. Shortest Path (Distance)
        for _ in range(counts['sp_dist']):
            s, t = self._get_rand_nodes(2)
            events.append({
                "type": "shortest_path", "id": eid, "source": s, "target": t, "mode": "distance"
            })
            eid += 1
            
        # 2. Shortest Path (Time)
        for _ in range(counts['sp_time']):
            s, t = self._get_rand_nodes(2)
            events.append({
                "type": "shortest_path", "id": eid, "source": s, "target": t, "mode": "time"
            })
            eid += 1

        # 3. KNN (Euclidean)
        for _ in range(counts['knn_euc']):
            n = random.choice(self.nodes)
            events.append({
                "type": "knn", "id": eid, "poi": random.choice(ALLOWED_POIS),
                "query_point": {"lat": n["lat"], "lon": n["lon"]},
                "k": random.randint(3, 10), "metric": "euclidean"
            })
            eid += 1

        # 4. KNN (Shortest Path)
        for _ in range(counts['knn_sp']):
            n = random.choice(self.nodes)
            events.append({
                "type": "knn", "id": eid, "poi": random.choice(ALLOWED_POIS),
                "query_point": {"lat": n["lat"], "lon": n["lon"]},
                "k": random.randint(3, 10), "metric": "shortest_path"
            })
            eid += 1

        # 5. Updates (Remove)
        for _ in range(counts['rem_edge']):
            e_id = self._get_rand_edge_id()
            if e_id is not None:
                events.append({"type": "remove_edge", "id": eid, "edge_id": e_id})
                eid += 1

        # 6. Updates (Modify)
        for _ in range(counts['mod_edge']):
            edge = self._get_rand_edge_data()
            if edge is not None:
                patch = {}
                if random.random() < 0.5: patch["length"] = round(edge["length"] * 1.5, 2)
                else: patch["road_type"] = "local" # distinct enough
                events.append({"type": "modify_edge", "id": eid, "edge_id": edge["id"], "patch": patch})
                eid += 1

        # SHUFFLE logic
        if shuffle:
            print("   -> Shuffling events to mix queries and updates.")
            random.shuffle(events)
            # Re-assign IDs after shuffle to ensure strict monotonicity if required?
            # Spec says IDs are unique integers. Usually good practice to keep them ordered 1..N
            # But for testing ID mapping, random order is fine. 
            # Let's re-index to be clean.
            for i, e in enumerate(events):
                e["id"] = i + 1

        return {"meta": {"id": "qset_phase1_custom"}, "events": events}

    # --------------------------------------------------------------------------
    # PHASE 2 GENERATION
    # --------------------------------------------------------------------------
    def generate_phase2(self, counts):
        print("Generating Phase 2 Queries...")
        events = []
        eid = 0

        # Exact KSP
        for _ in range(counts['ksp_exact']):
            s, t = self._get_rand_nodes(2)
            events.append({
                "type": "k_shortest_paths", "id": eid, "source": s, "target": t, 
                "k": random.randint(3, 10), "mode": "distance"
            })
            eid += 1

        # Heuristic KSP
        for _ in range(counts['ksp_heur']):
            s, t = self._get_rand_nodes(2)
            events.append({
                "type": "k_shortest_paths_heuristic", "id": eid, "source": s, "target": t, 
                "k": random.randint(3, 8), "overlap_threshold": random.randint(20, 80)
            })
            eid += 1

        # Approx Batch
        for _ in range(counts['approx']):
            batch_size = random.randint(5, 20)
            batch = [{"source": s, "target": t} for s, t in [self._get_rand_nodes(2) for _ in range(batch_size)]]
            events.append({
                "type": "approx_shortest_path", "id": eid, "queries": batch, 
                "time_budget_ms": 20 * batch_size, "acceptable_error_pct": 10.0
            })
            eid += 1

        return {"meta": {"id": "qset_phase2_custom"}, "events": events}

    # --------------------------------------------------------------------------
    # PHASE 3 GENERATION
    # --------------------------------------------------------------------------
    def generate_phase3(self, count):
        print("Generating Phase 3 Queries...")
        events = []
        
        for i in range(count):
            # Scale complexity with i
            complexity = i / max(1, count - 1)
            num_orders = int(2 + (complexity * 15))
            num_drivers = int(1 + (complexity * 3))
            
            orders = []
            for oid in range(num_orders):
                p, d = self._get_rand_nodes(2)
                orders.append({"order_id": oid + 1, "pickup": p, "dropoff": d})
                
            events.append({
                "orders": orders,
                "fleet": {
                    "num_delivery_guys": num_drivers,
                    "depot_node": self._get_rand_nodes(1)[0]
                }
            })
            
        return {"meta": {"id": "qset_phase3_custom"}, "events": events}

    # --------------------------------------------------------------------------
    # SAVER
    # --------------------------------------------------------------------------
    def save(self, data, filename):
        out_path = self.graph_path.parent / filename
        print(f"   -> Saving {out_path}")
        with open(out_path, 'w') as f:
            json.dump(data, f, indent=2)

# ==============================================================================
# MAIN CLI
# ==============================================================================
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate Custom Queries for an existing Graph")
    parser.add_argument("graph_file", help="Path to graph.json")
    
    # Phase 1 Args
    p1 = parser.add_argument_group('Phase 1 Customization')
    p1.add_argument("--sp_dist", type=int, default=10, help="Number of Shortest Path (Distance)")
    p1.add_argument("--sp_time", type=int, default=10, help="Number of Shortest Path (Time)")
    p1.add_argument("--knn_euc", type=int, default=5, help="Number of KNN (Euclidean)")
    p1.add_argument("--knn_sp", type=int, default=5, help="Number of KNN (Shortest Path)")
    p1.add_argument("--rem_edge", type=int, default=5, help="Number of Remove Edge updates")
    p1.add_argument("--mod_edge", type=int, default=5, help="Number of Modify Edge updates")
    p1.add_argument("--no_shuffle", action="store_true", help="Don't shuffle P1 events (keep updates at end?)")

    # Phase 2 Args
    p2 = parser.add_argument_group('Phase 2 Customization')
    p2.add_argument("--p2_exact", type=int, default=5, help="Exact KSP queries")
    p2.add_argument("--p2_heur", type=int, default=5, help="Heuristic KSP queries")
    p2.add_argument("--p2_approx", type=int, default=2, help="Approx Batch queries")

    # Phase 3 Args
    p3 = parser.add_argument_group('Phase 3 Customization')
    p3.add_argument("--p3_count", type=int, default=5, help="Number of Delivery Scheduling queries")

    args = parser.parse_args()

    gen = QueryGenerator(args.graph_file)

    # Phase 1
    p1_counts = {
        'sp_dist': args.sp_dist, 'sp_time': args.sp_time,
        'knn_euc': args.knn_euc, 'knn_sp': args.knn_sp,
        'rem_edge': args.rem_edge, 'mod_edge': args.mod_edge
    }
    q1_data = gen.generate_phase1(p1_counts, shuffle=not args.no_shuffle)
    gen.save(q1_data, "queries1.json")

    # Phase 2
    p2_counts = {'ksp_exact': args.p2_exact, 'ksp_heur': args.p2_heur, 'approx': args.p2_approx}
    q2_data = gen.generate_phase2(p2_counts)
    gen.save(q2_data, "queries2.json")

    # Phase 3
    q3_data = gen.generate_phase3(args.p3_count)
    gen.save(q3_data, "queries3.json")

    print("Done.")