#!/bin/bash
# Build a HARD STRATUM test set to close the paper's main weakness.
#
# Published Sudoku numbers are on 17-clue puzzles; ours are 20-30 (mean 24.4),
# so the comparison has had to be hedged as "different distribution". Our test
# set contains only 46 puzzles at <=22 clues - too few to report.
#
# Generating a large pool and keeping the hardest tail gives a stratum close to
# the published difficulty, on which the EXISTING trained models can be evaluated
# with no retraining. CPU only, so it costs nothing against the GPU allocation.
#SBATCH -J genh2
#SBATCH -p standard
#SBATCH -c 64
#SBATCH --mem=64g
#SBATCH -t 4:00:00
#SBATCH -o logs/genh2-%j.out
export SCR=/scratch/zt1/project/msml612/user/$USER
source /etc/profile
SCR=/scratch/zt1/project/msml612/user/$USER
module load python/3.10.10
export PYTHONPATH=
cd $SCR/pcmd/src
$SCR/hf_env/bin/python -u - <<'PY'
import numpy as np, os, sys, time
from multiprocessing import Pool
sys.path.insert(0, '.')
from gen_hard_sudoku import make_one, canon
from sudoku import CELLS

OUT = os.path.join(os.environ['SCR'], 'pcmd/data/sudoku_stratum')
os.makedirs(OUT, exist_ok=True)
N = 30000
t0 = time.time()
with Pool(64) as p:
    got = p.map(make_one, range(1_000_000, 1_000_000 + N), chunksize=32)
P = np.stack([g for g, _ in got]); S = np.stack([s for _, s in got])
clues = (P > 0).sum(1)
print(f"generated {N} minimal puzzles in {time.time()-t0:.0f}s", flush=True)
import collections
h = collections.Counter(clues.tolist())
for k in sorted(h): print(f"  {k} clues: {h[k]}", flush=True)

for cap, name in ((22, 'le22'), (23, 'le23')):
    m = clues <= cap
    if m.sum() < 50: continue
    np.save(os.path.join(OUT, f'{name}_puz.npy'), P[m])
    np.save(os.path.join(OUT, f'{name}_sol.npy'), S[m])
    print(f"  saved {name}: {int(m.sum())} puzzles, mean {clues[m].mean():.1f} clues",
          flush=True)

# these must not overlap the TRAINING solutions of the models we will evaluate
tr = np.load(os.path.join(os.environ['SCR'], 'pcmd/data/sudoku_hard/train_sol.npy'))
trc = {canon(r) for r in tr}
for cap, name in ((22, 'le22'), (23, 'le23')):
    f = os.path.join(OUT, f'{name}_sol.npy')
    if not os.path.exists(f): continue
    Ssub = np.load(f)
    leak = sum(canon(r) in trc for r in Ssub)
    print(f"  {name}: {leak}/{len(Ssub)} solutions also in training (must be 0)", flush=True)
    assert leak == 0
PY