"""LAB07: genuine small image VAE training, NumPy-only backpropagation.

Run: python lab07_image_vae.py --out results
Requires numpy, scipy, scikit-learn, matplotlib, threadpoolctl. No download, GPU or hidden state.
Uses sklearn's bundled 8x8 UCI digits, NOT MNIST. Pixel > 8 -> 1.
Primary references:
https://arxiv.org/html/1312.6114v11 (sections 2-3, appendices B/C)
https://scikit-learn.org/stable/modules/generated/sklearn.datasets.load_digits.html
"""
import os
for name in ("OPENBLAS_NUM_THREADS", "OMP_NUM_THREADS", "MKL_NUM_THREADS"):
    os.environ[name] = "1"
import argparse
import copy
import csv
import hashlib
import json
from pathlib import Path
import platform
import time
import numpy as np
from threadpoolctl import threadpool_limits
import scipy
from scipy.special import expit
import sklearn
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

CONFIG = dict(data="sklearn.datasets.load_digits (UCI 8x8 digits)",
              binarization="pixel > 8", train_fraction=0.8,
              split_seed=7301, init_seed=7302, train_seed=7303,
              evaluation_seed=7304, visual_seed=7305,
              architecture="64-tanh64-(mu2,logvar2); z2-tanh64-logits64",
              latent_dim=2, hidden_dim=64, epochs=150, batch_size=64,
              learning_rate=0.001, betas=[0.0, 1.0],
              reconstruction="Bernoulli negative log likelihood: sum 64 pixels, mean batch",
              kl="diagonal Gaussian to N(0,I): sum 2 latent dimensions, mean batch",
              train_mc_samples=1, evaluation_mc_samples=8,
              optimizer="Adam beta1=0.9 beta2=0.999 eps=1e-8",
              dtype="float64", version="lab07-v1")


def initialize(seed):
    rng = np.random.default_rng(seed)
    p = {}
    for name, a, b in [("e", 64, 64), ("mu", 64, 2), ("lv", 64, 2),
                       ("d", 2, 64), ("out", 64, 64)]:
        p[name+"W"] = rng.normal(0, np.sqrt(2/(a+b)), (a, b))
        p[name+"b"] = np.zeros(b)
    return p


def encode(p, x):
    h = np.tanh(x @ p["eW"] + p["eb"])
    return h @ p["muW"] + p["mub"], h @ p["lvW"] + p["lvb"]


def decode(p, z):
    return expit(np.tanh(z @ p["dW"] + p["db"]) @ p["outW"] + p["outb"])


def forward(p, x, eps, beta, gradient=False):
    h = np.tanh(x @ p["eW"] + p["eb"])
    mu, lv = h @ p["muW"] + p["mub"], h @ p["lvW"] + p["lvb"]
    sd = np.exp(0.5 * lv)
    z = mu + sd * eps                       # reparameterization, no detached mu/sd
    hd = np.tanh(z @ p["dW"] + p["db"])
    logits = hd @ p["outW"] + p["outb"]
    prob = expit(logits)
    recon = np.mean(np.sum(np.logaddexp(0, logits) - x * logits, axis=1))
    kl = np.mean(0.5 * np.sum(mu**2 + np.exp(lv) - 1 - lv, axis=1))
    result = {"objective": float(recon + beta * kl), "reconstruction_nll": float(recon),
              "kl": float(kl), "negative_elbo": float(recon + kl)}
    if not gradient:
        return result
    n = len(x)
    g = {}
    dl = (prob-x)/n                         # pixel SUM; batch MEAN
    g["outW"], g["outb"] = hd.T @ dl, dl.sum(0)
    dhd = (dl @ p["outW"].T) * (1-hd**2)
    g["dW"], g["db"] = z.T @ dhd, dhd.sum(0)
    dz = dhd @ p["dW"].T
    dmu = dz + beta * mu/n
    dlv = dz * eps * 0.5 * sd + beta * 0.5 * (np.exp(lv)-1)/n
    g["muW"], g["mub"] = h.T @ dmu, dmu.sum(0)
    g["lvW"], g["lvb"] = h.T @ dlv, dlv.sum(0)
    dh = (dmu @ p["muW"].T + dlv @ p["lvW"].T) * (1-h**2)
    g["eW"], g["eb"] = x.T @ dh, dh.sum(0)
    return result, g


