"""Aggregate additional archived text-to-protein benchmark metrics."""
import csv
import hashlib
import json
from collections import defaultdict
from pathlib import Path
from statistics import mean

MODEL_FILES = {
    'Ouroboros': ('large_newds_v5_esmif.json', 'ours_hard1_esmif.json'),
    'Pinal-1.2B': ('pinal1p2b_esmif.json', 'pinal1p2b_hard1_esmif.json'),
    'Pinal-15B': ('pinal15b_esmif.json', 'pinal15b_hard1_esmif.json'),
    'ESM-3': ('esm3_esmif.json', 'esm3_hard1_esmif.json'),
    'ProteinDT': ('proteindt_esmif.json', 'proteindt_hard1_esmif.json'),
    'Chroma': ('chroma_newds_esmif.json', 'chroma_hard1_esmif.json'),
    'ProtDAT': ('protdat_esmif.json', 'protdat_hard1_esmif.json'),
}


DATASET = {'id': 'swissprot_time_cutoff', 'name': 'Swiss-Prot time-cutoff',
           'pairs': 500, 'keywordScoredPairs': 476,
           'sourceSplits': ['hard1', 'hard']}


def time_cutoff_metrics(metrics):
    """Publish only the 500-pair temporal cohort; retain measured means and n."""
    return [dict(row, split=DATASET['id']) for row in metrics
            if row['split'] in {'hard', 'hard1', DATASET['id']}]


def augment(metrics, paper, project, provenance):
    def source(path):
        raw = path.read_bytes()
        provenance.append({'file': str(path.relative_to(project)), 'sha256': hashlib.sha256(raw).hexdigest()}
                          if path.is_relative_to(project) else
                          {'file': str(path.relative_to(paper)), 'sha256': hashlib.sha256(raw).hexdigest()})
        return raw

    metrics = [r for r in metrics if r['metric'] not in {'PAE', 'ProTrek score', 'Keyword recovery', 'Inverse folding recovery', 'Inverse folding SW'}]
    for filename in ('protrek_score.csv', 'keyword_recovery.csv'):
        path = paper / 'figures/text2prot/source_data' / filename
        grouped = defaultdict(list)
        for row in csv.DictReader(source(path).decode('utf-8-sig').splitlines()):
            if row['split'] in {'hard', 'hard1'} and row['model'] not in {'Ground Truth', 'Random'}:
                grouped[(row['metric'], row['split'], row['model'].replace('Ours (large)', 'Ouroboros'))].append(float(row['value']))
        for (metric, split, model), scores in grouped.items():
            metrics.append({'task': 't2p', 'metric': metric, 'split': split, 'model': model,
                            'n': len(scores), 'mean': mean(scores)})

    folders = {
        'hard1': project / 'prot_downstream/text2prot/results/inverse_folding',
    }
    for split, folder in folders.items():
        for model, filenames in MODEL_FILES.items():
            rows = json.loads(source(folder / filenames[split == 'hard1']))
            assert len(rows) == (500 if split == 'hard1' else 100)
            for label, key in [('Inverse folding recovery', 'recovery'), ('Inverse folding SW', 'sw_similarity')]:
                scores = [float(r[key]) for r in rows]
                assert len(scores) == len(rows)
                metrics.append({'task': 't2p', 'metric': label, 'split': split, 'model': model,
                                'n': len(scores), 'mean': mean(scores)})
    return time_cutoff_metrics(metrics)