"""Run the pinned author's create_model in a dedicated, serial Brian2 process.

This module never edits the author source or its parameters.  The caller supplies
an exclusive output directory and the already validated, signed connectivity
multipliers in ``raw_count``.  Global IDs are not FlyWire IDs and are mapped only
for this small-circuit diagnostic.  Brian stores voltages in SI volts internally;
the returned state and weights are converted to mV in the selected dtype.
"""
from __future__ import annotations

import hashlib
import importlib.metadata
import json
import platform
import sys
import types
from pathlib import Path

import numpy as np
import pandas as pd


AUTHOR_SHA256 = "fc45837d7122c6ce2a7f3f2f23c515992e4b232aadb919efabb72337fac88e4e"
AUTHOR_COMMIT = "91bdd1e7dcf193f3e7ca5a8933497fcef63b7960"
PACKAGE_PINS = {
    "brian2": "2.5.1", "numpy": "1.22.3", "pandas": "1.4.3",
    "scipy": "1.8.1", "pyarrow": "11.0.0", "joblib": "1.2.0",
    "cython": "0.29.33", "jinja2": "2.11.3", "markupsafe": "1.1.1",
    "sympy": "1.12", "mpmath": "1.3.0", "pyparsing": "3.0.9",
    "packaging": "23.1", "setuptools": "67.8.0",
    "python-dateutil": "2.8.2", "pytz": "2023.3", "six": "1.16.0",
    "wheel": "0.38.4",
}


def _json(path: Path, value) -> None:
    with path.open("x", encoding="utf-8") as output:
        json.dump(value, output, indent=2, sort_keys=True, allow_nan=False)
        output.write("\n")


def _steps(value: float, dt: float, label: str, *, positive=False) -> int:
    ratio = float(value) / dt
    if not np.isfinite(ratio) or ratio < 0:
        raise ValueError(f"{label} must be finite and nonnegative")
    result = int(round(ratio))
    if not np.isclose(ratio, result, rtol=0, atol=1e-9):
        raise ValueError(f"{label} must be an integer number of dt steps")
    if result > np.iinfo(np.int32).max or (positive and result == 0):
        raise ValueError(f"{label} is outside the supported step range")
    return result


def _int_vector(values, label: str, dtype=np.int64) -> np.ndarray:
    source = np.asarray(values)
    if source.ndim != 1 or source.dtype.kind not in "iu":
        raise ValueError(f"{label} must be a one-dimensional integer array")
    limits = np.iinfo(dtype)
    if np.any(source < limits.min) or np.any(source > limits.max):
        raise ValueError(f"{label} is outside {np.dtype(dtype).name} range")
    return source.astype(dtype, copy=True)


def _environment() -> dict:
    packages = {}
    for name, expected in PACKAGE_PINS.items():
        try:
            actual = importlib.metadata.version(name)
        except importlib.metadata.PackageNotFoundError:
            actual = None
        packages[name] = {"expected": expected, "actual": actual,
                          "matches": actual == expected}
    return {
        "python": platform.python_version(), "author_python": "3.10.11",
        "python_patch_matches": platform.python_version() == "3.10.11",
        "python_executable": sys.executable, "platform": platform.platform(),
        "author_platform": "Windows", "packages": packages,
        "all_package_pins_match": all(p["matches"] for p in packages.values()),
        "binary_identical_author_environment_claimed": False,
    }


def _code_text(value):
    if isinstance(value, dict):
        return {str(key): _code_text(item) for key, item in value.items()}
    return str(value)


def _code_records(network) -> list:
    records = []
    for obj in network.sorted_objects:
        record = {"name": obj.name, "type": type(obj).__name__,
                  "when": obj.when, "order": int(obj.order),
                  "active": bool(obj.active)}
        for key in ("abstract_code", "user_code"):
            if hasattr(obj, key):
                record[key] = _code_text(getattr(obj, key))
        record["code_objects"] = [
            {"name": code.name, "type": type(code).__name__,
             "generated_code": _code_text(code.code)}
            for code in getattr(obj, "code_objects", ())
        ]
        records.append(record)
    return records


