text:
import numpy as np
import numpy.matlib
from numpy import linalg as la
import sklearn.cluster as cluster
import probability_functions as prob
from joblib import Parallel, delayed
import multiprocessing
import coordinate_conversion as coconv

pi = np.pi

def find_maxima(cnn_response):

    maxima = np.zeros(cnn_response.shape)

    A = cnn_response.shape[1]
    B = cnn_response.shape[0]

    for b in range(B):
        for a in range(A):
            vm = cnn_response[b, a]
            vu = cnn_response[b, a+1] if a+1 < A else 0
            vd = cnn_response[b, a-1] if a-1 > 0 else 0
            vl = cnn_response[b-1, a] if b-1 > 0 else 0
            vr = cnn_response[b+1, a] if b+1 < B else 0

            if vm > vu and vm > vd and vm > vl and vm > vr:
                maxima[b, a] = 1

    return maxima

def line_rating_knn(lp, k1=10, k2=3, sigma=1):

    N = lp.shape[0]

    lscore = np.zeros(N)

    k1 = np.minimum(k1, N)
    k2 = np.minimum(k2, N)

    num_cores = multiprocessing.cpu_count()
    ldist = Parallel(n_jobs=num_cores)(delayed(calc_ldist_parfun)(i, lp) for i in range(N))
    ldist = np.stack(ldist)

    ldist_argsorted = np.argsort(ldist, axis=1)
    ldist_argbest = ldist_argsorted[:,0:k1]

    for li in range(N):
        lp1 = lp[li,:]

        cosphi = np.zeros(k1)
        for ki in range(k1):
            cosphi[ki] = lines_points_cosangle(lp[li, :], lp[ldist_argbest[li, ki], :], f=9)

        cosphi_argsorted = np.argsort(cosphi)
        cosphi_argsorted = cosphi_argsorted[::-1]
        cosphi_argbest = cosphi_argsorted[0:k2]

        lsim_temp = np.zeros(k2)
        for ki in range(k2):
            lj = ldist_argbest[li,cosphi_argbest[ki]]
            lp2 = lp[lj, :]
            prox = lines_proximity(lp1, lp2, sigma)
            lsim_temp[ki] = prox * cosphi[cosphi_argbest[ki]]

        lscore[li] = np.sum(lsim_temp)

    lscore /= k2

    return lscore

def calc_ldist_parfun(i, lp):
    N = lp.shape[0]
    ldist = np.zeros(N)
    for j in range(N):
        if i != j:
            ldist[j] = line_distance_closest(lp[i, :], lp[j, :])
        else:
            ldist[j] = 4

    return ldist

def calc_lsim(lp, sigma=0.1):

    num_cores = multiprocessing.cpu_count()
    N = lp.shape[0]

    lsim = Parallel(n_jobs=num_cores)(delayed(calc_lsim_parfun)(i, lp, sigma) for i in range(N))
    lsim = np.stack(lsim)

    for i in range(N):
        for j in range(i, N):
            lsim[i, j] = lsim[j, i]

    return lsim

def calc_lsim_parfun(i, lp, sigma):
    N = lp.shape[0]
    lsim = np.zeros(N)
    for j in range(i):
        lsim[j] = lines_similarity(lp[i, :], lp[j, :], sigma=sigma)

    return lsim

def find_initial_vps(sphere_image, cnn_response, num_max):

    sphere = sphere_image.copy()
    sphere = sphere[::-1, :]

    rA = cnn_response.shape[0]
    rB = cnn_response.shape[1]
    sA = sphere_image.shape[0]
    sB = sphere_image.shape[1]

    maxima = find_maxima(cnn_response).flatten()
    flat_cnn_response = cnn_response.flatten()
    best_maxima = np.argsort(flat_cnn_response[maxima == 1])
    best_maxima = best_maxima[::-1]
    maxima[np.where(maxima == 1)[0][best_maxima[num_max:]]] = 0
    maxima = np.reshape(maxima, cnn_response.shape)

    vps = []

    for ra in range(rA):
        for rb in range(rB):
            if maxima[ra,rb] == 1:
                sphere_slice = sphere[ra*sA/rA:(ra+1)*sA/rA, rb*sB/rB:(rb+1)*sB/rB]

                max_response = np.max(sphere_slice)
                sphere_slice_flat = sphere_slice.flatten()
                sphere_slice_flat[sphere_slice_flat < max_response] = 0
                maxed_idx = np.where(sphere_slice_flat>0)[0]
                unraveled_indices = []

                if maxed_idx.shape[0] == 0:
                    continue

                for i in range(maxed_idx.shape[0]):
                    unraveled = np.unravel_index(maxed_idx[i], sphere_slice.shape)
                    unraveled_indices.append(unraveled)

                average_index = np.zeros(2)
                for idx in unraveled_indices:
                    average_index += idx
                average_index /= len(unraveled_indices)

                max_response = average_index

                max_index = np.zeros(2)

                max_index[1] = max_response[0] + ra*sA/rA
                max_index[0] = max_response[1] + rb*sB/rB

                angle = coconv.index_to_angle(max_index, sphere_image.shape)
                vp = coconv.angle_to_point(angle)

                vps.append(vp)

    return np.vstack(vps)

def expectation_maximisation(l, lp, cnn_response, num_iter=100, sphere_image=None,
                             init_vp=None, do_merge=True, do_split=True, do_iterations=True,
                             distance_measure="angle", use_weights=True, wbias=1, num_init_vp=25, split_merge_freq=10,
                             merge_thresh=1e-3, outlier_thresh=1.96**2, final_convergence=5e-3,
                             s_thresh=1e-200, num_min_lines=3):

    N = l.shape[0]
    print "Number of line segments: ", N

    if use_weights:
        lsim = calc_lsim(lp, sigma=1)
    else:
        lsim = np.zeros((l.shape[0], l.shape[0]))

    lv = np.zeros((l.shape[0], 2))
    lm = np.zeros((l.shape[0], 2))

    for i in range(l.shape[0]):
        l[i,:] /= np.linalg.norm(l[i,:])
        lv[i,:] = lp[i,0:2] - lp[i,2:4]
        lm[i,:] = (lp[i,0:2] + lp[i,2:4])*0.5

    merge_thresh_final = merge_thresh*10
    merge_freq = split_merge_freq
    split_freq = split_merge_freq
    split_merge_it = 100
    splits = 1

    if distance_measure == "angle":
        max_stdd = 1e-6
        s_init_factor = 1e-6
    elif distance_measure == "dotprod":
        max_stdd = 1e-3
        s_init_factor = 1e-3
    else:
        assert False

    result = {'vp_assoc': None, 'vp': None, 'counts': None, 'count_id': None,
              'decision_metric': None, 'iterations': 0}

    v0 = find_initial_vps(sphere_image, cnn_response, num_init_vp)

    pdfpar = prob.pdf_params(cnn_response)

    if not (init_vp is None):
        v0 = init_vp.copy()
        for m in range(v0.shape[0]):
            v0[m,:] /= np.linalg.norm(v0[m,:])

    langles = lines_angles(lp)

    s_init = pdfpar.sigma * s_init_factor

    outlier_stdd = 1

    llen = np.ones(l.shape[0])

    for li in range(l.shape[0]):
        l[li,:] /= np.linalg.norm(l[li,:], ord=2)
        llen[li] = np.linalg.norm(np.array([lp[li,0]-lp[li,2], lp[li,1]-lp[li,3]]), ord=2)

    if use_weights:
        lscore = line_rating_knn(lp, k2=4)
        lscore = np.clip(lscore, 0.2, 1)
        lweight = llen
        lweight *= lscore
    else:
        lweight = np.ones(N)

    M = v0.shape[0]

    s = np.ones(M) * s_init

    v = np.zeros((num_iter+1, v0.shape[0], v0.shape[1]))

    v[0,:,:] = v0.copy()

    p = prob.calc_probabilities(0, pdfpar, v, l, lp, s, llen, distance_measure)
    w = weight_matrix(p.vl, lweight, lsim, bias=wbias)
    counts, counts_weighted, vp_assoc = calc_vp_line_counts(v[0,:,:], l, lp, s, w, lweight, distance_measure,
                                                            thresh=outlier_thresh)

    v = np.delete(v, np.where(counts < 3)[0], axis=1)
    s = np.delete(s, np.where(counts < 3)[0], axis=0)

    M = v.shape[1]
    print "Number of initial VPs: ", M

    for i in range(num_iter):

        if M == 0:
            print "No VPs left!"
            return result

        if np.mod(i, split_freq) == 0 and i > 0 and i < split_merge_it and do_split:
            for it in range(splits):
                p = prob.calc_probabilities(i, pdfpar, v, l, lp, s, llen, distance_measure)
                w = weight_matrix(p.vl, lweight, lsim, bias=wbias)
                split = split_best_vp(i, v, s, linePoints=lp, lines=l, weightMatrix=w, lineWeights=lweight,
                                      lineAngles=langles, min_diff=merge_thresh)
                v = split['v'].copy()
                s = split['s'].copy()

        M = v.shape[1]

        p = prob.calc_probabilities(i, pdfpar, v, l, lp, s, llen, distance_measure)

        max_err = 0
        max_id = 0

        to_be_removed = []

        lweight_temp = lweight.copy()

        w = weight_matrix(p.vl, lweight_temp, lsim, bias=wbias)

        for m in range(M):

            if not do_iterations:
                break

            wtemp = w[m,:]
            ltemp = l

            newVP = calc_new_vanishing_point(ltemp, wtemp)

            if newVP is None:
                to_be_removed.append(m)
                continue
            else:
                v[i + 1, m, :] = newVP

            try:
                p_vl_sum = np.sum(p.vl[m,:])

                s_log = np.log(np.sum(p.lvsq[:,m] * p.vl[m,:])) - np.log(p_vl_sum)
                s[m] = np.exp(s_log)

                s[m] = np.minimum(s[m], max_stdd)
                s[m] = np.maximum(s[m], s_thresh)

                if np.isnan(s[m]):
                    to_be_removed.append(m)
                else:
                    err = np.arccos(np.minimum(np.abs(np.dot(v[i,m,:], v[i+1,m,:])), 1.0))
                    max_err = np.maximum(max_err, err)
                    max_id = m if max_err == err else max_id

                    if err > 1.5:
                        to_be_removed.append(m)

            except np.linalg.linalg.LinAlgError as err:
                print err
                to_be_removed.append(m)
                continue

        if not do_iterations:
            v[i + 1, :, :] = v[i , :, :].copy()

        print "%03d - max. VP change: %.4f - VPs: %d" % (i, max_err, M)

        to_be_removed = np.array(to_be_removed)
        v = np.delete(v, to_be_removed, axis=1)
        s = np.delete(s, to_be_removed, axis=0)
        p = prob.calc_probabilities(i, pdfpar, v, l, lp, s, llen, distance_measure)
        M = v.shape[1]

        if max_err < final_convergence or i == num_iter-1 or not do_iterations:
            print "reached convergence"

            if do_merge:
                merged = merge_vps(i + 1, v, s, l, merge_thresh_final, lweight, lsim, wbias, pdfpar, lp, llen,
                                   distance_measure, outlier_stdd=outlier_stdd)
                v = merged['v']
                s = merged['s']

            p = prob.calc_probabilities(i, pdfpar, v, l, lp, s, llen, distance_measure)

            w = weight_matrix(p.vl, lweight_temp, lsim, bias=wbias)

            to_be_removed = []
            assoc = np.argmax(w, axis=0)

            M = v.shape[1]

            for m in range(M):

                if np.size(w[m, assoc == m]) == 0:
                    continue

                w[m, assoc == m] /= np.max(w[m, assoc == m])

                wtemp = w[m, assoc==m]
                ltemp = l[assoc==m,:]

                newVP = calc_new_vanishing_point(ltemp, wtemp)

                if newVP is None:
                    to_be_removed.append(m)
                    continue
                else:
                    v[i + 1, m, :] = newVP

                try:
                    p_vl_sum = np.sum(p.vl[m, :])

                    s_log = np.log(np.sum(p.lvsq[:, m] * p.vl[m, :])) - np.log(p_vl_sum)
                    s[m] = np.exp(s_log)

                    s[m] = np.minimum(s[m], max_stdd)

                    if np.isnan(s[m]) or s[m] < s_thresh:
                        to_be_removed.append(m)
                    else:

                        err = np.arccos(np.minimum(np.abs(np.dot(v[i, m, :], v[i + 1, m, :])), 1.0))
                        max_err = np.maximum(max_err, err)
                        max_id = m if max_err == err else max_id

                        if err > 1.5:
                            to_be_removed.append(m)

                except np.linalg.linalg.LinAlgError as err:
                    to_be_removed.append(m)
                    continue

            to_be_removed = np.array(to_be_removed)
            v = np.delete(v, to_be_removed, axis=1)
            s = np.delete(s, to_be_removed, axis=0)

            p = prob.calc_probabilities(i, pdfpar, v, l, lp, s, llen, distance_measure)

            decision_metric = weight_matrix(p.vl, lweight, lsim, bias=wbias)

            if decision_metric.size <= 0:
                print "decision metric is empty"
                return result

            max_decision = np.argmax(decision_metric, axis=0)

            good_vp = np.unique(max_decision)

            print "Number of VPs: ", good_vp.size

            v = v[:,good_vp,:]
            s = s[good_vp]

            p = prob.calc_probabilities(i+1, pdfpar, v, l, lp, s, llen, distance_measure)

            decision_metric = weight_matrix(p.vl, lweight, lsim, bias=wbias)
            counts, counts_weighted, vp_assoc = calc_vp_line_counts(v[i+1,:,:], l, lp, s, decision_metric, lweight,
                                                                    distance_measure, thresh=outlier_thresh)

            M = v.shape[1]

            vidx = 0
            while vidx < M:
                if counts[vidx] < num_min_lines:
                    v = np.delete(v, vidx, axis=1)
                    s = np.delete(s, vidx)
                    M = v.shape[1]

                    p = prob.calc_probabilities(i + 1, pdfpar, v, l, lp, s, llen, distance_measure)
                    decision_metric = weight_matrix(p.vl, lweight, lsim, bias=wbias)
                    counts, counts_weighted, vp_assoc = calc_vp_line_counts(v[i + 1, :, :], l, lp, s, decision_metric,
                                                                            lweight, distance_measure,
                                                                            thresh=outlier_thresh, vp_assoc=None)

                else:
                    vidx += 1

            vp = v[i + 1, :, :]

            return {'vp_assoc':vp_assoc,'vp':vp, 'counts':counts, 'counts_weighted':counts_weighted, 'count_id':None,
                    'decision_metric':decision_metric, 'iterations':i, 'distribution':p, 'sigma':s }

        if np.mod(i,merge_freq) == 0 and i > 0 and i <= split_merge_it+merge_freq and do_merge:
            merged = merge_vps(i+1, v, s, l, merge_thresh, lweight, lsim, wbias, pdfpar, lp, llen, distance_measure,
                               outlier_stdd=outlier_stdd)
            v = merged['v']
            s = merged['s']

    return result

