"""Q1-Q7 over the oxigraph triplestore: SPARQL retrieval (spec 08).

The store holds raw data only (no precomputed summaries), so queries that
need steady-state COF or run-in retrieve the raw cycles via SPARQL and
derive the values with the shared algorithm (common/track_summary.py) -
identical semantics to every other format. The SPARQL text of each query
is exported to out/bench/queries/rdf/q<N>.sparql by make_queries.py.
"""

from __future__ import annotations

from pathlib import Path

from pyoxigraph import Store

from common.queries import Q1_TRACK_CODE
from common.track_summary import summarize_track

PREFIX = "PREFIX tribo: <https://sandia.gov/ontology/tribology#>\n"
TRACK_IRI = f"<https://sandia.gov/ontology/tribology/id/track/{Q1_TRACK_CODE}>"

SPARQL = {
	"q1": PREFIX + f"""SELECT ?cycle ?cof WHERE {{
  {TRACK_IRI} tribo:hasCycle ?c .
  ?c tribo:cycle ?cycle ; tribo:cof ?cof .
}}""",
	"q2": PREFIX + """SELECT ?ccode ?env ?track ?cycle ?cof WHERE {
  ?coupon tribo:au_wtpct_mean ?au ; tribo:coupon_code ?ccode .
  FILTER(?au >= 9.5 && ?au <= 10.5)
  ?track tribo:partOf ?coupon ; tribo:environment ?env ; tribo:hasCycle ?c .
  ?c tribo:cycle ?cycle ; tribo:cof ?cof .
}""",
	"q3_hardness": PREFIX + """SELECT ?ccode ?bcode (AVG(?h) AS ?hardness) WHERE {
  ?coupon tribo:coupon_code ?ccode ; tribo:cutFrom ?wafer ; tribo:nanoindentation ?m .
  ?wafer tribo:partOf ?batch .
  ?batch tribo:batch_code ?bcode .
  ?m tribo:hasIndent ?i .
  ?i tribo:hardness_gpa ?h .
} GROUP BY ?ccode ?bcode""",
	"q3_cycles": PREFIX + """SELECT ?ccode ?track ?cycle ?cof WHERE {
  ?track a tribo:FrictionTest ; tribo:partOf ?coupon ; tribo:hasCycle ?c .
  ?coupon tribo:coupon_code ?ccode .
  ?c tribo:cycle ?cycle ; tribo:cof ?cof .
}""",
	"q4": PREFIX + """SELECT ?tcode ?track ?cycle ?cof WHERE {
  ?track tribo:environment "dry_n2" ; tribo:load_mn ?load ; tribo:track_code ?tcode ; tribo:hasCycle ?c .
  FILTER(?load = 100)
  ?c tribo:cycle ?cycle ; tribo:cof ?cof .
}""",
	"q5": PREFIX + """SELECT ?bcode ?track ?cycle ?cof WHERE {
  ?track a tribo:FrictionTest ; tribo:partOf ?coupon ; tribo:hasCycle ?c .
  ?coupon tribo:cutFrom ?wafer .
  ?wafer tribo:partOf ?batch .
  ?batch tribo:batch_code ?bcode .
  ?c tribo:cycle ?cycle ; tribo:cof ?cof .
}""",
	"q6": PREFIX + """SELECT ?ccode ?load (AVG(?wv) AS ?wear) WHERE {
  ?track tribo:partOf ?coupon ; tribo:load_mn ?load ; tribo:wear ?w .
  ?w tribo:wear_volume_um3 ?wv .
  ?coupon tribo:coupon_code ?ccode .
} GROUP BY ?ccode ?load""",
	"q7": PREFIX + """SELECT ?env ?speed ?load ?track ?cycle ?cof WHERE {
  ?track tribo:environment ?env ; tribo:speed_mm_s ?speed ; tribo:load_mn ?load ; tribo:hasCycle ?c .
  ?c tribo:cycle ?cycle ; tribo:cof ?cof .
}""",
}


def _store(store_path: Path) -> Store:
	try:
		return Store.read_only(str(store_path))
	except AttributeError:
		return Store(str(store_path))


