import csv
import gzip
import io
import struct
import sys
def main(src: str, dst: str) -> None:
particles: dict[str, int] = {}
neurons: dict[str, int] = {}
out = bytearray()
rows = 0
with gzip.open(src, "rb") as raw:
reader = csv.reader(io.TextIOWrapper(raw, encoding="utf-8"))
header = next(reader)
col = {name: i for i, name in enumerate(header)}
fi, ti, ni, hi = (col[k] for k in ("particle_from", "particle_to", "neuron", "height"))
for row in reader:
f = particles.setdefault(row[fi], len(particles))
t = particles.setdefault(row[ti], len(particles))
n = neurons.setdefault(row[ni], len(neurons))
out += struct.pack("<IIII", f, t, n, int(row[hi]))
rows += 1
with open(dst, "wb") as fh:
fh.write(b"TRUL" + struct.pack("<III", 1, rows, len(particles)))
fh.write(out)
print(f"{dst}: {rows} links, {len(particles)} particles, {len(neurons)} neurons")
if __name__ == "__main__":
if len(sys.argv) != 3:
sys.exit("usage: links_bin.py <cyberlinks_indexed.csv.gz> <out.links>")
main(sys.argv[1], sys.argv[2])