def calc_new_vanishing_point(l, w):

    try:
        if np.size(w) == 0:
            return None

        if np.max(w) == 0:
            return None

        W = np.diag(w / np.max(w))
        A = l
        Mat = np.dot(W, A)

        U, S, V = la.svd(Mat)

        V = V.T

        vp = np.squeeze(V[:, 2])

        vp /= np.linalg.norm(vp, ord=2)

        vp *= np.sign(vp[2])

    except np.linalg.linalg.LinAlgError as err:
        vp = None

    return vp

def calc_vp_line_counts(vp, l, lp, s, decision_metric, lweights, distance_measure, thresh=2.57, vp_assoc=None):

    N = l.shape[0]
    M = vp.shape[0]
    if vp_assoc is None:
        vp_assoc = np.argmax(decision_metric, axis=0)

    counts = np.zeros(M)
    counts_weighted = np.zeros(M)

    for n in range(N):
        m = vp_assoc[n]
        if m > -1:
            if distance_measure == "dotprod":
                dist = np.abs(np.dot(vp[m], l[n, :]))
            elif distance_measure == "angle":
                dist = prob.calc_lvsq_single(vp[m], l[n, :], lp[n, :])
            elif distance_measure == "area":
                dist = prob.calc_lvsq_area_single(vp[m], l[n, :], lp[n, :])
            else:
                assert False

            if dist > thresh * np.sqrt(s[m]):
                vp_assoc[n] = -1
            elif lweights[n] == 0:
                vp_assoc[n] = -1
            else:
                counts[m] += 1
                counts_weighted[m] += lweights[n]

    return counts, counts_weighted, vp_assoc

def weight_matrix(p_vl, lweight, lsim, bias=0.001):

    w = np.zeros(p_vl.shape)
    for m in range(w.shape[0]):
        w_ = p_vl[m, :] * lweight

        for k in range(w.shape[1]):
            w[m,k] = ( w_[k] + bias * lweight[k] * np.dot(w_, lsim[:, k]) ) / ( 1 + bias * lweight[k] * np.sum(lsim[:, k]) )

    return w

def split_best_vp(i, v, s, linePoints, lines, weightMatrix, lineWeights, lineAngles, numClusters=2, min_diff=0.0001):

    M = v.shape[1]
    N = lines.shape[0]

    mean_phi = np.zeros(M)
    stdd_phi = np.zeros(M)

    weightMatrixGreedy = np.zeros(weightMatrix.shape)
    weightIndices = weightMatrix.argmax(axis=0)
    for li in range(N):
        weightMatrixGreedy[weightIndices[li],li] = weightMatrix[weightIndices[li],li]
    weightMatrixGreedy /= weightMatrix.max()

    for m in range(M):

        mean_phi[m] = np.mean(lineAngles[weightMatrixGreedy[m,:]>0])
        stdd_phi[m] = np.std(lineAngles[weightMatrixGreedy[m,:]>0])

    worstVPs = np.argsort(stdd_phi)
    worstVPs = worstVPs[::-1]

    worstVP = None
    for m in range(M):
        vpAssoc = np.argmax(weightMatrix, axis=0)
        assocLines = np.where(vpAssoc == worstVPs[m])[0]
        lp = linePoints[assocLines]
        l = lines[assocLines]
        Nworst = lp.shape[0]

        vp = v[i, m, :].copy()
        vp /= vp[2]

        if  Nworst > numClusters*4 and ( ( vp[0] > -1 and vp[1] > -1 and vp[0] < 1 and vp[1] < 1 ) ):
            worstVP = worstVPs[m]
            break

    if not (worstVP is None):

        stdd = s[worstVP] / numClusters

        Ldist = np.zeros((Nworst,Nworst))
        for li in range(Nworst):
            for lj in range(Nworst):
                if lj != li:
                    Ldist[li,lj] = 1-lines_points_cosangle(lp[li,:], lp[lj,:], f=2)

        model = cluster.AgglomerativeClustering(linkage='average', connectivity=Ldist, n_clusters=numClusters,
                                                affinity='precomputed')
        model.fit_predict(Ldist)

        labels = model.labels_

        lw = lineWeights[assocLines]

        l[:,0] *= lw
        l[:,1] *= lw
        l[:,2] *= lw

        new_vps = []

        for c in range(numClusters):

            lineSet = l[labels == c]

            if lineSet.shape[0] < 3:
                continue

            U, S, V = la.svd(lineSet)
            V = V.T
            vp = np.squeeze(V[:, 2])
            vp /= np.linalg.norm(vp, ord=2)
            if vp[2] < 0:
                vp *= -1

            new_vps.append(vp)

        too_similar = True

        for c in range(len(new_vps)):
            for d in range(c+1, len(new_vps)):
                vp1 = new_vps[c]
                vp2 = new_vps[d]

                cosphi = np.clip(np.dot(vp1, vp2.T), -1, 1)
                angle = np.abs(np.arccos(np.clip(np.abs(cosphi), -1, 1)))

                if angle > min_diff:
                    too_similar = False

        if not too_similar:
            first = True
            for c in range(len(new_vps)):
                vp = new_vps[c]
                if first:
                    v[i, worstVP, :] = vp.copy()
                    s[worstVP] = stdd
                    first = False
                else:
                    v = np.append(v, np.zeros((v.shape[0],1,v.shape[2])), axis=1)
                    s = np.append(s,stdd)
                    v[i, -1, :] = vp.copy()

    return {'v': v, 's': s}

def merge_vps(i, v, s, l, thresh, lweight, lsim, wbias, pdfpar, lp, llen, distance_measure, max_stdd=0.01,
              outlier_stdd=1e-6):

    M = v.shape[1]

    num_cores = multiprocessing.cpu_count()

    tryAgain = True
    tries = 0

    while tryAgain and M > 1:

        tries += 1

        angles = Parallel(n_jobs=num_cores)(delayed(calc_angle_to_other_vp)(v, i, j) for j in range(M))
        angles = np.stack(angles)

        argmin_angle = numpy.unravel_index(angles.argmin(), angles.shape)
        j = argmin_angle[0]
        k = argmin_angle[1]
        min_angle = angles[j, k]

        if min_angle < thresh:

            try:
                p = prob.calc_probabilities(i, pdfpar, v, l, lp, s, llen, distance_measure)
                w = weight_matrix(p.vl, lweight, lsim, bias=wbias)

                newVP = calc_new_vanishing_point(l, w[j, :] + w[k, :])

                p_vl_sum = np.sum(p.vl[k, :] + p.vl[j, :])
                s_log = np.log(np.sum(0.5 * (p.lvsq[:, j] + p.lvsq[:, k]) * (p.vl[k, :] + p.vl[j, :]))) - \
                        np.log(p_vl_sum)
                s[k] = np.exp(s_log)

                if newVP is None or s[k] > max_stdd:
                    tryAgain = False
                    continue
                else:
                    v[i, k, :] = newVP

                v = np.delete(v, j, axis=1)
                s = np.delete(s, j, axis=0)

            except np.linalg.linalg.LinAlgError as err:
                continue
        else:
            tryAgain = False

        M = v.shape[1]

    return {'v': v, 's': s}

def calc_angle_to_other_vp(v, i, k):

    thisVP = np.squeeze(v[i,k,:])
    otherVPs = np.squeeze(v[i,:,:])
    cosphi = np.clip(np.dot(otherVPs, thisVP.T), -1, 1)
    angles = np.abs(np.arccos(np.clip(np.abs(cosphi), -1, 1)))
    if np.isscalar(angles):
        angles = pi
    else:
        angles[k] = pi
    return angles

def lines_similarity(lp1, lp2, sigma=0.1):
    cosphi = lines_points_cosangle(lp1, lp2, f=9)

    sim = cosphi * lines_proximity(lp1, lp2, sigma)

    return sim

def lines_proximity(lp1, lp2, sigma=0.1):
    sigma = sigma*np.minimum(line_length(lp1), line_length(lp2))
    d = line_distance_closest(lp1,lp2)
    prox = np.exp(-(d * d) / (2 * sigma * sigma))
    return prox

def lines_points_cosangle(lp1,lp2,f=1):
    v1 = lp1[0:2] - lp1[2:4]
    v2 = lp2[0:2] - lp2[2:4]

    cosdphi = np.abs(np.dot(v1,v2) / (np.linalg.norm(v1)*np.linalg.norm(v2)))

    dphi = np.abs(np.arccos(np.clip(cosdphi, -1, 1)))
    cosdphi = np.cos(np.clip(f*dphi,-pi/2,pi/2))

    return cosdphi

