#!/usr/bin/env python3
"""final_arcmix_pool.py — pula koncowych runow = caly arcmix (e87f594f) bez dokumentow z trafieniem wiki test/val lub ARC val/test
(skany scan_arcmix_valall_v6 + scan_arcmix_test_v6, ta sama regula co pula C), BEZ losowania, kolejnosc dokumentow zachowana.
Kompaktowanie W MIEJSCU (dysk 30 GB nie miesci dwoch kopii 18.8 GB): zapis zawsze <= odczyt, potem truncate. sha calosci w manifescie.
Uzycie: final_arcmix_pool.py ARCMIX.bin arcmix_index.npz arc_id_split.json scan_valall.json scan_test.json OUT.json [EXTRA.json EXP_REM EXP_TOK]
EXTRA.json (opcjonalnie): {"arcmix_doc_idx":[...]} z nc_scan.py (regula: tekst zawiera CC BY-NC-SA), wtedy asserty rem/tokenow z argumentow."""
import numpy as np, json, hashlib, os, sys
binp, idxp, splp, sv, stst, outj = sys.argv[1:7]
extra = sys.argv[7] if len(sys.argv)>7 else None
EXP_REM = int(sys.argv[8]) if extra else 4362
EXP_TOK = int(sys.argv[9]) if extra else 9_401_283_628
h=hashlib.sha256()
with open(binp,'rb') as f:
    for ch in iter(lambda: f.read(1<<24), b''): h.update(ch)
assert h.hexdigest().startswith('e87f594f'), 'arcmix sha'
idx=np.load(idxp); st,ln=idx['starts'],idx['lens']; split=json.load(open(splp))
assert int(st[-1]+ln[-1])*2==os.path.getsize(binp)
assert st[0]==0 and bool(np.all(st[1:]==st[:-1]+ln[:-1])), 'indeks nieciagly'
rem=set()
for f in (sv, stst):
    for off,n,src in json.load(open(f))['all_hits']:
        if any(k.startswith('wiki:') or (k.startswith('arc') and split.get(k.split(':',1)[1]) in ('validation','test')) for k in src):
            rem.add(int(np.searchsorted(st, off, 'right')-1))
nscan=len(rem)
if extra: rem |= set(int(i) for i in json.load(open(extra))['arcmix_doc_idx'])
assert len(rem)==EXP_REM, ('usuniete dok.', len(rem), EXP_REM)
keep=np.ones(len(st),bool); keep[sorted(rem)]=False
a=np.memmap(binp,dtype=np.uint16,mode='r+'); w=0; sha=hashlib.sha256(); nd=0
for i in range(len(st)):
    if not keep[i]: continue
    s,l=int(st[i]),int(ln[i])
    if s!=w: a[w:w+l]=a[s:s+l]      # w <= s zawsze, wiec nie nadpisujemy nieprzeczytanych danych
    sha.update(a[w:w+l].tobytes()); w+=l; nd+=1
a.flush(); del a
if w!=EXP_TOK: print('STOP: tokenow', w, '!=', EXP_TOK, 'bez truncate i manifestu', flush=True); sys.exit(4)
os.truncate(binp, w*2)
man={'pool':'final_arcmix','src':'arcmix train.bin (e87f594f)','removed_docs':len(rem),'removed_by_scan':nscan,'removed_by_nc_sa_rule':len(rem)-nscan,'extra':os.path.basename(extra) if extra else None,'docs':nd,'tokens':int(w),'sha256':sha.hexdigest(),
     'rule':'removed: wiki test/val and ARC validation/test hits (scan_arcmix_valall_v6 + scan_arcmix_test_v6)'+('; documents whose text contains CC BY-NC-SA (nc_scan)' if extra else '')+'; no sampling; order kept'}
json.dump(man, open(outj,'w'), indent=1); print('FINAL-POOL', json.dumps(man), flush=True)