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

import numpy as np
import torch
from torch import nn
from torch.nn import functional as F
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

torch.manual_seed(41)
torch.set_num_threads(1)
np.random.seed(41)
content = torch.randn(384, 2)
labels = (content[:, 0] + content[:, 1] > 0).long()
encoder = nn.Sequential(nn.Linear(8, 32), nn.ReLU(), nn.Linear(32, 12))
projector = nn.Sequential(nn.Linear(12, 16), nn.ReLU(), nn.Linear(16, 8))
optimizer = torch.optim.Adam(list(encoder.parameters()) + list(projector.parameters()), lr=0.006)

def view(latent):
    return torch.cat((latent + 0.04 * torch.randn_like(latent),
                      torch.randn(len(latent), 6)), dim=1)

def contrastive(a, b, temperature=0.2):
    n = len(a)
    if n < 2:
        raise ValueError("At least two paired examples are needed")
    z = F.normalize(torch.cat((a, b)), dim=-1)
    logits = z @ z.T / temperature
    logits = logits.masked_fill(torch.eye(2*n, dtype=torch.bool), -torch.inf)
    positive = (torch.arange(2*n) + n) % (2*n)
    return F.cross_entropy(logits, positive)

fixed_a, fixed_b = view(content[:128]), view(content[:128])
with torch.no_grad():
    initial = contrastive(projector(encoder(fixed_a)), projector(encoder(fixed_b))).item()
for _ in range(180):
    ids = torch.randperm(256)[:96]
    a, b = view(content[ids]), view(content[ids])
    optimizer.zero_grad(set_to_none=True)
    loss = contrastive(projector(encoder(a)), projector(encoder(b)))
    loss.backward()
    optimizer.step()
encoder.eval()
projector.eval()
with torch.no_grad():
    final = contrastive(projector(encoder(fixed_a)), projector(encoder(fixed_b))).item()
    features = encoder(torch.cat((content, torch.zeros(384, 6)), dim=1)).numpy()
probe = make_pipeline(StandardScaler(), LogisticRegression(C=1.0, max_iter=1000, random_state=41))
probe.fit(features[:128], labels[:128].numpy())
accuracy = probe.score(features[256:], labels[256:].numpy())
eigenvalues = np.linalg.eigvalsh(np.cov(features[:256], rowvar=False))
assert final < initial
assert np.isfinite(features).all() and eigenvalues[-1] > 1e-4
assert accuracy > 0.75
print({"initial_contrastive_loss": initial, "final_contrastive_loss": final,
       "probe_accuracy": accuracy, "largest_covariance_eigenvalue": eigenvalues[-1]})