"""Check complete matrices against reviewer_manifest.json using Python's standard library."""
import argparse
import csv
import gzip
import hashlib
import io
import json
import math
from pathlib import Path


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--directory', type=Path, default=Path('.'))
    args = parser.parse_args()
    root = args.directory
    manifest = json.loads((root / 'reviewer_manifest.json').read_text(encoding='utf-8'))
    sample_digest, gene_digest = None, None
    for matrix in manifest['matrices']:
        plain = root / matrix['name']
        path = plain if plain.exists() else root / (matrix['name'] + '.gz')
        digest = hashlib.sha256()
        with (gzip.open(path, 'rb') if path.suffix == '.gz' else path.open('rb')) as handle:
            size = 0
            for block in iter(lambda: handle.read(1024 * 1024), b''):
                digest.update(block)
                size += len(block)
        assert size == matrix['bytes'] and digest.hexdigest() == matrix['sha256'], matrix['kind']
        ids, seen = [], set()
        totals = [0.0] * manifest['samples']
        with (gzip.open(path, 'rt', encoding='utf-8', newline='') if path.suffix == '.gz' else path.open(encoding='utf-8', newline='')) as handle:
            reader = csv.reader(handle, delimiter='\t')
            header = next(reader)
            assert len(header) == manifest['samples'] + 1 and len(set(header[1:])) == manifest['samples']
            for row in reader:
                assert len(row) == len(header) and row[0] not in seen
                seen.add(row[0])
                ids.append(row[0])
                for index, cell in enumerate(row[1:]):
                    assert cell != ''
                    value = float(cell)
                    assert math.isfinite(value) and value >= 0
                    totals[index] += value
        assert len(ids) == manifest['genes']
        this_samples = hashlib.sha256('\n'.join(header[1:]).encode()).hexdigest()
        this_genes = hashlib.sha256('\n'.join(ids).encode()).hexdigest()
        if sample_digest is None:
            sample_digest, gene_digest = this_samples, this_genes
        assert sample_digest == this_samples and gene_digest == this_genes
        if matrix['kind'] == 'tpm':
            assert max(abs(total - 1000000) for total in totals) < 0.01
        print(matrix['kind'], 'PASS', len(ids), 'x', len(header) - 1, digest.hexdigest(), flush=True)
    print('All complete matrices passed checksum, shape, ID order and numeric checks.')


if __name__ == '__main__':
    main()