"""Recompute a metadata-matched analysis of MLPerf Inference v6.1.

Python 3 standard library only. No network, model inference, APIs, or subprocesses.
The source repository's summary row IDs are preserved as strings.
"""
from collections import Counter, defaultdict
from pathlib import Path
import ast
import csv
import hashlib
import json
import math
import re

ROOT = Path(__file__).resolve().parent
REVISION = '4bb63cd28eb136be221bc801a9e38c66c8a7553f'
FOCUS_IDS = ['285', '286', '287', '288', '289', '290', '300', '301']
VARIABLE_FIELDS = {'Scenario', 'Result', 'Accuracy', 'Location'}
EXPECTED_FIELDS = ['Organization', 'Availability', 'Division', 'SystemType', 'SystemName', 'Platform', 'Model', 'MlperfModel', 'Scenario', 'Result', 'Accuracy', 'number_of_nodes', 'host_processor_model_name', 'host_processors_per_node', 'host_processor_core_count', 'accelerator_model_name', 'accelerators_per_node', 'total_accelerators', 'Location', 'framework', 'operating_system', 'notes', 'compliance', 'errors', 'version', 'inferred', 'has_power', 'Units', 'weight_data_types', 'design_power_watts', 'PrivateID']
MATCH_FIELDS = [k for k in EXPECTED_FIELDS if k not in VARIABLE_FIELDS]


def digest(data):
    return hashlib.sha256(data).hexdigest()


def git_blob_sha1(data):
    return hashlib.sha1(b'blob ' + str(len(data)).encode() + b'\0' + data).hexdigest()


def check_provenance(root=ROOT):
    manifest = json.loads((root / 'provenance.json').read_text())
    for item in manifest['sources']:
        data = (root / item['local_path']).read_bytes()
        assert digest(data) == item['sha256'], item['local_path']
        assert git_blob_sha1(data) == item['git_blob_sha1'], item['local_path']
    return len(manifest['sources'])


def load_rows(root=ROOT):
    rows = json.loads((root / 'data/summary.json').read_text())
    assert list(rows) == [str(i) for i in range(492)], 'Unexpected source row IDs'
    csv_rows = list(csv.DictReader((root / 'data/summary.csv').open(newline='')))
    assert len(csv_rows) == len(rows)
    for i, c in enumerate(csv_rows):
        r = rows[str(i)]
        assert list(r) == EXPECTED_FIELDS, 'Schema changed'
        for key, value in r.items():
            if isinstance(value, bool):
                assert c[key] == str(value)
            elif isinstance(value, (float, int)):
                assert math.isclose(float(c[key]), value, rel_tol=1e-9, abs_tol=1e-9), (i, key)
            elif value is None:
                assert c[key] in ('', 'None', 'N/A'), (i, key)
            else:
                assert c[key] == value, (i, key)
        assert r['version'] == 'v6.1'
        assert math.isfinite(r['Result']) and r['Result'] > 0
        assert r['errors'] == 0
    return rows


def eligible(row):
    return (row['Division'] == 'closed' and row['compliance'] == 'closed'
            and row['SystemType'] == 'datacenter' and row['Availability'] == 'available'
            and row['inferred'] == 0 and row['errors'] == 0
            and row['Units'] == 'Tokens/s' and row['Scenario'] in ('Offline', 'Server')
            and row['has_power'] is False and row['design_power_watts'] is None)


