#!/usr/bin/env python3
"""Figures for site/pages/seismic.html (Claude Opus 5.5, 2026-10-17).

Reads the SU files the flow wrote on obsidian (copied to DIR) and the archive's 1981 stack
and scans; writes grey GIFs (680 px wide, the site's page width) into site/static/.
Display only: no processing here beyond gain for the raw shots and a per-panel clip.
"""
import os, sys
import numpy as np
from PIL import Image, ImageDraw, ImageFont

DIR = sys.argv[1] if len(sys.argv) > 1 else "/tmp/s3181"
OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "site", "static")
sys.path.insert(0, DIR)
from rd import load                       # IBM-float SEG-Y reader used for the quick look

try:
    FONT = ImageFont.truetype("DejaVuSans.ttf", 13)
except OSError:
    FONT = ImageFont.load_default()


def su(name, ntr=536):
    """SU file from obsidian: big-endian, 240-byte headers, 3000 samples at 2 ms."""
    d = np.fromfile(os.path.join(DIR, name), dtype=">f4").reshape(ntr, -1)
    return d[:, 60:].astype(np.float64)


def agc(x, win):
    e = np.sqrt(np.apply_along_axis(lambda t: np.convolve(t * t, np.ones(win) / win, "same"), 1, x)) + 1e-30
    return x / e


def panel(x, w, h, clip=98.0):
    """traces x samples -> w x h grey image, black = positive."""
    c = np.percentile(np.abs(x), clip) or 1.0
    g = np.clip(x / c, -1, 1)
    img = Image.fromarray(np.uint8(127.5 - 127.5 * g.T))
    return img.resize((w, h), Image.BILINEAR)


def sheet(panels, labels, w, h, gap=6, top=20):
    n = len(panels)
    out = Image.new("L", (n * w + (n - 1) * gap, h + top), 255)
    d = ImageDraw.Draw(out)
    for i, (p, lab) in enumerate(zip(panels, labels)):
        out.paste(p, (i * (w + gap), top))
        d.text((i * (w + gap) + 2, 3), lab, fill=0, font=FONT)
    return out


def save(img, name):
    img.convert("P", palette=Image.ADAPTIVE, colors=64).save(os.path.join(OUT, name))
    print(name, img.size)


# raw shots SP 109 (file 120) and SP 110 (file 121, "vented"), 0-4 s, t^2 gain + 250 ms AGC
data = np.load(os.path.join(DIR, "all.npy"), mmap_mode="r")
hdr = np.load(os.path.join(DIR, "hdr.npy"))
shots = []
for ff in (120, 121):
    tr = np.where((hdr[:, 0] == ff) & (hdr[:, 1] <= 96))[0]
    x = np.array(data[tr, :2000], dtype=np.float64) * (np.arange(2000) * 0.002) ** 2
    shots.append(panel(agc(x, 125), 337, 420))
save(sheet(shots, ["SP 109", "SP 110 (\"vented\")"], 337, 420), "seismic-shots.gif")

# ours (edited stack, Stolt) against the 1981 stack, 0-4 s
_, old = load(os.path.join(DIR, "31_81", "PROCESSED", "31_81_PR.SGY"), ns=1501)
old = old[:, :1000]
ours = su("stack2.su")[:, :2000:2]
stolt = su("mig_stolt.su")[:, :2000:2]
# The archive's stack is stored in the same CDP order and polarity as ours; it sits 24 ms earlier
# (best lag of the summed cross-correlation, 0.60; zero lag -0.15), which the display leaves alone.
save(sheet([panel(ours, 222, 400), panel(stolt, 222, 400), panel(old, 222, 400)],
           ["ours, stacked", "ours, Stolt migration", "TI ASC, 1981"], 222, 400), "seismic-compare.gif")

# spread as logged against reversed: CDP 1-300, 0.9-1.7 s (the dipping beds)
win = lambda d: agc(d[:300, 450:850], 100)
save(sheet([panel(win(su("stack.su")), 337, 260), panel(win(su("stack_rev.su")), 337, 260)],
           ["group 1 at the north end (as logged)", "groups reversed"], 337, 260), "seismic-reversed.gif")

# three migrations, CDP 150-350, 1.4-2.6 s
win = lambda d: d[150:350, 700:1300]
save(sheet([panel(win(su(f)), 222, 300) for f in ("mig_stolt.su", "mig_ps.su", "mig_tk.su")],
           ["Stolt, 33 s", "phase-shift, 868 s", "Kirchhoff, 274 s"], 222, 300), "seismic-migs.gif")

# the scans: a piece of the observer's log, and the whole 1981 paper section
log = Image.open(os.path.join(DIR, "o2crop.png")).convert("L")
save(log.resize((680, int(log.size[1] * 680 / log.size[0])), Image.LANCZOS), "seismic-log.gif")
sec = Image.open(os.path.join(DIR, "31_81", "IMAGES", "31_81_IM.JPG")).convert("L")
save(sec.resize((680, int(sec.size[1] * 680 / sec.size[0])), Image.LANCZOS), "seismic-1981.gif")

# explainer: one shot heard along the spread, and one midpoint heard from several shot-receiver pairs
def diagram(w=337, h=222, sy=40, ly=190):
    img = Image.new("L", (w, h), 255)
    d = ImageDraw.Draw(img)
    d.rectangle((0, 0, w, sy), fill=235)                  # air
    d.line((0, sy, w, sy), fill=0, width=2)               # ground
    d.line((0, ly, w, ly - 18), fill=0, width=2)          # a gently dipping layer boundary
    d.line((0, 110, w, 104), fill=150, width=1)           # a shallower one, faint
    for x in range(12, w, 22):                            # geophone groups
        d.polygon([(x - 4, sy), (x + 4, sy), (x, sy - 7)], fill=0)
    return img, d

def depth(x, w=337, ly=190):
    return ly - 18 * x / w

def star(d, x, y):
    d.ellipse((x - 5, y - 5, x + 5, y + 5), fill=0)

a, d = diagram()
sx = 166
star(d, sx, 52)
d.line((sx, 40, sx, 52), fill=0)
for rx in (34, 100, 232, 298):                            # rays shot -> boundary -> receiver
    mx = (sx + rx) / 2
    d.line((sx, 52, mx, depth(mx)), fill=0)
    d.line((mx, depth(mx), rx, 40), fill=0)
d.text((4, 4), "one shot, 96 groups listening", fill=0, font=FONT)
d.text((sx + 8, 52), "shot", fill=0, font=FONT)
d.text((6, 196), "a boundary between rock layers", fill=0, font=FONT)

b, d = diagram()
mx = 166
for off in (40, 90, 140):                                 # three pairs sharing one midpoint
    s, r = mx - off, mx + off
    star(d, s, 46)
    d.line((s, 46, mx, depth(mx)), fill=0)
    d.line((mx, depth(mx), r, 40), fill=0)
d.ellipse((mx - 4, depth(mx) - 4, mx + 4, depth(mx) + 4), outline=0, width=2)
d.text((4, 4), "different shots, same midpoint", fill=0, font=FONT)
d.text((mx - 60, depth(mx) + 10), "common midpoint", fill=0, font=FONT)
save(sheet([a, b], ["", ""], 337, 222, top=0), "seismic-how.gif")