def line_distance_closest(lp1, lp2):
    l1p1 = np.array([lp1[0], lp1[1], 1])
    l1p2 = np.array([lp1[2], lp1[3], 1])
    l2p1 = np.array([lp2[0], lp2[1], 1])
    l2p2 = np.array([lp2[2], lp2[3], 1])

    d1 = line_segment_point_distance(lp1, l2p1)
    d2 = line_segment_point_distance(lp1, l2p2)
    d4 = line_segment_point_distance(lp2, l1p1)
    d5 = line_segment_point_distance(lp2, l1p2)

    d = np.min(np.array([d1,d2,d4,d5]))

    return d

def line_segment_point_distance(lp, p):
    lp1 = np.array([lp[0], lp[1], 1])
    lp2 = np.array([lp[2], lp[3], 1])

    param = np.dot(p-lp1, lp2-lp1) / np.square(np.linalg.norm(lp2-lp1))

    if param < 0:
        pclosest = lp1
    elif param > 1:
        pclosest = lp2
    else:
        pclosest = lp1 + param*(lp2-lp1)

    d = np.linalg.norm(pclosest-p)

    return d

def line_length(lp):
    return np.linalg.norm(lp[0:2]-lp[2:4], ord=2)

def lines_angles(lp):
    N = lp.shape[0]
    angles = np.zeros(N)

    for i in range(N):
        v = np.array([lp[i,0] - lp[i,2], lp[i,1] - lp[i,3]])
        v /= np.linalg.norm(v)
        phi = np.abs(np.arccos(np.clip(v[0], -1, 1)))
        phi = pi-phi if phi > pi/2 else phi
        angles[i] = phi

    return angles

reason: import numpy,from numpy


text:
import numpy as np
import pandas as pd
import quantipy as qp
from quantipy.core.helpers.functions import emulate_meta
import savReaderWriter as srw
import copy
import json

def write_sav(path_sav, data, **kwargs):
    """
    Write the given records to a SAV file at path_sav.

    Using the various definitions indicated by the packed kwargs, write
    the given set of records to a SAV file at the location indicated by
    path_sav.

    For a full explanation of the kwargs used please see the
    savReaderWriter library documentation here:
    http://pythonhosted.org/savReaderWriter/

    Parameters
    ----------
    path_sav : str
        The full path, including extension, indicating where the output
        file should be saved.
    records : list
        A list of records (so a list of lists) holding the row data to
        be saved in the output SAV file.
    **kwargs : various
        Remaining keyword arguments passed to
        savReaderWriter.SavWriter().

    Returns
    -------
    None
    """
    with srw.SavWriter(path_sav, ioUtf8=True, **kwargs) as writer:
        records = data.fillna(writer.sysmis).values.tolist()
        for record in records:
            writer.writerow(record)

def split_series(series, sep, columns=None):
    """
    Splits all the items of a series using the given delimiter.

    Splits each item in series using the given delimiter and returns
    a DataFrame (as per Excel text-to-columns). Optionally, you can
    pass in a list of column names that should be used to name the
    resulting columns.

    Parameters
    ----------
    series : pandas.Series
        The series that should be split.
    sep : str
        The separator that should be used to split the series.
    columns : list-list, default=None
        A list of names that should be set into the resulting DataFrame
        columns.

    Returns
    -------
    df : pandas.DataFrame
        Series, split by sep, returned as a DataFrame.
    """

    df = pd.DataFrame(series.astype('str').str.split(sep).tolist())
    if not columns is None:
        df.columns = columns
    return df

def get_savwriter_integer_format(series):
    """
    Derive the required SAV format value for the given integer series.

    savReaderWriter requires specific instructions for each variable
    in order to correctly create target variables. This function
    determines the width of the maximum integer in the given series
    and constructs that format string accordingly.

    Parameters
    ----------
    series : pandas.Series
        The series for which the format string is being generated.

    Returns
    -------
    fmt : str
        The format string describing the given integer series.
    """

    fmt = 'F%s' % (
        len(str(series.dropna().astype('int').max()))
    )
    return fmt

def get_savwriter_float_format(series):
    """
    Derive the required SAV format value for the given float series.

    savReaderWriter requires specific instructions for each variable
    in order to correctly create target variables. This function
    determines the width of the maximum float in the given series
    and constructs that format string accordingly.

    Parameters
    ----------
    series : pandas.Series
        The series for which the format string is being generated.

    Returns
    -------
    fmt : str
        The format string describing the given float series.
    """

    if series.dropna().shape[0]==0:
        # If there's no data in the series it's impossible to predict
        # what sort of data may be expected later on, so a generic
        # interpretation of three whole numbers and two decimals is
        # applied.
        fmt = 'F5.2'
    else:
        df = split_series(
            series.dropna(),
            sep='.',
            columns=['int', 'dec']
        )
        w_int = len(str(df['int'].max()))
        w_dec = len(str(df['dec'].max()))
        if df['dec'].max()!='0':
            fmt = 'F%s.%s' %(
                w_int + w_dec,
                w_dec
            )
        else:
            fmt = 'F%s' %(
                w_int
            )
    return fmt

def get_value_text(values, value, text_key):
    """
    Get the text for the given value, using the given text_key.

    Values is a meta object in the form of a list of dicts. This
    function is used to find in that list the dict that has the given
    value, and from it return the text for that value using the given
    text_key.

    Parameters
    ----------
    values : list
        A Quantipy values meta object.
    value : int
        The value that should targeted in the list of values.
    text_key : str
        The key from which the the value text should be taken.

    Returns
    -------
    text : str
        The text associated with the targeted value and text_key.
    """

    for val in values:
        if val['value']==value:
            try:
                text = val['text'][text_key]
                return text
            except KeyError:
                print (
                    "The text key '%s' was not found in the text object: %s" % (
                        text_key,
                        json.dumps(val)
                    )
                )
    raise ValueError(
            "The value '%s' was not found in the values object: %s" % (
                value,
                json.dumps(values)
            )
        )

def list_known_columns(meta, from_set):
    """
    Get an ordered list of columns from meta's from_set set.

    The from_set set may include items that point to masks. This
    function will replace mask-references with the names of the columns
    those masks point to so that what you get in return is an ordered
    list of column names that definitely exist in meta['columns'].

    .. note:: The meta document must have a set called from_set.

    Parameters
    ----------
    meta : dict
        The meta document.
    from_set : str
        The set name from which the export should be drawn.

    Returns
    -------
    column_names : list
        The column names.
    """

    column_names = []
    for col in meta['sets'][from_set ]['items']:
        pointer, name = col.split('@')
        if pointer=='columns':
            if name in meta['columns'] and not name in column_names:
                column_names.append(name)
        elif pointer=='masks':
            for item in meta['masks'][name]['items']:
                name =  item['source'].split('@')[1]
                if name in meta['columns'] and not name in column_names:
                    column_names.append(name)
    return column_names

def stringify_dates(dates):
    """
    Convert a datetime64 series to string in the form 'YYYY-M-D'.

    Parameters
    ----------
    dates : pandas.Series
        The numpy.datetime64 dates.

    Returns
    -------
    series : pandas.Series
        The string dates.
    """

    def stringify_date(date):

        try:
            date = ' '.join([
                '-'.join([
                    str(date.year),
                    str(date.month).zfill(2),
                    str(date.day).zfill(2)]),
                ':'.join([
                    str(date.hour).zfill(2),
                    str(date.minute).zfill(2),
                    str(date.second).zfill(2)])])
        except:
            pass

        return date

    series = dates.apply(stringify_date)
    series = series.astype('str')

    return series

def fix_label(label):
    label = label.replace('\n', '')
    return label

def save_sav(path_sav, meta, data, index=False, text_key=None,
             mrset_tag_style='__', drop_delimited=True, from_set=None,
             verbose=True):
    """
    One-sentence description.

    More detailed description.

    .. note:: Important note (if any).

    Parameters
    ----------
    path_sav : str
        The full path, including extension, indicating where the output
        file should be saved.
    meta : dict
        The meta document to be converted to SAV.
    data : pandas.DataFrame
        The data to be converted to SAV.
    index : bool, default=False
        Should the index be inserted into the dataframe before the
        conversion happens?
    text_key : str, default=None
        The text_key that should be used when taking labels from the
        source meta. If the given text_key is not found for any
        particular text object, the default text key (as indicated under
        meta['lib']['default text']) will be used instead.
    mrset_tag_style : str, default='__'
        The delimiting character/string to use when naming dichotomous
        set variables. The mrset_tag_style will appear between the
        name of the variable and the dichotomous variable's value name,
        as taken from the delimited set value that dichotomous
        variable represents.
    drop_delimited : bool, default=True
        Should Quantipy's delimited set variables be dropped from
        the export after being converted to dichotomous sets/mrsets?
    from_set : str
        The set name from which the export should be drawn.

    Returns
    -------
    None
    """
    # This function will make edits to the meta and data objects, so
    # they should be copied first.
    meta = copy.deepcopy(meta)
    data = data.copy()

    if from_set is None:
        from_set = 'data file'
    if from_set not in meta['sets']:
        raise KeyError(
            "The set '{}' was not found in meta.".format(from_set)
        )

    # There is an issue converting numpy dates to SAV so dates
    # are currently being turned into strings in the form 'Y-M-D'
    date_cols = [
        col for col in meta['columns']
        if meta['columns'][col]['type'] == 'date']
    for date_col in date_cols:
        data[date_col] = stringify_dates(data[date_col])
        meta['columns'][date_col]['type'] = 'string'

        # This code can be used to instead simply remove all
        # dates from the dataset before conversion
