import sys, os, time, random, string
import Vmaps
Vm = Vmaps


# How many iterations of Operate() should be done
TotalIters = 1024 * 1024 * 1

# How many children to fork()
RunJobs = 3

# how much space (in bytes) to operate in
TestSize = 1024 * 8#1024 * 4

# how many columns per array element
TCols=32

# What to name the test file (it is deleted when we finish)
TestFileName = 'xxVmapdata.xxx'


# How many random numbers to generate
RandomSample1 = 1024 * 57
RandomSample2 = 1024 * 13

# -------------------------------------------------
def Operate( amap, rndx, cl ):
    Spun = 0
    while 1:
        dt = amap.atswap(rndx,cl,-1)
        if dt !=-1: break
        Spun +=1
        amap.schyield()
    # endw
    amap.atswap(rndx,cl,dt+1)
    return Spun

# -------------------------------------------------
def Summarize( amap, Workers ):
    Hlen=len(amap)+1
    print len(amap), 'elements in', amap.databytes, 'bytes,', amap.size, 'mapped space'
    print '----------------- sum retype'
    amap.astype(8)
    
    Mlen = len(amap) - Hlen
    todo=(TotalIters / Workers)
    # odd splits will result in "lost" iterations
    xrnd= ( TotalIters - (todo*Workers) ) 

    print Workers, 'kids with',  todo, 'each (',TotalIters,' total', xrnd, 'wksplit)'
    aa=ba=z=hw=c= 0
    avg = (TotalIters/ (Mlen-Hlen)) # one col per element should be 0
    spike=time.time()
    c = amap.sumrange()
    [mini,maxi] = amap.minmax()
    hw=amap[maxi]
    zc=amap.cntbndrange()
    print ("%8.2f" % (time.time()-spike)), 'sec to summarize'
    print 'total=',c
    print 'avg=',avg,
    print 'highwater=',hw, 'numzero=',zc,
    print
    if c != (TotalIters-xrnd):
        print ' OOOPS!! lost ', (TotalIters - c), 'iterations somewhere'
    #endif
    
    return
# -------------------------------------------------
def MakeMap(fd):
    vm_type = Vm.Int2d
    vm_cols = TCols
    vm_flags =  Vm.VM_STAYOPEN 
    mm_flags = Vm.MAP_SHARED | Vm.MAP_NORESERVE
    mm_prot = Vm.PROT_READ | Vm.PROT_WRITE
    amap = Vm.newmap( fd, TestSize, 0, mm_flags, mm_prot, vm_flags, vm_type, vm_cols )
    amap.open()
    #amap.astype(vm_type, vm_cols)
    #amap[0]=1
    #print 'data:', amap[0],amap[(len(amap)-1)]
    return amap
# -------------------------------------------------

