use BlockDist;
config var numTasks:int;
config var size:int;

iter BlockDom.withGranularity(userMaxTasks: int = 1) {
    for i in whole do yield i;
}

iter BlockDom.withGranularity(param tag: iterKind, userMaxTasks: int = 1) where tag == iterKind.leader {
    const maxTasks = userMaxTasks;
    const ignoreRunning = dist.dataParIgnoreRunningTasks;
    const minSize = dist.dataParMinGranularity;
    const wholeLow = whole.lowBound;

    // If this is the only task running on this locale, we don't want to
    // count it when we try to determine how many tasks to use. Here we
    // check if we are the only one running, and if so, use ignoreRunning=true 
    // for this locale only.  Obviously there's a bit of a race condition 
    // if some other task starts after we check, but in that case there is 
    // no correct answer anyways.
    //
    // Note that this code assumes that any locale will only be in the
    // targetLocales array once.  If this is not the case, then the
    // tasks on this locale will *all* ignoreRunning, which may have
    // performance implications.
    const hereId = here.id;
    const hereIgnoreRunning = if here.runningTasks() == 1 then true else ignoreRunning;
    
    coforall locDom in locDoms do on locDom {
        const myIgnoreRunning = if here.id == hereId then hereIgnoreRunning else ignoreRunning;
        
        // Use the internal function for untranslate to avoid having to do
        // extra work to negate the offset
        const tmpBlock = locDom.myBlock.chpl__unTranslate(wholeLow);
        var locOffset: rank*idxType;
        
        for param i in 0..tmpBlock.rank - 1 {
            const dim = tmpBlock.dim(i);
            const aStr = if dim.hasPositiveStride() then dim.stride else -dim.stride;
            locOffset(i) = dim.low / aStr:idxType;
        }

        writeln("Max tasks given on ", here.id , " = ", userMaxTasks);
        
        // Forward to defaultRectangular
        for followThis in tmpBlock.these(iterKind.leader, maxTasks, myIgnoreRunning, minSize, locOffset) do
            yield followThis;
    }
}

iter BlockDom.withGranularity(param tag: iterKind, userMaxTasks: int = 1, followThis) where tag == iterKind.follower {
    if chpl__testParFlag then chpl__testParWriteln("Block domain follower invoked on ", followThis);

    const myLocBlockDom = if _privatization then chpl_getPrivatizedCopy(this.type, pid) else this;

    if myLocBlockDom.locale != here then
      halt("Wait, what?!?");
    
    for i in myLocBlockDom.these(tag, followThis) do
      yield i;
}

iter BlockArr.withGranularity(userMaxTasks: int = 1) ref {
    foreach i in dom do yield dsiAccess(i);
}

iter BlockArr.withGranularity(param tag: iterKind, userMaxTasks: int = 1) where tag == iterKind.leader {
    for followThis in dom.withGranularity(tag, userMaxTasks) do yield followThis;
}

iter BlockArr.withGranularity(param tag: iterKind, userMaxTasks: int = 1, followThis, param fast: bool = false) 
                              ref where tag == iterKind.follower {
    const myLocBlockArr = if _privatization then chpl_getPrivatizedCopy(this.type, pid) else this;

    for x in myLocBlockArr.these(tag, followThis, fast) do
      yield x;
}

proc main() {
    var D = blockDist.createDomain({0..size-1});
    var D_localeSpace: domain(1) = {0..numLocales-1};
    var tasksPerLocale = blockDist.createArray(D_localeSpace, domain(int));
    
    writeln("#### TESTS WITH GRANULARITY FOR DOMAINS");
    writeln("The values per locale from here:");
    var tid: atomic int;
    forall d in D.withGranularity(numTasks) with (const myTid = tid.fetchAdd(1)) {
        tasksPerLocale[here.id] += myTid;
    }
    writeln();
    
    writeln("Should match the ones here:");
    coforall loc in Locales with (ref tasksPerLocale) do on loc {
        writeln("tasksPerLocale on ", here.id, " = ", tasksPerLocale[here.id].size);
        tasksPerLocale[here.id].clear();
    }
    writeln();
    writeln();

    writeln("#### TESTS WITH GRANULARITY FOR ARRAYS");
    writeln("The values per locale from here:");
    var A = blockDist.createArray(D, int);
    tid.write(0);
    forall a in A.withGranularity(numTasks) with (const myTid = tid.fetchAdd(1)) {
        tasksPerLocale[here.id] += myTid;
    }
    writeln();
    
    writeln("Should match the ones here:");
    coforall loc in Locales with (ref tasksPerLocale) do on loc {
        writeln("tasksPerLocale on ", here.id, " = ", tasksPerLocale[here.id].size);
        tasksPerLocale[here.id].clear();
    }
}