#         del meta['columns'][date_col]
#         mapper = 'columns@{}'.format(date_col)
#         try:
#             meta['sets'][from_set]['items'].remove(mapper)
#         except:
#             pass

    for key, val in meta['columns'].iteritems():
        if val['type'] == 'string':
            if key in data.columns:
                data[key].fillna('', inplace=True)

    if index:
        if data.index.name not in data.columns:
            # Put the index into the first column of data
            data.insert(0, data.index.name, data.index)
        mapper = 'columns@%s' % (data.index.name)
        if mapper not in meta['sets'][from_set]['items']:
            # Add the index meta-mapper to the set
            meta['sets'][from_set]['items'].insert(0, mapper)

    if text_key is None:
        # Get default text key instead
        text_key = meta['lib']['default text']

    # Remove columns from data not found in meta
    known_columns = list_known_columns(meta, from_set)
    for col in data.columns:
        if col not in known_columns:
            if col != '@1':
                if verbose:
                    print (
                        "Data column '{}' not included in"
                        " the '{}' set, it will be excluded"
                        " from the SAV file."
                    ).format(col, from_set)
            data.drop(col, axis=1, inplace=True)

    # Remove columns from meta not found in data
    for col in known_columns:
        if col not in data.columns:
            if verbose:
                print (
                    'Meta column "%s" not found in data, it will '
                    'be excluded from the SAV file.'
                ) % (col)
            if col in meta['columns']:
                del meta['columns'][col]

    # Create the varNames definition for the savWriter
    # Part of this process is to identify variable names that contain
    # illegal characters (as far as SPSS is concerned) and replace them
    # with an underscore ('_').
    varNames = list_known_columns(meta, from_set)
    new_names = []
    column_mapper = {}
    for varName in varNames:
        new_name = varName
        if varName[0] in [1, 2, 3, 4, 5, 6, 7, 8, 9, 0]:
            new_name = '_%s' % (new_name)
        for i, char in enumerate(new_name):
            if char in ['[', ']', '{', '}', '.']:
                new_name = '%s_%s' % (
                    new_name[:i],
                    new_name[i+1:]
                )
        new_names.append(new_name)
        column_mapper[varName] = new_name
    for old_name, new_name in column_mapper.iteritems():
        if new_name not in meta['columns']:
            meta['columns'][new_name] = meta['columns'].pop(old_name)
            idx = varNames.index(old_name)
            varNames = varNames[:idx] + [new_name] + varNames[idx+1:]
    data.columns = pd.Series(data.columns).map(column_mapper)

    # Create the multRespDefs definition for the savWriter
    delimited_sets = [
        c
        for c in varNames
        if meta['columns'][c]['type'] == 'delimited set'
    ]
    multRespDefs = {}
    for ds_name in delimited_sets:
        values = [
            val['value']
            for val in emulate_meta(meta, meta['columns'][ds_name]['values'])
        ]
        if data[ds_name].dtype == 'float':
            # The delimited set has no responses
            dichot = pd.DataFrame(np.NaN, index=data.index, columns=values)
        else:
            dichot = data[ds_name].str.get_dummies(sep=';')
            if '0' in dichot.columns and 0 not in values:
                # A '0' column here indicates rows that had no answers
                # in the delimited set, that column can be discarded.
                dichot.drop('0', axis=1, inplace=True)
            # Make sure all known values from the meta appear in the
            # dichotomous dataframe.
            for val in values:
                if val not in dichot.columns.astype('int'):
                    dichot[val] = 0
            no_responses = dichot.sum(axis=1) == 0
            dichot.loc[no_responses, :] = np.NaN

        dichot.columns = dichot.columns.astype(int)
        if pd.__version__ == '0.19.2':
            dichot.sort_index(axis=1, inplace=True)
        else:
            dichot.sort(axis=1, inplace=True)
        dsNames = ['%s%s%s' % (ds_name, mrset_tag_style, val) for val in values]
        ds_index = varNames.index(ds_name)
        varNames[ds_index+1:ds_index+1] = dsNames
        varNames.pop(ds_index)

        cols = [
            '%s%s%s' % (ds_name, mrset_tag_style, c) for c in dichot.columns]
        dichot.columns = cols
        # Synch dichotomous columns with varName order
        dichot = dichot[dsNames]
        # Find the position of the original delimited set in the source
        # dataframe's columns.
        ds_name_idx = data.columns.tolist().index(ds_name)
        # Insert the columns from the dichotomous dataframe after the
        # position of the delimited set.
        for i, col in enumerate(dichot.columns, start=1):
            data.insert(ds_name_idx+i, col, dichot[col])
        # Add the column metadata for each dichotomous column
        for dichName in dsNames:
            meta['columns'][dichName] = {
                'type': 'single',
                'values': [
                    {'value': 0, 'text': {text_key: 'No'}},
                    {'value': 1, 'text': {text_key: 'Yes'}}
                ],
                'text': {
                    text_key: get_value_text(
                        emulate_meta(meta, meta['columns'][ds_name]['values']),
                        int(dichName.split('_')[-1]),
                        text_key)
                }
            }

        # Add the savWriter-required definition of the mrset
        varLabel = fix_label(meta['columns'][ds_name]['text'].get(text_key, ''))
        if varLabel > 120:
            varLabel = varLabel[:120]
        multRespDefs[ds_name] = {
            'varNames': dsNames,
            'label': qp.core.tools.dp.io.unicoder(varLabel, like_ascii=True),
            'countedValue': 1,
            'setType': 'D'
        }

        if drop_delimited:
            data.drop(ds_name, axis=1, inplace=True)

    # Create the varLabels definition for the savWriter
    varLabels = {
        v: fix_label(meta['columns'][v]['text'].get(text_key, ''))
        for v in varNames
    }

    for v in varLabels:
        if len(varLabels[v]) > 120:
            varLabels[v] = varLabels[v][:120]

    # Create the valueLabels definition for the savWriter
    # This will now catch all of the newly added dichotomous set columns
    singles = [v for v in varNames if meta['columns'][v]['type'] == 'single']
    valueLabels = {
        var: {
            int(val['value']): fix_label(val['text'].get(text_key, ''))
            for val in emulate_meta(meta, meta['columns'][var]['values'])
        }
        for var in singles
    }

    # Create the varTypes definition for the savWriter
    varTypes = {
        v:
        0
        if meta['columns'][v]['type'] in ['single', 'int', 'float']
        else 1000
        for v in varNames
    }

    # Create the formats definition for the savWriter
    numerics = [v for v, t in varTypes.iteritems() if t == 0]
    strings = [
        v
        for v in varTypes.keys()
        if meta['columns'][v]['type'] in ['string', 'delimited set']
    ]
    dates = [
        v
        for v in varTypes.keys()
        if meta['columns'][v]['type'] in ['date']
    ]

    sav_formatter = {
        'single': get_savwriter_integer_format,
        'int': get_savwriter_integer_format,
        'float': get_savwriter_float_format
    }
    numeric_formats = {
        v: sav_formatter[meta['columns'][v]['type']](data[v])
        for v in numerics
    }
    string_formats = {
        s: 'A1000'
        for s in strings
    }
    date_formats = {
        d: 'EDATE40'
        for d in dates
    }
    formats = {}
    formats.update(numeric_formats)
    formats.update(string_formats)
    formats.update(date_formats)

    data = data[varNames]

    write_sav(
        path_sav,
        data,
        varNames=varNames,
        varTypes=varTypes,
        formats=formats,
        varLabels=varLabels,
        valueLabels=valueLabels,
        multRespDefs=multRespDefs
    )

ext: py

max_stars_repo_name: scottporter/quantipy

max_issues_repo_name: scottporter/quantipy

max_forks_repo_name: scottporter/quantipy

reason: import numpy


text:
import xgboost as xgb
import numpy as np

# test


text:
import cntk as C
import numpy as np

def lrelu(input, leak=0.2, name=""):
    return C.param_relu(C.constant((np.ones(input.shape)*leak).astype(np.float32)), input, name=name)

def bn(input, activation=None, name=""):
    if activation is not None:
        x = C.layers.BatchNormalization(map_rank=1, name=name+"_bn" if name else "")(input)
        x = activation(x, name=name)
    else:
        x = C.layers.BatchNormalization(map_rank=1, name=name)(input)
    return x

def bn_relu(input, name=""):
    return bn(input, activation=C.relu, name=name)

def bn_lrelu(input, name=""):
    return bn(input, activation=C.leaky_relu, name=name)

def conv(input, filter_shape, num_filters, strides=(1,1), init=C.he_normal(), activation=None, pad=True, name=""):
    return C.layers.Convolution(filter_shape, num_filters, strides=strides, pad=pad, activation=activation, init=init, bias=False, name=name)(input)

def conv_bn(input, filter_shape, num_filters, strides=(1,1), init=C.he_normal(), activation=None, name=""):
    x = conv(input, filter_shape, num_filters, strides, init, name=name+"_conv" if name else "")
    x = bn(x, activation, name=name)
    return x

def conv_bn_nopad(input, filter_shape, num_filters, strides=(1,1), init=C.he_normal(), activation=None, name=""):
    x = conv(input, filter_shape, num_filters, strides, init, pad=False, name=name+"_conv" if name else "")
    x = bn(x, activation, name=name)
    return x

def conv_bn_relu(input, filter_shape, num_filters, strides=(1,1), init=C.he_normal(), name=""):
    return conv_bn(input, filter_shape, num_filters, strides, init, activation=C.relu, name=name)

def conv_bn_lrelu(input, filter_shape, num_filters, strides=(1,1), init=C.he_normal(), name=""):
    return conv_bn(input, filter_shape, num_filters, strides, init, activation=C.leaky_relu, name=name)

def conv_bn_relu_nopad(input, filter_shape, num_filters, strides=(1,1), init=C.he_normal(), name=""):
    return conv_bn_nopad(input, filter_shape, num_filters, strides, init, activation=C.relu, name=name)

def conv_bn_lrelu_nopad(input, filter_shape, num_filters, strides=(1,1), init=C.he_normal(), name=""):
    return conv_bn_nopad(input, filter_shape, num_filters, strides, init, activation=C.leaky_relu, name=name)

def flatten(input, name=""):
    assert (len(input.shape) == 3)
    return C.reshape(input, input.shape[0]*input.shape[1]* input.shape[2], name=name)

def flatten_2D(input, name):
    assert (len(input.shape) >= 3)
    return C.reshape(input, (input.shape[-3], input.shape[-2]* input.shape[-1]), name=name)

def broadcast_xy(input_vec, h, w):
    """ broadcast input vector of length d to tensor (d x h x w) """
    assert(h > 0 and w > 0)
    d = input_vec.shape[0]
    # reshape vector to d x 1 x 1
    x = C.reshape(input_vec, (d, 1, 1))
    # create a zeros-like tensor of size (d x h x w)
    t = np.zeros((d, h, w), dtype=np.float32)
    y = C.constant(t)
    z = C.reconcile_dynamic_axes(y, x)
    z = z + x
    return z

def conv_from_weights(x, weights, bias=None, padding=True, name=""):
    """ weights is a numpy array """
    k = C.parameter(shape=weights.shape, init=weights)
    y = C.convolution(k, x, auto_padding=[False, padding, padding])
    if bias:
        b = C.parameter(shape=bias.shape, init=bias)
        y = y + bias
    y = C.alias(y, name=name)
    return y

# bi-directional recurrence function op
# fwd, bwd: a recurrent op, LSTM or GRU
def bi_recurrence(input, fwd, bwd, name=""):
    F = C.layers.Recurrence(fwd, go_backwards=False, name='fwd_rnn')(input)
    B = C.layers.Recurrence(bwd, go_backwards=True, name='bwd_rnn')(input)
    h = C.splice(F, B, name=name)
    return h


text:
# BSD 3-Clause License; see https://github.com/scikit-hep/uproot4/blob/master/LICENSE

from __future__ import absolute_import

import numpy
import pytest
import skhep_testdata

import uproot4

def test_fallback_reading():
    with uproot4.open(
        skhep_testdata.data_path("uproot-small-evnt-tree-fullsplit.root")
    ) as f:
        f["tree:evt/P3/P3.Py"]
        assert f.file._streamers is None

    with uproot4.open(skhep_testdata.data_path("uproot-demo-double32.root")) as f:
        f["T/fD64"]
        assert f.file._streamers is not None


text:
import math
import sc2
import random
from sc2 import Race, Difficulty
import os
import sys
from sc2.constants import *
from sc2.position import Pointlike, Point2
from sc2.player import Bot, Computer
from sc2.unit import Unit as sc2Unit
sys.path.insert(0, os.path.abspath('../'))
import torch
from agents.prolonet_agent import DeepProLoNet
from agents.py_djinn_agent import DJINNAgent
from agents.lstm_agent import LSTMNet
from agents.baseline_agent import FCNet
from ICML.runfiles import sc_helpers
import numpy as np
import time
import torch.multiprocessing as mp
import argparse

DEBUG = False
SUPER_DEBUG = False
if SUPER_DEBUG:
    DEBUG = True

FAILED_REWARD = -0.0
STEP_PENALTY = 0.01
SUCCESS_BUILD_REWARD = 0.
SUCCESS_TRAIN_REWARD = 0.
SUCCESS_SCOUT_REWARD = 0.
SUCCESS_ATTACK_REWARD = 0.
SUCCESS_MINING_REWARD = 0.

