# Source: content/notes/ml/types-of-ml.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import inspect
import numpy as np
from sklearn.datasets import make_moons
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import PolynomialFeatures, StandardScaler
from sklearn.semi_supervised import SelfTrainingClassifier

rng = np.random.default_rng(102)
X, y = make_moons(n_samples=700, noise=0.22, random_state=102)
pool_x, test_x, pool_y, test_y = train_test_split(
    X, y, test_size=300, stratify=y, random_state=103)
labeled_ids = np.concatenate([np.flatnonzero(pool_y == label)[:12] for label in (0, 1)])
unlabeled_ids = np.setdiff1d(np.arange(len(pool_x)), labeled_ids)
lx, ly = pool_x[labeled_ids], pool_y[labeled_ids]
ux = pool_x[unlabeled_ids]
def learner():
    return make_pipeline(PolynomialFeatures(3, include_bias=False), StandardScaler(),
                         LogisticRegression(C=1.0, max_iter=1000, random_state=102))
baseline = learner().fit(lx, ly)
print("Labeled-only accuracy:", accuracy_score(test_y, baseline.predict(test_x)))
parameter = "estimator" if "estimator" in inspect.signature(SelfTrainingClassifier).parameters else "base_estimator"
for name, unlabeled in (("matched", ux), ("shifted", ux + np.array([3.0, -2.0]))):
    features = np.vstack([lx, unlabeled])
    targets = np.concatenate([ly, np.full(len(unlabeled), -1, dtype=int)])
    model = SelfTrainingClassifier(**{parameter: learner()}, threshold=0.85, max_iter=10)
    model.fit(features, targets)
    accepted = np.sum(model.labeled_iter_[len(lx):] > 0)
    probability = model.predict_proba(test_x)
    assert probability.shape == (len(test_x), 2)
    np.testing.assert_allclose(probability.sum(axis=1), 1.0, atol=1e-12)
    assert 0 < accepted <= len(unlabeled)
    assert np.all(model.transduction_[:len(lx)] == ly)
    print(name, "pseudo-labels accepted:", accepted,
          "test accuracy:", accuracy_score(test_y, model.predict(test_x)))