"""Task 10 - Extrapolation to 600 GB / 1.2 TB / 6 TB.

Calibrates per-cell scaling laws on the measured 150 MB point (task 09
medians) and projects wall times and hardware needs per format x scale.
Model (spec 10, docs/rules/bench-methodology.md section 7):
- per-format fixed overhead = the format's cheapest measured median
  (interpreter start, imports, connection); the remainder is the
  data-dependent calibration constant;
- growth laws per access pattern: o1 (flat), log (index depth,
  1 + ln r / ln n0 with n0 = measured friction-cycle rows), logk (result
  set grows linearly with data), k_log_n (linear x index depth), scan
  (linear), scan_parallel (linear x base_cores / cores_at_scale; range
  partitions outnumber cores at every scale);
- RDF join penalties are already embedded in the measured constants;
- RDF q3/q5/q7 are scans because the implementation derives run-in /
  steady-state from raw cycles (no summary triples exist by design).
Every number written here is a PROJECTION, never a measurement; rows are
flagged IMPRACTICAL (> 1 h) and FAIL (> 24 h).
Spec: docs/specs/10_extrapolation_model.md.
"""

from __future__ import annotations

import csv
import logging
import math
import sys
from pathlib import Path

import psycopg
import yaml

from common.pg import get_dsn
from common.pipeline import (
	check_dependencies,
	load_lab_config,
	parse_task_args,
	remove_stale_marker,
	write_marker,
)
from common.relational import expected_counts

logger = logging.getLogger(__name__)

TASK_ID = "10"
DEPENDS_ON = ["09"]

SCALES_GB = (600, 1200, 6000)  # raw-CSV target scales (decimal GB), spec 10
IMPRACTICAL_S = 3600.0         # projection > 1 h, spec 10
FAIL_S = 24 * 3600.0           # projection > 24 h, spec 10
DISK_HEADROOM = 1.3            # 30% free-space headroom, spec 10
CACHE_INDEX_FRACTION = 0.10    # relational hot-set: 10% of index size, spec 10
WORKING_SET_GB = 4.0           # relational session working set, spec 10 default
RDF_HOT_FRACTION = 0.20        # dictionary+index working set, 15-25% midpoint, spec 10
STREAM_RAM_GB = 16.0           # flat streaming baseline (csv/json servers)
MIN_RAM_GB = 16.0              # smallest sensible server RAM
NODE_RAM_CEILING_GB = 1024.0   # single-node RAM threshold, spec 10 RDF note
BASE_CORES = 4                 # cores behind the measured PG parallel scans (task 09 host)
CORE_OPTIONS = (16, 32)        # priced server tiers, hw_prices.yaml

FORMATS = ("csv", "json", "sqlite", "pg", "rdf")
QUERY_NUMBERS = tuple(range(1, 8))

# working (queryable) footprint variant per format for storage scaling
SIZE_VARIANT = {
	"csv": ("csv", "tree"), "json": ("json", "full"), "sqlite": ("sqlite", "db"),
	"pg": ("pg", "live"), "rdf": ("rdf", "store"),
}

# access-pattern law per (format, query) - spec 10 table, adjusted to the
# ACTUAL implementations (rdf q3/q5/q7 read raw cycles; see module docstring)
LAWS = {
	("csv", 1): "o1", ("json", 1): "o1", ("sqlite", 1): "log", ("pg", 1): "log", ("rdf", 1): "log",
	("csv", 2): "scan", ("json", 2): "scan", ("sqlite", 2): "logk", ("pg", 2): "logk", ("rdf", 2): "logk",
	("csv", 3): "scan", ("json", 3): "scan", ("sqlite", 3): "k_log_n", ("pg", 3): "k_log_n", ("rdf", 3): "scan",
	("csv", 4): "scan", ("json", 4): "scan", ("sqlite", 4): "logk", ("pg", 4): "logk", ("rdf", 4): "logk",
	("csv", 5): "scan", ("json", 5): "scan", ("sqlite", 5): "scan", ("pg", 5): "scan_parallel", ("rdf", 5): "scan",
	("csv", 6): "scan", ("json", 6): "scan", ("sqlite", 6): "k_log_n", ("pg", 6): "k_log_n", ("rdf", 6): "k_log_n",
	("csv", 7): "scan", ("json", 7): "scan", ("sqlite", 7): "scan", ("pg", 7): "scan_parallel", ("rdf", 7): "scan",
}

