#!/usr/bin/env python3
"""Compile and census spatial passage candidates on the four owner strata."""
from __future__ import annotations

import argparse
import json
import re
import subprocess
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parents[4]
OUT = ROOT / "Saved/Linux/l90/candidate-census"
DATA = ROOT / "Standalone/data/owner_world.json"
BINARY = ROOT / "Saved/Linux/l90/bin/VoxelForgeCoreHeadless-l90-clang19"
STRATA = {
    "S0": {"id": 0, "camera": (8.0, 16.0, -372.2626953125), "near": 1408, "whole": 2850},
    "FI": {"id": 1, "camera": (25.8369140625, 3.19921875, -737.0), "near": 1216, "whole": 2249},
    "TN": {"id": 2, "camera": (19.166015625, -14.2001953125, -753.3740234375), "near": 1344, "whole": 2233},
    "S3": {"id": 3, "camera": (-20.0625, -5.068359375, -900.353515625), "near": 1344, "whole": 2329},
}
PIPELINE = re.compile(r"^\[StreamPipeline\] (.*)$", re.M)
FIELDS = re.compile(r"([a-zA-Z0-9_]+)=([^ ]+)")
NEAR = re.compile(
    r"^\[WorldAssemblerNearLODReady\] lods=0-2 tiles=(\d+)/(\d+) cold_ready_s=([0-9.]+)$", re.M)
STAMP = re.compile(r"^\[VoxelForgeSourceStamp\] sha256=(\S+)$", re.M)
BUILD = re.compile(r"^\[L90PassageCandidateBuild\] (.*)$", re.M)
CENSUS = re.compile(r"^\[L90PassageCandidateTile\] (.*)$", re.M)
HIST_NAMES = ("carve_hist", "tread_hist", "clearance_hist", "mouth_column_hist")


def command(stratum: str) -> list[str]:
    spec = STRATA[stratum]
    camera = ",".join(f"{value:.12f}" for value in spec["camera"]) + ",0,0"
    return [
        "taskset", "-c", "0-11", "prlimit", "--as=6442450944", "--",
        "env", "VOXELFORGE_L90_PASSAGE_CANDIDATE_CENSUS=1",
        "VOXELFORGE_L81_PASSAGE_BOUND_PADDING_VOXELS=0", "timeout", "1740s",
        str(BINARY), str(DATA), "--world-assembler-prototype", "--stream-benchmark",
        "--stream-pipeline", "--workers=8", "--clip-radius=3", "--max-clip-level=8",
        "--max-level=8", f"--camera={camera}", "--profile-lod-work",
        "--profile-passage-composer", "--stream-scheduler-policy=3",
    ]


def parse_hist(value: str) -> list[int]:
    items = [int(item) for item in value.split(",")]
    if len(items) != 17:
        raise ValueError(f"expected 17 histogram bins, got {len(items)}")
    return items


def parse_fields(line: str) -> dict[str, str]:
    return dict(FIELDS.findall(line))


def parse_tuple(value: str, expected: int) -> list[int]:
    items = [int(item) for item in value.strip("()").split(",")]
    if len(items) != expected:
        raise ValueError(f"expected {expected} values, got {value}")
    return items


