"""Read-only, float32 TemporalEligibility-v2 trace and fixed A/B identities.

This module consumes already observed full-graph spike events.  It neither
simulates neurons nor generates a weight delta.  A gate is a virtual readout
after the same-tick KC trace update, not evidence of an actual DAN event.
"""
from __future__ import annotations

import math

import numpy as np


N_GLOBAL = 166700
CATEGORY_CODES = {'inactive': 0, 'A_specific': 1, 'B_specific': 2, 'shared': 3}
MAJORITY = 3
REPLICATES = 5


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


def _vector(value, name, dtype):
    _require(isinstance(value, np.ndarray) and value.ndim == 1 and value.dtype == np.dtype(dtype),
             name + ' must be a one-dimensional ' + dtype + ' ndarray; no implicit dtype conversion')
    return value


def _frozen_scalar(value, expected, name):
    _require(not isinstance(value, (bool, np.bool_)) and isinstance(value, (int, float, np.number)) and
             np.ndim(value) == 0 and math.isfinite(float(value)) and float(value) == expected,
             name + ' must retain its frozen value ' + str(expected))


def replay_credit(step, neuron_idx, kc_globals, steps=1000, gate_tick=900,
                  dt_ms=.1, tau_ms=100, cap=1):
    """Replay the frozen trace without touching the caller's arrays.

    ``step`` and ``neuron_idx`` are equal-length int32 vectors for the complete
    observed graph (or its KC-only subset), ordered uniquely by (tick, global
    neuron). Non-KC events are validated then ignored. ``kc_globals`` is sorted
    unique int32. The trace starts at +0 and, on every tick including tick zero,
    performs separate float32 multiply, add current spikes, then clamp to 1.

    No 18-tick shift is applied: the recorded KC emission tick is the trace
    event tick. Gate counts include gate_tick; final counts cover [0, steps).
    dt/tau/cap are exposed for explicit provenance but may not be changed.
    Shorter steps/gates are allowed for synthetic contract tests.

    Returns owned trace_gate/trace_final float32[K] and
    spike_count_gate/spike_count_total int32[K], plus scalar audit metadata.
    Memory is O(events + KC + global-neuron count), never a dense T-by-K trace.
    """
    ticks = _vector(step, 'step', 'int32')
    nodes = _vector(neuron_idx, 'neuron_idx', 'int32')
    kc = _vector(kc_globals, 'kc_globals', 'int32')
    _require(type(steps) is int and 0 < steps <= np.iinfo(np.int32).max, 'steps must be a positive int32-range integer')
    _require(type(gate_tick) is int and 0 <= gate_tick < steps, 'gate_tick must lie in [0, steps)')
    _frozen_scalar(dt_ms, .1, 'dt_ms')
    _frozen_scalar(tau_ms, 100., 'tau_ms')
    _frozen_scalar(cap, 1., 'cap')
    _require(ticks.shape == nodes.shape, 'Event vectors have different lengths')
    _require(np.all((ticks >= 0) & (ticks < steps)), 'Event tick outside the observed window')
    _require(np.all((nodes >= 0) & (nodes < N_GLOBAL)), 'Event neuron outside the frozen global graph')
    _require(len(ticks) < 2 or np.all((ticks[1:] > ticks[:-1]) |
             ((ticks[1:] == ticks[:-1]) & (nodes[1:] > nodes[:-1]))),
             'Events must be sorted and unique by (tick, global neuron); duplicates are not extra spikes')
    _require(len(kc) > 0 and kc[0] >= 0 and kc[-1] < N_GLOBAL and np.all(kc[1:] > kc[:-1]),
             'kc_globals must be nonempty, sorted, unique, in-range global identities')

    lookup = np.full(N_GLOBAL, -1, np.int32)
    lookup[kc] = np.arange(len(kc), dtype=np.int32)
    local = lookup[nodes]
    selected = local >= 0
    kc_ticks, kc_local = ticks[selected], local[selected]
    counts_total = np.bincount(kc_local, minlength=len(kc)).astype(np.int32)
    counts_gate = np.bincount(kc_local[kc_ticks <= gate_tick], minlength=len(kc)).astype(np.int32)
    decay = np.float32(math.exp(-float(dt_ms) / float(tau_ms)))
    trace = np.zeros(len(kc), np.float32)
    trace_gate = None
    begin = 0
    for tick in range(steps):
        # Preserve the source's separate rounding operations; do not substitute
        # float64 exponentiation by elapsed time or shift spikes to arrivals.
        np.multiply(trace, decay, out=trace)
        end = int(np.searchsorted(kc_ticks, tick, side='right'))
        if end > begin:
            firing = kc_local[begin:end]
            trace[firing] = np.add(trace[firing], np.float32(1.0))
        np.minimum(trace, np.float32(1.0), out=trace)
        if tick == gate_tick:
            trace_gate = trace.copy()
        begin = end
    _require(begin == len(kc_ticks) and trace_gate is not None and np.isfinite(trace).all(), 'Incomplete or nonfinite trace replay')
    return {
        'trace_gate': trace_gate, 'trace_final': trace.copy(),
        'spike_count_gate': counts_gate, 'spike_count_total': counts_total,
        'steps': steps, 'gate_tick': gate_tick, 'dt_ms': float(dt_ms), 'tau_ms': float(tau_ms), 'cap': float(cap),
        'decay_f32_bits': f'{int(decay.view(np.uint32)):08x}',
        'total_input_events': int(len(ticks)), 'total_KC_events': int(len(kc_ticks)),
        'KC_events_through_gate': int(counts_gate.sum(dtype=np.int64)),
        'gate_includes_same_tick_spikes': True, 'event_shift_steps': 0,
        'gate_is_virtual_readout_not_DAN_activity': True,
        'weights_modified': False, 'weight_delta_generated': False,
    }


