"""Synthetic demonstration only: run after installing the package."""
import tempfile
from pathlib import Path
import numpy as np
import torch
from cfref import CFREF

torch.set_num_threads(1)  # Record thread count when reproducing experiments.
rng = np.random.RandomState(7)
X_train = rng.normal(size=(100, 10))
y_train = np.tile([0]*30 + [1]*20, 2)
cohorts = np.repeat(['source_A', 'source_B'], 50)
features = [f'bin_{i}' for i in range(10)]
X_train[:, 0] += y_train
model = CFREF(episodes=20, hidden_dim=16, embedding_dim=8)
model.fit(X_train, y_train, cohorts, feature_names=features)
# These are independent synthetic subjects, NOT a clinical validation dataset.
references = rng.normal(size=(5, 10))
queries = rng.normal(size=(3, 10))
print('Scores:', model.decision_function(queries, references, feature_names=features))
print('Labels:', model.predict(queries, references, feature_names=features))
print('Locked threshold:', model.threshold_)
with tempfile.TemporaryDirectory() as d:
    path = Path(d) / 'example.cfref'
    model.save(path)
    restored = CFREF.load(path)
    np.testing.assert_array_equal(model.predict(queries,references,feature_names=features),
                                  restored.predict(queries,references,feature_names=features))
print('Save/load check passed.')