def run(stratum: str) -> dict:
    OUT.mkdir(parents=True, exist_ok=True)
    log = OUT / f"{stratum.lower()}-w8-profile.log"
    started = time.monotonic()
    with log.open("w", encoding="utf-8") as stream:
        result = subprocess.run(command(stratum), cwd=ROOT, stdout=stream,
                                stderr=subprocess.STDOUT, check=False, timeout=1750)
    output = log.read_text(encoding="utf-8", errors="replace")
    near, stamp = NEAR.search(output), STAMP.search(output)
    if result.returncode or near is None or stamp is None:
        raise RuntimeError(f"run failed or required output is missing: {log}; exit={result.returncode}")
    expected = STRATA[stratum]["near"]
    if (int(near.group(1)), int(near.group(2))) != (expected, expected):
        raise RuntimeError(f"incomplete near set in {log}: {near.group(1)}/{near.group(2)}")
    builds = [parse_fields(match.group(1)) for match in BUILD.finditer(output)]
    if not builds:
        raise RuntimeError(f"candidate compiler build row missing: {log}")

    hist = [[0] * 17 for _ in range(4)]
    census_rows = 0
    near_queries = no_brick = fallback_iterations = fallback_ns = query_ns = 0
    fallback_reach = fallback_shortspan = fallback_invalid = fallback_proof = 0
    fallback_carve = fallback_tread = fallback_clearance = fallback_mouth = 0
    proof_checks = [0] * 6
    proof_misses = [0] * 6
    for match in CENSUS.finditer(output):
        fields = parse_fields(match.group(1))
        if int(fields["lod"]) > 2:
            continue
        census_rows += 1
        near_queries += int(fields["queries"])
        no_brick += int(fields["no_brick"])
        fallback_iterations += int(fields["fallback_iterations"])
        fallback_reach += int(fields["fallback_reach"])
        fallback_shortspan += int(fields["fallback_shortspan"])
        fallback_invalid += int(fields["fallback_invalid"])
        fallback_proof += int(fields["fallback_proof"])
        fallback_ns += int(fields["fallback_ns"])
        query_ns += int(fields["query_ns"])
        fallback_carve += int(fields["fallback_carve"])
        fallback_tread += int(fields["fallback_tread"])
        fallback_clearance += int(fields["fallback_clearance"])
        fallback_mouth += int(fields["fallback_mouth"])
        proof_checks = [a + b for a, b in zip(proof_checks, parse_tuple(fields["proof_checks"], 6))]
        proof_misses = [a + b for a, b in zip(proof_misses, parse_tuple(fields["proof_misses"], 6))]
        for semantic, name in enumerate(HIST_NAMES):
            values = parse_hist(fields[name])
            hist[semantic] = [a + b for a, b in zip(hist[semantic], values)]
    if census_rows == 0 or near_queries == 0:
        raise RuntimeError(f"near-LOD candidate census rows missing: {log}")
    pipeline = PIPELINE.search(output)
    if pipeline is None:
        raise RuntimeError(f"stream pipeline summary missing: {log}")
    pipeline_fields = parse_fields(pipeline.group(1))
    row = {
        "stratum": stratum, "workers": 8, "exit_code": result.returncode,
        "wall_s": round(time.monotonic() - started, 3),
        "lod02_ready_s": float(near.group(3)),
        "whole_view_ready_s": float(pipeline_fields["whole_view_ready"].removesuffix("s")),
        "source_stamp": stamp.group(1), "near_tiles": expected,
        "census_scopes_lod02": census_rows, "near_query_hits": near_queries,
        "no_brick_hits": no_brick, "fallback_iterations": fallback_iterations,
        "fallback_reasons": {"reach_overflow": fallback_reach, "short_span": fallback_shortspan,
                             "invalid_compiler": fallback_invalid, "proof_miss": fallback_proof},
        "fallback_ns": fallback_ns, "query_ns": query_ns,
        "fallback_share": fallback_ns / query_ns if query_ns else None,
        "fallback_work": {"carve_visits": fallback_carve, "tread_visits": fallback_tread,
                          "clearance_checks": fallback_clearance, "mouth_checks": fallback_mouth},
        "proof_checks": dict(zip(("carve", "tread", "clearance", "mouth_route", "columns", "bridges"), proof_checks)),
        "proof_misses": dict(zip(("carve", "tread", "clearance", "mouth_route", "columns", "bridges"), proof_misses)),
        "histograms": {name.removesuffix("_hist"): values
                       for name, values in zip(HIST_NAMES, hist)},
        "build_rows": builds, "log": str(log.relative_to(ROOT)),
    }
    print("[L90CandidateCensusRun] " + json.dumps(row, sort_keys=True), flush=True)
    return row


def aggregate(rows: list[dict]) -> dict:
    total_hist = [[0] * 17 for _ in range(4)]
    for row in rows:
        for semantic, name in enumerate(HIST_NAMES):
            values = row["histograms"][name.removesuffix("_hist")]
            total_hist[semantic] = [a + b for a, b in zip(total_hist[semantic], values)]
    weighted = {}
    for semantic, name in enumerate(HIST_NAMES):
        values = total_hist[semantic]
        total = sum(values)
        def percentile(p: float, first_bin: int = 0) -> int:
            population = sum(values[first_bin:])
            target = max(1, int((p * population) + 0.999999999))
            seen = 0
            for index in range(first_bin, len(values)):
                count = values[index]
                seen += count
                if seen >= target:
                    return index
            return 16
        nonempty = sum(values[1:])
        weighted[name.removesuffix("_hist")] = {
            "queries": total,
            "histogram_0_to_15_plus": values,
            "mean_lower_bound": (sum(i * n for i, n in enumerate(values)) / total
                                 if total else 0.0),
            "p50_lower_bound": percentile(0.50) if total else 0,
            "p90_lower_bound": percentile(0.90) if total else 0,
            "p95_lower_bound": percentile(0.95) if total else 0,
            "share_16_plus": values[16] / total if total else 0.0,
            "nonempty_share": nonempty / total if total else 0.0,
            "nonempty_mean_lower_bound": (
                sum(i * n for i, n in enumerate(values)) / nonempty if nonempty else 0.0),
            "nonempty_p90_lower_bound": percentile(0.90, 1) if nonempty else 0,
            "nonempty_p95_lower_bound": percentile(0.95, 1) if nonempty else 0,
        }
    return {"strata": rows, "weighted_semantics": weighted}


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--strata", nargs="+", choices=tuple(STRATA), default=tuple(STRATA))
    args = parser.parse_args()
    if not BINARY.is_file() or not DATA.is_file():
        raise SystemExit(f"missing input: {BINARY} or {DATA}")
    rows = [run(stratum) for stratum in args.strata]
    summary = aggregate(rows)
    (OUT / "summary.json").write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8")
    print("[L90CandidateCensusSummary] " + json.dumps(summary["weighted_semantics"], sort_keys=True), flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())