HW_PRICES_DEFAULT = """# Hardware cost parameters for extrapolation (task 10).
# Street prices retrieved 2026-07-11 from the sources below (server DRAM is
# in a documented 2025-2026 price surge - see market_context). The cost
# model buys whole components: 64 GB RDIMM modules and 7.68 TB NVMe drives.
# Operator-editable; re-run extrapolate.py after changes (code-config-yaml).
as_of: "2026-07-11"
ram_module_gb: 64
ram_module_usd: 1887.0       # A-Tech 64GB (2x32GB) DDR5-5600 ECC RDIMM
nvme_drive_tb: 7.68
nvme_drive_usd: 3995.0       # Cloud Ninjas NEW 7.68TB NVMe U.2 1DWPD (Dell 14-16G)
server_16_core_usd: 4618.0   # Supermicro AS-1015CS-TNR 1U (EPYC 9004/9005), Broadberry starting config
server_32_core_usd: 6272.0   # Supermicro AS-1115CS-TNR 1U, Broadberry starting config
extra_node_usd: 4618.0       # one additional entry 1U node (AS-1015CS-TNR chassis)
sources:
  ram: https://atechmemory.com/collections/ddr5-memory-ram
  nvme: https://cloudninjas.com/products/new-7-68tb-nvme-u-2-1dwpd-sie-2-5-enterprise-solid-state-drive-for-14th-15th-16th-gen-dell
  servers: https://www.broadberry.com/amd-epyc-9004-supermicro-servers
  market_context: https://www.techpowerup.com/342331/server-dram-pricing-jumps-50-only-70-of-orders-getting-filled
"""
HW_PRICE_KEYS = (
	"as_of", "ram_module_gb", "ram_module_usd", "nvme_drive_tb", "nvme_drive_usd",
	"server_16_core_usd", "server_32_core_usd", "extra_node_usd", "sources",
)


def read_medians(path: Path) -> dict[tuple[str, int], float]:
	"""Median walls of the 35 cells; abort unless every cell is OK + checksummed
	(projections built on invalid measurements are forbidden, bench-methodology).
	"""
	medians: dict[tuple[str, int], float] = {}
	with open(path, encoding="ascii", newline="") as fh:
		for row in csv.DictReader(fh):
			if row["status"] != "OK" or row["checksum_ok"] != "true":
				raise ValueError(f"{path}: cell {row['format']}/{row['query']} is {row['status']}/checksum={row['checksum_ok']} - cannot calibrate on it")
			medians[(row["format"], int(row["query"][1:]))] = float(row["wall_s_median"])
	expected_cells = {(fmt, qn) for fmt in FORMATS for qn in QUERY_NUMBERS}
	if set(medians) != expected_cells:
		raise ValueError(f"{path}: {len(medians)} cells, expected the full 35-cell matrix")
	return medians


def read_sizes(path: Path) -> dict[tuple[str, str], int]:
	sizes: dict[tuple[str, str], int] = {}
	with open(path, encoding="ascii", newline="") as fh:
		reader = csv.reader(fh)
		next(reader)
		for fmt, variant, nbytes in reader:
			sizes[(fmt, variant)] = int(nbytes)
	missing = set(SIZE_VARIANT.values()) - set(sizes)
	if missing:
		raise ValueError(f"{path}: missing rows {sorted(missing)}")
	return sizes


def pg_index_share(dsn: str) -> tuple[float, int]:
	"""Index share of the live PG footprint (read-only catalog query)."""
	with psycopg.connect(dsn) as conn:
		index_bytes, total_bytes = conn.execute(
			"SELECT COALESCE(sum(pg_indexes_size(c.oid)), 0)::bigint,"
			"       COALESCE(sum(pg_total_relation_size(c.oid)), 0)::bigint"
			" FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace"
			" WHERE n.nspname = 'public' AND c.relkind IN ('r', 'p', 'm')"
		).fetchone()
	if total_bytes == 0:
		raise ValueError("pg catalog reports zero relation bytes - is the database loaded?")
	return index_bytes / total_bytes, index_bytes


