"""Read the numbers behind three figures of Vasilyev et al. (2024) from the vector PDFs in their arXiv
source (2412.12265v1): the fraction of injected flares their search detected, by flare energy and
stellar temperature (detectability_of_flares.pdf); the average profile of the flares they found,
which they injected to measure it (mean_flare_profile.pdf); and their full sample's flare frequency
by energy before and after their correction for the flares their search misses
(pdf_flares_for_paper.pdf).

Each figure's content stream is interpreted (transformations, stroke colors, paths, text, and the
images and glyphs it places), its axes are calibrated from the tick marks and their labels, each
detectability curve's energy is read from where its color falls on the color bar's image, and the
frequency figure's two histograms are found through its legend.

Run once, outside the analysis's container, with pypdf 5.1.0:
python read_figures.py <arXiv source tarball> <output dir>
Writes vasilyev_detectability.csv (log_e, teff, r), vasilyev_profile.csv (cadence, flux), and
vasilyev_frequency.csv (log_e_lo, log_e_hi, observed, corrected: log10 flares per star per year per
erg).
"""
import io
import os
import sys
import tarfile

import numpy as np
from pypdf import PdfReader
from pypdf.generic import ContentStream


def matrix(a, b, c, d, e, f):
    return np.array([[a, b, 0.0], [c, d, 0.0], [e, f, 1.0]])


def drawing(pdf):
    """The page's stroked paths (page coordinates, with stroke color), text strings (with the
    position each starts at), images (with their transformation), and placed forms such as
    markers and glyphs (name and position)."""
    reader = PdfReader(io.BytesIO(pdf))
    page = reader.pages[0]
    xobjects = page['/Resources']['/XObject']
    ctm, color, stack, tm = np.eye(3), None, [], np.eye(3)
    subpaths, strokes, texts, images, forms = [], [], [], [], []
    for operands, op in ContentStream(page.get_contents(), reader).operations:
        op = op.decode()
        v = [float(x) for x in operands] if op in ('cm', 'RG', 'G', 'm', 'l', 'Td') else None
        if op == 'q':
            stack.append((ctm, color))
        elif op == 'Q':
            ctm, color = stack.pop()
        elif op == 'cm':
            ctm = matrix(*v) @ ctm
        elif op == 'RG':
            color = tuple(v)
        elif op == 'G':
            color = (v[0],) * 3
        elif op in ('m', 'l'):
            p = (np.array([v[0], v[1], 1.0]) @ ctm)[:2]
            if op == 'm':
                subpaths.append([p])
            else:
                subpaths[-1].append(p)
        elif op in ('S', 'B'):
            strokes += [(color, np.array(s)) for s in subpaths if len(s) > 1]
            subpaths = []
        elif op in ('f', 'n'):
            subpaths = []
        elif op == 'BT':
            tm = np.eye(3)
        elif op == 'Td':
            tm = matrix(1, 0, 0, 1, v[0], v[1]) @ tm
        elif op in ('Tj', 'TJ'):
            parts = operands[0] if op == 'TJ' else [operands[0]]
            text = ''.join(str(x) for x in parts if not isinstance(x, (int, float)))
            texts.append((text.strip(), *(np.array([0.0, 0.0, 1.0]) @ tm @ ctm)[:2]))
        elif op == 'Do':
            obj = xobjects[operands[0]].get_object() if operands[0] in xobjects else None
            if obj is not None and obj.get('/Subtype') == '/Image':
                images.append((obj, ctm.copy()))
            else:
                forms.append((str(operands[0]), ctm[2, 0], ctm[2, 1]))
    return strokes, texts, images, forms


def number(text):
    try:
        return float(text)
    except ValueError:
        return None


