"""Frozen eager Torch float32 APL component and binary oracle sidecar.

This is the Phase 4P operator sequence, not an author-neuron implementation.
No precision/determinism flags are changed here.  The full matrix dimensions
are retained even when the caller simulates a small neuronal fixture.
"""
import argparse
import hashlib
import json
import math
import os
from pathlib import Path
import struct
import sys
import traceback

import numpy as np
import torch


PROTOCOL = "phase4ag-apl-v1"
N_NEURONS, N_APL, N_KC = 166700, 2, 4064
BUNDLE_FILES = ("input_matrix.npy", "output_matrix.npy", "kc_globals.npy", "apl_globals.npy")
FRAME = struct.Struct("<qQ")
PAYLOAD_BYTES = (N_APL + N_APL + N_KC) * 4


def require(condition, message):
    if not condition:
        raise ValueError(message)


def file_sha256(path):
    digest = hashlib.sha256()
    with Path(path).open("rb") as stream:
        for block in iter(lambda: stream.read(8 << 20), b""):
            digest.update(block)
    return digest.hexdigest()


def array_sha256(value):
    return hashlib.sha256(np.ascontiguousarray(value).tobytes()).hexdigest()


def environment_metadata():
    device = torch.cuda.current_device()
    return {"torch": str(torch.__version__), "numpy": str(np.__version__),
            "cuda": torch.version.cuda, "cudnn": torch.backends.cudnn.version(),
            "device_name": torch.cuda.get_device_name(device), "device_index": device,
            "device_capability": list(torch.cuda.get_device_capability(device)),
            "deterministic_algorithms": torch.are_deterministic_algorithms_enabled(),
            "cudnn_deterministic": torch.backends.cudnn.deterministic,
            "cudnn_benchmark": torch.backends.cudnn.benchmark,
            "cuda_matmul_allow_tf32": torch.backends.cuda.matmul.allow_tf32,
            "cudnn_allow_tf32": torch.backends.cudnn.allow_tf32,
            "float32_matmul_precision": torch.get_float32_matmul_precision(),
            "CUBLAS_WORKSPACE_CONFIG": os.environ.get("CUBLAS_WORKSPACE_CONFIG"),
            "precision_configuration_changed_by_component": False}


