#!/usr/bin/env python3
"""real-but-hidden (BlackAlps 2026) -- recover the CKKS ternary key from real-part-only logs.

Re m(zeta_j) = m_0 + sum_{k=1}^{31} cos(pi(2j+1)k/n) * (m_k - m_{n-k})   (zeta^(n-k) = -conj(zeta^k), zeta^32 = +-i)
=> the 32 published reals give *exactly* the integers d_0 = m_0, d_k = m_k - m_{n-k}.
Each job: d(b + a*s) = d_published - d(smudge + e) mod q  -> 32 linear equations in s with ~2^17 error.
q ~ 2^40: Kannan embedding + LLL (fplll CLI) recovers s.  Usage: .venv/bin/python -I solve.py files/output.txt
"""
import ast, sys, subprocess, numpy as np

def load(p):
    d = {}
    for line in open(p):
        if '=' in line and not line.lstrip().startswith('#'):
            k, v = line.split('=', 1)
            d[k.strip()] = ast.literal_eval(v.strip())
    return d

D = load(sys.argv[1] if len(sys.argv) > 1 else 'files/output.txt')
n, q, delta = D['n'], D['q'], D['delta']
h = n // 2

def negacyclic_matrix(a):
    """M with (a*s) = M s in Z_q[X]/(X^n+1)."""
    M = [[0] * n for _ in range(n)]
    for i in range(n):
        for j in range(n):
            k = i + j
            if k < n: M[k][j] = (M[k][j] + a[i]) % q
            else:     M[k - n][j] = (M[k - n][j] - a[i]) % q
    return M

def dmap(v):  # v: length-n integer vector -> d (length h)
    return [v[0]] + [v[k] - v[n - k] for k in range(1, h)]

jj = np.arange(h)
R = np.cos(np.pi * np.outer(2 * jj + 1, jj) / n)

rows_A, rhs = [], []
t = 0
while f'job{t}_re' in D:
    b, a, re = D[f'job{t}_b'], D[f'job{t}_a'], D[f'job{t}_re']
    dpub = [int(round(x)) for x in np.linalg.solve(R, np.array(re) * delta)]
    M = negacyclic_matrix(a)
    # d(b + M s) ≡ dpub + err   ->  d(M) s ≡ dpub - d(b) + err
    dM = [M[0]] + [[(M[k][c] - M[n - k][c]) % q for c in range(n)] for k in range(1, h)]
    db = dmap(b)
    for i in range(h):
        rows_A.append(dM[i]); rhs.append((dpub[i] - db[i]) % q)
    t += 1

m = len(rows_A)
W, K = 1 << 16, 1 << 17  # weight s (ternary) to match error size ~2^17
B = []
for c in range(n):
    B.append([W if r == c else 0 for r in range(n)] + [rows_A[i][c] for i in range(m)] + [0])
for i in range(m):
    B.append([0] * n + [q if r == i else 0 for r in range(m)] + [0])
B.append([0] * n + [(-rhs[i]) % q for i in range(m)] + [K])

inp = '[' + '\n'.join('[' + ' '.join(map(str, r)) + ']' for r in B) + ']'
out = subprocess.run(['fplll', '-a', 'bkz', '-b', '20'], input=inp, capture_output=True, text=True, check=True).stdout
red = [list(map(int, l.strip(' []').split())) for l in out.strip().strip('[]').split('\n') if l.strip(' []')]

s = None
for r in red:
    if abs(r[-1]) == K and all(x % W == 0 and abs(x // W) <= 1 for x in r[:n]):
        sg = r[-1] // K
        s = [sg * x // W for x in r[:n]]
        break
assert s is not None, "key not found"
print("s =", s)

def mul(f, g):
    hh = [0] * n
    for i, fi in enumerate(f):
        for j, gj in enumerate(g):
            k = i + j
            if k < n: hh[k] += fi * gj
            else:     hh[k - n] -= fi * gj
    return [c % q for c in hh]

def center(v): return [c - q if c > q // 2 else c for c in v]

# sanity: residuals on jobs are small
for i in range(m):
    e = (sum(rows_A[i][c] * s[c] for c in range(n)) - rhs[i]) % q
    assert min(e, q - e) < 1 << 19

V = [[np.exp(1j * np.pi * (2 * k + 1) / n) ** p for p in range(n)] for k in range(h)]
flag = ''
t = 0
while f'flag{t}_b' in D:
    b, a = D[f'flag{t}_b'], D[f'flag{t}_a']
    mpoly = center([(x + y) % q for x, y in zip(b, mul(a, s))])
    z = [sum(V[j][k] * mpoly[k] for k in range(n)) / delta for j in range(h)]
    print('slots:', [round(c.real, 3) for c in z][:8], '...')
    flag += ''.join(chr(int(round(c.real))) for c in z if 0 < round(c.real) < 256)
    t += 1
print(flag)