def axis(strokes, labels, along, line):
    """pixel -> value along one axis, from labels [(position, value)] that sit in a row ('x') or a
    column ('y') at coordinate `line`. Tick marks are short segments across the axis; those whose
    nearer end lies closest to the labels sit on the labelled spine, the longest of them are the
    major ticks, and each label belongs to the major tick nearest it."""
    k = 0 if along == 'x' else 1   # the coordinate that places a tick mark
    ticks = [s for _, s in strokes if len(s) == 2 and abs(s[0, k] - s[1, k]) < 1e-6
             and 0 < abs(s[0, 1 - k] - s[1, 1 - k]) < 10]
    gap = np.array([np.abs(s[:, 1 - k] - line).min() for s in ticks])
    ticks = [s for s, g in zip(ticks, gap) if g < gap.min() + 1e-3]
    length = np.array([abs(s[0, 1 - k] - s[1, 1 - k]) for s in ticks])
    major = np.array(sorted(s[0, k] for s, n in zip(ticks, length) if n > length.max() - 1e-3))
    at = [major[np.argmin(np.abs(major - p))] for p, _ in labels]
    if len(set(at)) != len(at):
        raise ValueError(f'{along}: two labels share a tick')
    values = np.array([v for _, v in labels])
    fit = np.polyfit(at, values, 1)
    if np.max(np.abs(np.polyval(fit, at) - values)) > 1e-6 * np.ptp(values):
        raise ValueError(f'{along}: tick labels not evenly spaced')
    return lambda p: np.polyval(fit, p)


def linear_axis(strokes, texts, along):
    """An axis labelled with plain numbers: 'x' (the lowest row of numbers), 'y' (the leftmost
    column), or 'bar' (the rightmost column)."""
    labels = [(t, x, y) for t, x, y in texts if number(t) is not None]
    if along == 'x':
        line = min(y for _, _, y in labels)
        return axis(strokes, [(x, number(t)) for t, x, y in labels if abs(y - line) < 1], 'x', line)
    line = (min if along == 'y' else max)(x for _, x, _ in labels)
    return axis(strokes, [(y, number(t)) for t, x, y in labels if abs(x - line) < 1], 'y', line)


def log_axis(strokes, texts, forms, along):
    """An axis labelled as matplotlib labels a logarithmic one, 10 with a raised exponent, glyph by
    glyph: a base of "1" and "0" on one baseline and, just right of it and raised, the exponent's
    digits, after the glyph named minus when it is negative. Values are log10; 'x' takes the lowest
    row of labels and 'y' the leftmost column."""
    glyphs = [(x, y, t) for t, x, y in texts if t.isdigit()]
    glyphs += [(x, y, '-') for name, x, y in forms if name.endswith('-minus')]
    labels = []
    for x1, y1, t in glyphs:
        zero = [x for x, y, u in glyphs if t == '1' and u == '0' and abs(y - y1) < 0.5 and 0 < x - x1 < 10]
        if not zero:
            continue
        raised = sorted((x, u) for x, y, u in glyphs if 0 < y - y1 < 12 and 0 < x - min(zero) < 30)
        if raised:
            labels.append((x1, y1, (x1 + raised[-1][0]) / 2, int(''.join(u for _, u in raised))))
    if along == 'x':
        line = min(y for _, y, _, _ in labels)
        return axis(strokes, [(c, v) for x, y, c, v in labels if abs(y - line) < 1], 'x', line)
    line = min(x for x, _, _, _ in labels)
    return axis(strokes, [(y, v) for x, y, c, v in labels if abs(x - line) < 1], 'y', line)