def _traces(value, name):
    _require(isinstance(value, np.ndarray) and value.ndim == 2 and value.dtype == np.dtype('float32') and
             value.shape[0] == REPLICATES and value.shape[1] > 0,
             name + ' must have exact shape (5, KC) and float32 dtype')
    _require(np.isfinite(value).all() and np.all((value >= 0) & (value <= 1)),
             name + ' must contain finite frozen-cap traces in [0, 1]')
    return value


def _set_summary(active_a, active_b, category):
    intersection = int(np.count_nonzero(active_a & active_b))
    union = int(np.count_nonzero(active_a | active_b))
    return {
        'total': int(len(category)), 'A_active': int(np.count_nonzero(active_a)),
        'B_active': int(np.count_nonzero(active_b)), 'intersection': intersection, 'union': union,
        'active_set_jaccard': float(intersection / union) if union else 1.0,
        'empty_union': union == 0,
        'category_counts': {name: int(np.count_nonzero(category == code)) for name, code in CATEGORY_CODES.items()},
    }


def pair_identity(traces_a, traces_b, edge_source_local):
    """Classify fixed paired five-seed A/B traces, without installing plasticity.

    A KC is A-active/B-active if its gate trace is strictly positive in at
    least 3 of 5 corresponding trials. Categories are int8: inactive=0,
    A_specific=1, B_specific=2, shared=3. Every candidate edge inherits its
    source KC's category in the supplied edge order. Multiple edges sharing
    a KC are retained as separate rows, not independent neural observations.

    ``edge_source_local`` is int32[E] indexing the common KC axis of both
    trace matrices. The caller must bind that axis and the five paired seed
    rows to the frozen identities; this pure function cannot infer labels.
    Jaccard uses COMPLETE active A/B sets, not the disjoint specific sets.
    Empty/empty Jaccard is 1 with empty_union=True, never evidence of useful
    credit. Compartment stratification and any success gate belong to caller.
    """
    a, b = _traces(traces_a, 'traces_a'), _traces(traces_b, 'traces_b')
    _require(a.shape == b.shape, 'A and B trace matrices use different KC axes')
    source = _vector(edge_source_local, 'edge_source_local', 'int32')
    _require(np.all((source >= 0) & (source < a.shape[1])), 'Candidate edge source is not a valid KC-local index')
    positive_a = np.count_nonzero(a > np.float32(0), axis=0).astype(np.int32)
    positive_b = np.count_nonzero(b > np.float32(0), axis=0).astype(np.int32)
    active_a, active_b = positive_a >= MAJORITY, positive_b >= MAJORITY
    category = active_a.astype(np.int8) + np.int8(2) * active_b.astype(np.int8)
    edge_a, edge_b, edge_category = active_a[source].copy(), active_b[source].copy(), category[source].copy()
    return {
        'kc_category': category, 'edge_category': edge_category,
        'kc_active_a': active_a, 'kc_active_b': active_b,
        'edge_active_a': edge_a, 'edge_active_b': edge_b,
        'kc_positive_trials_a': positive_a, 'kc_positive_trials_b': positive_b,
        'summary': {
            'replicates_per_condition': REPLICATES, 'majority_required': MAJORITY,
            'positivity_rule': 'float32 trace > 0; no fitted epsilon', 'category_codes': dict(CATEGORY_CODES),
            'KC': _set_summary(active_a, active_b, category),
            'edges': _set_summary(edge_a, edge_b, edge_category),
            'jaccard_scope': 'complete majority-active A and B sets, including shared members',
            'empty_union_jaccard_convention': 1.0,
            'edge_counts_are_source_degree_weighted': True,
            'actual_DAN_credit_claimed': False, 'weights_modified': False, 'weight_delta_generated': False,
        },
    }