def load_hw_prices(config_dir: Path) -> dict:
	path = config_dir / "hw_prices.yaml"
	if not path.exists():
		path.write_text(HW_PRICES_DEFAULT, encoding="ascii", newline="\n")
		logger.info("seeded editable defaults: %s", path)
	with open(path, encoding="ascii") as fh:
		prices = yaml.safe_load(fh)
	missing = [k for k in HW_PRICE_KEYS if k not in prices]
	if missing:
		raise ValueError(f"{path}: missing keys {missing}")
	return prices


def growth(law: str, r: float, n0: int, cores: int) -> float:
	depth = 1.0 + math.log(r) / math.log(n0)
	if law == "o1":
		return 1.0
	if law == "log":
		return depth
	if law == "logk":
		return r
	if law == "k_log_n":
		return r * depth
	if law == "scan":
		return r
	if law == "scan_parallel":
		return r * BASE_CORES / cores
	logger.error("growth: unknown law=%r", law)
	raise ValueError(f"unknown scaling law: {law!r}")


def flag_of(seconds: float) -> str:
	if seconds > FAIL_S:
		return "FAIL"
	if seconds > IMPRACTICAL_S:
		return "IMPRACTICAL"
	return "OK"


def project(fmt: str, qn: int, medians: dict, overhead: dict, r: float, n0: int, cores: int) -> float:
	variable = max(0.0, medians[(fmt, qn)] - overhead[fmt])
	return overhead[fmt] + variable * growth(LAWS[(fmt, qn)], r, n0, cores)


def pick_pg_cores(medians: dict, overhead: dict, r: float, n0: int) -> int:
	"""Smallest priced tier; escalate only when it fixes an IMPRACTICAL flag
	on a parallel-scan cell (the only law that benefits from cores)."""
	for cores in CORE_OPTIONS:
		worst = max(project("pg", qn, medians, overhead, r, n0, cores) for qn in QUERY_NUMBERS if LAWS[("pg", qn)] == "scan_parallel")
		if worst <= IMPRACTICAL_S:
			return cores
	return CORE_OPTIONS[-1]


