#!/usr/bin/env python3
import MCPL as mcpl
import pathlib,os

def files_for_testing():
    for p in pathlib.Path(os.environ['SBLD_DATA_DIR']).glob('MCPLTests*/*.mcpl*'):
        if 'bad' in p.name or 'truncated' in p.name or 'crash' in p.name:
            continue
        yield p

for path in sorted(list(files_for_testing())):
    with mcpl.MCPLFile(path) as f:
        nparticles = f.nparticles
    for blocklength in sorted({max(1,nparticles),nparticles+1,max(1,nparticles-1)}):
        print('\n'*3+'='*100+'\n'*3)
        print ('Testing %s (nparticles=%i) with blocklength=%i'%(os.sep.join(path.parts[-2:]),nparticles,blocklength))
        with mcpl.MCPLFile(path,blocklength=blocklength) as f:
            #print (path,blocklength)
            assert f.blocklength == blocklength
            for p in f.particle_blocks:
                lens = [ len(p.x),len(p.y),len(p.z),
                         len(p.ux),len(p.uy),len(p.uz),
                         len(p.polx),len(p.poly),len(p.polz),
                         len(p.pdgcode),len(p.weight),len(p.userflags),
                         len(p.position),len(p.polarisation),len(p.direction),
                         len(p.time),len(p.ekin) ]
                #print (lens)
                assert len(set(lens))==1
                assert lens[0]<=blocklength
                assert blocklength>=1
            if nparticles>0:
                s=mcpl.collect_stats(f)
                mcpl.dump_stats(s)
