"""Train and inspect a 2 -> 3 -> 2 linear autoencoder. Requires only NumPy.

Run: python3 experiment.py
Try another initialization: python3 experiment.py --seed 11
Export the browser's snapshot: python3 experiment.py --export ../data/model.json

The export path is relative to the working directory. Training uses pixels only;
brightness/contrast labels are used afterward for representation analysis.
"""

import argparse
import json
from pathlib import Path

import numpy as np


def sample(rng, count):
    brightness = rng.uniform(0.25, 0.75, count)
    contrast = rng.uniform(-0.25, 0.25, count)
    pixels = np.column_stack((brightness + contrast, brightness - contrast))
    properties = np.column_stack((brightness, contrast))
    return pixels, properties


def gram_schmidt(vectors):
    """Modified Gram-Schmidt on columns; reject dependent directions."""
    q = np.zeros_like(vectors)
    r = np.zeros((vectors.shape[1], vectors.shape[1]))
    for j in range(vectors.shape[1]):
        residual = vectors[:, j].copy()
        for i in range(j):
            r[i, j] = q[:, i] @ residual
            residual -= r[i, j] * q[:, i]
        r[j, j] = np.linalg.norm(residual)
        if r[j, j] < 1e-12:
            raise ValueError("Feature directions are linearly dependent")
        q[:, j] = residual / r[j, j]
    return q, r


def run(seed=7):
    rng = np.random.default_rng(seed)
    x, properties = sample(rng, 2048)
    w = rng.normal(0, 0.4, (3, 2))
    d = rng.normal(0, 0.4, (2, 3))

    # Rows are examples. The loss is mean squared Euclidean distance.
    for _ in range(10_000):
        h = x @ w.T
        error = h @ d.T - x
        grad_d = 2 * error.T @ h / len(x)
        grad_w = 2 * (error @ d).T @ x / len(x)
        # Both gradients use the same, pre-update weights.
        w -= 0.08 * grad_w
        d -= 0.08 * grad_d

    # Freeze the network. Controlled changes identify the writing directions.
    v_b = w @ np.array([1.0, 1.0])
    v_c = w @ np.array([1.0, -1.0])
    v = np.column_stack((v_b, v_c))
    q, r = gram_schmidt(v)

    h = x @ w.T
    # Separate fit of a linear regression probe, without changing w or d.
    # h has rank 2: lstsq handles the non-unique 3-coordinate solutions.
    probe = np.linalg.lstsq(h, properties, rcond=None)[0]
    h_mean = h.mean(axis=0)
    _, singular_values, vt = np.linalg.svd(h - h_mean, full_matrices=False)
    pca = vt[:2]

    x_test, properties_test = sample(rng, 512)
    h_test = x_test @ w.T
    reconstructed = h_test @ d.T
    qr_properties = np.linalg.solve(r, (h_test @ q).T).T
    probe_properties = h_test @ probe

    # Check geometry and behavior on examples not used for training/fitting.
    np.testing.assert_allclose(reconstructed, x_test, atol=1e-9)
    np.testing.assert_allclose(qr_properties, properties_test, atol=1e-9)
    np.testing.assert_allclose(probe_properties, properties_test, atol=1e-9)
    np.testing.assert_allclose(q.T @ q, np.eye(2), atol=1e-12)
    np.testing.assert_allclose(q @ r, v, atol=1e-12)
    np.testing.assert_allclose(q @ q.T, pca.T @ pca, atol=1e-12)
    assert np.linalg.matrix_rank(h - h_mean) == 2

    alpha = 0.1
    brightness_shift = (h_test + alpha * v_b) @ d.T - reconstructed
    contrast_shift = (h_test + alpha * v_c) @ d.T - reconstructed
    np.testing.assert_allclose(brightness_shift, alpha * np.ones_like(x_test), atol=1e-9)
    np.testing.assert_allclose(
        contrast_shift, np.tile([alpha, -alpha], (len(x_test), 1)), atol=1e-9
    )

    projection_b = h_test @ v_b / (v_b @ v_b)
    overlap = (v_b @ v_c) / (v_b @ v_b)
    np.testing.assert_allclose(
        projection_b, properties_test[:, 0] + overlap * properties_test[:, 1], atol=1e-12
    )

    metrics = {
        "reconstruction_mse": float(np.mean((reconstructed - x_test) ** 2)),
        "probe_mse": float(np.mean((probe_properties - properties_test) ** 2)),
        "qr_mse": float(np.mean((qr_properties - properties_test) ** 2)),
        "projection_b_mse": float(np.mean((projection_b - properties_test[:, 0]) ** 2)),
        "cosine": float((v_b @ v_c) / (np.linalg.norm(v_b) * np.linalg.norm(v_c))),
        "rank": int(np.linalg.matrix_rank(h - h_mean)),
    }
    return {
        "seed": seed,
        "steps": 10_000,
        "learning_rate": 0.08,
        "train_count": len(x),
        "test_count": len(x_test),
        "W": w.tolist(),
        "D": d.tolist(),
        "probe": probe.tolist(),
        "mean": h_mean.tolist(),
        "pca": pca.tolist(),
        "singular_values": singular_values.tolist(),
        "explained_variance": (singular_values**2 / (singular_values**2).sum()).tolist(),
        "metrics": metrics,
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--seed", type=int, default=7)
    parser.add_argument("--export", type=Path, help="Optional path for the browser snapshot")
    args = parser.parse_args()
    result = run(args.seed)
    print(json.dumps(result, indent=2, allow_nan=False))
    if args.export:
        args.export.parent.mkdir(parents=True, exist_ok=True)
        args.export.write_text(json.dumps(result, indent=2, allow_nan=False) + "\n")
        print(f"Saved {args.export}")


if __name__ == "__main__":
    main()