#!/usr/bin/env python3
"""General-purpose compressor baselines for tick files (see README.md).

usage: baselines.py ZSTD_BINARY TMP_DIR file.ticks:digits ...
Prints bytes/tick for xz -9e and zstd -19 --long on
  raw24              the raw 24-byte records,
  delta-rows12       per-tick (dt, dbid, dask) as 3 x int32,
  delta-varint-cols  columnar zig-zag varints of dt, dbid, dask.
"""
import struct, sys, subprocess, os
def varint(x, out):
    while True:
        b = x & 0x7f; x >>= 7
        if x: out.append(b | 0x80)
        else: out.append(b); return
def zz(x): return (x << 1) ^ (x >> 63) if x < 0 else x << 1
ZSTD=sys.argv[1]
tmp=sys.argv[2]
for p, dg in [(a.split(':')[0], int(a.split(':')[1])) for a in sys.argv[3:]]:
    d = open(p, 'rb').read(); n = len(d)//24; s = 10**dg
    T = [struct.unpack_from('<qdd', d, i*24) for i in range(n)]
    B = [round(t[1]*s) for t in T]; A = [round(t[2]*s) for t in T]; TM = [t[0] for t in T]
    st = [bytearray(), bytearray(), bytearray()]
    for i in range(1, n):
        varint(TM[i]-TM[i-1], st[0]); varint(zz(B[i]-B[i-1]), st[1]); varint(zz(A[i]-A[i-1]), st[2])
    col = bytes(st[0]+st[1]+st[2])
    rows = bytearray()
    for i in range(1,n):
        rows += struct.pack('<Iii', TM[i]-TM[i-1], B[i]-B[i-1], A[i]-A[i-1])
    res = {}
    for name, blob in [('raw24', d), ('delta-rows12', bytes(rows)), ('delta-varint-cols', col)]:
        fn = os.path.join(tmp, 'b.bin'); open(fn,'wb').write(blob)
        xz = len(subprocess.run(['xz','-9e','-c',fn],capture_output=True).stdout)
        zs = len(subprocess.run([ZSTD,'-19','--long=27','-q','-c',fn],capture_output=True).stdout)
        res[name] = (xz/n, zs/n)
    print(os.path.basename(p), ' '.join('%s: xz=%.3f zstd=%.3f'%(k,v[0],v[1]) for k,v in res.items()), 'B/tick')