"""Learn 3 -> 3 -> k -> 1 maps and audit their representations.

Run with Python 3.10+ and NumPy. --export writes the article's frozen data;
--figures additionally requires Matplotlib. All paths are relative to this file.
"""

import argparse
import json
from pathlib import Path
import platform

import numpy as np

ROOT = Path(__file__).resolve().parents[1]
BETA = np.array([2.0, 1.0, -0.2])
SEEDS = [7, 11, 23, 42, 101]


def sample(rng, n):
    # Independent continuous feature scores: population mean 0, variance 1.
    return rng.uniform(-np.sqrt(3), np.sqrt(3), (n, 3))


def loss_and_grad(weights, x, y):
    w1, w2, w3 = weights
    h1 = x @ w1.T
    h2 = h1 @ w2.T
    error = h2 @ w3.T - y
    delta = 2 * error / len(x)
    gradients = [(delta @ w3 @ w2).T @ x,
                 (delta @ w3).T @ h1, delta.T @ h2]
    return float(np.mean(error**2)), gradients


def gradient_check():
    rng = np.random.default_rng(99)
    x = sample(rng, 9)
    y = (x @ BETA)[:, None]
    for width in (1, 2, 3):
        weights = [rng.normal(0, 0.4, shape)
                   for shape in [(3, 3), (width, 3), (1, width)]]
        _, analytic = loss_and_grad(weights, x, y)
        for w, g in zip(weights, analytic):
            numeric = np.empty_like(w)
            for index in np.ndindex(w.shape):
                original = w[index]
                w[index] = original + 1e-6
                plus = loss_and_grad(weights, x, y)[0]
                w[index] = original - 1e-6
                minus = loss_and_grad(weights, x, y)[0]
                w[index] = original
                numeric[index] = (plus - minus) / 2e-6
            np.testing.assert_allclose(g, numeric, atol=2e-9, rtol=2e-6)


def r2(actual, predicted):
    return 1 - ((actual - predicted)**2).sum(axis=0) / (
        (actual - actual.mean(axis=0))**2).sum(axis=0)


def analyze(mapping, x_train, x_test):
    h_train, h_test = x_train @ mapping.T, x_test @ mapping.T
    mean = h_train.mean(axis=0)
    _, singular_values, vt = np.linalg.svd(h_train - mean, full_matrices=False)
    # Fix otherwise arbitrary SVD signs for stable plots and exported numbers.
    for axis in vt:
        if axis[np.argmax(np.abs(axis))] < 0:
            axis *= -1
    design = np.column_stack((h_train, np.ones(len(h_train))))
    probe = np.linalg.lstsq(design, x_train, rcond=None)[0]
    recovered = np.column_stack((h_test, np.ones(len(h_test)))) @ probe
    scores = (h_test - mean) @ vt.T
    correlations = np.corrcoef(scores.T, x_test.T)[:len(vt), len(vt):]
    projector = np.linalg.pinv(mapping) @ mapping
    return {
        "rank": int(np.linalg.matrix_rank(mapping)),
        "mean": mean.tolist(), "pca": vt.tolist(),
        "singular_values": singular_values.tolist(),
        "variance_share": (singular_values**2 / np.sum(singular_values**2)).tolist(),
        "pc_feature_correlations": correlations.tolist(),
        "probe": probe.tolist(), "probe_r2": r2(x_test, recovered).tolist(),
        # Isotropic population benchmark: optimal linear R² = diag(projector).
        "population_linear_probe_r2": np.diag(projector).tolist(),
    }


