import numpy as np
from mifs import mifs
from tqdm import tqdm

from sklearn.metrics import balanced_accuracy_score

from keras.models import Sequential
from keras.layers import Dense, Activation
from keras.optimizers import RMSprop

x = np.load("ANI_Training.Input.npy")
y = np.load("ANI_Training.Label.npy")

ixx = np.load("ixx.npy")
ixy = np.load("ixy.npy")

x_train = x[:2700]
x_test = x[2700:]

y_train = y[:2700]
y_test = y[2700:]


def minus1(x):
    if x == 0:
        return -1
    else:
        return 1

def divide(x):
    if x < 0:
        return -1
    else:
        return 1

def ber(y_true, y_pred):
    return 1 - balanced_accuracy_score(y_true, y_pred)

def main():
    sel = 9
    pop = []

    selection = mifs(ixx, ixy, sel)

    for _ in range(4):
        pop.append(make_gene(selection))

    bers = []


    while(1):

        for gene in tqdm(pop):
            model = Sequential()
            model.add(Dense(sel, input_dim=len(gene_to_selection(gene))))
            model.add(Activation('tanh'))

            model.add(Dense(sel))
            model.add(Activation('tanh'))

            model.add(Dense(sel))
            model.add(Activation('tanh'))

            model.add(Dense(1))
            model.add(Activation('sigmoid'))

            gene = np.array(gene)

            x_train_sel = x_train[:,gene_to_selection(gene)]
            x_test_sel = x_test[:,gene_to_selection(gene)]

            model.compile(optimizer=RMSprop(lr=4e-3),
                        loss='mse')


            for _ in tqdm(range(20)):
                model.fit(x_train_sel, y_train, epochs=50, batch_size=len(x_train), verbose=0)


            y_pred = np.round(model.predict(x_test_sel)).astype(int)

            bers.append(ber(y_test, y_pred))

        print(bers)
        ber_ranks = np.argsort(bers)
        print(ber_ranks)
        best_ber_ranks = ber_ranks[:2]
        print(best_ber_ranks)
        best_bers = bers[best_ber_ranks]
        print(best_bers)
        best_genes = pop[best_ber_ranks]
        print(best_genes)

        new_genes = []

        for g1, ber1 in zip(best_genes, best_bers):
            for g2, ber2 in zip(best_genes, best_bers):
                new_genes.append(crossover(selection, g1, 1-ber1, g2, 1-ber2))

        pop = new_genes
        print("BEST:", best_bers)


def make_gene(selection, size=250, high_bits=10):
    gene = np.array([0] * size)
    gene[selection] = 1

    choices = np.random.choice(range(size), size=high_bits)

    gene[choices] = 1

    return gene

def crossover(selection, gene_1, score_1, gene_2, score_2, size=250, mutation_rate=0.05):
    gene = []
    weight = score_1 / (score_1 + score_2)

    randoms = np.random.random(size=size)
    mutate = np.random.random(size=size)

    for g1, g2, r, m in zip(gene_1, gene_2, randoms, mutate):
        if m < mutation_rate:
            gene.append(np.random.choice([0,1]))
        else:
            if r > weight:
                gene.append(g1)
            else:
                gene.append(g2)

    gene[selection] = 1

    return gene

def gene_to_selection(gene):
    selection = []
    for g, i in zip(gene, range(len(gene))):
        if g == 1:
            selection.append(i)

    return np.array(selection)




if __name__ == '__main__':

    main()