# Stage 37 (read-only): compose before/after review sheets.
#
#   blender --background --python 37_sheets.py -- <out_dir>
#
# "Before" is the body as it was handed to this lane: decimated v01, Tripo's 5,870-chart atlas,
# single-sided (so the inward-wound patches are culled into grey tears). "After" is v04.
import bpy, sys, os
import numpy as np

OUT = os.path.abspath(sys.argv[sys.argv.index("--") + 1])
os.makedirs(OUT, exist_ok=True)
HERE = os.path.dirname(os.path.abspath(__file__))


def load(p):
    im = bpy.data.images.load(p)
    w, h = im.size
    b = np.empty(w * h * 4, dtype=np.float32)
    im.pixels.foreach_get(b)
    a = b.reshape(h, w, 4)[:, :, :3].copy()
    bpy.data.images.remove(im)
    return a[::-1]


def fit(a, w, h):
    ys = (np.arange(h) * (a.shape[0] / h)).astype(int).clip(0, a.shape[0] - 1)
    xs = (np.arange(w) * (a.shape[1] / w)).astype(int).clip(0, a.shape[1] - 1)
    return a[ys][:, xs]


def save(a, name):
    h, w = a.shape[:2]
    img = bpy.data.images.new(name, w, h, alpha=False)
    buf = np.ones((h, w, 4), dtype=np.float32)
    buf[:, :, :3] = np.clip(a, 0, 1)[::-1]
    img.pixels.foreach_set(buf.reshape(-1))
    p = os.path.join(OUT, name)
    img.file_format = 'JPEG'
    img.filepath_raw = p
    img.save(filepath=p)
    print(f"wrote {p}  {w}x{h}  {os.path.getsize(p)/1e6:.2f} MB", flush=True)


GAP = 10
BG = 0.10


def grid(rows, cell_w, cell_h):
    nrow, ncol = len(rows), max(len(r) for r in rows)
    W = ncol * cell_w + (ncol + 1) * GAP
    H = nrow * cell_h + (nrow + 1) * GAP
    sheet = np.full((H, W, 3), BG)
    for r, row in enumerate(rows):
        for c, p in enumerate(row):
            if p is None:
                continue
            y = GAP + r * (cell_h + GAP)
            x = GAP + c * (cell_w + GAP)
            sheet[y:y + cell_h, x:x + cell_w] = fit(load(p), cell_w, cell_h)
    return sheet


B4 = os.path.join(HERE, "beauty_input_v01")     # as handed over: Tripo atlas, single-sided
AF = os.path.join(HERE, "beauty_refill_v04")    # now

save(grid([[os.path.join(B4, "full_front.png"), os.path.join(AF, "full_front.png")],
           [os.path.join(B4, "chest.png"), os.path.join(AF, "chest.png")],
           [os.path.join(B4, "hip.png"), os.path.join(AF, "hip.png")]],
          900, 900), "sheet_body_before_after.jpg")

save(grid([[os.path.join(HERE, "probe_atlas", "basecolor.png"),
            os.path.join(HERE, "refill_v04", "lena_nude_refill_base.jpg")]],
          1000, 1000), "sheet_atlas_before_after.jpg")

save(grid([[os.path.join(HERE, "probe_atlas", "islands.png"),
            os.path.join(HERE, "atlas_v03", "charts.png")]],
          1000, 1000), "sheet_charts_before_after.jpg")
print("SHEETS_DONE")