class StarmniBot(sc2.BotAI):
    def __init__(self, rl_agent):
        super(StarmniBot, self).__init__()
        self.player_unit_tags = []
        self.agent = rl_agent
        self.corners = None
        self.action_buffer = []
        self.prev_state = None
        self.last_known_enemy_units = []
        self.itercount = 0
        self.last_scout_iteration = -100
        self.scouted_bases = []
        self.last_reward = 0
        # self.num_nexi = 1
        self.mining_reward = SUCCESS_MINING_REWARD
        self.last_sc_action = 43  # Nothing
        self.positions_for_pylons = []
        self.positions_for_buildings = []
        self.army_below_half = 0
        self.last_attack_loop = 0

    async def on_step(self, iteration):
        if self.units(UnitTypeId.NEXUS).amount < 1 and self.workers.amount < 1:
            await self._client.leave()
        if iteration == 0:
            base_height = self._game_info.terrain_height[self.start_location.rounded]
            max_x = self._game_info.map_size.width
            max_y = self._game_info.map_size.height
            min_pos = self.state.mineral_field.closer_than(8, self.start_location).center
            vector = min_pos.direction_vector(self.start_location)
            if vector.x == 0 or vector.y == 0:
                vector = min_pos.direction_vector(self.game_info.map_center)
            p0 = self.start_location + (vector * Point2((5.5, 5.5))).rounded
            p1 = p0 + vector * Point2((4, 6))
            b0 = p0 + vector * Point2((0.5, 2.5))
            b1 = b0 + vector * Point2((0, 3))
            b2 = p0 + vector * Point2((4.5, 0.5))
            b3 = b2 + vector * Point2((0, 3))
            for y in range(-3, 3):
                for x in range(-3, 3):
                    if x < 0 and y < 0:
                        continue
                    for pos in [p0 + vector * Point2((x * 8, y * 9)), p1 + vector * Point2((x * 8, y * 9))]:
                        if 0 < pos.x < max_x and 0 < pos.y < max_y and \
                                self._game_info.terrain_height[pos.rounded] == base_height:
                            self.positions_for_pylons.append(pos)
                    for pos in [b0 + vector * Point2((x * 8, y * 9)), b1 + vector * Point2((x * 8, y * 9)),
                                b2 + vector * Point2((x * 8, y * 9)), b3 + vector * Point2((x * 8, y * 9))]:
                        if 0 < pos.x < max_x and 0 < pos.y < max_y and \
                                self._game_info.terrain_height[pos.rounded] == base_height:
                            self.positions_for_buildings.append(pos)
            self.positions_for_pylons.sort(key=lambda a: self.start_location.distance_to(a))
            self.positions_for_buildings.sort(key=lambda a: self.start_location.distance_to(a))

            await self.chat_send("ProLo")
            self.corners = [Point2(Pointlike([0, self.game_info.map_size[1]])),
                            Point2(Pointlike([0, 0])),
                            Point2(Pointlike([self.game_info.map_size[0], self.game_info.map_size[1]])),
                            Point2(Pointlike([self.game_info.map_size[0], 0]))]
            closest_dist = self.units(UnitTypeId.NEXUS).first.distance_to(self.corners[0])
            closest = 0
            for corner in range(1, len(self.corners)):
                corner_dist = self.units(UnitTypeId.NEXUS).first.distance_to(self.corners[corner])
                if corner_dist < closest_dist:
                    closest = corner
                    closest_dist = corner_dist
            del self.corners[closest]

        self.itercount += 1
        # if self.itercount % 10 != 0:
        #     return
        # Get current state (minerals, gas, idles, etc.)
        current_state = sc_helpers.get_player_state(self.state)
        current_state = np.array(current_state)
        my_unit_type_arr = sc_helpers.my_units_to_type_count(self.units)
        enemy_unit_type_arr = sc_helpers.enemy_units_to_type_count(self.known_enemy_units)
        # Get pending
        pending = []
        for unit_type in sc_helpers.MY_POSSIBLES:
            if self.already_pending(unit_type):
                pending.append(1)
            else:
                pending.append(0)

        # Reshape all into batch of 1
        my_unit_type_arr = my_unit_type_arr.reshape(-1)  # batch_size, len
        current_state = current_state.reshape(-1)
        enemy_unit_type_arr = enemy_unit_type_arr.reshape(-1)
        pending = np.array(pending).reshape(-1)
        last_act = np.array([0]).reshape(-1)
        self.prev_state = np.concatenate((current_state,
                                          my_unit_type_arr,
                                          enemy_unit_type_arr,
                                          pending,
                                          last_act))

        action = self.agent.get_action(self.prev_state)
        self.last_reward = await self.activate_sub_bot(action)
        self.last_reward -= STEP_PENALTY
        self.agent.save_reward(self.last_reward)
        self.last_sc_action = action
        try:
            await self.do_actions(self.action_buffer)
        except sc2.protocol.ProtocolError:
            print("Not in game?")
            self.action_buffer = []

            return
        self.action_buffer = []

    async def activate_sub_bot(self, general_bot_out):
        if general_bot_out == 43:
            if SUPER_DEBUG:
                print("Nothing")
            return FAILED_REWARD*0.05
        elif general_bot_out == 42:
            if SUPER_DEBUG:
                print("Scouting")
            try:
                return await self.send_scout()
            except Exception as e:
                print("Scouting Exception", e)
                return FAILED_REWARD
        elif general_bot_out == 41:
            if SUPER_DEBUG:
                print("Defending")
            try:
                return await self.defend()
            except Exception as e:
                print("Defending exception", e)
                return FAILED_REWARD
        elif general_bot_out == 40:
            if SUPER_DEBUG:
                print("Mining")
            try:
                # return await self.back_to_mining()
                if self.prev_state[4] > 0 or self.prev_state[9] > self.prev_state[28]*20:
                    await self.distribute_workers()
                    return SUCCESS_MINING_REWARD
                else:
                    return FAILED_REWARD
                # return self.mining_reward
            except Exception as e:
                print("Mining Exception", e)
                return FAILED_REWARD
            # return await self.distribute_workers()
        elif general_bot_out == 39 or general_bot_out == 13:
            if SUPER_DEBUG:
                print("Attacking")
            # return await self.activate_combat()
            try:
                return await self.army_attack()
            except Exception as e:
                print("Attacking Exception", e)
                return FAILED_REWARD
        else:
            if SUPER_DEBUG:
                print("Building")
            try:
                return await self.activate_builder(general_bot_out)
            except Exception as e:
                print("Builder Exception", e)
                return FAILED_REWARD

    async def activate_builder(self, hard_coded_choice):
        """
        call on the builder_bot to choose a unit to create, and select a target location if applicable
        :param hard_coded_choice: output from the RL agent
        :return: reward for action based on legality
        """
        unit_choice = hard_coded_choice
        if SUPER_DEBUG:
            print(unit_choice)

        # Need to know unit type to know if placement is valid.
        success = FAILED_REWARD
        index_to_unit = {
            0: UnitTypeId.NEXUS,
            1: UnitTypeId.PYLON,
            2: UnitTypeId.ASSIMILATOR,
            3: UnitTypeId.GATEWAY,
            4: UnitTypeId.WARPGATE,
            5: UnitTypeId.FORGE,
            6: UnitTypeId.CYBERNETICSCORE,
            7: UnitTypeId.PHOTONCANNON,
            8: UnitTypeId.SHIELDBATTERY,
            9: UnitTypeId.ROBOTICSFACILITY,
            10: UnitTypeId.STARGATE,
            11: UnitTypeId.TWILIGHTCOUNCIL,
            12: UnitTypeId.ROBOTICSBAY,
            13: UnitTypeId.FLEETBEACON,
            14: UnitTypeId.TEMPLARARCHIVE,
            15: UnitTypeId.DARKSHRINE
        }
        if unit_choice < 16:
            # If the builder wants to build a building, we need to place it:
            if unit_choice == 0:
                # Nexuses will automatically move into expansion slots
                target_pt = None
            elif unit_choice == 2:
                # Assimilators will automatically move onto vespene geysers
                target_pt = None
            elif unit_choice == 4:
                # Warpgates are upgrades of existing Gateways
                target_pt = None
            elif unit_choice == 1:
                target_pt = self.positions_for_pylons[0]
                pos_ind = 0
            else:
                target_pt = self.positions_for_buildings[0]
                pos_ind = 0
            if target_pt is None:
                try:
                    success = await self.build_building(unit_choice, None)
                except Exception as e:
                    if DEBUG:
                        print("exception", e)
                    success = FAILED_REWARD
            elif unit_choice == 1 or self.state.psionic_matrix.covers(target_pt):
                while True:
                    if unit_choice == 1 and await self.can_place(index_to_unit[unit_choice], target_pt):
                        break  # TODO: Eventually I'd like to pop(pos_ind) off of the list...
                    elif self.state.psionic_matrix.covers(target_pt) and await self.can_place(index_to_unit[unit_choice], target_pt):
                        break  # TODO: Eventually I'd like to pop(pos_ind) off of the list...
                    else:
                        pos_ind += 1
                        if unit_choice == 1:
                            if pos_ind >= len(self.positions_for_pylons):
                                target_pt = await self.find_placement(index_to_unit[unit_choice], self.units(UnitTypeId.NEXUS).random.position)
                                break
                            target_pt = self.positions_for_pylons[pos_ind]
                        else:
                            if pos_ind >= len(self.positions_for_buildings):
                                target_pt = await self.find_placement(index_to_unit[unit_choice], self.units(UnitTypeId.PYLON).random.position)
                                break
                            target_pt = self.positions_for_buildings[pos_ind]
                try:
                    # If can place and psionic covers (or is pylon)
                    success = await self.build_building(unit_choice, target_pt)
                except Exception as e:
                    if DEBUG:
                        print("Exception", e)
                    success = FAILED_REWARD
            else:
                success = FAILED_REWARD
        elif unit_choice < 34:
            try:
                success = await self.train_unit(unit_choice)
            except Exception as e:
                if DEBUG:
                    print("Exception", e)
                success = FAILED_REWARD
        elif unit_choice < 39:
            try:
                success = await self.research_upgrade(unit_choice)
            except Exception as e:
                if DEBUG:
                    print("Exception", e)
                success = FAILED_REWARD
        return success

    async def build_building(self, action_building_index, placement_location=None):
        """
            This function takes in an action index from the RL agent and maps it to a building to build. For now, these are
            totally hand engineered. Hash. Tag. Deal.
            :param action_building_index: index of action from the RL agent
            :param placement_location: where to drop the new building
            :return: reward for action based on legality
            """
        action_indices_to_buildings = {
            0: UnitTypeId.NEXUS,
            1: UnitTypeId.PYLON,
            2: UnitTypeId.ASSIMILATOR,
            3: UnitTypeId.GATEWAY,
            4: UnitTypeId.WARPGATE,
            5: UnitTypeId.FORGE,
            6: UnitTypeId.CYBERNETICSCORE,
            7: UnitTypeId.PHOTONCANNON,
            8: UnitTypeId.SHIELDBATTERY,
            9: UnitTypeId.ROBOTICSFACILITY,
            10: UnitTypeId.STARGATE,
            11: UnitTypeId.TWILIGHTCOUNCIL,
            12: UnitTypeId.ROBOTICSBAY,
            13: UnitTypeId.FLEETBEACON,
            14: UnitTypeId.TEMPLARARCHIVE,
            15: UnitTypeId.DARKSHRINE,
            # 16: UnitTypeId.ORACLESTASISTRAP
        }
        building_target = action_indices_to_buildings[action_building_index]
        if self.workers.amount < 1:
            return FAILED_REWARD
        if DEBUG:
            if placement_location is not None:
                print("Building a", building_target.name, "at (", placement_location.position.x, ",", placement_location.position.y, ")")
            else:
                print("Building a", building_target.name)

        if self.can_afford(building_target):
            if building_target == UnitTypeId.ASSIMILATOR:
                nexi = self.units(UnitTypeId.NEXUS).ready
                for nexus in nexi:
                    geysers_nearby = self.state.vespene_geyser.closer_than(20.0, nexus)

                    for geyser in geysers_nearby:
                        worker = self.select_build_worker(geyser.position, force=True)
                        if not self.units(UnitTypeId.ASSIMILATOR).closer_than(1.0, geyser).exists and worker is not None:
                            self.action_buffer.append(worker.build(UnitTypeId.ASSIMILATOR, geyser))
                            return SUCCESS_BUILD_REWARD
                return FAILED_REWARD
            elif building_target == UnitTypeId.WARPGATE:
                for gateway in self.units(UnitTypeId.GATEWAY).ready:
                    abilities = await self.get_available_abilities(gateway)
                    if AbilityId.MORPH_WARPGATE in abilities:
                        if self.can_afford(AbilityId.MORPH_WARPGATE):
                            self.action_buffer.append(gateway(AbilityId.MORPH_WARPGATE))
                            return SUCCESS_BUILD_REWARD
                    else:
                        if self.units(UnitTypeId.CYBERNETICSCORE).exists:
                            core = self.units(UnitTypeId.CYBERNETICSCORE).random
                            abilities = await self.get_available_abilities(core)
                            if AbilityId.RESEARCH_WARPGATE in abilities:
                                    self.action_buffer.append(core(AbilityId.RESEARCH_WARPGATE))
                                    return SUCCESS_BUILD_REWARD
                return FAILED_REWARD
            elif building_target == UnitTypeId.NEXUS:
                location = await self.get_next_expansion()
                location = await self.find_placement(UnitTypeId.NEXUS, location, placement_step=1)
                worker = self.select_build_worker(location, force=True)
                if worker is not None:
                    self.action_buffer.append(worker.build(UnitTypeId.NEXUS, location))
                    if self.units(NEXUS).amount > 4:
                        return FAILED_REWARD
                    return SUCCESS_BUILD_REWARD
                else:
                    return FAILED_REWARD
            else:
                # pos_dist = random.random()*2 + 3
                pos_dist = 0.5
                pos = placement_location.position.to2.towards(random.choice(self.corners), pos_dist)

                worker = self.select_build_worker(pos, force=True)
                if worker is not None:
                    self.action_buffer.append(worker.build(building_target, pos))
                    # await self.build(building_target, near=pos)
                    if building_target == UnitTypeId.PYLON:
                        return SUCCESS_BUILD_REWARD*0.1
                    if self.units(building_target).amount > 3:
                        return FAILED_REWARD
                    return SUCCESS_BUILD_REWARD
                else:
                    return FAILED_REWARD
        return FAILED_REWARD

    async def train_unit(self, action_unit_index):
        action_to_unit = {
            16: UnitTypeId.PROBE,
            17: UnitTypeId.ZEALOT,
            18: UnitTypeId.STALKER,
            19: UnitTypeId.SENTRY,
            20: UnitTypeId.ADEPT,
            21: UnitTypeId.HIGHTEMPLAR,
            22: UnitTypeId.DARKTEMPLAR,
            23: UnitTypeId.OBSERVER,
            24: UnitTypeId.WARPPRISM,
            25: UnitTypeId.IMMORTAL,
            26: UnitTypeId.COLOSSUS,
            27: UnitTypeId.DISRUPTOR,
            28: UnitTypeId.PHOENIX,
            29: UnitTypeId.VOIDRAY,
            30: UnitTypeId.ORACLE,
            31: UnitTypeId.TEMPEST,
            32: UnitTypeId.CARRIER,
            33: UnitTypeId.MOTHERSHIP,
            # 34: UnitTypeId.INTERCEPTOR,  # TRAIN BY DEFAULT, DONT NEED TO TRAIN
            # 35: UnitTypeId.ARCHON,  # CURRENT IMPOSSIBLE WITH THE API
        }
        unit_to_build = action_to_unit[action_unit_index]
        if DEBUG:
            print("Training a", unit_to_build.name)
        warpgate_gateway_units = [UnitTypeId.ZEALOT, UnitTypeId.STALKER, UnitTypeId.ADEPT, UnitTypeId.HIGHTEMPLAR,
                                  UnitTypeId.DARKTEMPLAR,  UnitTypeId.SENTRY]
        robotics_units = [UnitTypeId.OBSERVER, UnitTypeId.IMMORTAL, UnitTypeId.WARPPRISM, UnitTypeId.COLOSSUS,
                          UnitTypeId.DISRUPTOR]
        stargate_units = [UnitTypeId.PHOENIX, UnitTypeId.ORACLE, UnitTypeId.VOIDRAY, UnitTypeId.TEMPEST,
                          UnitTypeId.CARRIER]
        if unit_to_build == UnitTypeId.PROBE:
            nexus = self.units(UnitTypeId.NEXUS).ready.random
            if nexus.noqueue and self.can_afford(unit_to_build):
                self.action_buffer.append(nexus.train(unit_to_build))
                return SUCCESS_TRAIN_REWARD
            else:
                return FAILED_REWARD
        elif unit_to_build in warpgate_gateway_units:
            # Look for a warpgate first, since those are apparently way better
                if self.units(UnitTypeId.WARPGATE).ready.exists:
                    if self.units(UnitTypeId.PYLON).ready.exists:
                        warp_at = self.units(UnitTypeId.PYLON).ready.random
                    else:
                        warp_at = self.units(UnitTypeId.NEXUS).ready.random
                    abilities_necessary = [AbilityId.WARPGATETRAIN_ZEALOT, AbilityId.WARPGATETRAIN_STALKER,
                                           AbilityId.TRAINWARP_ADEPT, AbilityId.WARPGATETRAIN_HIGHTEMPLAR,
                                           AbilityId.WARPGATETRAIN_DARKTEMPLAR, AbilityId.WARPGATETRAIN_SENTRY]
                    ability_req = abilities_necessary[warpgate_gateway_units.index(unit_to_build)]
                    building = self.units(UnitTypeId.WARPGATE).ready.random
                    abilities = await self.get_available_abilities(building)
                    # all the units have the same cooldown anyway so let's just look at ZEALOT
                    if ability_req in abilities and self.can_afford(unit_to_build):
                        pos = warp_at.position.to2.random_on_distance(4)
                        placement = await self.find_placement(ability_req, pos, placement_step=1)
                        if placement is None:
                            # Can't place
                            return FAILED_REWARD
                        self.action_buffer.append(building.warp_in(unit_to_build, placement))
                        return SUCCESS_TRAIN_REWARD
                    else:
                        return FAILED_REWARD
                elif self.units(UnitTypeId.GATEWAY).ready.exists:
                    building = self.units(UnitTypeId.GATEWAY).ready.random
                    abilities_necessary = [AbilityId.GATEWAYTRAIN_ZEALOT, AbilityId.GATEWAYTRAIN_STALKER,
                                           AbilityId.TRAIN_ADEPT, AbilityId.GATEWAYTRAIN_HIGHTEMPLAR,
                                           AbilityId.GATEWAYTRAIN_DARKTEMPLAR, AbilityId.GATEWAYTRAIN_SENTRY]
                    ability_req = abilities_necessary[warpgate_gateway_units.index(unit_to_build)]
                    abilities = await self.get_available_abilities(building)
                    if ability_req in abilities and self.can_afford(unit_to_build):
                        self.action_buffer.append(building.train(unit_to_build))
                        return SUCCESS_TRAIN_REWARD*3
                    else:
                        # Can't afford or can't make unit type
                        return FAILED_REWARD
        elif unit_to_build in robotics_units:
            if self.units(UnitTypeId.ROBOTICSFACILITY).ready.exists:

                building = self.units(UnitTypeId.ROBOTICSFACILITY).ready.random
                abilities_necessary = [AbilityId.ROBOTICSFACILITYTRAIN_OBSERVER,
                                       AbilityId.ROBOTICSFACILITYTRAIN_IMMORTAL,
                                       AbilityId.ROBOTICSFACILITYTRAIN_WARPPRISM,
                                       AbilityId.ROBOTICSFACILITYTRAIN_COLOSSUS, AbilityId.TRAIN_DISRUPTOR]
                ability_req = abilities_necessary[robotics_units.index(unit_to_build)]
                abilities = await self.get_available_abilities(building)
                if ability_req in abilities and self.can_afford(unit_to_build):
                    self.action_buffer.append(building.train(unit_to_build))
                    return SUCCESS_TRAIN_REWARD
                else:
                    return FAILED_REWARD
            else:
                return FAILED_REWARD
        elif unit_to_build in stargate_units:
            if self.units(UnitTypeId.STARGATE).ready.exists:
                building = self.units(UnitTypeId.STARGATE).ready.random
                abilities_necessary = [AbilityId.STARGATETRAIN_PHOENIX, AbilityId.STARGATETRAIN_ORACLE,
                                       AbilityId.STARGATETRAIN_VOIDRAY,
                                       AbilityId.STARGATETRAIN_CARRIER, AbilityId.STARGATETRAIN_TEMPEST]
                ability_req = abilities_necessary[stargate_units.index(unit_to_build)]
                abilities = await self.get_available_abilities(building)
                if ability_req in abilities and self.can_afford(unit_to_build):
                    self.action_buffer.append(building.train(unit_to_build))
                    return SUCCESS_TRAIN_REWARD*5
                else:
                    return FAILED_REWARD
            else:
                return FAILED_REWARD
        elif unit_to_build == UnitTypeId.MOTHERSHIP:
            nexus = self.units(UnitTypeId.NEXUS).ready.random
            nexus_abilities = await self.get_available_abilities(nexus)
            if self.can_afford(unit_to_build) and AbilityId.NEXUSTRAINMOTHERSHIP_MOTHERSHIP in nexus_abilities:
                self.action_buffer.append(nexus.train(unit_to_build))
                return SUCCESS_TRAIN_REWARD
            else:
                return FAILED_REWARD
        return FAILED_REWARD

    async def army_attack(self):
        """
        :return: reward based on if we sent out attackers
        """
        def in_range(unit, target, gap=0.5):
            # if gap > 0, in_range will be true if unit is closer than maximum fire range
            return unit.distance_to(target) + gap < unit.radius + target.radius + max(unit.air_range, unit.ground_range)

        def half_health(unit):
            return unit.health + unit.shield < min(unit.health_max, unit.shield_max)

        attacking_classes = [UnitTypeId.ZEALOT, UnitTypeId.STALKER, UnitTypeId.SENTRY, UnitTypeId.ADEPT,
                             UnitTypeId.HIGHTEMPLAR, UnitTypeId.DARKTEMPLAR, UnitTypeId.IMMORTAL, UnitTypeId.COLOSSUS,
                             UnitTypeId.PHOENIX, UnitTypeId.VOIDRAY, UnitTypeId.ORACLE, UnitTypeId.TEMPEST,
                             UnitTypeId.CARRIER, UnitTypeId.INTERCEPTOR, UnitTypeId.MOTHERSHIP]

        all_targets = self.known_enemy_units.exclude_type([UnitTypeId.LARVA, UnitTypeId.EGG])
        ground_targets = all_targets.not_flying
        flying_targets = all_targets.flying
        attacking_ground_targets = ground_targets.filter(lambda x: (x.can_attack_ground or x.can_attack_air))
        attacking_flying_targets = flying_targets.filter(lambda x: (x.can_attack_ground or x.can_attack_air))
        my_army = self.units.of_type(attacking_classes)
        army_half_health = len([unit for unit in my_army if half_health(unit)])
        # if army_half_health <= self.army_below_half and self.state.game_loop - self.last_attack_loop > 120:  # If nobody got hurt and we did this recently, save time and don't do it again.
        #     self.last_attack_loop = self.state.game_loop
        #     self.army_below_half = army_half_health
        #     return SUCCESS_ATTACK_REWARD
        self.last_attack_loop = self.state.game_loop
        self.army_below_half = army_half_health

        for unit in my_army:
            if unit.type_id == UnitTypeId.SENTRY and await self.can_cast(unit, AbilityId.GUARDIANSHIELD_GUARDIANSHIELD):
                await self.do(unit(AbilityId.GUARDIANSHIELD_GUARDIANSHIELD))

            target = None
            nearest_target = None
            possible_targets = all_targets.filter(lambda x: unit.target_in_range(x))
            if unit.type_id == UnitTypeId.STALKER and attacking_flying_targets.exists:
                target = attacking_flying_targets.closest_to(unit)
            elif possible_targets:
                    nearest_target = possible_targets.closest_to(unit)
                    target = min(possible_targets, key=lambda t: (t.health + t.shield) /
                                                                 (t.ground_dps + t.air_dps + 1))
            treat_outrange = -100
            if all_targets:
                if unit.type_id is UnitTypeId.COLOSSUS:
                    treat = max(all_targets, key=lambda e: max(e.ground_range, e.air_range)** 2 -
                                                                e.position._distance_squared(unit.position))
                    treat_outrange = unit.radius + treat.radius + max(treat.ground_range, treat.air_range) + 1 - \
                                     unit.distance_to(treat)
                elif unit.is_flying:
                    treat = max(all_targets, key=lambda e: e.air_range ** 2 -
                                                                e.position._distance_squared(unit.position))
                    treat_outrange = unit.radius + treat.radius + treat.air_range + 1 - \
                                     unit.distance_to(treat)
                else:
                    treat = max(all_targets, key=lambda e: e.ground_range ** 2 -
                                                                e.position._distance_squared(unit.position))
                    treat_outrange = unit.radius + treat.radius + treat.ground_range + 1 - \
                                     unit.distance_to(treat)

            if target is None:
                if attacking_flying_targets.exists and unit.can_attack_air:
                    target = attacking_flying_targets.closest_to(unit)
                elif attacking_ground_targets.exists:
                    target = attacking_ground_targets.closest_to(unit)
                elif flying_targets.exists and unit.can_attack_air:
                    target = flying_targets.closest_to(unit)
                elif ground_targets.exists:
                    target = ground_targets.closest_to(unit)
                else:
                    target = random.choice(self.enemy_start_locations).position
            retreat = False
            if unit.type_id is not UnitTypeId.ZEALOT:
                if (unit.type_id is not UnitTypeId.VOIDRAY and unit.weapon_cooldown > 2 and
                    possible_targets.filter(lambda x: (x.can_attack_air or x.can_attack_ground)).exists
                    and nearest_target is not None and (in_range(unit, nearest_target) or treat_outrange >= 0)) or \
                        (army_half_health < my_army.amount / 2 and half_health(unit) and treat_outrange+unit.radius*2>=0
                        or unit.type_id is UnitTypeId.SENTRY):
                    # target = self.start_location
                    retreat = True
                    if unit.distance_to(self.start_location) < 5:
                        target = all_targets.closest_to(unit)
            unit_abs = await self.get_available_abilities(unit)
            if AbilityId.EFFECT_BLINK_STALKER in unit_abs and await self.can_cast(unit, AbilityId.EFFECT_BLINK_STALKER)\
                    and unit.distance_to(target) > 6 and \
                    target is not UnitTypeId.ZERGLING:  # don't jump into zerglings
                if not retreat:
                    towards = -unit.radius + unit.ground_range  # one corpus closer
                    if isinstance(target, sc2Unit):
                        towards += target.radius
                        target = target.position.towards(unit.position, towards)
                    else:
                        target = target.towards(unit.position, towards)
                await self.do(unit(AbilityId.EFFECT_BLINK_STALKER, target))
                if retreat and possible_targets.exists:
                    # blink in other direction if can't blink back to base
                    alt_target = unit.position.towards(possible_targets.center, -8)
                    self.action_buffer.append(unit(AbilityId.EFFECT_BLINK_STALKER, alt_target, queue=True))
            if retreat:
                self.action_buffer.append(unit.move(target))
            else:
                self.action_buffer.append(unit.attack(target))
        attack_reward = min((len(my_army)-5)*SUCCESS_ATTACK_REWARD, 10)
        return attack_reward

    async def defend(self):
        """
        This function find enemy army centroids, and the number of attackers in each army
        so if we're being hit on 3 fronts, it'll attempt to distribute to each.
        #TODO: intelligently grab first available attackers (like.. grab nearest)
        :return: success for a defensive mobilization, fail for nobody to defend or nothing to defend
        """
        enemy_positions = []
        num_attackers = []
        for depo in self.units(UnitTypeId.NEXUS):
            num_attackers.append(0)
            have_position = False
            for unit in self.known_enemy_units.not_structure:
                if unit.position.to2.distance_to(depo.position.to2) < 15:
                    if not have_position:
                        enemy_positions.append(unit.position)
                        have_position = True
                    num_attackers[-1] += 1
        defensive_orders = []
        for position, army_size in zip(enemy_positions, num_attackers):
            attacking_classes = [UnitTypeId.ZEALOT, UnitTypeId.STALKER, UnitTypeId.SENTRY, UnitTypeId.ADEPT,
                                 UnitTypeId.HIGHTEMPLAR, UnitTypeId.DARKTEMPLAR, UnitTypeId.IMMORTAL,
                                 UnitTypeId.COLOSSUS,
                                 UnitTypeId.PHOENIX, UnitTypeId.VOIDRAY, UnitTypeId.ORACLE, UnitTypeId.TEMPEST,
                                 UnitTypeId.CARRIER, UnitTypeId.INTERCEPTOR, UnitTypeId.MOTHERSHIP]
            for unit_type in attacking_classes:
                for unit in self.units(unit_type).idle:
                    defensive_orders.append(unit.attack(position))
                    if len(defensive_orders) > army_size:
                        break
            if len(defensive_orders) < army_size:
                for worker in self.workers:
                    defensive_orders.append(worker.attack(position))
                    if len(defensive_orders) > army_size:
                        break

        self.action_buffer.extend(defensive_orders)
        if len(defensive_orders) > 0:
            return SUCCESS_ATTACK_REWARD
        else:
            return FAILED_REWARD

    async def send_scout(self):
        """
        :return: fail if we scouted too recently, success otherwise
        """
        # TODO: add global list of recently checked base locations so that I can prioritize based on recency of check
        # TODO: use global list of recency / discovery to reward bot for scouting that pays off
        # If I have scouted recently (within 200 iterations?), don't re-scout
        if self.itercount - self.last_scout_iteration < 200:
            return FAILED_REWARD
        a_scout = None
        places_to_look = []
        scout_loc = random.choice(self.enemy_start_locations)

        for base_location in self.expansion_locations.keys():
            if base_location.position in self.scouted_bases:
                continue
            else:
                places_to_look.append(base_location)
        if len(places_to_look) == 0:
            self.scouted_bases = []

        if len(self.known_enemy_structures) < 1:
            # If we haven't found their main yet, prioritize enemy base locs
            enemy_base_locs = self.enemy_start_locations
            for enemy_base_loc in enemy_base_locs:
                invalid = False
                # This one is a candidate
                if enemy_base_loc.position in self.scouted_bases:
                    invalid = True
                    # If we've been nearby, it's no longer valid
                if not invalid:
                    # If we haven't been nearby, this is valid. break search
                    scout_loc = enemy_base_loc
                    break
        elif len(places_to_look) > 0:
            scout_loc = random.choice(places_to_look).position
        else:
            scout_loc = random.choice(self.enemy_start_locations)

        if self.units(UnitTypeId.OBSERVER).exists:
            a_scout = self.units(UnitTypeId.OBSERVER).closest_to(scout_loc)
        elif self.workers.exists:
            if self.workers.idle.exists:
                a_scout = self.workers.idle.closest_to(scout_loc)
            else:
                a_scout = self.workers.closest_to(scout_loc)

        if a_scout is not None:
            self.action_buffer.append(a_scout.move(scout_loc))
            self.last_scout_iteration = self.itercount  # Track when we sent out our last scout
            self.scouted_bases.append(scout_loc)
            return SUCCESS_SCOUT_REWARD
        else:
            return FAILED_REWARD

    async def research_upgrade(self, research_index):
        index_to_upgrade = {
            34: "ground_attacks",
            35: "air_attacks",
            36: "ground_armor",
            37: "air_armor",
            38: "shields",
            # 39: "speed",
            # 40: "range",
            # 41: "spells",
            # 42: "misc"
        }
        research_topic = index_to_upgrade[research_index]
        if research_topic in ["ground_attacks", "ground_armor", "shields"]:
            if self.units(UnitTypeId.FORGE).exists:
                forge = self.units(UnitTypeId.FORGE).random
                abilities = await self.get_available_abilities(forge)
                if research_topic == "ground_attacks":
                    all_topics = [AbilityId.FORGERESEARCH_PROTOSSGROUNDWEAPONSLEVEL1,
                                  AbilityId.FORGERESEARCH_PROTOSSGROUNDWEAPONSLEVEL2,
                                  AbilityId.FORGERESEARCH_PROTOSSGROUNDWEAPONSLEVEL3]
                elif research_topic == "ground_armor":
                    all_topics = [AbilityId.FORGERESEARCH_PROTOSSGROUNDARMORLEVEL1,
                                  AbilityId.FORGERESEARCH_PROTOSSGROUNDARMORLEVEL2,
                                  AbilityId.FORGERESEARCH_PROTOSSGROUNDARMORLEVEL3]
                elif research_topic == "shields":
                    all_topics = [AbilityId.FORGERESEARCH_PROTOSSSHIELDSLEVEL1,
                                  AbilityId.FORGERESEARCH_PROTOSSSHIELDSLEVEL2,
                                  AbilityId.FORGERESEARCH_PROTOSSSHIELDSLEVEL3]
                for topic in all_topics:
                    if topic in abilities:
                        self.action_buffer.append(forge(topic))
                        return SUCCESS_BUILD_REWARD
            return FAILED_REWARD
        elif research_topic in ["air_attacks", "air_armor"]:
            if self.units(UnitTypeId.CYBERNETICSCORE).exists:
                core = self.units(UnitTypeId.CYBERNETICSCORE).random
                abilities = await self.get_available_abilities(core)
                if research_topic == "air_attacks":
                    air_research = [AbilityId.CYBERNETICSCORERESEARCH_PROTOSSAIRWEAPONSLEVEL1,
                                    AbilityId.CYBERNETICSCORERESEARCH_PROTOSSAIRWEAPONSLEVEL2,
                                    AbilityId.CYBERNETICSCORERESEARCH_PROTOSSAIRWEAPONSLEVEL3]
                elif research_topic == "air_armor":
                    air_research = [AbilityId.CYBERNETICSCORERESEARCH_PROTOSSAIRARMORLEVEL1,
                                    AbilityId.CYBERNETICSCORERESEARCH_PROTOSSAIRARMORLEVEL1,
                                    AbilityId.CYBERNETICSCORERESEARCH_PROTOSSAIRARMORLEVEL1]
                for topic in air_research:
                    if topic in abilities:
                        self.action_buffer.append(core(topic))
                        return SUCCESS_BUILD_REWARD
            return FAILED_REWARD

    async def back_to_mining(self):
        reassigned_miners = 0
        for a in self.units(UnitTypeId.ASSIMILATOR):
            if a.assigned_harvesters < a.ideal_harvesters:
                w = self.workers.closer_than(20, a)
                if w.exists:
                    self.action_buffer.append(w.random.gather(a))
                    reassigned_miners += 1
        for idle_worker in self.workers.idle:
            mf = self.state.mineral_field.closest_to(idle_worker)
            self.action_buffer.append(idle_worker.gather(mf))
            reassigned_miners += 1
        if reassigned_miners > 0:
            return SUCCESS_MINING_REWARD
        else:
            return FAILED_REWARD

    def finish_episode(self, game_result):
        print("Game over!")
        if game_result == sc2.Result.Defeat:
            reward = -250  # + self.itercount/500.0 + self.units.amount
        elif game_result == sc2.Result.Tie:
            reward = 50
        elif game_result == sc2.Result.Victory:
            reward = 250  # - min(self.itercount/500.0, 900) + self.units.amount
        else:
            # ???
            return -13
        bot_fn = '../txts/' + self.agent.bot_name + '_victories.txt'
        with open(bot_fn, "a") as myfile:
            myfile.write(str(reward) + '\n')
        self.agent.save_reward(reward)
        R = 0
        rewards = []
        all_rewards = self.agent.replay_buffer.rewards_list
        reward_sum = sum(all_rewards)
        all_values = self.agent.replay_buffer.value_list
        deeper_all_values = self.agent.replay_buffer.deeper_value_list
        # Discount future rewards back to the present using gamma
        advantages = []
        deeper_advantages = []

        for r, v, d_v in zip(all_rewards[::-1], all_values[::-1], deeper_all_values[::-1]):
            R = r + 0.99 * R
            rewards.insert(0, R)
            advantages.insert(0, R - v)
            if d_v is not None:
                deeper_advantages.insert(0, R - d_v)
        advantages = torch.Tensor(advantages)
        rewards = torch.Tensor(rewards)

        if len(deeper_advantages) > 0:
            deeper_advantages = torch.Tensor(deeper_advantages)
            deeper_advantages = (deeper_advantages - deeper_advantages.mean()) / (
                    deeper_advantages.std() + torch.Tensor([np.finfo(np.float32).eps]))
            self.agent.replay_buffer.deeper_advantage_list = deeper_advantages.detach().clone().cpu().numpy().tolist()
        else:
            self.agent.replay_buffer.deeper_advantage_list = [None] * len(all_rewards)
        # Scale rewards
        rewards = (rewards - rewards.mean()) / (rewards.std() + torch.Tensor([np.finfo(np.float32).eps]))
        advantages = (advantages - advantages.mean()) / (advantages.std() + torch.Tensor([np.finfo(np.float32).eps]))
        self.agent.replay_buffer.rewards_list = rewards.detach().clone().cpu().numpy().tolist()
        self.agent.replay_buffer.advantage_list = advantages.detach().clone().cpu().numpy().tolist()
        return reward_sum