def _track_cofs(solutions, track_var: str, keep: tuple[str, ...]) -> dict[str, dict]:
	"""Group cycle solutions by track: ordered cof list + kept literals."""
	tracks: dict[str, dict] = {}
	for sol in solutions:
		key = str(sol[track_var])
		entry = tracks.get(key)
		if entry is None:
			entry = {"pairs": []}
			for name in keep:
				entry[name] = sol[name].value
			tracks[key] = entry
		entry["pairs"].append((int(sol["cycle"].value), float(sol["cof"].value)))
	for entry in tracks.values():
		entry["pairs"].sort()
		entry["cofs"] = [cof for _, cof in entry["pairs"]]
	return tracks


def _mean(values: list[float]) -> float:
	total = 0.0
	for v in values:
		total += v
	return total / len(values)


def q1(store_path: Path) -> list[tuple]:
	store = _store(store_path)
	return [(int(s["cycle"].value), float(s["cof"].value)) for s in store.query(SPARQL["q1"])]


def q2(store_path: Path) -> list[tuple]:
	store = _store(store_path)
	tracks = _track_cofs(store.query(SPARQL["q2"]), "track", ("ccode", "env"))
	coupons: dict[tuple[str, str], list[float]] = {}
	for entry in sorted(tracks.values(), key=lambda e: e["ccode"]):
		coupons.setdefault((entry["ccode"], entry["env"]), []).append(summarize_track(entry["cofs"])[0])
	return [(code, env, _mean(values)) for (code, env), values in sorted(coupons.items())]


def q3(store_path: Path) -> list[tuple]:
	store = _store(store_path)
	hardness = {
		s["ccode"].value: (s["bcode"].value, float(s["hardness"].value))
		for s in store.query(SPARQL["q3_hardness"])
	}
	tracks = _track_cofs(store.query(SPARQL["q3_cycles"]), "track", ("ccode",))
	coupons: dict[str, list[float]] = {}
	for entry in sorted(tracks.values(), key=lambda e: e["ccode"]):
		coupons.setdefault(entry["ccode"], []).append(summarize_track(entry["cofs"])[0])
	return [
		(code, hardness[code][0], hardness[code][1], _mean(values))
		for code, values in sorted(coupons.items())
	]


def q4(store_path: Path) -> list[tuple]:
	store = _store(store_path)
	tracks = _track_cofs(store.query(SPARQL["q4"]), "track", ("tcode",))
	out = []
	for entry in sorted(tracks.values(), key=lambda e: e["tcode"]):
		cof_ss = summarize_track(entry["cofs"])[0]
		if cof_ss > 0.20:
			out.append((entry["tcode"], cof_ss))
	return out


def q5(store_path: Path) -> list[tuple]:
	store = _store(store_path)
	tracks = _track_cofs(store.query(SPARQL["q5"]), "track", ("bcode",))
	sums: dict[str, list[float]] = {}
	for entry in sorted(tracks.values(), key=lambda e: e["bcode"]):
		acc = sums.setdefault(entry["bcode"], [0.0, 0.0])
		acc[0] += summarize_track(entry["cofs"])[2]
		acc[1] += 1
	return [(batch, acc[0] / acc[1]) for batch, acc in sorted(sums.items())]


def q6(store_path: Path) -> list[tuple]:
	store = _store(store_path)
	return [
		(s["ccode"].value, float(s["load"].value), float(s["wear"].value))
		for s in store.query(SPARQL["q6"])
	]


def q7(store_path: Path) -> list[tuple]:
	store = _store(store_path)
	tracks = _track_cofs(store.query(SPARQL["q7"]), "track", ("env", "speed", "load"))
	groups: dict[tuple[str, int], list[float]] = {}
	for key in sorted(tracks):
		entry = tracks[key]
		bucket = int(round(float(entry["speed"]) * float(entry["load"])))
		run_in = summarize_track(entry["cofs"])[2]
		acc = groups.setdefault((entry["env"], bucket), [0.0, 0.0])
		for cof in entry["cofs"][run_in:]:
			acc[0] += cof
			acc[1] += 1
	return [(env, bucket, acc[0] / acc[1]) for (env, bucket), acc in sorted(groups.items())]