def train(seed, width, x_train, x_test):
    rng = np.random.default_rng(seed)
    weights = [rng.normal(0, 0.4, shape)
               for shape in [(3, 3), (width, 3), (1, width)]]
    y = (x_train @ BETA)[:, None]
    initial_w1 = weights[0].copy()
    history = []
    for step in range(10_001):
        loss, gradients = loss_and_grad(weights, x_train, y)
        if step % 25 == 0 or loss < 1e-20:
            history.append([step, loss])
        if loss < 1e-20:
            break
        if step == 10_000 or not np.isfinite(loss):
            raise RuntimeError(f"Training failed: seed={seed}, width={width}, loss={loss}")
        for w, grad in zip(weights, gradients):
            w -= 0.03 * grad
    w1, w2, w3 = weights
    b = w2 @ w1
    prediction = (x_test @ b.T @ w3.T).ravel()
    target = x_test @ BETA
    np.testing.assert_allclose((w3 @ b).ravel(), BETA, atol=1e-9)
    np.testing.assert_allclose(prediction, target, atol=1e-9)
    first, second = analyze(w1, x_train, x_test), analyze(b, x_train, x_test)
    np.testing.assert_allclose(first["probe_r2"], np.ones(3), atol=1e-10)
    np.testing.assert_allclose(sum(second["population_linear_probe_r2"]), width, atol=1e-10)
    result = {
        "seed": seed, "width": width, "steps": step, "train_mse": loss,
        "test_mse": float(np.mean((prediction - target)**2)),
        "price_r2": float(r2(target, prediction)),
        "initial_W1": initial_w1.tolist(),
        "W1": w1.tolist(), "W2": w2.tolist(), "W3": w3.tolist(),
        "B": b.tolist(), "effective_weights": (w3 @ b).ravel().tolist(),
        "h1": first, "h2": second, "history": history,
    }
    if width == 2:
        # Full SVD retains the third right singular vector: the lost direction.
        _, _, vt = np.linalg.svd(b, full_matrices=True)
        null = vt[-1]
        if null[np.argmax(np.abs(null))] < 0:
            null *= -1
        np.testing.assert_allclose(b @ null, 0, atol=1e-12)
        np.testing.assert_allclose((x_test + 0.4 * null) @ b.T, x_test @ b.T, atol=1e-12)
        np.testing.assert_allclose(BETA @ null, 0, atol=1e-9)
        h = x_test @ b.T
        effects = []
        for j in range(3):
            shift = ((h + 0.25 * b[:, j]) @ w3.T).ravel() - prediction
            np.testing.assert_allclose(shift, 0.25 * BETA[j], atol=1e-9)
            effects.append(float(shift.mean()))
        # Reparameterize a trained representation; no retraining is involved.
        transform = np.array([[1.0, 2.0], [0.0, 0.3]])
        b_new = transform @ b
        w3_new = w3 @ np.linalg.inv(transform)
        changed_prediction = (x_test @ b_new.T @ w3_new.T).ravel()
        np.testing.assert_allclose(changed_prediction, prediction, atol=1e-12)
        transformed = analyze(b_new, x_train, x_test)
        np.testing.assert_allclose(transformed["probe_r2"], second["probe_r2"], atol=1e-10)
        result.update({
            "null_direction": null.tolist(), "intervention_delta": 0.25,
            "intervention_effects": effects, "basis_transform": transform.tolist(),
            "transformed": transformed,
            "basis_change_max_prediction_error": float(np.max(np.abs(changed_prediction - prediction))),
        })
    return result


def figures(report):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    plt.rcParams.update({"svg.fonttype": "none", "font.size": 10})
    output = ROOT / "figures"
    output.mkdir(exist_ok=True)
    run = next(r for r in report["runs"] if r["width"] == 2 and r["seed"] == 7)
    fig, axes = plt.subplots(1, 2, figsize=(10, 3.6), layout="constrained")
    for axis, layer in zip(axes, ["h1", "h2"]):
        corr = np.array(run[layer]["pc_feature_correlations"])
        im = axis.imshow(corr, vmin=-1, vmax=1, cmap="RdBu_r", aspect="auto")
        axis.set(xticks=range(3), xticklabels=report["features"],
                 yticks=range(len(corr)), yticklabels=[f"PC {i+1}" for i in range(len(corr))],
                 title=f"{layer}: test correlations with known features")
        for (i, j), value in np.ndenumerate(corr):
            axis.text(j, i, f"{value:.2f}", ha="center", va="center",
                      color="white" if abs(value) > 0.6 else "black")
    fig.colorbar(im, ax=axes, label="Pearson correlation", shrink=0.85)
    fig.savefig(output / "pca-correlations.svg", metadata={"Date": None})
    plt.close(fig)
    fig, ax = plt.subplots(figsize=(8, 4), layout="constrained")
    for j, feature in enumerate(report["features"]):
        means = [np.mean([r["h2"]["probe_r2"][j] for r in report["runs"] if r["width"] == w]) for w in (1, 2, 3)]
        ax.plot([1, 2, 3], means, "o-", label=feature)
        for r in report["runs"]:
            ax.scatter(r["width"], r["h2"]["probe_r2"][j], alpha=0.4, s=14, color=f"C{j}")
    ax.set(xticks=[1, 2, 3], xlabel="Bottleneck width", ylabel="Held-out linear probe R²",
           title="Perfect price prediction does not preserve every input", ylim=(-0.08, 1.06))
    ax.legend(title="Lines: means over five seeds; dots: individual runs")
    fig.savefig(output / "feature-recovery.svg", metadata={"Date": None})
    plt.close(fig)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--export", action="store_true", help="Write ../data/results.json")
    parser.add_argument("--figures", action="store_true", help="Write SVG figures")
    args = parser.parse_args()
    gradient_check()
    rng = np.random.default_rng(20260925)
    x_train, x_test = sample(rng, 2048), sample(rng, 4096)
    runs = [train(seed, width, x_train, x_test) for width in (1, 2, 3) for seed in SEEDS]
    report = {"features": ["Size", "Rooms", "Age"], "beta": BETA.tolist(),
              "data_seed": 20260925, "train_count": len(x_train), "test_count": len(x_test),
              "learning_rate": 0.03, "init_std": 0.4, "seeds": SEEDS,
              "python": platform.python_version(), "numpy": np.__version__,
              "samples": x_test[:160].tolist(), "runs": runs}
    for r in runs:
        print(f"width={r['width']} seed={r['seed']:3d} steps={r['steps']:4d} "
              f"test MSE={r['test_mse']:.2e} probe R²={np.round(r['h2']['probe_r2'], 4)}")
    print("Passed: finite-difference gradients, held-out predictions, probes, nullspace, and interventions.")
    if args.export:
        (ROOT / "data").mkdir(exist_ok=True)
        (ROOT / "data/results.json").write_text(json.dumps(report, indent=2, allow_nan=False) + "\n")
    if args.figures:
        figures(report)


if __name__ == "__main__":
    main()