def run_episode(q, main_agent):
    result = None
    agent_in = main_agent.duplicate()

    bot = StarmniBot(rl_agent=agent_in)
    opponents = [
        Computer(Race.Zerg, Difficulty.Medium),
        # Computer(Race.Protoss, Difficulty.VeryEasy),
        # Computer(Race.Terran, Difficulty.VeryEasy)
    ]
    enemy = random.choice(opponents)

    try:
        result = sc2.run_game(sc2.maps.get("Acid Plant LE"),
                              [Bot(Race.Protoss, bot), enemy],
                              realtime=False)
    except KeyboardInterrupt:
        result = [-1, -1]
    except Exception as e:
        print(str(e))
        print("No worries", e, " carry on please")
    if type(result) == list and len(result) > 1:
        result = result[0]

    reward_sum = bot.finish_episode(result)
    if q is not None:
        try:
            q.put([reward_sum, bot.agent.replay_buffer.__getstate__()])
        except RuntimeError as e:
            print(e)
            return [reward_sum, bot.agent.replay_buffer.__getstate__()]
    return [reward_sum, bot.agent.replay_buffer.__getstate__()]

# def main(episodes, agent_in, num_processes, reset_on_fail=False):
#     running_reward_array = []
#     # lowered = False
#     agent = agent_in.duplicate()
# #     mp.set_start_method('spawn')
#     for episode in range(episodes):
#         successful_runs = 0
#         master_reward, reward, running_reward = 0, 0, 0
#         processes = []
#         queueue = mp.Manager().Queue()
#         for proc in range(num_processes):
#             p = mp.Process(target=run_episode, args=(queueue, agent))
#             p.start()
#             processes.append(p)
#         for p in processes:
#             p.join()
#         while not queueue.empty():
#             try:
#                 fake_out = queueue.get()
#             except MemoryError as e:
#                 print(e)
#                 fake_out = [-13, None]
#             if fake_out[0] != -13:
#                 master_reward += fake_out[0]
#                 running_reward_array.append(fake_out[0])
#                 agent.replay_buffer.extend(fake_out[1])
#                 successful_runs += 1
#
#         if successful_runs > 0:
#             if (reset_on_fail and master_reward > 0) or not reset_on_fail:
#                 reward = master_reward / float(successful_runs)
#                 agent.end_episode(reward, num_processes)
#                 running_reward = sum(running_reward_array[-100:]) / float(min(100.0, len(running_reward_array)))
#                 agent.save('../models/last')
#             else:
#                 agent = agent_in.duplicate()
#                 print("Resetting...")
#
#         if episode % 50 == 0:
#             print(f'Episode {episode}  Last Reward: {reward}  Average Reward: {running_reward}')
#             # print(f"Running {num_processes} concurrent simulations per episode")
#     return running_reward_array

