"""Check frozen predictions through the downloaded release; never select models."""
import argparse
import csv
import hashlib
import json
from pathlib import Path
import os
import sys

os.environ['USE_TF'] = '0'
os.environ['HF_HUB_OFFLINE'] = '1'
os.environ['TRANSFORMERS_OFFLINE'] = '1'


def compare(rows, gold, predictions, expected_correct):
    assert len(rows) == len(gold) == len(predictions)
    assert [int(r['artifact_row_index']) for r in rows] == list(range(len(rows)))
    assert [r['gold_label'] for r in rows] == gold, 'Dataset row order or labels differ'
    assert all(int(r['correct']) == int(r['gold_label'] == r['predicted_label']) for r in rows)
    mismatches = sum(p != r['predicted_label'] for p, r in zip(predictions, rows))
    correct = sum(p == y for p, y in zip(predictions, gold))
    result = {'n': len(rows), 'correct': correct, 'accuracy_percent': correct / len(rows) * 100,
              'prediction_mismatches': mismatches, 'expected_correct': expected_correct,
              'passed': mismatches == 0 and correct == expected_correct,
              'scope': 'Frozen retrospective reproduction; no tuning or model selection'}
    return result


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--package', required=True, choices=['banking_sam13', 'clinc_single', 'clinc_ensemble'])
    parser.add_argument('--checkpoint', type=Path, required=True)
    parser.add_argument('--data', type=Path, required=True)
    parser.add_argument('--output', type=Path, required=True)
    args = parser.parse_args()
    folder = args.checkpoint.resolve()
    sys.path.insert(0, str(folder))
    rows = list(csv.DictReader((folder / 'evaluation/predictions.csv').open(encoding='utf-8-sig')))
    if args.package == 'banking_sam13':
        import torch
        from transformers import AutoTokenizer, AutoModelForSequenceClassification, DataCollatorWithPadding
        from torch.utils.data import DataLoader
        torch._dynamo.config.disable = True
        torch.set_num_threads(4)
        data = [json.loads(line) for line in args.data.read_text().splitlines()]
        gold = [r['label_text'] for r in data]
        assert len(data) == 3080 and len(set(gold)) == 77
        model = AutoModelForSequenceClassification.from_pretrained(folder, local_files_only=True,
                                                                 attn_implementation='sdpa').cuda().eval()
        tok = AutoTokenizer.from_pretrained(folder, local_files_only=True)
        tokens = tok([r['text'] for r in data], truncation=True, max_length=128)
        items = [{k: v[i] for k, v in tokens.items()} for i in range(len(data))]
        predictions = []
        with torch.inference_mode():
            for batch in DataLoader(items, batch_size=32, collate_fn=DataCollatorWithPadding(tok)):
                logits = model(**{k: v.cuda() for k, v in batch.items()}).logits.float()
                assert torch.isfinite(logits).all()
                predictions.extend(model.config.id2label[i] for i in logits.argmax(-1).tolist())
        expected = 2899
    else:
        data = json.loads(args.data.read_text())['test']
        gold = [y for _, y in data]
        assert len(data) == 4500 and len(set(gold)) == 150
        if args.package == 'clinc_single':
            import torch
            from safetensors.torch import load_file
            from transformers import PreTrainedTokenizerFast
            from torch.utils.data import DataLoader
            from train_clinc import IntentModel, TextData, infer
            expected = 4399
            cfg = json.loads((folder / 'config.json').read_text())
            assert json.loads((folder / 'inference_config.json').read_text())['variant'] == 'classifier_only'
            tok = PreTrainedTokenizerFast.from_pretrained(folder / 'base/tokenizer', local_files_only=True)
            model = IntentModel(str(folder / 'base'))
            model.load_state_dict(load_file(str(folder / 'model.safetensors')), strict=True)
            model.cuda().eval()
            loader = DataLoader(TextData(data, cfg['labels'], tok, cfg['max_len']), batch_size=64)
            logits, _, _ = infer(model, loader)
            assert torch.isfinite(logits).all()
            answers = [{'intent': cfg['labels'][i]} for i in logits.argmax(-1).tolist()]
        else:
            from infer_hybrid import predict
            expected = 4411
            answers = predict(str(folder), [t for t, _ in data])
        predictions = [r['intent'] for r in answers]
    result = compare(rows, gold, predictions, expected)
    result['input_sha256'] = hashlib.sha256(args.data.read_bytes()).hexdigest()
    args.output.write_text(json.dumps(result, indent=2) + '\n')
    print(json.dumps(result), flush=True)
    if not result['passed']:
        raise SystemExit('Frozen predictions differ: publication stopped; do not tune on this test')


if __name__ == '__main__':
    main()