def match_pairs(rows):
    """Match all summary fields except scenario, value, accuracy result and log path.

    Duplicate scenarios fail instead of silently choosing the fastest value.
    Keeping Model and MlperfModel preserves the benchmark/accuracy-target label.
    This is a metadata match, not proof of identical scenario-specific tuning.
    """
    groups = defaultdict(dict)
    for row_id, row in rows.items():
        if not eligible(row):
            continue
        key = tuple(row[field] for field in MATCH_FIELDS)
        scenario = row['Scenario']
        assert scenario not in groups[key], 'Ambiguous duplicate scenario'
        groups[key][scenario] = (row_id, row)
    pairs, unmatched = [], []
    for group in groups.values():
        if set(group) != {'Offline', 'Server'}:
            unmatched.extend(row_id for row_id, _ in group.values())
            continue
        offline_id, offline = group['Offline']
        server_id, server = group['Server']
        assert all(offline[k] == server[k] for k in MATCH_FIELDS)
        ratio = server['Result'] / offline['Result']
        pairs.append({
            'offline_row_id': offline_id, 'server_row_id': server_id,
            'metadata': {k: offline[k] for k in MATCH_FIELDS},
            'offline_tokens_per_second': offline['Result'],
            'server_tokens_per_second': server['Result'],
            'server_over_offline': ratio,
            'server_over_offline_percent': ratio * 100,
            'difference_percent_of_offline': (ratio - 1) * 100,
            'offline_accuracy': offline['Accuracy'], 'server_accuracy': server['Accuracy'],
            'offline_location': offline['Location'], 'server_location': server['Location'],
        })
    return pairs, unmatched


def _constant(node):
    """Read the pinned constant dictionaries without executing upstream code."""
    if isinstance(node, ast.Constant):
        return node.value
    if isinstance(node, ast.Dict):
        return {_constant(k): _constant(v) for k, v in zip(node.keys, node.values)}
    if isinstance(node, (ast.Tuple, ast.List)):
        return [_constant(n) for n in node.elts]
    if isinstance(node, ast.BinOp) and isinstance(node.op, (ast.Mult, ast.Add, ast.Sub, ast.Div)):
        a, b = _constant(node.left), _constant(node.right)
        if isinstance(node.op, ast.Mult): return a * b
        if isinstance(node.op, ast.Add): return a + b
        if isinstance(node.op, ast.Sub): return a - b
        return a / b
    if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub):
        return -_constant(node.operand)
    raise ValueError('Unsupported constant expression: ' + type(node).__name__)


def accuracy_and_log_checks(rows, root=ROOT):
    tree = ast.parse((root / 'sources/submission-constants.py').read_text())
    configs = {}
    for node in tree.body:
        if isinstance(node, ast.Assign):
            name = getattr(node.targets[0], 'id', None)
            if name in ('MODEL_CONFIG', 'LLM_LATENCY_LIMITS'):
                configs[name] = _constant(node.value)
    targets = configs['MODEL_CONFIG']['v6.1']['accuracy-target']
    checks = []
    for row_id in FOCUS_IDS:
        row = rows[row_id]
        model = row['MlperfModel']
        values = {key: float(value) for key, value in re.findall(r'([A-Za-z0-9_]+):\s*([0-9.]+)', row['Accuracy'])}
        target = targets[model]
        thresholds = dict(zip(target[::2], target[1::2]))
        for metric, minimum in thresholds.items():
            assert values[metric] >= minimum, (row_id, metric, values[metric], minimum)
        log = (root / f'sources/logs/{row_id}-mlperf_log_summary.txt').read_text()
        assert 'Result is : VALID' in log
        assert f'Scenario : {row["Scenario"]}' in log
        pattern = r'^Tokens per second:\s*([0-9.]+)' if row['Scenario'] == 'Offline' else r'^Completed tokens per second:\s*([0-9.]+)'
        rate = float(re.search(pattern, log, re.M).group(1))
        # The official summary rounds some log rates; half a token/s covers
        # the displayed integer precision of these highlighted rows.
        assert abs(rate - row['Result']) <= 0.5
        check = {'row_id': row_id, 'model': model, 'scenario': row['Scenario'],
                 'thresholds': thresholds, 'reported_accuracy': values,
                 'log_tokens_per_second': rate, 'summary_tokens_per_second': row['Result'],
                 'log_result': 'VALID'}
        if row['Scenario'] == 'Server':
            check['limits_ns'] = configs['LLM_LATENCY_LIMITS'][model]['Server']
            for field, expected in check['limits_ns'].items():
                actual = int(re.search(r'^' + field + r'_latency \(ns\):\s*(\d+)', log, re.M).group(1))
                assert actual == expected, (row_id, field)
            assert 'Performance constraints satisfied : Yes' in log
        checks.append(check)
    return checks


