"""Lab053: новый профиль Lab052 и точный файловый источник для конвейера."""

from __future__ import annotations

import argparse
from contextlib import contextmanager
from dataclasses import replace
import hashlib
import json
from pathlib import Path
import sys
import subprocess
import time
from types import SimpleNamespace

if __package__ in (None, ""):
    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import numpy as np

from experiments import lab052_candidate_detection as detection
from experiments import lab052_detector_study as study
from experiments.lab051_receiver_profile import StageProfiler

ROOT = Path(__file__).resolve().parents[1]
OUTPUT = ROOT / "data/processed/lab053"
SOURCE = ROOT / "data/raw/lab053/source.c64"


def write_result(name: str, value: dict) -> None:
    OUTPUT.mkdir(parents=True, exist_ok=True)
    (OUTPUT / name).write_text(json.dumps(value, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def distribution(values) -> dict:
    values = np.asarray(values, dtype=np.float64)
    if not values.size:
        return {"count": 0, "sum": 0.0, "mean": None, "median": None, "p95": None, "max": None}
    return {"count": len(values), "sum": float(values.sum()), "mean": float(values.mean()),
            "median": float(np.median(values)), "p95": float(np.quantile(values, .95)),
            "max": float(values.max())}


class DetailedProfiler(StageProfiler):
    """Вложенные времена не складываются повторно; только отдельный замер."""

    @contextmanager
    def measure(self, *, optimized=True):
        with super().measure(optimized=optimized):
            patches = []

            def wrap(function, name):
                def measured(*args, **kwargs):
                    scope = self.shared if self.current is None else self.current
                    stage = scope[name]
                    started = time.perf_counter()
                    entry = [0.0]
                    self.stack.append(entry)
                    try:
                        return function(*args, **kwargs)
                    finally:
                        elapsed = time.perf_counter() - started
                        self.stack.pop()
                        if self.stack:
                            self.stack[-1][0] += elapsed
                        stage.calls += 1
                        stage.seconds += elapsed
                        stage.exclusive_seconds += elapsed - entry[0]
                return measured

            targets = ((detection.CandidateFrontEnd, "coarse_scores", "coarse_detection"),
                       (detection.CandidateFrontEnd, "split_scores", "cheap_lab052"),
                       (detection.previous.DecimationCache, "get", "shared_decimation"),
                       (detection.previous.BlockFrontEnd, "extract", "soft_extraction"))
            try:
                for owner, name, label in targets:
                    original = getattr(owner, name)
                    patches.append((owner, name, original))
                    setattr(owner, name, wrap(original, label))
                yield self
            finally:
                for owner, name, original in reversed(patches):
                    setattr(owner, name, original)

    def detailed_summary(self):
        output = self.summary()
        for kind in ("valid", "false"):
            rows = [row for row in self.rows if row["kind"] == kind]
            names = sorted({name for row in rows for name in row["stages"]})
            output[kind]["packet_total"] = distribution([row["seconds"] for row in rows])
            output[kind]["per_packet_stages"] = {
                name: {field: distribution([row["stages"].get(name, {}).get(field, 0.0) for row in rows])
                       for field in ("seconds", "exclusive_seconds")} for name in names}
        return output


def baseline_profile() -> None:
    sequence = study.sequence_fixture()
    profiler = DetailedProfiler()
    started = time.perf_counter()
    with profiler.measure():
        stream, receiver = study.measured_stream(sequence, detection.StreamingReceiver)
    full_experiment_seconds = time.perf_counter() - started
    result = {"stream": stream, "profile": profiler.detailed_summary(), "packets": profiler.rows,
              "experiment_wall_seconds": full_experiment_seconds,
              "source_and_diagnostics_seconds": full_experiment_seconds - stream["processing_seconds"],
              "timing_includes_instrumentation": True, "rf_tx": 0}
    write_result("baseline_profile.json", result)
    print(json.dumps({"profile": result["profile"], "processing": stream["processing_seconds"],
                      "rx_packets": receiver.packet_count}), flush=True)


def prepare_source(frame_count=30) -> None:
    """Точные блоки прежнего генератора; конечный файл принадлежит источнику.

    RX далее читает ограниченные блоки обычным read, без mmap всего файла.
    Стоимость синтеза передаваемой формы не объявляется стоимостью live RX.
    """
    source_path = SOURCE if frame_count == 30 else SOURCE.with_name(f"source_{frame_count}.c64")
    source_path.parent.mkdir(parents=True, exist_ok=True)
    sequence = study.sequence_fixture()
    sequence = [replace(sequence[number % len(sequence)], image_id=number + 1) for number in range(frame_count)]
    checksum = hashlib.sha256()
    counts = []
    started = time.perf_counter()
    with source_path.open("wb") as output:
        class Recorder(detection.StreamingReceiver):
            def push_iq_block(self, samples):
                data = samples.tobytes()
                output.write(data)
                checksum.update(data)
                counts.append(samples.size)
                self.total_samples += samples.size
                return []

            def finish(self):
                self.closed = True
                return []

        original = study.performance.fast
        try:
            study.performance.fast = SimpleNamespace(StreamingReceiver=Recorder)
            study.performance.run_stream(sequence)
        finally:
            study.performance.fast = original
    result = {"path": str(source_path.relative_to(ROOT)), "dtype": "complex64", "blocks": len(counts),
              "input_frames": frame_count, "sample_count": sum(counts), "bytes": source_path.stat().st_size, "sha256": checksum.hexdigest(),
              "block_samples": 65536, "last_block_samples": counts[-1],
              "iq_seconds": sum(counts) / 8000000, "preparation_seconds": time.perf_counter() - started,
              "source_fixture": "data/processed/lab052/source_fixture.json", "rf_tx": 0}
    write_result("source.json" if frame_count == 30 else f"source_{frame_count}.json", result)
    print(json.dumps(result), flush=True)


def optimized_stream(workers=1) -> dict:
    from experiments.lab053_streaming_receiver import StreamingReceiver
    rows = []
    result, receiver = study.measured_stream(study.sequence_fixture(), lambda: StreamingReceiver(
        workers=workers, packet_observer=lambda row, outcome: rows.append(row)))
    result.update(workers=workers, packet_rows=rows, candidate_queue_peak=receiver.candidate_queue_peak,
                  candidate_overflow=receiver.candidate_overflow,
                  speculative_checks=receiver.speculative_checks,
                  discarded_speculative_checks=receiver.discarded_speculative_checks,
                  discarded_speculative_seconds=receiver.discarded_speculative_seconds)
    valid = [row for row in rows if row["accepted"]]
    result["valid_packet_total"] = distribution([row["seconds"] for row in valid])
    result["valid_packet_stages"] = {name: distribution([row["exclusive_stages"].get(name, 0) for row in valid])
                                    for name in sorted({name for row in valid for name in row["stages"]})}
    write_result(f"initial_workers_{workers}.json", result)
    print(json.dumps({key: result[key] for key in ("workers", "processing_seconds", "processing_per_rf_second",
          "rx_physical_packets", "received_images", "object_jpeg_integrity", "fresh_display_fps",
          "image_age_min_seconds", "image_age_median_seconds", "image_age_max_seconds", "candidate_count",
          "valid_packet_total", "candidate_queue_peak", "discarded_speculative_checks")}), flush=True)
    return result


def summarize_results() -> None:
    """Выбрать целые медианные проходы и пересчитать критерии из их строк."""
    def read(name):
        return json.loads((OUTPUT / name).read_text(encoding="utf-8"))

    selected = {}
    repeats = {}
    for mode in ("paced", "throughput"):
        runs = [read(f"{mode}_1_repeat_{number}.json") for number in (1, 2, 3)]
        chosen = sorted(range(3), key=lambda number: runs[number]["processing_seconds"])[1]
        result = dict(runs[chosen])
        result["selection"] = {"method": "whole median processing-time run", "repeat": chosen + 1,
                               "source": f"{mode}_1_repeat_{chosen + 1}.json"}
        selected[mode], repeats[mode] = result, runs
        write_result("stream.json" if mode == "paced" else "throughput.json", result)
        write_result(f"{mode}_workers_1.json", result)
    main, throughput = selected["paced"], selected["throughput"]
    comparison = [{"workers": number,
        "throughput_processing_per_rf_second": read(f"throughput_workers_{number}.json")["processing_per_rf_second"],
        **{key: read(f"paced_workers_{number}.json")[key] for key in (
            "processing_per_rf_second", "fresh_display_fps", "image_age_max_seconds", "candidate_queue_peak",
            "iq_queue_peak", "speculative_checks", "discarded_speculative_checks")}} for number in (1, 2, 4)]
    write_result("worker_comparison.json", {"selected_workers": 1, "comparisons": comparison,
        "reason": "single worker is sufficient; additional workers give no measured throughput or latency benefit"})
    long = read("long_stream.json")
    equivalence, sensitivity = read("equivalence.json"), read("sensitivity.json")
    regression = read("regression20k/result.json")
    source = read("source.json")
    overflow_probe = read("overflow_probe.json")
    with (ROOT / source["path"]).open("rb") as source_handle:
        source_sha256 = hashlib.file_digest(source_handle, "sha256").hexdigest()
    cases = repeats["paced"] + repeats["throughput"] + [read(f"{mode}_workers_{worker}.json")
        for mode in ("paced", "throughput") for worker in (2, 4)]
    checks = {
        "all_main_runs_240_packets_30_exact_jpegs": all(row["rx_physical_packets"] == 240 and row["received_images"] == 30
            and row["object_jpeg_integrity"] for row in cases),
        "all_main_runs_255_strict_15_false": all(row["candidate_count"] == 255 for row in cases),
        "all_paced_runs_5_fresh_fps": all(row["fresh_display_fps"] == 5 for row in cases if row["mode"] == "PACED_WALL_CLOCK"),
        "all_main_age_at_most_200ms": all(row["image_age_max_seconds"] <= .2 for row in cases if row["mode"] == "PACED_WALL_CLOCK"),
        "all_main_processing_ratio_below_08": all(row["processing_per_rf_second"] <= .8 for row in cases),
        "all_main_sources_complete_and_no_overflow": all(row["source_consumed_completely"] and not row["errors"] and
            row["iq_overflow"] == row["candidate_overflow"] == row["output_overflow"] == 0 for row in cases),
        "all_main_queues_within_caps": all(row["iq_queue_peak"] <= row["iq_queue_capacity"] and
            row["candidate_queue_peak"] <= row["candidate_queue_capacity"] and row["output_queue_peak"] <= row["output_queue_capacity"] for row in cases),
        "same_240_phy_crc_timing_all_worker_counts": equivalence["passed"],
        "sensitivity_37_of_45_identical": sensitivity["same_verdicts"] and sensitivity["optimized_correct"] == 37 and sensitivity["lost_baseline_packets"] == 0,
        "long_960_packets_120_jpeg_5_fps": long["rx_physical_packets"] == 960 and long["received_images"] == 120 and
            long["object_jpeg_integrity"] and long["fresh_display_fps"] == 5 and long["image_age_max_seconds"] <= .2,
        "long_queues_within_caps_no_overflow": long["iq_queue_peak"] <= long["iq_queue_capacity"] and
            long["candidate_queue_peak"] <= long["candidate_queue_capacity"] and not long["errors"] and
            long["iq_overflow"] == long["candidate_overflow"] == long["output_overflow"] == 0,
        "long_complete_source_and_cleared_state": long["consumed_iq_seconds"] == long["iq_seconds"] and
            long["retained_after_close_bytes"] == long["pending_jobs_after_close"] == 0,
        "old20k_36_crc_sha": regression["strict"]["crc_pass"] == 36 and regression["object"]["valid_fragments"] == 36 and
            regression["after_independent_verdict"]["sha_match"],
        "source_exact_duration": source["sample_count"] == 48355328 and source["iq_seconds"] == 6.044416,
        "source_file_size": (ROOT / source["path"]).stat().st_size == source["bytes"],
        "source_file_sha256": source_sha256 == source["sha256"],
        "forced_overflow_is_reported_and_state_cleared": overflow_probe["iq_overflow"] == 1 and
            overflow_probe["iq_queue_peak"] == overflow_probe["iq_queue_capacity"] and
            not overflow_probe["source_consumed_completely"] and
            any(row["reason"] == "IQ_QUEUE_OVERFLOW" for row in overflow_probe["errors"]) and
            overflow_probe["retained_after_close_bytes"] == overflow_probe["pending_jobs_after_close"] == 0,
        "protocol_unchanged": not subprocess.check_output(["git", "status", "--porcelain", "--untracked-files=all", "--", "protocol"], cwd=ROOT).strip(),
        "lab052_detector_unchanged": not subprocess.check_output(["git", "diff", "--", "experiments/lab052_candidate_detection.py"], cwd=ROOT).strip(),
        "lab053_index_empty": not subprocess.check_output(["git", "diff", "--cached", "--name-only"], cwd=ROOT).strip(),
    }
    fec_crc = sum(main["valid_packet_stages"][name]["sum"] for name in ("viterbi", "protocol_crc"))
    partition = {"A_iq_read_and_input_preparation": main["input_statistics"]["read_seconds"] + main["iq_input_preparation_seconds"],
        "B_common_decimation": main["shared_front_end_seconds"], "C_coarse_detection": main["coarse_detection_exclusive_seconds"],
        "D_cheap_lab052": main["cheap_seconds"], "E_true_packet_without_G": main["true_strict_seconds"] - fec_crc,
        "F_false_strict": main["false_strict_seconds"], "G_viterbi_crc": fec_crc,
        "H_fragment_assembly": main["assembly_seconds"], "I_jpeg_decode": main["jpeg_decode_seconds"],
        "J_freshness": main["output_statistics"]["freshness_seconds"]}
    partition["other_orchestration_and_instrumentation"] = main["processing_seconds"] - sum(partition.values())
    rss = long["rss_samples"]
    late = rss[len(rss)//2:]
    # К старой оценке временных массивов добавлены ЯВНЫЕ резервы постоянных кэшей.
    ring_samples = main["allocated_iq_ring_bytes"] // 16
    maximum_waveform = ring_samples - 2*65536 - 8000
    rotation_cache = 16 * maximum_waveform * 8
    detection_rotation_cache = 8 * ((ring_samples + 7)//8) * 8
    constants_reserve = 1024*1024
    memory = {"retained_iq_peak_bytes": main["retained_iq_peak_bytes"], "allocated_ring_bytes": main["allocated_iq_ring_bytes"],
        "iq_queue_peak_bytes": main["iq_queue_peak"] * 65536*8, "iq_queue_limit_bytes": main["iq_queue_limit_bytes"],
        "temporary_working_buffers_estimate_bytes": main["working_arrays_engineering_estimate_bytes"],
        "maximum_frame_waveform_samples": maximum_waveform, "rotation_cache_maxsize": 16,
        "maximum_rotation_cache_bytes": rotation_cache, "maximum_detection_rotation_cache_bytes": detection_rotation_cache,
        "constant_templates_and_trellis_reserve_bytes": constants_reserve,
        "working_arrays_and_constant_caches_estimate_bytes": main["working_arrays_engineering_estimate_bytes"] + rotation_cache + detection_rotation_cache + constants_reserve,
        "estimate_is_not_a_universal_c_allocation_bound": True,
        "process_memory": main["process_memory"], "long_process_memory": long["process_memory"],
        "long_late_rss_min_bytes": min(row["rss_bytes"] for row in late), "long_late_rss_max_bytes": max(row["rss_bytes"] for row in late),
        "long_late_rss_slope_bytes_per_second": float(np.polyfit([row["iq_seconds"] for row in late], [row["rss_bytes"] for row in late], 1)[0]),
        "receiver_state_bounded": checks["all_main_queues_within_caps"] and checks["long_queues_within_caps_no_overflow"]
            and checks["long_complete_source_and_cleared_state"], "experiment_diagnostic_history_is_external": True}
    write_result("memory.json", memory)
    write_result("acceptance.json", {"checks": checks, "required_checks_pass": all(checks.values()),
        "real_time_capable": all(checks.values()), "software_5fps_target_pass": all(checks.values()),
        "hardware_validated": False, "head": subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=ROOT, text=True).strip(),
        "source_code_sha256": {str(path.relative_to(ROOT)): hashlib.sha256(path.read_bytes()).hexdigest()
            for path in sorted((ROOT/"experiments").glob("lab053_*.py"))},
        "paced_service_partition_seconds": partition, "paced_service_processing_per_rf_second": main["processing_per_rf_second"],
        "saturated_wall_processing_per_rf_second": throughput["processing_per_rf_second"],
        "estimated_single_worker_service_for_40_packets_seconds": 40*main["valid_packet_total"]["mean"],
        "budget_note": "perf_counter service durations; Windows thread_time has 15.625ms accounting quantization", "rf_tx": 0})
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    figure, axes = plt.subplots(2, 2, figsize=(12, 8), constrained_layout=True)
    for number, row in enumerate(repeats["paced"], 1):
        axes[0, 0].plot([image["image_id"] for image in row["images"]],
            [1000*image["image_age_seconds"] for image in row["images"]], label=f"Проход {number}")
    axes[0, 0].axhline(200, color="red", linestyle="--", label="Предел 200 мс")
    axes[0, 0].set(xlabel="Номер изображения", ylabel="Возраст JPEG, мс", title="Подача IQ по часам: 5 свежих кадров/с")
    axes[0, 0].legend(fontsize=8)
    axes[0, 1].bar(["Lab052", "Lab053\nработа стадий", "Lab053\nбез ожидания"], [1.2192732766011487, main["processing_per_rf_second"], throughput["processing_per_rf_second"]])
    axes[0, 1].axhline(1, color="red", linestyle="--")
    axes[0, 1].axhline(.8, color="green", linestyle=":")
    axes[0, 1].set(ylabel="Время обработки / длительность IQ", title="Пропускная способность")
    axes[1, 0].plot([row["capture_timestamp"] for row in long["images"]], [1000*row["image_age_seconds"] for row in long["images"]])
    axes[1, 0].axhline(200, color="red", linestyle="--")
    axes[1, 0].set(xlabel="Время входной сессии, с", ylabel="Возраст JPEG, мс", title="Длинная сессия: 120 точных JPEG")
    axes[1, 1].plot([row["iq_seconds"] for row in rss], [row["rss_bytes"]/2**20 for row in rss], marker=".")
    axes[1, 1].set(xlabel="Принятая длительность IQ, с", ylabel="RSS процесса, МиБ", title="Память в длинной сессии")
    for axis in axes.flat:
        axis.grid(alpha=.2)
    figure.savefig(OUTPUT / "performance_and_age.png", dpi=150)
    figure.savefig(OUTPUT / "performance_and_age.svg")
    plt.close(figure)
    svg = OUTPUT / "performance_and_age.svg"
    svg.write_text("\n".join(line.rstrip() for line in svg.read_text(encoding="utf-8").splitlines()) + "\n", encoding="utf-8")
    print(json.dumps({"checks": checks, "memory": memory, "partition": partition}), flush=True)


def run_functional_tests() -> None:
    assert distribution([1, 2, 3])["median"] == 2


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--stage", choices=("profile", "source", "optimized", "report"), default="profile")
    parser.add_argument("--workers", type=int, default=1)
    parser.add_argument("--frames", type=int, default=30)
    arguments = parser.parse_args()
    if arguments.stage == "optimized":
        optimized_stream(arguments.workers)
    elif arguments.stage == "source":
        prepare_source(arguments.frames)
    elif arguments.stage == "report":
        summarize_results()
    else:
        baseline_profile()


if __name__ == "__main__":
    main()