#!/usr/bin/python3
import gzip, tqdm

halters = set()
with open('bb6.halters.20M.txt') as f:
    for i, line in enumerate(f):
        if i % 2 == 0:
            seed, desc = line.split()
            assert desc.startswith('1RB')
            halters.add(int(seed))

non_bouncers = set()
with open('BB6_100K.umf', 'rb') as idx:
    while (seed := idx.read(4)):
        non_bouncers.add(int.from_bytes(seed, 'big'))

with (open('bbchallenge-deciders/decider-finite-automata-reduction/output/finite_automata_reduction.index', 'rb') as idx,
      gzip.open('halters.txt.gz', 'wt') as hlt, gzip.open('bouncers.txt.gz', 'wt') as bnc, gzip.open('ctl_decided.txt.gz', 'wt') as dec, gzip.open('ctl_undecided.txt.gz', 'wt') as und, gzip.open('bb.6x2.unknown.txt.gz', 'rt') as dsc,
      open('ctl_eligible.index', 'wb') as ctl_index):
    ctl_decided = set()
    while (seed := idx.read(4)):
        ctl_decided.add(int.from_bytes(seed, 'big'))
    for i, line in enumerate(tqdm.tqdm(dsc, unit=' beavers')):
        if i in halters:
            hlt.write(line)
        elif i not in non_bouncers:
            bnc.write(line)
        elif i in ctl_decided:
            dec.write(line)
            ctl_index.write(i.to_bytes(4, 'big'))
        else:
            und.write(line)
            ctl_index.write(i.to_bytes(4, 'big'))
print('Wrote halters.txt.gz, bouncers.txt.gz, ctl_decided.txt.gz, ctl_undecided.txt.gz -- and index ctl_eligible.index as decider input')
