from __future__ import annotations

import argparse

import json

import math

import sqlite3

import statistics

import time

from dataclasses import dataclass

from pathlib import Path

from typing import Any, Dict, List, Tuple

import pandas as pd

ROOT = Path(__file__).resolve().parent

DB_PATH = ROOT / "data" / "milano_places.sqlite"

QUERY_SETS_PATH = ROOT / "reports" / "query_sets_allthemes.json"

REPORTS = ROOT / "reports"

REPORTS.mkdir(parents=True, exist_ok=True)

DB_BUILD_CSV = REPORTS / "db_build_times_allthemes.csv"

DB_QUERY_CSV = REPORTS / "db_query_times_allthemes.csv"

DB_SUMMARY_CSV = REPORTS / "db_query_summary_allthemes.csv"

EARTH_R = 6371000.0

@dataclass(frozen=True)

class Projection:

lat0_rad: float

lon0_rad: float

def project_m(lat: float, lon: float, prj: Projection) -> Tuple[float, float]:

latr = math.radians(lat)

lonr = math.radians(lon)

x = EARTH_R * (lonr - prj.lon0_rad) * math.cos(prj.lat0_rad)

y = EARTH_R * (latr - prj.lat0_rad)

return x, y

def ensure_schema(conn: sqlite3.Connection) -> None:

conn.executescript(

"""

CREATE TABLE IF NOT EXISTS places_xy (

id INTEGER PRIMARY KEY,

fsq_place_id TEXT UNIQUE,

name TEXT,

top_category TEXT,

themes_json TEXT,

lat REAL,

lon REAL,

x REAL,

y REAL

);

CREATE VIRTUAL TABLE IF NOT EXISTS places_rtree USING rtree(

id,

minX, maxX,

minY, maxY

);

CREATE INDEX IF NOT EXISTS idx_places_xy_themes ON places_xy(themes_json);

CREATE INDEX IF NOT EXISTS idx_places_xy_topcat ON places_xy(top_category);

"""

)

conn.commit()

def rebuild_xy_and_rtree_for_theme(conn: sqlite3.Connection, theme: str) -> Tuple[Projection, float, int]:

"""

Rebuilds places_xy and places_rtree for a single theme (restaurants/theaters/libraries).

Assumes a base table 'places' exists with (fsq_place_id, name, lat, lon, categories_json, themes_json).

"""

cur = conn.cursor()

# Load only theme rows to compute projection + insert

cur.execute(

"""

SELECT fsq_place_id, name, lat, lon, categories_json, themes_json

FROM places

WHERE themes_json LIKE ?

AND lat IS NOT NULL AND lon IS NOT NULL

""",

(f"%{theme}%",),

)

rows = cur.fetchall()

if not rows:

raise RuntimeError(f"No rows found in DB for theme={theme}. Check places.themes_json content.")

lats = [float(r[2]) for r in rows]

lons = [float(r[3]) for r in rows]

lat0 = sum(lats) / len(lats)

lon0 = sum(lons) / len(lons)

prj = Projection(lat0_rad=math.radians(lat0), lon0_rad=math.radians(lon0))

# clear previous

conn.execute("DELETE FROM places_xy")

conn.execute("DELETE FROM places_rtree")

conn.commit()

ins_xy = """

INSERT INTO places_xy (id, fsq_place_id, name, top_category, themes_json, lat, lon, x, y)

VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)

"""

ins_rt = "INSERT INTO places_rtree (id, minX, maxX, minY, maxY) VALUES (?, ?, ?, ?, ?)"

t0 = time.perf_counter()

i = 1

for pid, name, lat, lon, cats_json, themes_json in rows:

# parse top category

try:

cats = json.loads(cats_json) if cats_json else []

top_cat = cats[0].get("name") if cats and isinstance(cats[0], dict) else "UNKNOWN"

except Exception:

top_cat = "UNKNOWN"

x, y = project_m(float(lat), float(lon), prj)

conn.execute(ins_xy, (i, pid, name or "", str(top_cat), themes_json or "[]", float(lat), float(lon), float(x), float(y)))

conn.execute(ins_rt, (i, float(x), float(x), float(y), float(y)))

i += 1

conn.commit()

t1 = time.perf_counter()

return prj, (t1 - t0), len(rows)

def nn_expand(cur: sqlite3.Cursor, qx: float, qy: float) -> List[int]:

"""

SQLite RTree does not provide a native NN operator, so we expand a window until candidates exist,

then pick closest by squared distance.

"""

r = 50.0

for _ in range(18):

cur.execute(

"""

SELECT p.id, p.x, p.y

FROM places_rtree r

JOIN places_xy p ON p.id=r.id

WHERE r.minX <= ? AND r.maxX >= ?

AND r.minY <= ? AND r.maxY >= ?

""",

(qx + r, qx - r, qy + r, qy - r),

)

cand = cur.fetchall()

if cand:

best_id = cand[0][0]

best_d2 = float("inf")

for pid, x, y in cand:

d2 = (x - qx) * (x - qx) + (y - qy) * (y - qy)

if d2 < best_d2:

best_d2 = d2

best_id = pid

return [int(best_id)]

r *= 2.0

return [1]

def bbox_query(cur: sqlite3.Cursor, minx: float, miny: float, maxx: float, maxy: float) -> List[int]:

