# Figures for zond6.n4ze3m.com, drawn from the same tone tracks as zond6.py (correct time base).
# usage: python make_figures.py <zond6.mp3> <img_dir> <summary.json>
import json
import os
import sys

import numpy as np
from PIL import Image, ImageDraw, ImageFont

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import zond6 as z

x = z.load(sys.argv[1])
L = z.tone_tracks(x)
S = json.load(open(sys.argv[3]))
SLOT = S['slot_s']
CELL = SLOT / 8.5
FIRST = CELL / 2 + 0.003          # first full cell starts ~half a cell after the marker onset
DT = z.DT
TR, CO = '/System/Library/Fonts/Supplemental/Times New Roman.ttf', '/System/Library/Fonts/Supplemental/Courier New.ttf'
F, C = (lambda s: ImageFont.truetype(TR, s)), (lambda s: ImageFont.truetype(CO, s))
COL = {'mark': (30, 90, 200), 'sync': (20, 140, 70), 'bit': (200, 30, 30), 'base': (200, 200, 200), 'none': (245, 245, 245)}


def state(i):
    for k, thr in (('mark', 0.8), ('sync', 0.8), ('bit', 0.8)):
        if L[k][i] > thr:
            return k
    return 'base' if L['base'][i] > 0.4 else 'none'


def marker_near(t):
    """Exact onset of the data marker closest to tape time t."""
    i0, i1 = int((t - 0.3) / DT), int((t + 0.3) / DT)
    on = L['mark'][i0:i1] > 1.0
    j = np.flatnonzero(np.diff(on.astype(int)) == 1)
    return (i0 + j[0] + 1) * DT if len(j) else t


def strip(dr, t0, x0, y0, w, h, label, word):
    for j in range(w):
        dr.line([(x0 + j, y0), (x0 + j, y0 + h)], fill=COL[state(int((t0 + j / w * SLOT) / DT))])
    dr.rectangle([x0, y0, x0 + w, y0 + h], outline=(0, 0, 0))
    for k in range(9):
        X = x0 + int((FIRST + k * CELL) / SLOT * w)
        if X <= x0 + w:
            dr.line([(X, y0 + h), (X, y0 + h + 8)], fill=(0, 0, 0))
        if k < 8:
            Xm = x0 + int((FIRST + (k + 0.5) * CELL) / SLOT * w)
            dr.text((Xm - 8, y0 + h + 10), f'b{k}', font=C(15), fill=(0, 0, 0))
            dr.text((Xm - 6, y0 - 26), word[k], font=C(20), fill=(200, 30, 30) if word[k] == '1' else (0, 0, 0))
    dr.text((x0, y0 - 52), label, font=F(19), fill=(0, 0, 0))


W, H = 1100, 430
im = Image.new('RGB', (W, H), (255, 255, 255))
dr = ImageDraw.Draw(im)
strip(dr, marker_near(175.0), 40, 70, 1000, 60, 'Tape 2:55, about 03:26 UT, 14 Nov 1968 (just after the Moon): 11010000', '11010000')
strip(dr, marker_near(28.78), 40, 250, 1000, 60, 'Tape 0:29, about 01:51 UT (approaching the Moon): 11111000', '11111000')
dr.text((40, H - 40), f'blue = data marker, red = 978 Hz pulse, grey = 700 Hz base tone, pale = quiet. One strip = one second '
        f'({SLOT:.3f} s on this tape). Ticks = {CELL * 1000:.1f} ms cells.', font=F(15), fill=(60, 60, 60))
im.save(f'{sys.argv[2]}/one-second.png')


def raster(t0, t1, title):
    on = L['mark'] > 1.0
    tt = np.arange(len(on)) * DT
    r = tt[1:][np.diff(on.astype(int)) == 1]
    sy = tt[1:][np.diff((L['sync'] > 1.0).astype(int)) == 1]
    M = np.sort(np.r_[r, sy]); M = M[(M >= t0) & (M <= t1)]
    k = np.round((M - M[0]) / SLOT); good = np.ones(len(M), bool)
    for _ in range(6):
        b, a = np.polyfit(k[good], M[good], 1); res = M - (a + b * k); new = np.abs(res) < 0.05
        if (new == good).all():
            break
        good = new
    first, last = int(np.ceil((t0 - a) / b)), int((t1 - a) / b)
    w, rowh, x0 = 760, 9, 110
    img = Image.new('RGB', (900, 60 + (last - first) * rowh + 10), (255, 255, 255)); g = ImageDraw.Draw(img)
    g.text((10, 8), title, font=F(19), fill=(0, 0, 0))
    for i in range(first, last):
        s = a + i * b; y = 50 + (i - first) * rowh
        g.text((10, y - 3), f'{int(s // 60)}:{s % 60:04.1f}', font=C(11), fill=(80, 80, 80))
        for j in range(w):
            g.line([(x0 + j, y), (x0 + j, y + rowh - 2)], fill=COL[state(int((s + j / w * b) / DT))])
    return img


p1 = raster(10, 63, 'Approaching the Moon (tape 0:10-1:03, about 01:51 UT). One row = one second.')
p2 = raster(135, 212, 'Just out from behind the Moon (tape 2:15-3:32, about 03:25 UT). One row = one second.')
out = Image.new('RGB', (900, p1.size[1] + p2.size[1] + 20), (255, 255, 255))
out.paste(p1, (0, 0)); out.paste(p2, (0, p1.size[1] + 20))
out.save(f'{sys.argv[2]}/raster.png')
print('figures saved')