class LegacyAPL:
    """Independent zero-initialized APL state with original full-width matrices."""

    def __init__(self, bundle_dir, dt_ms):
        require(float(dt_ms) == 0.1, "Only frozen dt_ms=0.1 is allowed")
        require(torch.cuda.is_available(), "CUDA is required for the legacy Torch component")
        self.bundle_dir = Path(bundle_dir)
        expected = {"input_matrix.npy": (np.float32, (N_APL, N_NEURONS)),
                    "output_matrix.npy": (np.float32, (N_APL, N_KC)),
                    "kc_globals.npy": (np.int32, (N_KC,)),
                    "apl_globals.npy": (np.int32, (N_APL,))}
        arrays = {}
        for name, (dtype, shape) in expected.items():
            value = np.load(self.bundle_dir / name, allow_pickle=False)
            require(value.dtype == dtype and value.shape == shape and value.flags.c_contiguous,
                    "APL bundle dtype/shape/layout mismatch: " + name)
            require(bool(np.isfinite(value).all()), "Nonfinite APL bundle: " + name)
            arrays[name] = value
        for name in ("kc_globals.npy", "apl_globals.npy"):
            ids = arrays[name]
            require(bool(np.all((ids >= 0) & (ids < N_NEURONS))) and
                    bool(np.all(np.diff(ids) > 0)), "APL/KC globals must be sorted unique: " + name)
        require(not np.intersect1d(arrays["kc_globals.npy"], arrays["apl_globals.npy"]).size,
                "APL and KC identities overlap")
        require(bool(np.all(arrays["input_matrix.npy"] >= 0)) and
                bool(np.all(arrays["output_matrix.npy"] >= 0)), "APL raw-contact matrices must be nonnegative")
        self.device = torch.device("cuda", torch.cuda.current_device())
        self.input_gpu = torch.from_numpy(arrays["input_matrix.npy"].copy()).to(self.device)
        self.output_gpu = torch.from_numpy(arrays["output_matrix.npy"].copy()).to(self.device)
        self.kc_globals = arrays["kc_globals.npy"].copy()
        self.apl_globals = arrays["apl_globals.npy"].copy()
        self.state = torch.zeros(N_APL, dtype=torch.float32, device=self.device)
        self.spikes = torch.zeros(N_NEURONS, dtype=torch.uint8, device=self.device)
        self.decay = np.float32(math.exp(-float(dt_ms) / 20.0))
        self.gain = np.float32(0.01)
        self.metadata = {"ok": True, "protocol": PROTOCOL, "n_neurons": N_NEURONS,
            "n_apl": N_APL, "n_kc": N_KC, "dt_ms": float(dt_ms),
            "files_sha256": {name: file_sha256(self.bundle_dir / name) for name in BUNDLE_FILES},
            "array_sha256": {name: array_sha256(value) for name, value in arrays.items()},
            "decay_f32_bits": f"{int(self.decay.view(np.uint32)):08x}",
            "gain_f32_bits": f"{int(self.gain.view(np.uint32)):08x}",
            "environment": environment_metadata(),
            "operator_sequence": "state.mul_(decay32); drive=mv(I32,spikes32); state.add_(drive); h=matmul(state,O32)*gain32",
            "h_units": "legacy mV", "initial_state": "float32 zeros"}

    @torch.no_grad()
    def advance(self, spikes_u8):
        require(spikes_u8.device == self.device and spikes_u8.dtype == torch.uint8 and
                tuple(spikes_u8.shape) == (N_NEURONS,) and spikes_u8.is_contiguous(),
                "APL expects full-width contiguous CUDA uint8 spikes")
        # Preserve separate eager operations, including Python evaluation order
        # in the frozen state.mul_(...).add_(torch.mv(...)) expression.
        self.state.mul_(float(self.decay))
        drive = torch.mv(self.input_gpu, spikes_u8.to(torch.float32))
        self.state.add_(drive)
        h = torch.matmul(self.state, self.output_gpu) * self.gain
        return drive, self.state, h

    def advance_global_fired(self, fired):
        require(fired.dtype == np.int32 and fired.ndim == 1, "Global spikes must be int32 vector")
        require(bool(np.all((fired >= 0) & (fired < N_NEURONS))) and
                bool(np.all(np.diff(fired) > 0)), "Global spikes must be sorted unique and in range")
        self.spikes.zero_()
        if fired.size:
            indices = torch.from_numpy(fired.copy()).to(device=self.device, dtype=torch.int64)
            self.spikes[indices] = 1
        return self.advance(self.spikes)


def read_exact(stream, count, boundary_eof=False):
    result = bytearray()
    while len(result) < count:
        block = stream.read(count - len(result))
        if not block:
            if boundary_eof and not result:
                return None
            raise EOFError("Truncated APL request frame")
        result.extend(block)
    return bytes(result)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--bundle-dir", required=True)
    parser.add_argument("--dt-ms", required=True, type=float)
    args = parser.parse_args()
    # stdout is exclusively the handshake line and binary response frames.
    output, source = sys.stdout.buffer, sys.stdin.buffer
    component = LegacyAPL(args.bundle_dir, args.dt_ms)
    output.write((json.dumps(component.metadata, allow_nan=False) + "\n").encode("utf-8"))
    output.flush()
    expected_tick = 0
    while True:
        header = read_exact(source, FRAME.size, boundary_eof=True)
        if header is None:
            return
        tick, count = FRAME.unpack(header)
        require(tick == expected_tick and count <= N_NEURONS, "Invalid APL request tick/count")
        fired = np.frombuffer(read_exact(source, count * 4), dtype="<i4").copy()
        drive, state, h = component.advance_global_fired(fired)
        arrays = [value.detach().cpu().numpy().astype("<f4", copy=False) for value in (drive, state, h)]
        require(all(bool(np.isfinite(value).all()) for value in arrays), "Nonfinite APL output")
        output.write(FRAME.pack(tick, PAYLOAD_BYTES))
        for value in arrays:
            output.write(value.tobytes())
        output.flush()
        expected_tick += 1


if __name__ == "__main__":
    try:
        main()
    except Exception:
        traceback.print_exc(file=sys.stderr)
        sys.exit(1)