# AUTOGENERATED FILE! PLEASE DON'T EDIT
from k1lib.imports import *
def getIr(base): return None | applyCl.aS(lambda: ls(base) | iden() & apply(os.path.getsize) | transpose() | deref()) | ungroup(begin=True) | insertIdColumn(True) | deref()
def normalize(d):
    d = d | deref(); s = d | cut(1) | toSum()
    return d | apply(op()/s, 1) | sort(0, False) | deref()
@lru_cache
def statsCpu():
    cpu = None | applyCl.aS(applyCl.cpu) | sort(0, False) | deref()
    cpuF = None | applyCl.aS(applyCl.cpu) | aS(normalize); cpuF # "cpuF" = cpu fraction. List[nodeId, cpu fraction]
    return [cpu, cpuF]
@lru_cache
def statsNodeId(): return applyCl.nodeIds()
_statsS1 = statsNodeId() | apply(wrapList() | insert(0, False)) | toDict()
def stats(ir):
    cpu, cpuF = statsCpu()
    # size fraction. List[nodeId, size fraction]
    sizeF = normalize({**_statsS1, **ir | groupBy(1, True) | apply(cut(2) | toSum(), 1) | toDict()}.items()); sizeF
    return *statsCpu(), sizeF
def scores(ir): cpu, cpuF, sizeF = stats(ir); return [cpuF, sizeF] | cut(1).all() | transpose() | ~apply(lambda c,s: s-c) | deref() # negpos
def score(ir): return scores(ir) | apply(lambda x: abs(x)**2) | toSum() # pos
def move(ir, nA:str, nB:str, idx:int):
    ir1 = ir | deref(); ir1[idx][1] = nB; a = ir1 | pretty() | join('\n'); return ir1
def optimize(ir):
    cpu, cpuF, sizeF = stats(ir); scs = scores(ir); a = np.argmax(scs); b = np.argmin(scs)
    # print(scs, a, b)
    nodeIds = applyCl.nodeIds() | sort(None, False) | deref()
    files = {**nodeIds | apply(wrapList() | insert([], False)) | toDict(), **ir | groupBy(1, True) | toDict()}.items() | sort(0, False) | deref()
    fA = files[a]; fB = files[b] # files A. [nodeId, List[idx, url, size]]
    # print(fA); print(fB)
    nA = nodeIds[a]; nB = nodeIds[b] # A node id
    sA, sB = [fA, fB] | apply(op()[1] | cut(2) | deref() | aS(np.array, dtype=int)) # file sizes A
    # print(f"sA: {sA} {sB}")
    spA = sA.sum() - sA; spB = sB.sum() + sA # sum prime A, array[files]
    sp = spA + spB; sfA = spA/sp; sfB = spB/sp
    cA = cpu[a][1]; cB = cpu[b][1] # cpu A
    c = cA + cB; cfA = cA/c; cfB = cB/c # cpu fraction A
    # print(f"sfA: {sfA}, cfA: {cfA}")
    exp = 5 # intuition says that exp should be even. But that doesn't work. Odd values work tho, but I have no idea why
    idx = fA[1][((sfA-cfA)**exp + (sfB-cfB)**exp).argmin()][0]
    ir2 = move(ir, nA, nB, idx)
    return ir2, [nA, nB, idx, score(ir2)]
def traj(ir, maxSteps=20):
    sc = score(ir); aux = None; auxs = []; irs = []
    for i in range(maxSteps):
        ir, aux = optimize(ir)
        if aux[3] > sc:
            #irs.append(ir); auxs.append(aux)
            break
        irs.append(ir); auxs.append(aux); sc = aux[3]
    return irs, auxs
def collapse(it):
    a, b = it | rows(0, -1)
    c = [a[0], b[1], a[2], b[3]]
    if c[0] == c[1]: return []
    return [c]
def traj2(ir):
    idx2FileName = ir | apply(lambda arr: [arr[0], arr[2]]) | toDict()
    a = traj(ir)[1] | groupBy(2) | filt(lambda x: len(x) > 1).split() | (apply(collapse)) + iden() | joinStreams(2) | deref()
    return a | lookup(idx2FileName, 2) | deref()
def moveFile(fileName:str, destNodeId:str): # executed on remotes
    fn = os.path.expanduser(fileName); dirname = os.path.dirname(fn)
    [destNodeId] | applyCl.aS(lambda: None | cmd(f"mkdir -p {dirname}; rm {fn}") | deref()) | deref()
    for chunk in cat(fn, False, True): [destNodeId] | applyCl.aS(lambda: chunk >> file(fn))
    None | cmd(f"rm {fn}") | deref()
def balance(base): traj2(getIr(base)) | apply(lambda arr: [arr[0], arr]) | ~applyCl(lambda a,b,fn,sc: moveFile(fn, b), pre=True) | deref()