"""autoencrypt (BlackAlps 2026) -- solve.

Flaw: the repetition code makes the encoder input x = ecc(m) live in a 64-dim
space.  W·x + b = W'·m + b where W' (96x64) sums each triple of columns of W.
So the "18432-weight" encoder is really an affine map 64 -> 96 with 65 unknowns
per output row.  70 known (plaintext, ciphertext) pairs >= 65 unknowns, so
least squares recovers W' and b exactly (up to 1e-6 rounding).  Then W' is a
tall full-rank 96x64 matrix: c - b = W'·m is injective on bits, so the
"bottleneck" loses nothing -- solve by least squares and round to {0,1}.

Usage: .venv/bin/python -I solve.py files/output.txt
"""
import ast
import sys

import numpy as np


def bits(data):
    return [(byte >> (7 - i)) & 1 for byte in data for i in range(8)]


def unbits(b):
    return bytes(sum(bit << (7 - i) for i, bit in enumerate(b[j:j + 8]))
                 for j in range(0, len(b), 8))


def parse(path):
    vals = {}
    for line in open(path):
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        name, rhs = line.split("=", 1)
        vals[name.strip()] = ast.literal_eval(rhs.strip())
    return vals


def main(path):
    v = parse(path)
    logs = sorted(k for k in v if k.startswith("log_"))
    flags = sorted(k for k in v if k.startswith("flag_"))

    # design matrix: [m bits | 1]
    A = np.array([bits(bytes.fromhex(v[k][0])) + [1] for k in logs], float)
    C = np.array([v[k][1] for k in logs], float)
    print(f"{len(logs)} queries, design rank {np.linalg.matrix_rank(A)} / {A.shape[1]}")

    X, *_ = np.linalg.lstsq(A, C, rcond=None)  # (65, 96)
    Wp, b = X[:64].T, X[64]
    print("fit residual max:", np.abs(A @ X - C).max())

    out = b""
    for k in flags:
        c = np.array(v[k], float)
        m, *_ = np.linalg.lstsq(Wp, c - b, rcond=None)
        mb = [int(round(t)) for t in m]
        assert all(t in (0, 1) for t in mb) and np.abs(m - mb).max() < 0.1, m
        out += unbits(mb)
    print(out.rstrip(b"\0").decode())


if __name__ == "__main__":
    main(sys.argv[1] if len(sys.argv) > 1 else "files/output.txt")