def self_tests(rows):
    """Adversarial fixtures prevent silent mixing of scenarios and conditions."""
    a, b = rows['285'].copy(), rows['286'].copy()
    assert len(match_pairs({'a': a, 'b': b})[0]) == 1
    for field, value in [('Division', 'open'), ('Availability', 'preview'), ('Units', 'Samples/s'),
                         ('inferred', 1), ('has_power', True), ('Scenario', 'Interactive')]:
        altered = dict(b, **{field: value})
        assert not match_pairs({'a': a, 'b': altered})[0], field
    for field, value in [('MlperfModel', 'llama2-70b-99'), ('total_accelerators', 72),
                         ('framework', 'different version'), ('weight_data_types', 'fp16'),
                         ('notes', 'different power condition')]:
        assert not match_pairs({'a': a, 'b': dict(b, **{field: value})})[0], field
    try:
        match_pairs({'a': a, 'b': b, 'duplicate': dict(b)})
    except AssertionError as error:
        assert 'duplicate' in str(error)
    else:
        raise AssertionError('Duplicate row was accepted')
    return 13


def analyze(root=ROOT):
    verified_sources = check_provenance(root)
    rows = load_rows(root)
    pairs, unmatched = match_pairs(rows)
    assert len(pairs) == 152 and unmatched == ['264']
    raw_csv = list(csv.DictReader((root / 'data/summary.csv').open(newline='')))
    for pair in pairs:
        off, srv = raw_csv[int(pair['offline_row_id'])], raw_csv[int(pair['server_row_id'])]
        assert all(off[k] == srv[k] for k in MATCH_FIELDS), 'CSV metadata mismatch'
    assert sum(eligible(r) for r in rows.values()) == 305
    focus = [p for p in pairs if p['offline_row_id'] in FOCUS_IDS]
    assert len(focus) == 4
    assert len({p['metadata']['Platform'] for p in focus[:3]}) == 1
    percentages = [p['server_over_offline_percent'] for p in focus[:3]]
    assert [round(x, 2) for x in percentages] == [85.36, 94.37, 97.87]
    assert round(focus[3]['server_over_offline_percent'], 2) == 101.72
    above = sum(p['server_over_offline'] > 1 for p in pairs)
    assert above == 11
    assert all(r['has_power'] is False and r['design_power_watts'] is None for r in rows.values())
    return {
        'release': 'MLPerf Inference v6.1', 'results_revision': REVISION,
        'method': 'Within-submission metadata-matched Server/Offline ratios; no system ranking.',
        'source_rows': len(rows), 'closed_rows': sum(r['Division'] == 'closed' for r in rows.values()),
        'eligible_rows': 305, 'complete_pairs': len(pairs), 'unmatched_row_ids': unmatched,
        'pairs_with_server_greater_than_offline': above,
        'rows_with_measured_power': sum(r['has_power'] is True for r in rows.values()),
        'match_fields': MATCH_FIELDS, 'excluded_match_fields': sorted(VARIABLE_FIELDS),
        'focus_pairs': focus, 'b300_ratio_span_percentage_points': max(percentages) - min(percentages),
        'accuracy_and_log_checks': accuracy_and_log_checks(rows, root),
        'negative_and_positive_test_cases': self_tests(rows), 'verified_source_files': verified_sources,
        'all_pairs': pairs,
    }


if __name__ == '__main__':
    results = analyze()
    (ROOT / 'results.json').write_text(json.dumps(results, indent=2) + '\n')
    rows = load_rows()
    (ROOT / 'data/selected_rows.json').write_text(json.dumps({i: rows[i] for i in FOCUS_IDS}, indent=2) + '\n')
    print(json.dumps({k: results[k] for k in ['source_rows', 'eligible_rows', 'complete_pairs',
          'pairs_with_server_greater_than_offline', 'rows_with_measured_power',
          'b300_ratio_span_percentage_points', 'negative_and_positive_test_cases', 'verified_source_files']}, indent=2))