import hashlib
import json
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
AA = dict(zip('ALA ARG ASN ASP CYS GLN GLU GLY HIS ILE LEU LYS MET PHE PRO SER THR TRP TYR VAL UNK'.split(), 'ARNDCQEGHILKMFPSTWYVX'))


def test_public_case_sequences_match_the_actual_structure_files():
    cases = json.loads((ROOT / 'public/data/petase.json').read_text())
    assert cases[0]['id'] == 'de2'
    assert cases[0]['length'] == 279
    assert cases[0]['sequenceSha256'] == '421800d7574578af255aa700b2f346c284af2f05c1f7e58ef210d17093ee0ccd'
    assert len(cases) == 10
    assert {r['id'] for r in cases} == {f'de{i}' for i in range(1, 10)} | {'IsPETase'}
    for case in cases:
        pdb = ROOT / 'public' / case['pdbUrl'].lstrip('/')
        ca = [s for s in pdb.read_text().splitlines() if s.startswith('ATOM  ') and s[12:16].strip() == 'CA']
        assert ''.join(AA.get(s[17:20], 'X') for s in ca) == case['sequence']
        assert len(ca) == case['length'] == len(case['residues'])
        assert hashlib.sha256(pdb.read_bytes()).hexdigest() == case['pdbSha256']
        assert abs(sum(float(s[60:66]) for s in ca)/len(ca) - case['plddt']) < .001
        fasta = ROOT / 'public' / case['fastaUrl'].lstrip('/')
        assert ''.join(fasta.read_text().splitlines()[1:]) == case['sequence']


def test_benchmark_has_both_directions_and_correctly_scoped_samples():
    data = json.loads((ROOT / 'public/data/evidence.json').read_text())
    assert {m['task'] for m in data['metrics']} == {'t2p', 'p2t'}
    assert data['dataset']['name'] == 'Swiss-Prot time-cutoff'
    assert len(data['metrics']) == 63
    assert {m['split'] for m in data['metrics']} == {'swissprot_time_cutoff'}
    for row in data['metrics']:
        expected = 476 if row['metric'] == 'Keyword recovery' else 500
        assert row['n'] == expected
    assert not any(row['metric'] == 'PAE' for row in data['metrics'])
    ours = {(row['metric'], row['split']): row for row in data['metrics'] if row['model'] == 'Ouroboros'}
    assert abs(ours['ProTrek score', 'swissprot_time_cutoff']['mean'] - 11.4167) < 0.0001
    assert abs(ours['Keyword recovery', 'swissprot_time_cutoff']['mean'] - 0.5917) < 0.0001
    assert abs(ours['Inverse folding recovery', 'swissprot_time_cutoff']['mean'] - 0.3039) < 0.0001
    assert abs(ours['Inverse folding SW', 'swissprot_time_cutoff']['mean'] - 128.686) < 0.0001
    assert len(data['scaling']) == 4
    assert all(r['n'] == 500 for r in data['scaling'])


def test_wetlab_values_and_screened_outcome_preserve_supplied_report():
    import csv
    from decimal import Decimal
    data = json.loads((ROOT / 'public/data/wetlab.json').read_text())
    rows = {r['id']: r for r in data['rows']}
    assert set(rows) == {f'de{i}' for i in range(1, 10)} | {'Y1'}
    assert data['submittedDesigns'] == 12 and data['reportedDesigns'] == 9
    assert data['activeDesigns'] == 9 and data['inactiveDesigns'] == 3
    assert data['assayDate'] is None
    assert rows['de2']['total'] == 1514.1531
    assert rows['Y1']['name'] == 'IsPETase'
    assert rows['Y1']['total'] == 1074.8149
    assert 1.408 < rows['de2']['total'] / rows['Y1']['total'] < 1.409
    assert all(r['total'] < rows['Y1']['total'] for r in rows.values() if r['id'] not in ['de2', 'Y1'])
    expected = {
        'Y1': ('1034.0479', '40.7670', '1074.8149'),
        'de1': ('256.5466', '185.2372', '441.7838'),
        'de2': ('752.5710', '761.5821', '1514.1531'),
        'de3': ('16.1355', '4.8610', '20.9965'),
        'de4': ('19.6139', '1.5279', '21.1418'),
        'de5': ('25.0867', '18.7743', '43.8610'),
        'de6': ('121.7478', '1.8291', '123.5769'),
        'de7': ('114.5323', '11.4136', '125.9460'),
        'de8': ('37.7339', '7.6817', '45.4156'),
        'de9': ('72.1781', '2.0316', '74.2097'),
    }
    csv_rows = {r['id']: r for r in csv.DictReader((ROOT/'public/data/petase-assay.csv').open())}
    assert set(csv_rows) == set(expected)
    for row in rows.values():
        for key, supplied in zip(['tpa','mhet','total'], expected[row['id']]):
            assert Decimal(str(row[key])) == Decimal(supplied)
        raw=csv_rows[row['id']]
        assert tuple(raw[k] for k in ['TPA_uM','MHET_uM','total_uM']) == expected[row['id']]
        assert abs(Decimal(raw['TPA_uM']) + Decimal(raw['MHET_uM']) - Decimal(raw['total_uM'])) <= Decimal('0.0001')
        fasta = ROOT / 'public' / row['fastaUrl'].lstrip('/')
        sequence = ''.join(fasta.read_text().splitlines()[1:])
        assert hashlib.sha256(sequence.encode()).hexdigest() == row['sequenceSha256']
    assert data['generation']['predictionIndex0'] == 2890
    assert data['structurePrediction']['experimentalStructure'] is False
    assert data['structurePrediction']['designCount'] == 9


def test_annotation_statistics_preserve_csv_and_multilabel_counts():
    import csv
    original=list(csv.DictReader((ROOT/'public/data/annotation_key_counts.csv').open(encoding='utf-8-sig')))
    data=json.loads((ROOT/'public/data/annotation-stats.json').read_text())
    assert data['keys']==len(original)==29
    assert data['recordKeyPairs']==sum(int(r['record_count']) for r in original)==2078270
    for raw,row in zip(original,data['rows']):
        assert row['key']==raw['key']
        assert row['count']==int(raw['record_count'])
        assert row['coverage']==float(raw['coverage_percent'])
    assert sum(row['coverage'] for row in data['rows'])>100
    assert data['rows'][-1]['key']=='interaction' and data['rows'][-1]['count']==2
    assert abs(data['beyondTopFiveShare']-21.3992407146)<1e-8