#!/usr/bin/env python3
"""Post-run Phase 4L audit corrections; does not execute or alter brain dynamics."""
from __future__ import annotations

import argparse
import json
import shutil
import time
from pathlib import Path

import numpy as np
import pandas as pd
import yaml

from phase4h_audit import sha256, write_json
from phase4l import load_external, pn_drive, seal


def main() -> None:
    parser = argparse.ArgumentParser(); parser.add_argument("--config", type=Path, required=True); args = parser.parse_args()
    cfg = yaml.safe_load(args.config.read_text()); exp = Path(cfg["paths"]["experiment_root"]); model = Path(cfg["paths"]["phase4l_model_root"])
    summary = json.loads((exp / "audit/final-summary.json").read_text()); identity = summary["eligible_edges"]
    a, b, shared = (int(identity[k]) for k in ("A_specific_eligible_edges", "B_specific_eligible_edges", "shared_eligible_edges"))
    identity["eligible_edge_Jaccard"] = shared / (a + b + shared)
    identity["specific_edge_intersection"] = 0
    identity["jaccard_definition"] = "(A-eligible intersection B-eligible)/(A-eligible union B-eligible), including shared edges"
    write_json(exp / "eligible-edges/summary.json", identity)

    # This is a separate read-only timing pass, not the timing of the original run.
    mapping = pd.read_parquet(exp / "receptor-mapping/olfactory-mapping.parquet")
    panel = pd.read_parquet(exp / "odor-encoder/frozen-odor-panel.parquet")
    vectors_npz = np.load(exp / "odor-encoder/mapped-glomerular-vectors.npz")
    vectors = {oid: vectors_npz[oid] for oid in panel.odor_id}; gloms = summary["formal_odors"]["glomeruli_order"]
    median = float(np.median(panel.mapped_response_sum)); start = time.perf_counter()
    for mode in cfg["odor_encoder"]["modes"]:
        for oid in panel.odor_id: pn_drive(oid, vectors[oid], mapping, gloms, mode, median, cfg)
    elapsed = time.perf_counter() - start
    perf = summary["performance"]; perf["odor_encoder_preprocessing_seconds"] = elapsed
    perf["odor_encoder_preprocessing_mean_seconds_per_odor_mode"] = elapsed / (len(panel) * len(cfg["odor_encoder"]["modes"]))
    perf["preprocessing_timing_method"] = "separate read-only repetition of frozen 12-odor, 2-mode response-to-PN transformation; not original simulation wall interval"
    perf["total_control_latency_mean_seconds_per_100ms"] = perf["odor_encoder_preprocessing_mean_seconds_per_odor_mode"] + perf["PN_event_generation_seconds_mean"] + perf["brain_wall_seconds_mean_per_100ms"]
    write_json(exp / "performance/summary.json", perf)

    write_json(exp / "audit/final-summary.json", summary)
    version = json.loads((model / "VERSION.json").read_text()); version["summary_sha256"] = sha256(exp / "audit/final-summary.json"); write_json(model / "VERSION.json", version)
    write_json(exp / "audit/post-run-corrections.json", {"only_phase4l_reports_modified": True, "brain_rerun": False,
        "eligible_Jaccard_corrected_to_include_shared_edges": identity["eligible_edge_Jaccard"],
        "encoder_preprocessing_seconds_separate_timing": elapsed, "run_log_sealed_after_process_exit": True})
    shutil.copy2(Path(__file__).resolve(), exp / "audit/phase4l_finalize.py")
    shutil.copy2(Path(__file__).resolve(), model / "phase4l_finalize.py")
    seal(model, model / "MANIFEST.json", model / "checksums.sha256")
    seal(exp, exp / "manifests/manifest.json", exp / "manifests/checksums.sha256")
    print(json.dumps({"eligible_Jaccard": identity["eligible_edge_Jaccard"], "preprocess_seconds": elapsed,
                      "experiment_checksum_count": len((exp / "manifests/checksums.sha256").read_text().splitlines())}))


if __name__ == "__main__": main()