"""Learned-feature dataset for the cats_dogs KNP binary (2048 inputs = K 512 × 4 quadrants).

Features: kmeans_features without contrast division (linear patch centring, so layer 1 folds into
neuron weights), alpha 0.25. Each pooled feature is scaled by its 99th percentile on the training set
to 0–255 (a fixed gain per pooling neuron). Training part = 9000 images + their mirror copies, shuffled;
then the 1000 validation images. Writes <out>/kmeans_validation.{bin,target} and the model.

usage: ~/anvenv/bin/python make_knp_features.py --out /home/int/run/cats-dogs/features-v3
"""
import argparse
import json
from pathlib import Path

import numpy as np

import kmeans_features as kf


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default="/home/int/run/cats-dogs/data-v3")
    ap.add_argument("--out", required=True)
    ap.add_argument("--k", type=int, default=512)
    ap.add_argument("--alpha", type=float, default=0.25)
    ap.add_argument("--seed", type=int, default=20261014)
    ap.add_argument("--model", help="use a learned model (npz with mean, zca, d) instead of learning K prototypes")
    ap.add_argument("--name", default="kmeans", help="output file prefix")
    a = ap.parse_args()
    out = Path(a.out)
    out.mkdir(exist_ok=True)
    kf.CONTRAST_DIVIDE = False
    tr, va = np.load(f"{a.data}/train.npz"), np.load(f"{a.data}/val.npz")
    if a.model:
        m = np.load(a.model)
        model = {key: m[key] for key in ("mean", "zca", "d")}
        a.k = len(model["d"])
    else:
        model = kf.learn(tr["rgb"], k=a.k)
    ft = np.vstack([kf.encode(tr["rgb"], model, alpha=a.alpha), kf.encode(tr["rgb"][:, :, :, ::-1], model, alpha=a.alpha)])
    yt = np.r_[tr["labels"], tr["labels"]]
    fv = kf.encode(va["rgb"], model, alpha=a.alpha)
    gain = np.percentile(ft, 99, axis=0) + 1e-6
    q = lambda f: np.clip(np.rint(f / gain * 255), 0, 255).astype(np.uint8)
    order = np.random.default_rng(a.seed).permutation(len(ft))
    x = np.vstack([q(ft)[order], q(fv)])
    y = np.r_[yt[order], va["labels"]]
    assert x.shape[1] == 4 * a.k
    x.tofile(out / f"{a.name}_validation.bin")
    (out / f"{a.name}_validation.target").write_text("\n".join(map(str, y)) + "\n")
    np.savez(out / f"{a.name}_model.npz", gain=gain, **model)
    info = {"k": a.k, "alpha": a.alpha, "train": len(ft), "validation": len(fv), "contrast_divide": False,
            "mean_value": float(x.mean()), "spikes_per_image": float(x.mean() / 255 * 10 * x.shape[1])}
    (out / f"{a.name}_manifest.json").write_text(json.dumps(info, indent=2) + "\n")
    print(info)


if __name__ == "__main__":
    main()