#!/usr/bin/env python3
"""Compare dense vs sparse obikmer query output for one specimen.

Both files are `obikmer query --count-missing` output: OBITools4-style
FASTA, one record per read — `>id {"kmer_count":N,"kmer_missing":M,
"kmer_strict_matches":{"label":count,...}}`. Packing format (dense vs
sparse presence matrix) must not change query results — only I/O access
pattern differs. Matched by read id rather than by stream position: the
query pipeline processes input in chunks across worker threads and does
not guarantee output order matches input order.

Output to stdout: one CSV row
  species, strain, n_reads, n_common, missing_in_dense, missing_in_sparse, mismatched, mismatch_pct
"""
import argparse
import gzip
import json
import sys


def load(path):
    records = {}
    opener = gzip.open if path.endswith('.gz') else open
    with opener(path, 'rt') as fh:
        for line in fh:
            if not line.startswith('>'):
                continue
            header = line[1:].rstrip('\n')
            read_id, _, json_part = header.partition(' ')
            records[read_id] = json.loads(json_part) if json_part else {}
    return records


def annotations_equal(a, b):
    return (
        a.get('kmer_count') == b.get('kmer_count')
        and a.get('kmer_missing') == b.get('kmer_missing')
        and a.get('kmer_strict_matches', {}) == b.get('kmer_strict_matches', {})
    )


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--species', required=True)
    ap.add_argument('--strain', required=True)
    ap.add_argument('dense_fasta')
    ap.add_argument('sparse_fasta')
    args = ap.parse_args()

    dense = load(args.dense_fasta)
    sparse = load(args.sparse_fasta)

    dense_ids, sparse_ids = set(dense), set(sparse)
    common = dense_ids & sparse_ids
    missing_in_dense = len(sparse_ids - dense_ids)
    missing_in_sparse = len(dense_ids - sparse_ids)

    mismatched = sum(1 for rid in common if not annotations_equal(dense[rid], sparse[rid]))

    n_reads = len(dense_ids | sparse_ids)
    mismatch_pct = 100.0 * (mismatched + missing_in_dense + missing_in_sparse) / n_reads if n_reads else 0.0

    print(f'{args.species},{args.strain},{n_reads},{len(common)},'
          f'{missing_in_dense},{missing_in_sparse},{mismatched},{mismatch_pct:.6f}')


if __name__ == '__main__':
    main()