#!/usr/bin/env python3
"""Read-only source audit; stop Phase 3B if its neuron-count gate fails."""
import collections
import hashlib
import json
import platform
from pathlib import Path

import pyarrow as pa
import pyarrow.feather as feather
import pyarrow.parquet as pq

ROOT = Path('/data/digital-drosophila/datasets/male-cns-v1.0')
OUT = ROOT / 'processed/neuron-only/audit'


def digest(path):
    h = hashlib.sha256()
    with path.open('rb') as f:
        for block in iter(lambda: f.read(8 * 1024 * 1024), b''):
            h.update(block)
    return h.hexdigest()


def main():
    OUT.mkdir(parents=True, exist_ok=True)
    if any(OUT.iterdir()):
        raise RuntimeError('Refusing to overwrite an existing audit; use a fresh audit directory.')
    source = ROOT / 'raw/annotations/body-annotations-male-cns-v1.0-minconf-0.5.feather'
    before = digest(source)
    table = feather.read_table(source)
    fields = [n for n in table.column_names if n.casefold() == 'superclass']
    if len(fields) != 1:
        raise RuntimeError(f'Ambiguous superclass field: {fields}')
    name = fields[0]
    if not pa.types.is_string(table.schema.field(name).type):
        raise RuntimeError(f'Uninspected superclass dtype: {table.schema.field(name).type}')
    values = table[name].to_pylist()
    counts = collections.Counter(values)
    valid = [v is not None and bool(v.strip()) for v in values]
    neurons = table.filter(pa.array(valid))
    types = {}
    for field in table.schema:
        if 'type' not in field.name.casefold():
            continue
        vals = neurons[field.name].to_pylist()
        good = [v for v in vals if v is not None and (not isinstance(v, str) or bool(v.strip()))]
        freq = collections.Counter(good)
        types[field.name] = dict(dtype=str(field.type), nulls=sum(v is None for v in vals),
                                blank_strings=sum(isinstance(v, str) and not v.strip() for v in vals),
                                populated_rows=len(good), unique_valid_values=len(freq))
        (OUT / f'{field.name}-value-counts.json').write_text(json.dumps(freq, indent=2, sort_keys=True)+'\n')
    # This is a classification audit, not a new neuron index or metadata product.
    pq.write_table(pa.table({'body_id': table['bodyId'].cast(pa.uint64()),
                            'superclass': table[name], 'passes_superclass_rule': valid}),
                   OUT / 'annotation-classification-audit.parquet', compression='zstd')
    report = dict(status='BLOCKED_NEURON_COUNT_MISMATCH', annotation_rows=table.num_rows,
                  superclass_field=name, superclass_dtype=str(table.schema.field(name).type),
                  actual_neurons=neurons.num_rows, expected_neurons=166691,
                  neuron_difference=neurons.num_rows-166691, fragments=len(values)-neurons.num_rows,
                  null_count=counts[None], empty_string_count=counts[''],
                  whitespace_only_count=sum(isinstance(v,str) and v!='' and not v.strip() for v in values),
                  nonnull_values_with_surrounding_whitespace=sum(isinstance(v,str) and v!=v.strip() for v in values),
                  superclass_value_counts={str(k): v for k,v in counts.items()},
                  type_fields=types, expected_type_count=11691,
                  actual_primary_type_count=types['type']['unique_valid_values'],
                  type_difference=types['type']['unique_valid_values']-11691,
                  expected_counts_source='User-specified formal paper reference counts',
                  reason='166700 != 166691: stopped before generating mappings, unified metadata or CSR.',
                  specific_nine_ids='Undetermined: a paper-matched reference ID set is required; count alone cannot identify the nine IDs.',
                  software=dict(python=platform.python_version(),pyarrow=pa.__version__))
    (OUT/'difference-audit.json').write_text(json.dumps(report,indent=2,sort_keys=True)+'\n')
    (OUT/'annotation-schema.txt').write_text(str(table.schema)+'\n')
    (OUT/'README.md').write_text(
        '# Phase 3B stopped at count gate\n\n'
        'Rule: non-null, non-blank superclass; original values are not rewritten.\n'
        'Actual neurons: 166700; expected: 166691; difference: +9. Fragments: 44877.\n'
        'Primary type: 11751 unique non-blank values; expected: 11691; difference: +60; null rows: 2194.\n'
        'No neuron index, unified metadata, connectivity or CSR was generated.\n'
        'All type-like columns are reported separately. No type fallback or subclass exclusions were applied.\n'
        'The nine differing body IDs cannot be identified without a paper-matched reference ID list.\n'
        'To reproduce, run the archived script with the dataset processed/.venv/bin/python; '
        'choose a fresh OUT directory because the script refuses existing audit contents.\n')
    (OUT/'audit-malecns-phase3b.py').write_bytes(Path(__file__).read_bytes())
    assert digest(source)==before
    inventory=[dict(path=p.name,bytes=p.stat().st_size,sha256=digest(p)) for p in sorted(OUT.iterdir()) if p.is_file()]
    manifest=dict(status=report['status'], source=str(source),source_sha256=before,
                  source_unchanged_after_audit=True,outputs=inventory,
                  downstream_generation_performed=False)
    (OUT/'manifest.json').write_text(json.dumps(manifest,indent=2,sort_keys=True)+'\n')
    (OUT/'checksums.sha256').write_text(''.join(f'{digest(p)}  {p.name}\n' for p in sorted(OUT.iterdir()) if p.is_file()))
    print(json.dumps(report,indent=2))


if __name__=='__main__':
    main()