from __future__ import annotations

import json

from model import ContrastiveEncoder, parameter_count
from safetensors.torch import load_file
from train import (
    ARTIFACT_DIR,
    embed,
    invariance_score,
    load_split,
    probe_suite,
    seed_everything,
)


def main() -> None:
    seed_everything(2035)
    train_pixels, train_labels = load_split("train")
    test_pixels, test_labels = load_split("test")
    learned = ContrastiveEncoder()
    learned.load_state_dict(load_file(ARTIFACT_DIR / "model.safetensors"))
    random_encoder = ContrastiveEncoder()
    learned_train = embed(learned, train_pixels)
    learned_test = embed(learned, test_pixels)
    random_train = embed(random_encoder, train_pixels)
    random_test = embed(random_encoder, test_pixels)
    raw_train = train_pixels.reshape(len(train_pixels), -1).numpy()
    raw_test = test_pixels.reshape(len(test_pixels), -1).numpy()
    results = {
        "model": "Contrastive Pocket",
        "parameters": parameter_count(learned),
        "unlabeled_pretraining_examples": len(train_pixels),
        "epochs": 220,
        "embedding_dimensions": learned_train.shape[1],
        "linear_probe_accuracy_by_examples_per_class": {
            "contrastive_encoder": probe_suite(
                learned_train,
                train_labels,
                learned_test,
                test_labels,
            ),
            "random_encoder": probe_suite(
                random_train,
                train_labels,
                random_test,
                test_labels,
            ),
            "raw_pixels": probe_suite(
                raw_train,
                train_labels,
                raw_test,
                test_labels,
            ),
        },
        "augmentation_invariance": {
            "contrastive_encoder": invariance_score(learned, train_pixels),
            "random_encoder": invariance_score(random_encoder, train_pixels),
        },
        "embedding_standard_deviation": float(learned_train.std(axis=0).mean()),
    }
    (ARTIFACT_DIR / "evaluation.json").write_text(
        json.dumps(results, indent=2),
        encoding="utf-8",
    )
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()