"""fits.py -- the fit bot: offline, free, no Opus. It loads boards Pin really played, tries pieces from Pin's bin on their open
spots, and writes down what actually fits (Ivy, Oct 23 2026: "We should have a bot that try fitting stuff together randomly from
played boards to try to which parts fit together best... Someway to offline learn for free what pieces Fit with which so it can
get handed a piece that will at least fit when it's stuck... so it can choose from legal smart choices without having to even
think about the dumb ones").

Every try goes through the game's own engine (play.Reef.try_move: real connectors, the sampled-surface collision check, the
invisible baseplate), so "fits" here means exactly what it means in a game. Results append to fits.jsonl one line per
(board, open spot, piece), so a run can be stopped and resumed; table() folds them into fits.json:
    host part -> connector kind -> piece -> [fitted, tried]

    python3 fits.py --boards 1 --spots 6        # a timing run
    python3 fits.py                              # every saved board not yet swept (run it at idle priority)
    python3 fits.py --table                      # just rebuild fits.json
Written by Iris."""
import os, sys, json, glob, time, argparse
import numpy as np
HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import play
from play import G
import build as B

C = B.C; Ci = np.linalg.inv(C)
LOG = os.path.join(HERE, 'fits.jsonl'); TABLE = os.path.join(HERE, 'fits.json')

def load_board(ldr, seed=0):
    """a saved .ldr back into a live board (the inverse of kit.write_ldr: it stores M = C^-1 T C, so T = C M C^-1)"""
    r = play.Reef(seed, 'tidy', os.path.join(HERE, 'runs', '_fits'))
    r.g.asm.parts.clear()
    for line in open(ldr):
        f = line.split()
        if len(f) < 15 or f[0] != '1': continue
        name = f[14][:-4] if f[14].lower().endswith('.dat') else f[14]
        if name not in play.PALETTE and name not in play.BASEPLATES: continue
        M = np.eye(4); M[:3, 3] = [float(v) for v in f[2:5]]; M[:3, :3] = np.array([float(v) for v in f[5:14]]).reshape(3, 3)
        try: p = G.fresh(name)
        except Exception: continue
        p.T = C @ M @ Ci; p.colour = int(f[1]); p.attached = True; r.g.asm.parts.append(p)
    if not r.g.asm.parts: return None
    r.g.reconcile(); r.floor_off = play.floor_grid(r.bodies())[1]
    return r

def sweep(r, board, spots=12, ways=4, rng=None, out=None):
    """for up to `spots` open spots on the board (never the bare floor: a piece on the floor fits trivially and joins nothing),
    try every piece in the palette up to `ways` random ways (connector index, spin) and record the first that fits"""
    rng = rng or np.random.default_rng(0)
    cs = [c for c in r.conns() if c[0] != 'floor']
    if len(cs) > spots: cs = [cs[i] for i in rng.choice(len(cs), spots, replace=False)]
    rows = []
    for c in cs:
        host = c[1].name
        for name in play.PALETTE:
            ok, n = False, 0
            for _ in range(ways):
                n += 1
                res = r.try_move(c, name, int(rng.integers(8)), float(rng.choice(G.SPINS)))
                if res: r.g.undo(*res); ok = True; break
            row = {'board': board, 'host': host, 'kind': c[0], 'piece': name, 'ok': ok, 'tries': n}
            rows.append(row)
            if out: out.write(json.dumps(row) + '\n')
        if out: out.flush()
    return rows

def table():
    t = {}
    if not os.path.exists(LOG): return t
    for l in open(LOG):
        try: x = json.loads(l)
        except ValueError: continue
        e = t.setdefault(x['host'], {}).setdefault(x['kind'], {}).setdefault(x['piece'], [0, 0])
        e[0] += int(x['ok']); e[1] += 1
    json.dump(t, open(TABLE, 'w'))
    return t

def fits_onto(host, kind=None, t=None):
    """pieces that have fitted onto this host part (optionally on one connector kind), best first: [(piece, rate, tried)]"""
    t = t if t is not None else (json.load(open(TABLE)) if os.path.exists(TABLE) else {})
    agg = {}
    for k, d in (t.get(host) or {}).items():
        if kind and k != kind: continue
        for p, (a, n) in d.items(): e = agg.setdefault(p, [0, 0]); e[0] += a; e[1] += n
    return sorted(((p, a / n, n) for p, (a, n) in agg.items() if a), key=lambda x: -x[1])

if __name__ == '__main__':
    ap = argparse.ArgumentParser(); ap.add_argument('--boards', type=int, default=0); ap.add_argument('--spots', type=int, default=12)
    ap.add_argument('--ways', type=int, default=4); ap.add_argument('--table', action='store_true'); a = ap.parse_args()
    if a.table: t = table(); print(len(t), 'host parts in fits.json'); sys.exit()
    done = set()
    if os.path.exists(LOG):
        for l in open(LOG):
            try: done.add(json.loads(l)['board'])
            except ValueError: pass
    boards = sorted(glob.glob(os.path.join(HERE, 'runs', '*', 'build.ldr')), key=os.path.getmtime, reverse=True)
    boards = [b for b in boards if os.path.basename(os.path.dirname(b)) not in done and not os.path.basename(os.path.dirname(b)).startswith('_')]
    if a.boards: boards = boards[:a.boards]
    with open(LOG, 'a') as out:
        for i, b in enumerate(boards):
            name = os.path.basename(os.path.dirname(b)); t0 = time.time()
            r = load_board(b, seed=i)
            if r is None: print(name, 'nothing loadable'); continue
            rows = sweep(r, name, a.spots, a.ways, np.random.default_rng(i), out)
            print(f"{name}: {len(r.bodies())} pieces, {len({(x['host'], x['kind']) for x in rows})} spots, "
                  f"{sum(x['ok'] for x in rows)}/{len(rows)} fit, {time.time() - t0:.0f} s", flush=True)
    table()