def gradient_check(p, x):
    """Finite differences with fixed epsilon; check both beta paths, every tensor."""
    rng = np.random.default_rng(7306)
    eps = rng.normal(size=(len(x), 2))
    checks = []
    for beta in [0.0, 1.0]:
        _, analytic = forward(p, x, eps, beta, True)
        for name, tensor in p.items():
            for flat in rng.choice(tensor.size, min(4, tensor.size), replace=False):
                index = np.unravel_index(flat, tensor.shape)
                original = tensor[index]
                tensor[index] = original+1e-5
                up = forward(p, x, eps, beta)["objective"]
                tensor[index] = original-1e-5
                down = forward(p, x, eps, beta)["objective"]
                tensor[index] = original
                numerical = (up-down)/2e-5
                actual = analytic[name][index]
                scaled = abs(numerical-actual)/max(1, abs(numerical), abs(actual))
                checks.append({"beta": beta, "parameter": name, "index": list(index),
                               "analytic": float(actual), "finite_difference": float(numerical),
                               "scaled_error": float(scaled)})
    assert max(c["scaled_error"] for c in checks) < 1e-6
    return checks


def evaluate(p, x, eps_samples, beta):
    records = [forward(p, x, eps, beta) for eps in eps_samples]
    return {key: float(np.mean([v[key] for v in records])) for key in records[0]}


def train(initial, xtrain, xtest, test_eps, beta, cfg):
    p = copy.deepcopy(initial)
    m, v = ({k: np.zeros_like(a) for k, a in p.items()} for _ in range(2))
    rng = np.random.default_rng(cfg["train_seed"])
    log = [{"epoch": 0, **evaluate(p, xtest, test_eps, beta)}]
    tick, step = time.perf_counter(), 0
    for epoch in range(1, cfg["epochs"]+1):
        permutation = rng.permutation(len(xtrain))
        for start in range(0, len(xtrain), cfg["batch_size"]):
            batch = xtrain[permutation[start:start+cfg["batch_size"]]]
            eps = rng.normal(size=(len(batch), cfg["latent_dim"]))
            _, grad = forward(p, batch, eps, beta, True)
            step += 1
            for k in p:
                m[k] = .9*m[k] + .1*grad[k]
                v[k] = .999*v[k] + .001*grad[k]**2
                p[k] -= cfg["learning_rate"] * (m[k]/(1-.9**step))/(np.sqrt(v[k]/(1-.999**step))+1e-8)
        if epoch % 5 == 0 or epoch == cfg["epochs"]:
            log.append({"epoch": epoch, **evaluate(p, xtest, test_eps, beta)})
    return p, log, time.perf_counter()-tick, step


def image_rows(rows, labels, target, titles=None, title="", scale=1):
    nrows, ncols = len(rows), len(rows[0])
    fig, axes = plt.subplots(nrows, ncols, figsize=(ncols*1.15*scale, nrows*1.25*scale),
                             squeeze=False)
    for i, row in enumerate(rows):
        for j, value in enumerate(row):
            ax = axes[i,j]
            ax.imshow(value.reshape(8,8), cmap="gray", vmin=0, vmax=1, interpolation="nearest")
            ax.set_xticks([]); ax.set_yticks([])
            if j == 0:
                ax.set_ylabel(labels[i], fontsize=9)
            if i == 0 and titles:
                ax.set_title(titles[j], fontsize=9)
    fig.suptitle(title, fontsize=12)
    fig.tight_layout()
    fig.savefig(target, dpi=160, bbox_inches="tight")
    plt.close(fig)


def json_dump(path, value):
    path.write_text(json.dumps(value, ensure_ascii=False, indent=2, default=lambda x: int(x) if isinstance(x,np.integer) else str(x)))