# -------------------------------------------------
# -------------------------------------------------
KidSeedRange= 1024 * 1024 * 1024
KidWaitSleep = 0.1
BlockSize= min(1024 * 512, TestSize)
# -------------------------------------------------
def main():
    Spinners=Spun=0
    Workers = RunJobs
    TestFile = TestFileName+str(os.getpid())
    # init file
    fo = open(TestFile,'w+')
    wrote=0
    bk=chr(0) * BlockSize
    while wrote<TestSize:
        fo.write( bk )
        wrote += BlockSize
    # nedwhi
    fo.seek(0)
    fd=fo.fileno()
    # do map
    amap = MakeMap(fd)
    Mlen=len(amap)
    todo = TotalIters / Workers
    todone = todo
    print Mlen, 'elements of', amap.elsize(),'bytes in',
    print amap.databytes, 'bytes,', amap.size, 'mapped space'
    print 'spawning', Workers, 'to do', todo, 'each (',TotalIters,' total)'
    kids=[]
    parent=1
    child=0
    # spawn children
    for n in range(Workers):
        kseed=random.randrange(KidSeedRange)
        child +=1
        pid=os.fork()
        if pid==0:
            # iam child
            parent=0
            break
        else:
            # parent
            kids.append( pid )
        # endif
    # endfor
    # this results in a different random seed for each child
    random.seed( (os.getpid() ^ kseed) )
    if parent:
        # Parent 
        print 'Spawned:', kids
        # parent work happens after this if
    else:
        # Child
        # build random numbers 
        randNdx=[0]
        randNdx2=[0]
        randNdx *= RandomSample1
        randNdx2 *= RandomSample2
        randNdx = map( lambda x: random.randrange(1,Mlen), randNdx )
        randNdx2 = map( lambda x: random.randrange(1,TCols), randNdx2 )
        # checkin (atomic increment of element 0)
        while 1:
            dt = amap.atswap(0,0,-1)
            if dt !=-1: break
            amap.schyield()
        amap.atswap(0,0,dt+1)
        # <------ ---------------------------- kid checked in
        # sleep rather than schyield while waiting for other children
        #  (takes less CPU)
        while amap[0][0] < Workers :
            time.sleep(KidWaitSleep)
            # amap.schyield()
        # <-------------------------------------  sync; this kid may go
        print '+',
        sys.stdout.flush()
        spike=time.time()
        # ACTUAL OPERATION LOOP
        while todo:
            todo -= 1
            el=randNdx[(todo % RandomSample1)]
            cl=randNdx2[(todo % RandomSample2)]
            x=Operate( amap, el,cl )
            if x:
                Spinners +=1
                Spun+=x
            # endif
        # endwhile
        print '++xit',os.getpid(), 'spun', Spun, '\ton', Spinners,
        print '\t', ("%8.2f" % (time.time()-spike)), 'sec'
        sys.stdout.flush()
        # checkout (atomic decrement element 0)
        while 1:
            dt = amap.atswap(0,0,-1)
            if dt !=-1: break
            amap.schyield()
        amap.atswap(0,0,dt-1)
        # <------ ---------------------------- kid checked out
        # wait for the number of children working to hit 0
        #  (or say -2 if you want the parent to do the msync() )
        while amap[0][0]!=0:
            time.sleep(KidWaitSleep)
            # amap.schyield()
        # <-------------------------------------  sync; this kid close
        amap.close()
        # child done
        os._exit(0)
    #endif
    # children don't reach this point.
    #go signal kids
    print '...waiting for sync',
    sys.stdout.flush()
    spike=time.time()
    lcnt=0
    # Simply watch element 0 when the children all run
    while 1:
        cnt=amap[0][0]
        if cnt!=lcnt and cnt>0:
            print cnt,("%5.2f" % (time.time()-spike)), 'sec;  ', 
            sys.stdout.flush()
            lcnt=cnt
        # nedif
        if cnt >= Workers: break
        time.sleep(KidWaitSleep)
    # endw
    print
    print ("%5.2f" % (time.time()-spike)), 'sec; Go!'
    sys.stdout.flush()
    # <------ children are operating; wait for them all to signal done
    print
    spike=time.time()
    while amap[0][0]!=0:
        time.sleep(KidWaitSleep)
        # amap.schyield()
    # <-------------------------------------  sync; all kids done now
    print '---------------------------------- done'
    print 'wait',("%8.2f" % (time.time()-spike)), 'sec'
    spike=time.time()
    amap.raw_msync(Vm.MS_SYNC)
    print 'sync',("%8.2f" % (time.time()-spike)), 'sec'
    # signal children we have done msync()
    amap.atswap(0,0,-2)
    # <-------------------------------------  sync; kids may close and exit
    # wait for children to exit
    while len(kids):
        nk=[]
        for pid in kids:
            (rp,st) = os.waitpid(pid,os.WNOHANG)
            if not rp: nk.append( pid )
        # endfor
        kids=nk
    # endwhile
    # summarize 
    amap.atswap(0,0,0)
    #amap[0] = ([0] * TCols)
    print 'kwait', ("%8.2f" % (time.time()-spike)), 'sec'
    # print repr(amap[0:4])
    Summarize( amap, Workers )
    # cleanup
    amap.close()
    fo.close()
    os.unlink(TestFile)
    return
# -------------------------------------------------
# -------------------------------------------------
# -------------------------------------------------
if '__main__'==__name__: main()