"""Pure CPU contact accounting and immutable diagnostic weight operations."""
import numpy as np
import pandas as pd


def contact_scope(apl_pn, pn_kc, rois):
    """Only exact, observed PN+postsynaptic-ROI pairs; no side/name inference."""
    required = {'source_neuron_idx', 'target_neuron_idx', 'primary_post'}
    if not required.issubset(apl_pn) or not required.issubset(pn_kc):
        raise ValueError('Actual schema lacks required identifiers/primary_post')
    permitted = apl_pn.loc[apl_pn.primary_post.isin(rois),
                           ['target_neuron_idx', 'primary_post']].drop_duplicates()
    permitted = permitted.rename(columns={'target_neuron_idx': 'source_neuron_idx'})
    keys = pd.MultiIndex.from_frame(permitted.astype({'primary_post': 'string'}))
    observed = pd.MultiIndex.from_frame(pn_kc[['source_neuron_idx', 'primary_post']]
                                        .astype({'primary_post': 'string'}))
    keep = observed.isin(keys)
    pairs = ['source_neuron_idx', 'target_neuron_idx']
    total = pn_kc.groupby(pairs, observed=True).size().rename('raw_partner_count')
    ca = pn_kc.loc[keep].groupby(pairs, observed=True).size().rename('blocked_contact_count')
    frame = total.to_frame().join(ca).fillna(0).reset_index()
    for key in pairs + ['raw_partner_count', 'blocked_contact_count']:
        frame[key] = frame[key].astype(np.int64)
    frame['retained_contact_count'] = frame.raw_partner_count - frame.blocked_contact_count
    assert (frame.retained_contact_count >= 0).all()
    return frame, permitted, keep


def edge_ids_for_pairs(frame, ptr, dst):
    out = []
    for s, t in frame[['source_neuron_idx', 'target_neuron_idx']].itertuples(index=False, name=None):
        ids = np.arange(ptr[s], ptr[s + 1], dtype=np.int64)
        found = ids[dst[ids] == t]
        if len(found) != 1:
            raise ValueError(f'Expected exactly one existing CSR edge for {s}->{t}')
        out.append(found[0])
    return np.asarray(out, np.int64)


def blocked_weights(base, raw, contacts):
    base = np.asarray(base)
    raw = np.asarray(raw)
    contacts = np.asarray(contacts)
    if base.dtype != np.float32 or raw.dtype != np.int64 or contacts.dtype != np.int64:
        raise TypeError('Expected float32 effective weight and int64 anatomical counts')
    if base.shape != raw.shape or raw.shape != contacts.shape:
        raise ValueError('Shape mismatch')
    if not np.isfinite(base).all() or np.any(raw <= 0) or np.any(contacts < 0) or np.any(contacts > raw):
        raise ValueError('Invalid contact partition or weights')
    out = base.copy()
    m = contacts > 0  # No-op copies original bits; no rounding at strength zero.
    out[m] = (base[m].astype(np.float64) * (raw[m] - contacts[m]) / raw[m]).astype(np.float32)
    return out