import numpy as np
from sklearn.linear_model import LogisticRegression

#shuffles the data according to the distance to decision boundary, rotating between tp,fp,tn,fn
def shuffle_rotate(X, y0, y1):
    
    clf = LogisticRegression(penalty="none").fit(X, y0)
    
    scores = clf.decision_function(X)
    scores = np.abs(scores)

    indices = np.argsort(scores)

    X = X[indices]
    y0 = y0[indices]
    y1 = y1[indices]

    preds = clf.predict(X)
    tp = ((preds == 1) & (y0 == 1)).nonzero()[0]
    fp = ((preds == 1) & (y0 == 0)).nonzero()[0]
    tn = ((preds == 0) & (y0 == 0)).nonzero()[0]
    fn = ((preds == 0) & (y0 == 1)).nonzero()[0]

    indices = []
    rotation = [tp, fp, tn, fn]

    for i in range(X.shape[0]):
        for arr in rotation:
            if i >= arr.shape[0]:
                continue
            indices.append(arr[i])

    # print(tp, fp, tn, fn)
    # print(indices)
    
    X = X[indices]
    y0 = y0[indices]
    y1 = y1[indices]

    return X, y0, y1