def detectability(pdf):
    strokes, texts, images, _ = drawing(pdf)
    to_teff, to_r = linear_axis(strokes, texts, 'x'), linear_axis(strokes, texts, 'y')
    to_bar = linear_axis(strokes, texts, 'bar')
    (bar, place), = images
    width, height = int(bar['/Width']), int(bar['/Height'])
    _, _, _, palette = bar['/ColorSpace']
    palette = np.frombuffer(palette, np.uint8).reshape(-1, 3) / 255.0
    rows = np.frombuffer(bar.get_data(), np.uint8).reshape(height, width)[:, width // 2]
    row_y = place[2, 1] + place[1, 1] * (1 - (np.arange(height) + 0.5) / height)  # image rows run top down
    out = []
    for c, s in strokes:
        dist = np.linalg.norm(palette[rows] - np.array(c), axis=1)
        if len(s) < 3 or dist.min() > 0.01:   # not a curve drawn in the color bar's colors
            continue
        log_e = float(to_bar(row_y[dist <= dist.min() + 1e-9].mean()))
        out += [(log_e, float(to_teff(x)), float(to_r(y))) for x, y in s]
    return out


def profile(pdf):
    strokes, texts, _, _ = drawing(pdf)
    to_flux = linear_axis(strokes, texts, 'y')
    line = max((s for _, s in strokes if len(s) > 2), key=len)
    flux = to_flux(line[:, 1])
    step = np.median(np.diff(line[:, 0]))
    at = np.argmax(flux)                     # the profile is normalized at its start
    offset = (line[:, 0] - line[at, 0]) / step
    keep = np.abs(offset - np.round(offset)) < 1e-3   # drops the line's end clipped at the frame
    return [(int(round(c)), float(f)) for c, f, k in zip(offset, flux, keep) if k]


def frequency(pdf, observed='Full sample', corrected='Corrected full sample'):
    """Each energy bin of the two histograms the legend names: (log_e_lo, log_e_hi, observed,
    corrected), the frequencies as log10."""
    strokes, texts, _, forms = drawing(pdf)
    to_log_e, to_log_f = log_axis(strokes, texts, forms, 'x'), log_axis(strokes, texts, forms, 'y')

    def histogram(name):
        (_, tx, ty), = [t for t in texts if t[0] == name]
        key = [c for c, s in strokes if len(s) == 2 and abs(s[0, 1] - s[1, 1]) < 1e-6 and c != (0.0, 0.0, 0.0)
               and 0 < tx - s[:, 0].max() < 20 and abs(s[0, 1] - ty) < 8]
        (color,) = set(key)
        (path,) = [s for c, s in strokes if c == color and len(s) > 10]
        return {(round(float(to_log_e(a[0])), 4), round(float(to_log_e(b[0])), 4)): float(to_log_f(a[1]))
                for a, b in zip(path[:-1], path[1:]) if abs(a[1] - b[1]) < 1e-6 and b[0] > a[0]}

    seen, fixed = histogram(observed), histogram(corrected)
    return [(lo, hi, f, fixed.get((lo, hi), np.nan)) for (lo, hi), f in sorted(seen.items())]


def main():
    source, out = sys.argv[1:3]
    with tarfile.open(source) as tf:
        read = lambda name: tf.extractfile(name).read()  # noqa: E731
        det = detectability(read('detectability_of_flares.pdf'))
        prof = profile(read('mean_flare_profile.pdf'))
        freq = frequency(read('pdf_flares_for_paper.pdf'))
    with open(os.path.join(out, 'vasilyev_detectability.csv'), 'w') as fh:
        fh.write('log_e,teff,r\n')
        fh.writelines(f'{log_e:.4f},{teff:.1f},{r:.5f}\n' for log_e, teff, r in sorted(det))
    with open(os.path.join(out, 'vasilyev_profile.csv'), 'w') as fh:
        fh.write('cadence,flux\n')
        fh.writelines(f'{c},{f:.5f}\n' for c, f in prof)
    with open(os.path.join(out, 'vasilyev_frequency.csv'), 'w') as fh:
        fh.write('log_e_lo,log_e_hi,observed,corrected\n')
        fh.writelines(f'{lo:.4f},{hi:.4f},{f:.4f},{"" if np.isnan(g) else f"{g:.4f}"}\n' for lo, hi, f, g in freq)


if __name__ == '__main__':
    main()