cur.execute(

"""

SELECT p.id

FROM places_rtree r

JOIN places_xy p ON p.id=r.id

WHERE r.minX <= ? AND r.maxX >= ?

AND r.minY <= ? AND r.maxY >= ?

""",

(maxx, minx, maxy, miny),

)

return [int(x[0]) for x in cur.fetchall()]

def radius_query(cur: sqlite3.Cursor, qx: float, qy: float, rr: float) -> List[int]:

cur.execute(

"""

SELECT p.id, p.x, p.y

FROM places_rtree r

JOIN places_xy p ON p.id=r.id

WHERE r.minX <= ? AND r.maxX >= ?

AND r.minY <= ? AND r.maxY >= ?

""",

(qx + rr, qx - rr, qy + rr, qy - rr),

)

r2 = rr * rr

out = []

for pid, x, y in cur.fetchall():

if (x - qx) * (x - qx) + (y - qy) * (y - qy) <= r2:

out.append(int(pid))

return out

def composite_query(cur: sqlite3.Cursor, qx: float, qy: float, rr: float, cat_contains: str) -> List[int]:

needle = cat_contains.lower()

cur.execute(

"""

SELECT p.id, p.x, p.y, p.top_category

FROM places_rtree r

JOIN places_xy p ON p.id=r.id

WHERE r.minX <= ? AND r.maxX >= ?

AND r.minY <= ? AND r.maxY >= ?

""",

(qx + rr, qx - rr, qy + rr, qy - rr),

)

r2 = rr * rr

out = []

for pid, x, y, top_cat in cur.fetchall():

if (x - qx) * (x - qx) + (y - qy) * (y - qy) <= r2:

if needle in (top_cat or "").lower():

out.append(int(pid))

return out

def main() -> None:

ap = argparse.ArgumentParser()

ap.add_argument("--repeats", type=int, default=1)

args = ap.parse_args()

if not DB_PATH.exists():

raise FileNotFoundError(f"Missing DB: {DB_PATH}")

if not QUERY_SETS_PATH.exists():

raise FileNotFoundError(f"Missing query sets: {QUERY_SETS_PATH} (run 03_indexes_benchmark_allthemes.py first)")

query_sets_all = json.loads(QUERY_SETS_PATH.read_text(encoding="utf-8"))

conn = sqlite3.connect(DB_PATH)

ensure_schema(conn)

build_rows: List[Dict[str, Any]] = []

query_rows: List[Dict[str, Any]] = []

for theme, blob in query_sets_all.items():

qset = blob["queries"]

prj, build_s, n_points = rebuild_xy_and_rtree_for_theme(conn, theme)

build_rows.append({"theme": theme, "method": "sqlite_rtree", "build_s": build_s, "n_points": n_points})

cur = conn.cursor()

def bench(qtype: str, qid: int, fn):

times = []

res_len = None

for _ in range(args.repeats):

t0 = time.perf_counter()

res = fn()

t1 = time.perf_counter()

times.append((t1 - t0) * 1000.0)

res_len = len(res) if isinstance(res, list) else 1

query_rows.append(

{

"theme": theme,

"method": "sqlite_rtree",

"qtype": qtype,

"qid": qid,

"elapsed_ms_mean": float(statistics.mean(times)),

"elapsed_ms_median": float(statistics.median(times)),

"result_count": int(res_len),

}

)

for i in range(len(qset["nn"])):

qnn = qset["nn"][i]

qb = qset["bbox"][i]

qr = qset["radius"][i]

qc = qset["composite"][i]

qx, qy = float(qnn["qx"]), float(qnn["qy"])

minx, miny, maxx, maxy = map(float, (qb["minx"], qb["miny"], qb["maxx"], qb["maxy"]))

rr = float(qr["r"])

cat_f = str(qc["category_contains"])

bench("nn", i, lambda: nn_expand(cur, qx, qy))

bench("bbox", i, lambda: bbox_query(cur, minx, miny, maxx, maxy))

bench("radius", i, lambda: radius_query(cur, qx, qy, rr))

bench("composite", i, lambda: composite_query(cur, qx, qy, rr, cat_f))

print(f"[OK] theme={theme} points={n_points} | db rtree built + benchmarked")

pd.DataFrame(build_rows).to_csv(DB_BUILD_CSV, index=False, encoding="utf-8")

dfq = pd.DataFrame(query_rows)

dfq.to_csv(DB_QUERY_CSV, index=False, encoding="utf-8")

summ_rows = []

for (theme, method, qtype), g in dfq.groupby(["theme", "method", "qtype"]):

times = list(g["elapsed_ms_mean"].astype(float))

summ_rows.append(

{

"theme": theme,

"method": method,

"qtype": qtype,

"n": int(len(times)),

"mean_ms": float(statistics.mean(times)),

"median_ms": float(statistics.median(times)),

"min_ms": float(min(times)),

"max_ms": float(max(times)),

}

)

pd.DataFrame(summ_rows).sort_values(["theme", "qtype", "mean_ms"]).to_csv(DB_SUMMARY_CSV, index=False, encoding="utf-8")

print("Wrote:", DB_BUILD_CSV)

print("Wrote:", DB_QUERY_CSV)

print("Wrote:", DB_SUMMARY_CSV)

conn.close()

if __name__ == "__main__":

main()