#!/usr/bin/env python3
"""crew.py -- the demolition crew: after the big brush fills a shape, a few critics with ONE taste each look at it and mark the sections
that should come out, the crew tears them out (and anything left hanging), and Pin gets the gaps (Ivy, Oct 23 2026: "Could have a demo
crew that has taste about what isn't working, what doesn't read, what's boring and repetitive. It colors inside the lines. Then erases,
then fills in with pin.").

Sections are clusters of pieces (k-means on where they sit), each painted one loud colour on a MAP render beside the real one, so a
critic can name a section by its colour; the prompt says where each one is (high up, low down -- from the geometry, because a critic
left to read the map called the base 'the crown'). The critics are Sonnet through `claude -p`, one call each. A section comes out when
two critics vote for it. Erasing a middle leaves the top floating, so after the votes everything with no stud-and-socket
path down to the bottom layer comes out too.

    python3 crew.py BUILD.ldr --what "the Statue of Liberty" [--k 8]
-> runs/crew_<time>/: map.png (real | sections, front and side), votes.json, erased.ldr, after.png
Written by Iris."""
import argparse, json, os, re, subprocess, sys, time
import numpy as np
from scipy.cluster.vq import kmeans2
from scipy.spatial import cKDTree
HERE = os.path.dirname(os.path.abspath(__file__)); sys.path.insert(0, HERE)
LEGO = os.path.expanduser('~/lego')

# loud, far apart, and every one a real Lego colour the critics can name (no black: LeoCAD's background is dark grey and the black
# section vanished into it, Oct 24 2026)
MAP = [('red', 4), ('blue', 1), ('yellow', 14), ('green', 2), ('orange', 25), ('white', 15), ('magenta', 26), ('medium azure', 322),
       ('lime', 27), ('reddish brown', 70)]

TASTES = {
    'reads': "You only care whether it READS as {what}. Walk round it in your head: which sections make it look like something "
             "else, or like nothing -- a blob, a column, a heap? A section that reads as the right thing stays, however rough.",
    'boring': "You only care about BORING AND REPETITIVE. Which sections are the same piece over and over in a flat run, with no "
              "change of shape, rhythm or direction -- filler that a builder would have done something with?",
    'working': "You only care whether it's WORKING as a build: ragged or thin stacks, bits sticking out at random, a section that "
               "looks loose, collapsed or unfinished. Not whether it's pretty -- whether it holds together and looks built on purpose.",
}

def parts(ldr):
    """[(colour, M 4x4 in LDraw units, name, line)] from an .ldr"""
    out = []
    for line in open(ldr):
        f = line.split()
        if len(f) < 15 or f[0] != '1': 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)
        out.append((int(f[1]), M, f[14], line))
    return out

def write(path, ps, colours=None):
    with open(path, 'w') as fh:
        fh.write(f'0 {os.path.basename(path)}\n')
        for i, (c, M, n, _) in enumerate(ps):
            c = colours[i] if colours is not None else c
            fh.write(f'1 {c} ' + ' '.join(f'{v:.3f}' for v in M[:3, 3]) + ' ' + ' '.join(f'{v:.4f}' for v in M[:3, :3].ravel()) + f' {n}\n')

def render(ldr, png, view='home', w=360, h=520):
    subprocess.run(['flatpak', 'run', '--filesystem=home', 'org.leocad.LeoCAD', '-l', os.path.join(LEGO, 'ldraw'), '-i', png,
                    '-w', str(w), '-h', str(h), '--viewpoint', view, ldr], capture_output=True, timeout=300)

def graph(ldr):
    """which pieces are joined (a stud in a socket), from the game's own parts: {i: set(j)} and each piece's bottom z (mm)"""
    import fits, play
    r = fits.load_board(ldr); bs = r.bodies()
    S = [(i, b.world(t['c'])) for i, b in enumerate(bs) for t in b.studs]
    K = [(i, b.world(k['c'])) for i, b in enumerate(bs) for k in b.sockets]
    nb = {i: set() for i in range(len(bs))}
    if S and K:
        ts = cKDTree(np.array([p for _, p in S]))
        for j, p in K:
            for q in ts.query_ball_point(p, 0.6):
                i = S[q][0]
                if i != j: nb[i].add(j); nb[j].add(i)
    return bs, nb

def sections(bs, k):
    X = np.array([b.mesh.vertices.mean(0) for b in bs])
    X = X / [8.0, 8.0, 9.6]                                  # in studs and bricks, so a section is about as tall as it is wide
    _, lab = kmeans2(X, k, seed=1, minit='++')
    return lab

def where(bs, lab, ids):
    """{colour: 'high up, 14 pieces'} from the pieces' heights -- a critic told where things are stops misreading the map (flash-lite called
    the base 'the crown', Oct 24 2026)"""
    z = np.array([b.mesh.vertices.mean(0)[2] for b in bs]); lo, hi = z.min(), z.max()
    words = [(0.15, 'at the very bottom'), (0.35, 'low down'), (0.6, 'in the middle'), (0.85, 'high up'), (1.01, 'at the very top')]
    out = {}
    for i in ids:
        f = (z[lab == i].mean() - lo) / max(hi - lo, 1e-6)
        out[MAP[i][0]] = f"{next(w for lim, w in words if f < lim)}, {int((lab == i).sum())} pieces"
    return out