def bernoulli_main(episodes, agent_in, num_processes):
    def bernoulli_test(p, n, k, alpha):
        coef = math.factorial(n) / (math.factorial(k) * math.factorial(n - k))
        p_to_k = p ** k
        q_to_n_k = (1 - p) ** (n - k)
        test_res = coef * p_to_k * q_to_n_k
        if test_res < alpha:
            return True
        else:
            return False

    win_prob = 0.25
    min_games = 15
    alpha = 0.05
    k, n, successful_runs, master_reward, reward, running_reward = 0, 0, 0, 0, 0, 0
    find_new_step = True
    running_reward_array = []
    # lowered = False
    agent = agent_in.duplicate()
    # mp.set_start_method('spawn')
    for episode in range(6, episodes):
        try:
            last_agent = agent.duplicate()
            if find_new_step:
                successful_runs = 0
                master_reward, reward, running_reward = 0, 0, 0
                tuple_out = run_episode(None,
                                        agent)
                if tuple_out[0] != -13:
                    master_reward += tuple_out[0]
                    running_reward_array.append(tuple_out[0])
                    agent.replay_buffer.extend(tuple_out[1])
                    successful_runs += 1
            if successful_runs > 0 or not find_new_step:  # if we have a non-empty replay buffer
                reward = master_reward
                agent.end_episode(reward, num_processes)  # take a gradient step
                running_reward = sum(running_reward_array[-100:]) / float(min(100.0, len(running_reward_array)))
                while True:  # Do the bernoulli testing

                    tuple_out = run_episode(None, agent)
                    if tuple_out[0] != -13:
                        if tuple_out[1]['rewards'][-1] > 0:
                            k += 1
                        n += 1
                        if n < num_processes:
                            master_reward += tuple_out[0]
                            agent.replay_buffer.extend(tuple_out[1])
                    if n >= min_games:
                        new_win_prob = float(k)/float(n)
                        if bernoulli_test(p=win_prob, n=n, k=k, alpha=alpha):
                            if new_win_prob > win_prob:
                                win_prob = new_win_prob
                                k = 0
                                n = 0
                                find_new_step = False
                                agent.save(f'../models/{episode}')
                                break
                            else:
                                agent = last_agent
                                k = 0
                                n = 0
                                find_new_step = True
                                break
                    print(f"After this episode, n={n}")
                    if n > 100 and k > n*.95:
                        print("Finishing this model")
                        agent.save('FINAL')
                        return

            if episode % 50 == 0:
                print(f'Episode {episode}  Last Reward: {reward}  Average Reward: {running_reward}')
                # print(f"Running {num_processes} concurrent simulations per episode")
        except RuntimeError as e:
            print(e)
            find_new_step = True
            time.sleep(5)
            continue
    return running_reward_array

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-a", "--agent_type", help="architecture of agent to run", type=str, default='djinn')
    parser.add_argument("-e", "--episodes", help="how many episodes", type=int, default=1000)
    parser.add_argument("-p", "--processes", help="how many processes?", type=int, default=1)
    parser.add_argument("-env", "--env_type", help="environment to run on", type=str, default='cart')
    parser.add_argument("-gpu", help="run on GPU?", action='store_true')
    parser.add_argument("-vec", help="Vectorized ProLoNet?", action='store_true')
    parser.add_argument("-adv", help="Adversarial ProLoNet?", action='store_true')
    parser.add_argument("-rand", help="Random ProLoNet?", action='store_true')
    parser.add_argument("-deep", help="Deepening?", action='store_true')
    parser.add_argument("-s", "--sl_init", help="sl to rl for fc net?", action='store_true')

    args = parser.parse_args()
    AGENT_TYPE = args.agent_type  # 'shallow_prolo', 'prolo', 'random', 'fc', 'lstm'
    ADVERSARIAL = args.adv  # Adversarial prolo, applies for AGENT_TYPE=='shallow_prolo'
    SL_INIT = args.sl_init  # SL->RL fc, applies only for AGENT_TYPE=='fc'
    NUM_EPS = args.episodes  # num episodes Default 1000
    NUM_PROCS = args.processes  # num concurrent processes Default 1
    ENV_TYPE = args.env_type  # 'cart' or 'lunar' Default 'cart'
    USE_GPU = args.gpu  # Applies for 'prolo' only. use gpu? Default false
    VECTORIZED = args.vec  # Applies for 'prolo' vectorized or no? Default false
    RANDOM = args.rand  # Applies for 'prolo' random init or no? Default false
    DEEPEN = args.deep  # Applies for 'prolo' deepen or no? Default false
    # torch.set_num_threads(NUM_PROCS)
    dim_in = 194
    dim_out = 44
    bot_name = AGENT_TYPE + 'SC_Macro'+'Medium'
    mp.set_sharing_strategy('file_system')
    if AGENT_TYPE == 'prolo':
        policy_agent = DeepProLoNet(distribution='one_hot',
                                    bot_name=bot_name,
                                    input_dim=dim_in,
                                    output_dim=dim_out,
                                    use_gpu=USE_GPU,
                                    vectorized=VECTORIZED,
                                    randomized=RANDOM,
                                    adversarial=ADVERSARIAL,
                                    deepen=DEEPEN)
    elif AGENT_TYPE == 'fc':
        policy_agent = FCNet(input_dim=dim_in,
                             bot_name=bot_name,
                             output_dim=dim_out,
                             sl_init=SL_INIT)
    elif AGENT_TYPE == 'lstm':
        policy_agent = LSTMNet(input_dim=dim_in,
                               bot_name=bot_name,
                               output_dim=dim_out)
    elif AGENT_TYPE == 'djinn':
        policy_agent = DJINNAgent(bot_name=bot_name,
                                  input_dim=dim_in,
                                  output_dim=dim_out)
    else:
        raise Exception('No valid network selected')
    start_time = time.time()

    # main(episodes=NUM_EPS, agent_in=policy_agent, num_processes=NUM_PROCS, reset_on_fail=True)
    policy_agent.load('../models/5')
    bernoulli_main(episodes=NUM_EPS, agent_in=policy_agent, num_processes=NUM_PROCS)
    with open(bot_name+'time.txt', 'a') as writer:
        writer.write(str(time.time()-start_time))

max_stars_repo_name: CORE-Robotics-Lab/ProLoNets

max_issues_repo_name: CORE-Robotics-Lab/ProLoNets

max_forks_repo_name: CORE-Robotics-Lab/ProLoNets


text:
import numpy as np

w1 = np.array([[0.2, 0.2, 0.2], [0.4, 0.4, 0.4], [0.6, 0.6, 0.6]])

w2 = np.zeros((1, 3))

w2[0, :] = np.array([0.5, 0.5, 0.5])

b1 = np.array([0.8, 0.8, 0.8])

b2 = np.array([0.2])

def f(x):
    return 1 / (1 + np.exp(-x))

def simple_looped_nn_calc(n_layers, x, w, b):
    for l in range(n_layers - 1):
        if l == 0:
            node_in = x
        else:
            node_in = h

        z = w[l].dot(node_in) + b[l]

        h = f(z)

    return h

w = [w1, w2]
b = [b1, b2]

x = [1.5, 2.0, 3.0]

simple_looped_nn_calc(3, w, w, b)

max_stars_repo_name: llGuy/tf-tests

max_issues_repo_name: llGuy/tf-tests

max_forks_repo_name: llGuy/tf-tests