def run_brian(author_path: Path, io_dir: Path, nodes: np.ndarray,
              edges: pd.DataFrame, config: dict, precision: str,
              seed_node: int | None, external_arrivals=None) -> dict:
    """Execute an unchanged author model and return canonical end-of-tick arrays.

    ``external_arrivals`` is either None or a (T, N) array in mV.  Every nonzero
    entry gets its own SpikeGeneratorGroup source and zero-delay conditional
    synaptic write, ordered after the anatomical synapses.  It is not direct
    Python assignment to g.  ``seed_margin_mv`` may be zero for a strict-threshold
    fixture; positive margins must remain representably above the threshold.

    This function controls process-global Brian preferences and is not thread
    safe.  Use only the isolated oracle environment, serially or in subprocesses.
    A failed run keeps its exclusive inputs and an error record for diagnosis.
    """
    if precision not in ("float64", "float32"):
        raise ValueError("precision must be float64 or float32")
    dtype = np.dtype(precision)
    nodes = _int_vector(nodes, "nodes", np.int32)
    if not len(nodes) or np.any(np.diff(nodes.astype(np.int64)) <= 0):
        raise ValueError("nodes must be nonempty, sorted, and unique")
    required = ["edge_id", "source_neuron_idx", "target_neuron_idx", "raw_count"]
    if not set(required).issubset(edges.columns):
        raise ValueError(f"edges require columns {required}")
    edge_ids = _int_vector(edges.edge_id, "edge_id")
    if np.any(np.diff(edge_ids) <= 0):
        raise ValueError("edge_id order must be canonical, increasing, and unique")
    sources = _int_vector(edges.source_neuron_idx, "source_neuron_idx", np.int32)
    targets = _int_vector(edges.target_neuron_idx, "target_neuron_idx", np.int32)
    multipliers = _int_vector(edges.raw_count, "raw_count")
    if not np.isin(sources, nodes).all() or not np.isin(targets, nodes).all():
        raise ValueError("all edge endpoints must occur in nodes")
    lookup = {int(node): i for i, node in enumerate(nodes)}
    if seed_node is not None and int(seed_node) not in lookup:
        raise ValueError("seed_node must occur in nodes")

    dt_ms = float(config["dt_ms"])
    if not np.isfinite(dt_ms) or dt_ms <= 0:
        raise ValueError("dt_ms must be finite and positive")
    total_steps = _steps(config["duration_ms"], dt_ms, "duration_ms", positive=True)
    model = config["model"]
    refractory_steps = _steps(model["refractory_ms"], dt_ms, "refractory_ms")
    delay_steps = _steps(model["synaptic_delay_ms"], dt_ms, "synaptic_delay_ms")
    margin_mv = float(config["seed_margin_mv"])
    if not np.isfinite(margin_mv) or margin_mv < 0:
        raise ValueError("seed_margin_mv must be finite and nonnegative")
    if "dt_ms" in model and float(model["dt_ms"]) != dt_ms:
        raise ValueError("top-level and model dt_ms disagree")
    arrivals = None
    if external_arrivals is not None:
        arrivals = np.asarray(external_arrivals, dtype=np.float64)
        if arrivals.shape != (total_steps, len(nodes)) or not np.isfinite(arrivals).all():
            raise ValueError("external_arrivals must be finite with shape (T, N)")

    author_path = Path(author_path).resolve(strict=True)
    author_bytes = author_path.read_bytes()
    author_hash = hashlib.sha256(author_bytes).hexdigest()
    if author_hash != AUTHOR_SHA256:
        raise ValueError(f"Author source hash mismatch: {author_hash}")
    io_dir = Path(io_dir).resolve()
    io_dir.mkdir(parents=True, exist_ok=False)
    environment = _environment()
    _json(io_dir / "environment.json", environment)
    if environment["packages"]["brian2"]["actual"] != "2.5.1":
        raise RuntimeError("The author oracle requires exactly Brian2 2.5.1")

    import brian2 as b2

    saved_preferences = {
        "codegen.target": b2.prefs["codegen.target"],
        "core.default_float_dtype": b2.prefs["core.default_float_dtype"],
        "legacy.refractory_timing": b2.prefs["legacy.refractory_timing"],
    }
    saved_dt = b2.defaultclock.dt
    try:
        b2.start_scope()
        b2.prefs["codegen.target"] = "numpy"
        b2.prefs["core.default_float_dtype"] = dtype.type
        b2.prefs["legacy.refractory_timing"] = False
        b2.defaultclock.dt = dt_ms * b2.ms
        # Compile the checked bytes themselves: no source rewrite or __pycache__
        # is possible, and a subsequent file change cannot alter imported code.
        author = types.ModuleType("phase4ad_pinned_author_" + author_hash[:12])
        author.__file__ = str(author_path)
        exec(compile(author_bytes, str(author_path), "exec"), author.__dict__)
        parameters = dict(author.default_params)
        checks = {}
        for config_key, author_key, unit in (
            ("membrane_tau_ms", "t_mbr", b2.ms),
            ("synaptic_tau_ms", "tau", b2.ms),
            ("refractory_ms", "t_rfc", b2.ms),
            ("synaptic_delay_ms", "t_dly", b2.ms),
            ("resting_potential_mv", "v_0", b2.mV),
            ("reset_potential_mv", "v_rst", b2.mV),
            ("threshold_mv", "v_th", b2.mV),
            ("w_syn_mv_per_anatomical_synapse", "w_syn", b2.mV),
        ):
            expected = float(model[config_key])
            actual = float(parameters[author_key] / unit)
            if not np.isfinite(expected) or not np.isclose(actual, expected, rtol=0, atol=1e-12):
                raise ValueError(f"config {config_key}={expected} differs from author {actual}")
            checks[config_key] = {"config": expected, "author": actual}

        completeness = pd.DataFrame({"local_index": np.arange(len(nodes), dtype=np.int32)},
                                    index=pd.Index(nodes, name="global_neuron_idx"))
        connections = pd.DataFrame({
            "Presynaptic_Index": [lookup[int(x)] for x in sources],
            "Postsynaptic_Index": [lookup[int(x)] for x in targets],
            "Excitatory x Connectivity": multipliers,
            "edge_id": edge_ids,
            "source_neuron_idx": sources, "target_neuron_idx": targets,
        }).astype({"Presynaptic_Index": "int32", "Postsynaptic_Index": "int32"})
        comp_path = io_dir / "completeness.csv"
        con_path = io_dir / "connectivity.parquet"
        completeness.to_csv(comp_path)
        connections.to_parquet(con_path, index=False, engine="pyarrow")
        neu, syn, spikes = author.create_model(comp_path, con_path, parameters)
        if np.dtype(neu.variables["v"].dtype) != dtype or np.dtype(syn.variables["w"].dtype) != dtype:
            raise AssertionError("Brian did not allocate the requested floating-point dtype")
        np.testing.assert_array_equal(np.asarray(syn.i[:]), connections.Presynaptic_Index)
        np.testing.assert_array_equal(np.asarray(syn.j[:]), connections.Postsynaptic_Index)

        objects = [neu, syn, spikes]
        seed_state = {"requested_global_node": None if seed_node is None else int(seed_node),
                      "margin_mv": margin_mv, "injections": 0}
        if seed_node is not None:
            seed_local = lookup[int(seed_node)]

            def inject_seed():
                if seed_state["injections"] == 0:
                    if float(b2.defaultclock.t / b2.ms) != 0:
                        raise AssertionError("Seed injection must occur at tick zero")
                    neu.v[seed_local] = parameters["v_th"] + margin_mv * b2.mV
                    applied = float(neu.v[seed_local] / b2.mV)
                    if margin_mv > 0 and not applied > float(parameters["v_th"] / b2.mV):
                        raise ValueError("The positive seed margin rounded away in the selected dtype")
                    seed_state["applied_mv"] = applied
                    seed_state["injections"] = 1

            objects.append(b2.NetworkOperation(inject_seed, clock=neu.clock,
                                               when="before_thresholds", order=0,
                                               name="oracle_seed_injection"))

        external_metadata = {"events": 0, "ordering": "after default_synapses.pre"}
        if arrivals is not None:
            event_ticks, event_targets = np.nonzero(arrivals)
            event_values = arrivals[event_ticks, event_targets]
            external_metadata["events"] = int(len(event_ticks))
            if len(event_ticks):
                event_ids = np.arange(len(event_ticks), dtype=np.int32)
                external_source = b2.SpikeGeneratorGroup(
                    len(event_ticks), indices=event_ids,
                    times=event_ticks.astype(np.float64) * dt_ms * b2.ms,
                    sorted=True, clock=neu.clock, name="zz_external_sources")
                external_syn = b2.Synapses(
                    external_source, neu, model="w : volt", on_pre="g_post += w",
                    delay=0 * b2.ms, name="zz_external_arrivals")
                external_syn.connect(i=event_ids, j=event_targets)
                external_syn.w = event_values * b2.mV
                external_syn.pre.order = int(syn.pre.order) + 1
                if external_syn.pre.when != syn.pre.when:
                    raise AssertionError("External and anatomical pathways must share a schedule slot")
                objects.extend([external_source, external_syn])
                external_metadata["pathway_order"] = int(external_syn.pre.order)
                external_metadata["actual_weights_mv"] = np.asarray(external_syn.w[:] / b2.mV, dtype=dtype).tolist()
            pd.DataFrame({"tick": event_ticks, "neuron_idx": nodes[event_targets],
                          "requested_mv": event_values}).to_csv(io_dir / "external-arrivals.csv", index=False)

        trace = b2.StateMonitor(neu, ["v", "g", "lastspike", "not_refractory"],
                                record=True, when="end", name="oracle_end_state")
        objects.append(trace)
        network = b2.Network(*objects)
        _json(io_dir / "run-spec.json", {
            "author_sha256": author_hash, "author_commit": AUTHOR_COMMIT,
            "precision": precision, "config": config, "parameter_checks": checks,
            "dt_ms": dt_ms, "steps": total_steps, "refractory_steps": refractory_steps,
            "delay_steps": delay_steps, "nodes": nodes.tolist(), "edge_ids": edge_ids.tolist(),
            "seed": seed_state.copy(), "external_arrivals": external_metadata,
            "preferences": {"codegen.target": "numpy", "legacy.refractory_timing": False,
                            "core.default_float_dtype": precision},
            "weight_rule": "raw_count multiplied by unchanged author w_syn; no sign inference",
            "internal_voltage_unit": "volt", "returned_voltage_unit": "mV",
        })
        _json(io_dir / "schedule.json", {"schedule": list(network.schedule),
                                        "summary": str(network.scheduling_summary())})
        network.run(float(config["duration_ms"]) * b2.ms, namespace={})
        _json(io_dir / "generated-code.json", _code_records(network))

        ticks_float = np.asarray(spikes.t[:] / b2.ms, dtype=np.float64) / dt_ms
        ticks = np.rint(ticks_float).astype(np.int32)
        if not np.allclose(ticks_float, ticks, rtol=0, atol=1e-6):
            raise AssertionError("Spike timestamps do not lie on the configured integer clock")
        event_nodes = nodes[np.asarray(spikes.i[:], dtype=np.int32)]
        order = np.lexsort((event_nodes, ticks))
        ticks, event_nodes = ticks[order], event_nodes[order]
        counts = np.asarray(spikes.count[:], dtype=np.int32).copy()
        if int(counts.sum()) != len(ticks):
            raise AssertionError("SpikeMonitor counts disagree with event records")
        trace_ticks = np.asarray(trace.t[:] / b2.ms, dtype=np.float64) / dt_ms
        if trace_ticks.shape != (total_steps,) or not np.allclose(
                trace_ticks, np.arange(total_steps), rtol=0, atol=1e-6):
            raise AssertionError("End-state trace has the wrong tick range")
        voltage = np.asarray(trace.v / b2.mV, dtype=dtype).T.copy()
        conductance = np.asarray(trace.g / b2.mV, dtype=dtype).T.copy()
        last_spike_ms = np.asarray(trace.lastspike / b2.ms, dtype=np.float64).T
        seen = np.isfinite(last_spike_ms) & (last_spike_ms >= 0)
        # Brian's never-spiked sentinel can be far outside int32.  Convert only
        # real nonnegative timestamps, then compute an end-of-tick countdown.
        last_ticks = np.zeros(last_spike_ms.shape, dtype=np.int64)
        last_ticks[seen] = np.rint(last_spike_ms[seen] / dt_ms).astype(np.int64)
        elapsed = np.arange(total_steps, dtype=np.int64)[:, None] - last_ticks
        ref = np.where(seen, np.maximum(refractory_steps - elapsed, 0), 0).astype(np.int32)
        available = np.asarray(trace.not_refractory, dtype=bool).T
        if not np.array_equal(available, ref == 0):
            raise AssertionError("Recorded refractory flag disagrees with the integer countdown")
        history = np.stack([voltage, conductance, ref.astype(dtype)], axis=1)
        if not np.isfinite(history).all():
            raise AssertionError("Nonfinite author state")
        if seed_node is not None:
            seed_state["spiked_at_zero"] = bool(np.any((ticks == 0) & (event_nodes == seed_node)))
            if seed_state["injections"] != 1 or (margin_mv > 0 and not seed_state["spiked_at_zero"]):
                raise AssertionError("The common initial seed event was not produced exactly once")
        _json(io_dir / "seed.json", seed_state)
        weights = np.asarray(syn.w[:] / b2.mV, dtype=dtype).copy()
        _json(io_dir / "actual-weights.json", {
            "edge_id": edge_ids.tolist(), "weights_mv": weights.tolist(),
            "weights_volt": np.asarray(syn.w_[:], dtype=np.float64).tolist(),
        })
        if hashlib.sha256(author_path.read_bytes()).hexdigest() != author_hash:
            raise AssertionError("Author source changed during execution")
        return {"step": ticks, "neuron_idx": event_nodes, "counts": counts,
                "v": voltage[-1].copy(), "g": conductance[-1].copy(),
                "ref": ref[-1].copy(), "history": history, "weights": weights}
    except Exception as error:
        _json(io_dir / "error.json", {"type": type(error).__name__, "message": str(error)})
        raise
    finally:
        for name, value in saved_preferences.items():
            b2.prefs[name] = value
        b2.defaultclock.dt = saved_dt