Jyotisha/scripts/research/report_reference_perf_ab.py

"""A/B timing for the professional long report.

A is both switches off. B is both on. Same process, fictional charts only.
A field sha mismatch exits 1 and prints a short diff context, not the report.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import platform
import sys
import tempfile
import time
import tracemalloc
from pathlib import Path
from types import SimpleNamespace


ROOT = Path(__file__).resolve().parents[2]
for _folder in (ROOT / "scripts", ROOT / "jyotish_vedic", ROOT):
    _text = str(_folder)
    if _text not in sys.path:
        sys.path.insert(0, _text)

PRODUCTION_BODY = {
    "format": "markdown",
    "edition": "full_data",
    "include_fact_tables": True,
    "languages": ["zh", "en"],
}
SECTION2_CHART = {
    "year": 1990,
    "month": 1,
    "day": 1,
    "hour": 12,
    "minute": 0,
    "second": 0,
    "lat": 28.6139,
    "lon": 77.2090,
    "tz": 5.5,
    "node_mode": "mean",
    "ayanamsa": "lahiri",
    "today": "2026-10-13",
    "transit_date": "2026-10-13",
    "target_year": 2026,
    "age": 36,
    "birth_time_accuracy": "confirmed",
}
FIELDS = (
    "response",
    "markdown",
    "markdown_en",
    "fact_table_packet",
    "reader_dasha_applicability",
)


def _mask_pyjhora() -> str:
    root = Path(tempfile.mkdtemp(prefix="nojhora-"))
    package = root / "jhora"
    package.mkdir()
    (package / "__init__.py").write_text('raise ImportError("masked")\n', encoding="utf-8")
    sys.path.insert(0, str(root))
    prior = os.environ.get("PYTHONPATH", "")
    os.environ["PYTHONPATH"] = str(root) + (os.pathsep + prior if prior else "")
    for name in list(sys.modules):
        if name == "jhora" or name.startswith("jhora."):
            del sys.modules[name]
    return str(root)


def _sha(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def _canonical(value) -> str:
    if isinstance(value, str):
        return value
    return json.dumps(value, ensure_ascii=False, default=str)


def _field_texts(result: dict) -> dict[str, str]:
    whole = json.dumps(result, ensure_ascii=False, default=str)
    return {
        "response": whole,
        "markdown": _canonical(result.get("markdown", "")),
        "markdown_en": _canonical(result.get("markdown_en", "")),
        "fact_table_packet": _canonical(result.get("fact_table_packet")),
        "reader_dasha_applicability": _canonical(result.get("reader_dasha_applicability", "")),
    }


def _field_shas(result: dict) -> dict[str, str]:
    return {name: _sha(text) for name, text in _field_texts(result).items()}


def _diff_context(left: str, right: str) -> str:
    limit = min(len(left), len(right))
    index = next((i for i in range(limit) if left[i] != right[i]), limit)
    start = max(0, index - 40)
    stop = index + 40
    return (
        f"at {index} len {len(left)} vs {len(right)} "
        f"A={left[start:stop]!r} B={right[start:stop]!r}"
    )


def _median(values: list[float]) -> float:
    ordered = sorted(values)
    return ordered[len(ordered) // 2]


class _Bucket:
    def __init__(self) -> None:
        self.cmd = 0.0
        self.packet = 0.0
        self.zh = 0.0
        self.en = 0.0
        self.hash = 0.0

    def reset(self) -> None:
        self.cmd = self.packet = self.zh = self.en = self.hash = 0.0

    def as_dict(self) -> dict[str, float]:
        return {
            "cmd_full_reading_s": round(self.cmd, 4),
            "packet_s": round(self.packet, 4),
            "zh_render_s": round(self.zh, 4),
            "en_render_s": round(self.en, 4),
            "hash_s": round(self.hash, 4),
        }


def _install_timers(bucket: _Bucket) -> None:
    import scripts.calculation_profile_contract as contract
    import scripts.full_report_quality_gate as gate
    import scripts.jyotish_engine as engine
    import scripts.professional_report_reference as reference

    orig_packet = engine.build_professional_report_reference_packet

    def timed_packet(*args, **kwargs):
        started = time.perf_counter()
        try:
            return orig_packet(*args, **kwargs)
        finally:
            bucket.packet += time.perf_counter() - started

    engine.build_professional_report_reference_packet = timed_packet
    orig_render = reference.render_full_data_markdown

    def timed_render(packet, *args, **kwargs):
        started = time.perf_counter()
        try:
            return orig_render(packet, *args, **kwargs)
        finally:
            elapsed = time.perf_counter() - started
            if isinstance(packet, dict) and packet.get("report_language") == "en":
                bucket.en += elapsed
            else:
                bucket.zh += elapsed

    reference.render_full_data_markdown = timed_render
    orig_hash = contract.hash_bound_result

    def timed_hash(*args, **kwargs):
        started = time.perf_counter()
        try:
            return orig_hash(*args, **kwargs)
        finally:
            bucket.hash += time.perf_counter() - started

    contract.hash_bound_result = timed_hash
    gate.hash_bound_result = timed_hash


def _set_flags(fast: bool, cache: bool) -> None:
    os.environ["JYOTISH_REPORT_RENDER_FAST"] = "1" if fast else "0"
    os.environ["JYOTISH_REPORT_REFERENCE_CACHE"] = "1" if cache else "0"


class _Handler:
    def __init__(self, birth: dict, bucket: _Bucket) -> None:
        self.birth = birth
        self.bucket = bucket

    def _high_rigor_birth_payload(self, _body: dict) -> dict:
        return dict(self.birth)

    def _compute_full_reading_for_thematic(self, birth: dict):
        from scripts.jyotish_engine import cmd_full_reading

        started = time.perf_counter()
        try:
            return cmd_full_reading(SimpleNamespace(**birth))
        finally:
            self.bucket.cmd += time.perf_counter() - started


def _run_once(birth: dict, body: dict, bucket: _Bucket, handler=None) -> dict:
    from scripts.professional_report_reference import build_professional_report_reference

    bucket.reset()
    tracemalloc.start()
    started = time.perf_counter()
    result = build_professional_report_reference(handler or _Handler(birth, bucket), body)
    elapsed = time.perf_counter() - started
    _current, peak = tracemalloc.get_traced_memory()
    tracemalloc.stop()
    texts = _field_texts(result)
    return {
        "elapsed_s": round(elapsed, 4),
        "peak_bytes": peak,
        "response_bytes": len(texts["response"].encode("utf-8")),
        "segments": bucket.as_dict(),
        "sha": _field_shas(result),
        "_texts": texts,
    }


def _repeat(birth: dict, body: dict, bucket: _Bucket, times: int) -> list[dict]:
    return [_run_once(birth, body, bucket) for _ in range(times)]


def _public_row(row: dict) -> dict:
    return {key: value for key, value in row.items() if key != "_texts"}


def _compare(label: str, left: dict, right: dict, problems: list[str]) -> None:
    left_texts = left["_texts"]
    right_texts = right["_texts"]
    for name in FIELDS:
        if left["sha"][name] != right["sha"][name]:
            problems.append(f"{label} {name} {_diff_context(left_texts[name], right_texts[name])}")


def _freeze_clocks():
    """Pin datetime.now / utcnow in every already imported module.

    ``format=json`` and the reference edition embed wall-clock fields. Freezing
    only the engine module leaves other imports on the real clock.
    """
    import datetime as dt_mod

    real = dt_mod.datetime
    fixed_utc = real(2026, 10, 13, 0, 0, 0)
    fixed_local = real(2026, 10, 13, 8, 0, 0)

    class Frozen(real):
        @classmethod
        def now(cls, tz=None):
            if tz is not None:
                return fixed_utc.replace(tzinfo=tz)
            return fixed_local

        @classmethod
        def utcnow(cls):
            return fixed_utc

    saved = []
    for module in list(sys.modules.values()):
        if module is None:
            continue
        current = getattr(module, "datetime", None)
        if isinstance(current, type) and current.__name__ == "datetime":
            saved.append((module, current))
            module.datetime = Frozen
    return saved


def _restore_clocks(saved) -> None:
    for module, current in saved:
        module.datetime = current


def _strip_generated_at(value):
    if isinstance(value, dict):
        return {key: _strip_generated_at(item) for key, item in value.items() if key != "generated_at"}
    if isinstance(value, list):
        return [_strip_generated_at(item) for item in value]
    return value


def _sha_pair(birth: dict, body: dict, bucket: _Bucket) -> dict:
    """Compare one frozen clock on one reused full reading.

    A second ``cmd_full_reading`` appends onto shared trait lists, so two live
    engine calls are not a byte test of the render switch.
    """
    import copy

    from scripts.jyotish_engine import cmd_full_reading

    _set_flags(False, False)
    reading = cmd_full_reading(SimpleNamespace(**birth))

    class _Reuse:
        def _high_rigor_birth_payload(self, _body: dict) -> dict:
            return dict(birth)

        def _compute_full_reading_for_thematic(self, _birth: dict):
            return copy.deepcopy(reading)

    _set_flags(False, False)
    left = _run_once(birth, body, bucket, _Reuse())
    _set_flags(True, False)
    right = _run_once(birth, body, bucket, _Reuse())
    problems = []
    _compare("side", left, right, problems)
    compared = "full"
    if problems:
        left_stripped = _sha(_canonical(_strip_generated_at(json.loads(left["_texts"]["response"]))))
        right_stripped = _sha(_canonical(_strip_generated_at(json.loads(right["_texts"]["response"]))))
        if left_stripped == right_stripped:
            problems.clear()
            compared = "without_generated_at"
        else:
            compared = "mismatch"
    return {
        "compared": compared,
        "sha_a": left["sha"]["response"],
        "sha_b": right["sha"]["response"],
        "problems": problems,
    }


def _charts() -> dict[str, dict]:
    from tests.test_report_english_edition import CASES

    return {
        "day": {key: value for key, value in CASES["day"].items() if key != "name"},
        "night": {key: value for key, value in CASES["night"].items() if key != "name"},
        "south": {key: value for key, value in CASES["south"].items() if key != "name"},
        "section2": dict(SECTION2_CHART),
    }


def _clear_report_cache() -> None:
    folder = ROOT / "scratch" / "local" / "report_reference_cache"
    if not folder.is_dir():
        return
    for path in folder.glob("*.json"):
        path.unlink(missing_ok=True)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--mask-pyjhora", action="store_true")
    parser.add_argument("--only-a", action="store_true")
    parser.add_argument("--side-only", action="store_true")
    args = parser.parse_args()
    mask_path = _mask_pyjhora() if args.mask_pyjhora else ""
    bucket = _Bucket()
    _install_timers(bucket)
    os.environ.pop("GITHUB_SHA", None)
    report: dict = {
        "python": platform.python_version(),
        "platform": platform.platform(),
        "mask_pyjhora": bool(args.mask_pyjhora),
        "mask_path": mask_path,
        "only_a": args.only_a,
        "charts": {},
    }
    problems: list[str] = []
    chart_items = () if args.side_only else _charts().items()
    for label, birth in chart_items:
        _set_flags(False, False)
        arm_a = _repeat(birth, PRODUCTION_BODY, bucket, 3)
        entry: dict = {"A": [_public_row(row) for row in arm_a]}
        if not args.only_a:
            _set_flags(True, False)
            arm_b = _repeat(birth, PRODUCTION_BODY, bucket, 3)
            for row in arm_b:
                _compare(f"{label} B-miss", arm_a[0], row, problems)
            os.environ["GITHUB_SHA"] = "report-reference-perf-ab"
            _set_flags(True, True)
            _clear_report_cache()
            miss = _run_once(birth, PRODUCTION_BODY, bucket)
            _compare(f"{label} cache-miss", arm_a[0], miss, problems)
            hits = _repeat(birth, PRODUCTION_BODY, bucket, 3)
            for row in hits:
                _compare(f"{label} cache-hit", miss, row, problems)
            entry["B_miss"] = [_public_row(row) for row in arm_b]
            entry["cache_miss"] = _public_row(miss)
            entry["B_hit"] = [_public_row(row) for row in hits]
            os.environ.pop("GITHUB_SHA", None)
        report["charts"][label] = entry
        print(f"{label} done", file=sys.stderr)
    if not args.only_a:
        saved_clocks = _freeze_clocks()
        try:
            sample = _charts()["section2"]
            report["format_json"] = _sha_pair(sample, {"format": "json", "edition": "full_data"}, bucket)
            report["edition_reference"] = _sha_pair(
                sample, {"format": "markdown", "edition": "reference"}, bucket,
            )
        finally:
            _restore_clocks(saved_clocks)
        problems.extend(report["format_json"]["problems"])
        problems.extend(report["edition_reference"]["problems"])
        _set_flags(False, False)
        from scripts.calculation_profile_contract import hash_bound_result
        from scripts.jyotish_engine import build_professional_report_reference_packet, cmd_full_reading
        from scripts.professional_report_reference import _export_args

        reading = cmd_full_reading(SimpleNamespace(**sample))
        packet = build_professional_report_reference_packet(reading, _export_args(sample), [])
        input_hash = packet["calculation_profile"]["input_hash"]

        def _time_hash(fast: bool) -> float:
            _set_flags(fast, False)
            started = time.perf_counter()
            hash_bound_result(packet, input_hash)
            return time.perf_counter() - started

        off = [_time_hash(False) for _ in range(3)]
        on = [_time_hash(True) for _ in range(3)]
        report["hash_once_s"] = {"off": [round(item, 4) for item in off], "on": [round(item, 4) for item in on]}
    report["problems"] = problems
    print(json.dumps(report, ensure_ascii=False))
    return 1 if problems else 0


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