def main() -> int:
	args = parse_task_args("Task 10: extrapolate measurements to 600 GB / 1.2 TB / 6 TB")
	out_root: Path = args.out_root
	bench_dir = out_root / "bench"
	try:
		check_dependencies(out_root, DEPENDS_ON)
		remove_stale_marker(out_root, TASK_ID)
		medians = read_medians(bench_dir / "results_median.csv")
		sizes = read_sizes(bench_dir / "storage_sizes.csv")
		prices = load_hw_prices(out_root / "config")
		n0 = expected_counts(load_lab_config(out_root))["friction_cycles"]
		csv_bytes = sizes[SIZE_VARIANT["csv"]]
		coeff = {fmt: sizes[SIZE_VARIANT[fmt]] / csv_bytes for fmt in FORMATS}
		index_share, pg_index_bytes = pg_index_share(get_dsn())
		logger.info("storage coefficients vs csv: %s; pg index share %.3f",
			", ".join(f"{fmt}={coeff[fmt]:.2f}" for fmt in FORMATS), index_share)
		overhead = {fmt: min(medians[(fmt, qn)] for qn in QUERY_NUMBERS) for fmt in FORMATS}

		flags = {"OK": 0, "IMPRACTICAL": 0, "FAIL": 0}
		with open(bench_dir / "extrapolation.csv", "w", encoding="ascii", newline="") as fh:
			writer = csv.writer(fh)
			writer.writerow(["format", "query", "scale_gb", "projected_wall_s", "flag"])
			for scale_gb in SCALES_GB:
				r = scale_gb * 1e9 / csv_bytes
				cores = {fmt: CORE_OPTIONS[0] for fmt in FORMATS}
				cores["pg"] = pick_pg_cores(medians, overhead, r, n0)
				for fmt in FORMATS:
					for qn in QUERY_NUMBERS:
						seconds = project(fmt, qn, medians, overhead, r, n0, cores[fmt])
						flag = flag_of(seconds)
						flags[flag] += 1
						writer.writerow([fmt, f"q{qn}", scale_gb, round(seconds, 3), flag])

		with open(bench_dir / "hardware_sizing.csv", "w", encoding="ascii", newline="") as fh:
			writer = csv.writer(fh)
			writer.writerow(["format", "scale_gb", "disk_tb", "ram_gb", "cores", "nodes",
				"est_cost_usd", "ram_modules", "nvme_drives", "base_usd", "ram_usd", "disk_usd", "extra_nodes_usd"])
			for scale_gb in SCALES_GB:
				r = scale_gb * 1e9 / csv_bytes
				for fmt in FORMATS:
					fmt_bytes = coeff[fmt] * scale_gb * 1e9
					disk_tb = fmt_bytes * DISK_HEADROOM / 1e12
					if fmt in ("csv", "json"):
						ram_gb = STREAM_RAM_GB
					elif fmt in ("sqlite", "pg"):
						# same index share assumed for sqlite: identical rows,
						# keys, and indexes as the PG schema (variant B)
						ram_gb = max(MIN_RAM_GB, CACHE_INDEX_FRACTION * index_share * fmt_bytes / 1e9 + WORKING_SET_GB)
					else:
						ram_gb = max(MIN_RAM_GB, RDF_HOT_FRACTION * fmt_bytes / 1e9)
					cores = pick_pg_cores(medians, overhead, r, n0) if fmt == "pg" else CORE_OPTIONS[0]
					nodes = max(1, math.ceil(ram_gb / NODE_RAM_CEILING_GB))
					# bill of materials: whole RDIMM modules and whole NVMe drives
					ram_modules = math.ceil(ram_gb / prices["ram_module_gb"])
					nvme_drives = math.ceil(disk_tb / prices["nvme_drive_tb"])
					base_usd = prices["server_16_core_usd"] if cores == 16 else prices["server_32_core_usd"]
					ram_usd = ram_modules * prices["ram_module_usd"]
					disk_usd = nvme_drives * prices["nvme_drive_usd"]
					extra_nodes_usd = (nodes - 1) * prices["extra_node_usd"]
					cost = base_usd + ram_usd + disk_usd + extra_nodes_usd
					writer.writerow([fmt, scale_gb, round(disk_tb, 3), round(ram_gb, 1), cores, nodes,
						round(cost), ram_modules, nvme_drives,
						round(base_usd), round(ram_usd), round(disk_usd), round(extra_nodes_usd)])

		entries = {
			"extrapolation_rows": len(SCALES_GB) * len(FORMATS) * len(QUERY_NUMBERS),
			"sizing_rows": len(SCALES_GB) * len(FORMATS),
			"flags_ok": flags["OK"],
			"flags_impractical": flags["IMPRACTICAL"],
			"flags_fail": flags["FAIL"],
			"pg_index_share": round(index_share, 3),
			"pg_index_mb": round(pg_index_bytes / 1048576, 1),
			"prices_as_of": prices["as_of"],
			"note_pandas": "pandas CSV variant not measured; its RAM is O(n) - infeasible above ~10 GB (spec 10)",
		}
		for fmt in FORMATS:
			entries[f"coeff_{fmt}"] = round(coeff[fmt], 3)
		marker_path = write_marker(out_root, TASK_ID, entries)
	except Exception:
		logger.critical("task %s failed", TASK_ID, exc_info=True)
		return 1

	print(f"task 10 ok: {entries['extrapolation_rows']} projections "
		f"(ok={flags['OK']}, impractical={flags['IMPRACTICAL']}, fail={flags['FAIL']}), "
		f"{entries['sizing_rows']} sizing rows")
	print(f"outputs: {bench_dir / 'extrapolation.csv'}, {bench_dir / 'hardware_sizing.csv'}")
	print(f"marker: {marker_path}")
	return 0


if __name__ == "__main__":
	sys.exit(main())