@threadpool_limits.wrap(limits=1, user_api='blas')
def run(out, epochs=None):
    out = Path(out); out.mkdir(parents=True, exist_ok=True)
    cfg = copy.deepcopy(CONFIG)
    if epochs is not None: cfg["epochs"] = epochs
    d = load_digits()
    # Fixed binary observations make the product Bernoulli likelihood explicit.
    x = (d.data > 8).astype(np.float64)
    train_id, test_id = train_test_split(np.arange(len(x)), test_size=.2,
        stratify=d.target, random_state=cfg["split_seed"])
    xtrain, xtest = x[train_id], x[test_id]
    cfg.update(train_size=len(train_id), test_size=len(test_id), parameters=0)
    initial = initialize(cfg["init_seed"])
    cfg["parameters"] = sum(a.size for a in initial.values())
    json_dump(out/"config.json", cfg)
    checks = gradient_check(initial, xtrain[:5])
    json_dump(out/"gradient_check.json", {"passed": True, "checks": checks,
        "max_scaled_error": max(c["scaled_error"] for c in checks)})
    np.savez_compressed(out/"data.npz", gray=d.data/16, binary=x, labels=d.target,
                        train_indices=train_id, test_indices=test_id)
    (out/"DATASET_DESCRIPTION.txt").write_text(d.DESCR)
    test_eps = np.random.default_rng(cfg["evaluation_seed"]).normal(
        size=(cfg["evaluation_mc_samples"], len(xtest), cfg["latent_dim"]))
    visual_rng = np.random.default_rng(cfg["visual_seed"])
    fixed_prior = visual_rng.normal(size=(25,2))
    fixed_uniform = visual_rng.uniform(size=(25,64))
    selected = np.array([np.flatnonzero(d.target[test_id] == digit)[0] for digit in range(10)])
    selected_x = xtest[selected]
    post_eps = visual_rng.normal(size=(4,8,2))
    np.savez_compressed(out/"fixed_inputs.npz", prior=fixed_prior, observation_uniform=fixed_uniform,
        selected_global_indices=test_id[selected], posterior_epsilon=post_eps, test_epsilon=test_eps)
    np.savez_compressed(out/"weights_initial.npz", **initial)
    models, histories, summaries = {}, {}, {}
    for beta in cfg["betas"]:
        tag = f"beta{beta:g}"
        p, log, seconds, steps = train(initial, xtrain, xtest, test_eps, beta, cfg)
        models[tag], histories[tag] = p, log
        np.savez_compressed(out/f"weights_{tag}.npz", **p)
        mu, lv = encode(p, xtest)
        summaries[tag] = {"beta": beta, "elapsed_training_seconds": seconds, "updates": steps,
                          "initial_test": log[0], "final_test": log[-1],
                          "posterior_mean_radius_mean": float(np.linalg.norm(mu,axis=1).mean()),
                          "posterior_std_mean": float(np.exp(lv/2).mean())}
        print(json.dumps(summaries[tag]), flush=True)
    json_dump(out/"summary.json", {"config":cfg,"runs":summaries,
        "environment":{"python":platform.python_version(),"numpy":np.__version__,
          "scipy":scipy.__version__,"sklearn":sklearn.__version__,"matplotlib":matplotlib.__version__},
        "gradient_check_max_scaled_error":max(c["scaled_error"] for c in checks)})
    with (out/"learning_curves.csv").open("w", newline="") as f:
        writer = csv.DictWriter(f,fieldnames=["beta","epoch","objective","reconstruction_nll","kl","negative_elbo"])
        writer.writeheader()
        for tag, log in histories.items():
            for row in log: writer.writerow({"beta": summaries[tag]["beta"],**row})
    # Fixed test input; using z=mu is deterministic reconstruction, not E_q[decoder(z)].
    recon_rows = [d.data[test_id[selected]]/16, selected_x, decode(initial, encode(initial,selected_x)[0])]
    for p in models.values(): recon_rows.append(decode(p, encode(p,selected_x)[0]))
    model_labels = [f"beta={b:g}" for b in cfg["betas"]]
    image_rows(recon_rows,["original\ngray","binary\ninput","initial\ndecoder(mu)"]+[label+"\ndecoder(mu)" for label in model_labels],
               out/"reconstruction.png",[str(i) for i in range(10)],"Same held-out digits, before and after learning")
    # Exactly the same z vectors in all rows; second figure uses same observation uniforms.
    prior_rows = [decode(p,fixed_prior) for p in [initial,*models.values()]]
    image_rows([r[:10] for r in prior_rows],["initial",*model_labels],out/"prior_same_seed.png",
               title="Same ten z ~ N(0,I): decoder Bernoulli probabilities")
    observation_tag = "beta1" if "beta1" in models else list(models)[-1]
    p = models[observation_tag]
    mean = decode(p,fixed_prior)
    image_rows([mean[:10],(fixed_uniform<mean).astype(float)[:10]],
               ["probability\n(mean)","Bernoulli\nobservation"],out/"mean_vs_observation.png",
               title="Same z, two different objects: probabilities and binary draws")
    for tag,p in models.items():
        means, logs = encode(p,xtest)
        mu, lv = encode(p,selected_x[[0,3,6,9]])
        zs = mu[:,None,:]+np.exp(lv[:,None,:]/2)*post_eps
        post = decode(p,zs.reshape(-1,2)).reshape(4,8,64)
        image_rows([np.concatenate([selected_x[[0,3,6,9]][i:i+1],post[i]],axis=0) for i in range(4)],
                   ["digit 0","digit 3","digit 6","digit 9"],out/f"posterior_{tag}.png",
                   ["input"]+[f"q draw {j+1}" for j in range(8)],title=f"{tag}: one input -> different posterior z -> mean images")
        ends = encode(p,selected_x[[3,8]])[0]
        alpha = np.linspace(0,1,11)
        interp = (1-alpha[:,None])*ends[0]+alpha[:,None]*ends[1]
        image_rows([decode(p,interp)],["linear z\npath"],out/f"interpolation_{tag}.png",
                   [f"{a:.1f}" for a in alpha],title=f"{tag}: interpolation from digit 3 to digit 8 (not prior draws)")
        grid_axis = np.linspace(-3,3,15)
        gx,gy = np.meshgrid(grid_axis,grid_axis[::-1])
        grid = decode(p,np.column_stack([gx.ravel(),gy.ravel()]))
        mosaic = grid.reshape(15,15,8,8).transpose(0,2,1,3).reshape(120,120)
        fig,axs=plt.subplots(1,2,figsize=(11,5),constrained_layout=True)
        axs[0].scatter(means[:,0],means[:,1],c=d.target[test_id],cmap="tab10",s=13,alpha=.65)
        angle=np.linspace(0,2*np.pi,200)
        axs[0].plot(2*np.cos(angle),2*np.sin(angle),color="black",ls="--",lw=1,label="radius 2 reference")
        axs[0].scatter(fixed_prior[:,0],fixed_prior[:,1],marker="x",s=35,c="black",label="25 fixed prior draws")
        axs[0].set(xlabel="z1",ylabel="z2")
        axs[0].set_title(f"{tag}: held-out posterior means",fontsize=11,pad=10)
        axs[0].legend(fontsize=7);axs[0].set_aspect("equal",adjustable="box")
        axs[1].imshow(mosaic,cmap="gray",vmin=0,vmax=1,extent=(-3,3,-3,3),interpolation="nearest")
        axs[1].set(xlabel="z1",ylabel="z2")
        axs[1].set_title("Uniform grid in [-3,3]^2 (not prior samples)",fontsize=11,pad=10)
        fig.savefig(out/f"latent_{tag}.png",dpi=160,bbox_inches="tight");plt.close(fig)
        np.savez_compressed(out/f"observations_{tag}.npz",test_mu=means,test_logvar=logs,
            selected_reconstruction=decode(p,encode(p,selected_x)[0]),prior_mean=decode(p,fixed_prior),
            prior_observation=(fixed_uniform<decode(p,fixed_prior)).astype(np.uint8),
            posterior_z=zs,posterior_mean_images=post,interpolation_z=interp,
            interpolation_mean=decode(p,interp),grid_z=np.column_stack([gx.ravel(),gy.ravel()]),grid_mean=grid)
    fig,axes=plt.subplots(1,2,figsize=(10,3.8))
    for tag,log in histories.items():
        axes[0].plot([v["epoch"] for v in log],[v["reconstruction_nll"] for v in log],label=tag)
        axes[1].plot([v["epoch"] for v in log],[v["kl"] for v in log],label=tag)
    axes[0].set(title="Held-out expected reconstruction NLL",xlabel="epoch",ylabel="nats/image; sum 64 pixels")
    axes[1].set(title="Held-out KL(q(z|x) || N(0,I))",xlabel="epoch",ylabel="nats/image; sum 2 dimensions")
    for ax in axes: ax.legend();ax.grid(alpha=.2)
    fig.tight_layout();fig.savefig(out/"learning_curves.png",dpi=170);plt.close(fig)
    manifest=[]
    for path in sorted(out.iterdir()):
        if path.is_file() and path.name!="manifest.json":
            manifest.append({"file":path.name,"bytes":path.stat().st_size,"sha256":hashlib.sha256(path.read_bytes()).hexdigest()})
    json_dump(out/"manifest.json",manifest)
    return summaries