def sheet(out, real_f, map_f, real_s, map_s, used):
    from PIL import Image, ImageDraw
    ims = [Image.open(p).convert('RGB') for p in (real_f, map_f, real_s, map_s)]
    W, H = ims[0].size; S = Image.new('RGB', (W * 4, H + 60), 'white'); d = ImageDraw.Draw(S)
    for i, (im, cap) in enumerate(zip(ims, ['the build, front', 'its sections, front', 'the build, side', 'its sections, side'])):
        S.paste(im, (W * i, 30)); d.text((W * i + 10, 8), cap, fill='black')
    d.text((10, H + 38), 'sections: ' + ', '.join(used), fill='black')
    p = os.path.join(out, 'map.png'); S.save(p); return p

def critic(taste, what, img, used, wh, model='sonnet'):
    """one taste, one look -- through `claude -p` like Pin (the direct Haiku path was refused nearly every call on Oct 24 2026 with the
    week's meter at 1%; Gemini flash-lite answered, but flipped between 'tear out the torso' and 'nothing' on the same picture, and with
    the map's locations given it only ever saw 'a messy base' where Pin saw a thin loose torso and no head)"""
    import bigbrush
    q = (f"This is a big Lego build of {what}, filled fast by a machine; a builder will rebuild whatever you tear out, by hand. "
         f"Read the picture {img}. Left pair: the build and the same build with each SECTION painted one colour (front view). Right "
         f"pair: the same from the side. The sections, and where they are: " + '; '.join(f"{c} ({wh[c]})" for c in used) + ".\n\n" +
         TASTES[taste].format(what=what) +
         "\n\nTear out at most three sections, or none if nothing fails your test. Answer ONLY with JSON: "
         '{"tear_out": ["<section colour>", ...], "why": "<one sentence>"}')
    try:
        d = bigbrush.pin(q, read_dir=os.path.dirname(img), model=model)
        return {'tear_out': [c.lower().strip() for c in d.get('tear_out') or [] if c.lower().strip() in used][:3], 'why': d.get('why', ''), 'by': model}
    except Exception as e:
        return {'tear_out': [], 'why': f'(no answer: {str(e)[:80]})', 'by': model}

def main():
    ap = argparse.ArgumentParser(); ap.add_argument('ldr'); ap.add_argument('--what', required=True); ap.add_argument('--k', type=int, default=8)
    ap.add_argument('--out'); a = ap.parse_args()
    out = os.path.abspath(a.out or os.path.join(HERE, 'runs', 'crew_' + time.strftime('%m%d_%H%M%S'))); os.makedirs(out, exist_ok=True)
    ps = parts(a.ldr); bs, nb = graph(a.ldr)
    assert len(ps) == len(bs), (len(ps), len(bs))            # load_board keeps every line it can read; these must line up
    lab = sections(bs, min(a.k, len(MAP)))
    ids = sorted(set(lab)); used = [MAP[i][0] for i in ids]
    src = os.path.join(out, 'before.ldr'); write(src, ps)
    mp = os.path.join(out, 'sections.ldr'); write(mp, ps, [MAP[l][1] for l in lab])
    for f, v in ((src, 'front'), (mp, 'front'), (src, 'right'), (mp, 'right')):
        render(f, f[:-4] + f'_{v}.png', v)
    img = sheet(out, src[:-4] + '_front.png', mp[:-4] + '_front.png', src[:-4] + '_right.png', mp[:-4] + '_right.png', used)
    wh = where(bs, lab, ids)
    from concurrent.futures import ThreadPoolExecutor
    with ThreadPoolExecutor(3) as ex: votes = dict(zip(TASTES, ex.map(lambda t: critic(t, a.what, img, used, wh), TASTES)))
    for t, v in votes.items(): print(f'{t:8s} tears out {v["tear_out"] or "nothing"} -- {v["why"]}', flush=True)
    n = {c: sum(c in v['tear_out'] for v in votes.values()) for c in used}
    size = {MAP[i][0]: int((lab == i).sum()) for i in ids}
    gone = [c for c in used if n[c] >= 2]           # two critics agree (one alone would have taken the torch: a 1-stud stack reads as 'precarious')
    kill = {i for i, l in enumerate(lab) if MAP[l][0] in gone}
    # whatever has no path down to the bottom any more comes out with them -- a torso torn out takes the head with it
    z = np.array([b.mesh.bounds[0][2] for b in bs]); ground = {i for i in range(len(bs)) if z[i] < z.min() + 0.5 and i not in kill}
    keep, todo = set(ground), list(ground)
    while todo:
        for j in nb[todo.pop()] - keep - kill: keep.add(j); todo.append(j)
    hanging = set(range(len(bs))) - keep - kill
    left = [p for i, p in enumerate(ps) if i in keep]
    write(os.path.join(out, 'erased.ldr'), left)
    render(os.path.join(out, 'erased.ldr'), os.path.join(out, 'after.png'))
    json.dump({'what': a.what, 'from': a.ldr, 'sections': size, 'where': wh, 'votes': votes, 'torn_out': gone, 'pieces': len(bs),
               'voted_out': len(kill), 'hanging': len(hanging), 'left': len(left)}, open(os.path.join(out, 'votes.json'), 'w'), indent=1)
    print(f'torn out: {gone or "nothing"} -- {len(kill)} voted, {len(hanging)} left hanging, {len(left)}/{len(bs)} pieces stay -> {out}')

if __name__ == '__main__':
    main()