# 若重跑，请用新目录保留本册原始结果。
# 单变量对照任选一个：CONFIG["betas"] = [0.0, 0.5, 1.0]
# 或 run("lab07_shorter", epochs=50)，不要同时改变两个设置。
published_repeat = run("lab07_rerun")

# 这些是本次重算目录里的新图，和上方保存的基线分别读取。
from IPython.display import display, Image
for figure in ["reconstruction.png", "prior_same_seed.png", "latent_beta0.png", "latent_beta1.png"]:
    display(Image(filename="lab07_rerun/" + figure))

{"beta": 0.0, "elapsed_training_seconds": 1.5911131510001724, "updates": 3450, "initial_test": {"epoch": 0, "objective": 44.874893795622526, "reconstruction_nll": 44.874893795622526, "kl": 0.45410676032218933, "negative_elbo": 45.32900055594472}, "final_test": {"epoch": 150, "objective": 16.06052023596692, "reconstruction_nll": 16.06052023596692, "kl": 103.37481381849595, "negative_elbo": 119.43533405446287}, "posterior_mean_radius_mean": 12.323355857715377, "posterior_std_mean": 0.009828530272238265}
{"beta": 1.0, "elapsed_training_seconds": 1.5743949399984558, "updates": 3450, "initial_test": {"epoch": 0, "objective": 45.32900055594472, "reconstruction_nll": 44.874893795622526, "kl": 0.45410676032218933, "negative_elbo": 45.32900055594472}, "final_test": {"epoch": 150, "objective": 20.89589535072323, "reconstruction_nll": 18.637906550351587, "kl": 2.2579888003716375, "negative_elbo": 20.89589535072323}, "posterior_mean_radius_mean": 1.2849481784600931, "posterior_std_mean": 0.3496167331762039}