#!/usr/bin/env python3

#               ------------USP - São Carlos------------
#        David Cairuz da Silva - 10830061 - davidcairuz@gmail.com
#       João Guilherme Madeira Araújo - 9725165 - joaogui1@usp.br
#          Luísa Souza Moura - 10692179 - luisamoura@usp.br
#  Daniel Sá Barretto Prado Garcia - 10374344 - danielbarretto@usp.br

import os
from os import system
from copy import deepcopy
import pandas as pd
from sklearn import metrics
from sklearn import preprocessing
import matplotlib.pyplot as plt

import numpy as np

def mutate(ind, mutation_rate):
    '''Mutates individuals'''
    ret = np.asarray(deepcopy(ind))
    ret = ret + np.random.normal(0, mutation_rate, list(ret.shape))
    return ret


def calinskiHarabaz(points, centers):
    '''Calculates fitness function'''
    ret = 0.0
    clusters = [[c] for c in centers]
    label = [0 for _ in enumerate(points)]

    for pos, p in enumerate(points):
        closest = np.argmin([np.linalg.norm(p - c) for c in centers])
        clusters[closest].append(p)
        label[pos] = closest

    if len(set(label)) == 1:
        return 0.001

    ret = metrics.calinski_harabaz_score(points, label)
    return ret


def helper():
    '''Prints help menu'''
    print("\033[94m---------Darwin commands---------")
    print(">Load [file]: Loads data from [file]", ">print data: prints loaded data",\
     ">plot data: plots data", sep="\n")
    print(">Params: Print the enviroment parameters")
    print(">Set clusters x: sets the ammount of clusters to x", \
    ">Set population x: \
    sets the number of individuals to x", ">Set selection x: Sets selection type to x", sep='\n')
    print(">Set generations x: sets the number of generations to x", \
    ">Set mutation x: sets mutation rate to x", sep='\n')
    print(">init: initializes population", ">evolve: evolves the population", \
    ">plot clusters: Prints data, now clustered", sep="\n")
    print(">normalize: scales data to achieve zero mean and unit variance\
     (use before loading data)", sep="\n")
    print(">Compare: Plots comparison graph between the different selection methods", \
    ">ignore lines x: skips x lines from the csv", sep="\n")
    print(">ignore columns x: ignores the x last columns", \
    ">quit: quits the program\033[0m", sep="\n")


def init_pop(points, num_coord, num_genes, pop_size):
    '''Initializes population with random values for its genes'''
    pop = []
    mean = sum(points)/len(points)
    aux = np.asarray([mean for _ in range(num_genes)])
    for _ in range(pop_size):
        pop.append(aux + np.random.normal(0, 2, list(aux.shape)))
    return pop


def crossover(ind1, ind2):
    '''Mates two individuals'''
    ret = deepcopy(ind1)
    for pos, a in enumerate(ret):
        closest = np.argmin([np.linalg.norm(a - b) for b in ind2])
        ret[pos] = (a + ind2[closest])/2.0
    return ret


def elitism(fitness, pop, points, generations, mutation_rate, in_compare = 0):
    '''Selection method: elitism'''
    '''A new popoulation is generated by cloning the fittest individual and mutating it'''
    history = []
    for i in range(generations):
        fit = np.array([fitness(points, p) for p in pop])
        pop[0], pop[np.argmax([fit])] = pop[np.argmax([fit])], pop[0]
        ubermensch = deepcopy(pop[0]) #fittest individual
        if len(set(history[-8:])) == 1:
            pop = [mutate(ubermensch, 5*mutation_rate) for _ in pop[1:]]
        else:
            pop = [mutate(ubermensch, mutation_rate) for _ in pop[1:]]

        pop.append(ubermensch)
        history.append(max(fit))

        with open('graphs/log.txt', 'a') as log_file:
            print("generation ", i, " fitness ", history[-1])
            log_file.write("generation {} fitness {}\n".format(i, history[-1]))

    plot_clusters(points, ubermensch, "elitism", in_compare)
    return pop, ubermensch, history


def harem(fitness, pop, points, generations, mutation_rate, in_compare=0):
    '''Selection method: harem
    Generates a new population by crossing the fittest individual with
    each member of the old generation and mutating them'''
    history = []
    for i in range(generations):
        fit = np.array([fitness(points, p) for p in pop])
        pop[0], pop[np.argmax([fit])] = pop[np.argmax([fit])], pop[0]
        ubermensch = deepcopy(pop[0]) #fittest individual

        if len(set(history[-8:])) == 1 :
            pop = [mutate(crossover(ubermensch, p), 5*mutation_rate) for p in pop[1:]]
        else:
            pop = [mutate(crossover(ubermensch, p), mutation_rate) for p in pop[1:]]

        pop.append(ubermensch)
        history.append(max(fit))
        #Populate log with the fittest individual's info from each generation
        with open('graphs/log.txt', 'a') as log_file:
            print("generation ", i, " fitness ", history[-1])
            log_file.write("generation {} fitness {}\n".format(i, history[-1]))

    plot_clusters(points, ubermensch, "harem", in_compare)
    return pop, ubermensch, history


def roulette(fitness, pop, points, generations, mutation_rate, in_compare = 0):
    '''Selection method: roulette
    Generates a new population by choosing two random idividuals whose
    pobability of being choosen is proportional to their fitness'''
    history = []
    for i in range(generations):
        new_pop = []
        fit = np.asarray([fitness(points, p) for p in pop])
        pop[0], pop[np.argmax([fit])] = pop[np.argmax([fit])], pop[0]
        ubermensch = deepcopy(pop[0]) #fittest individual

        history.append(max(fit))
        total_fit = sum(fit)
        fit /= total_fit

        for _ in pop[1:]:
            selected = np.random.multinomial(2, fit)
            parent1 = pop[selected[0]]
            parent2 = pop[selected[1]]

            if len(set(history[-8:])) == 1:
                new_pop.append(mutate(crossover(parent1, parent2), 5*mutation_rate))
            else:
                new_pop.append(mutate(crossover(parent1, parent2), mutation_rate))

        new_pop.append(ubermensch)

        with open('graphs/log.txt', 'a') as log_file:
            print("generation ", i, " fitness ", history[-1])
            log_file.write("generation {} fitness {}\n".format(i, history[-1]))

        pop = deepcopy(new_pop)

    plot_clusters(points, ubermensch, "roulette", in_compare)
    return pop, ubermensch, history


def tournanment(fitness, pop, points, generations, mutation_rate, in_compare = 0):
    '''Selection method: tournanment
    Generates a new population by comparing two individuals and choosing
    the best, this is done twice and both winners mate'''
    history = []
    for i in range(generations):
        new_pop = []
        fit = np.array([fitness(points, p) for p in pop])
        pop[0], pop[np.argmax([fit])] = pop[np.argmax([fit])], pop[0]
        ubermensch = deepcopy(pop[0]) #fittest individual

        for _ in pop[1:]:
            id1 = np.random.randint(0, len(pop))
            id2 = np.random.randint(0, len(pop))
            contender1 = pop[id1]
            contender2 = pop[id2]
            parent1 = contender1 if fit[id1] > fit[id2] else contender2

            id1 = np.random.randint(0, len(pop))
            id2 = np.random.randint(0, len(pop))
            contender1 = pop[id1]
            contender2 = pop[id2]
            parent2 = contender1 if fit[id1] > fit[id2] else contender2

            if len(set(history[-8:])) == 1:
                new_pop.append(mutate(crossover(parent1, parent2), 5*mutation_rate))
            else:
                new_pop.append(mutate(crossover(parent1, parent2), mutation_rate))

        new_pop.append(ubermensch)
        history.append(max(fit))

        with open('graphs/log.txt', 'a') as log_file:
            print("generation ", i, " fitness ", history[-1])
            log_file.write("generation {} fitness {}\n".format(i, history[-1]))

        pop = deepcopy(new_pop)

    plot_clusters(points, ubermensch, "tournanment", in_compare)
    return pop, ubermensch, history


def compare(f, pop, points, generations, mutation_rate):
    '''Evolves population with all selection methods and plots a comparison graph
    it also saves the clusters graph in the 'graph' directory'''
    history = []
    with open('graphs/log.txt', 'a') as log_file:
        print("-------elitism-------")
        log_file.write("-------elitism-------\n")

    pop_ = deepcopy(pop)
    pop_, _, hist = elitism(f, pop_, points, generations, mutation_rate, 1)
    history.append(hist)

    with open('graphs/log.txt', 'a') as log_file:
        print("-------harem-------")
        log_file.write("-------harem-------\n")

    pop_ = deepcopy(pop)
    pop_, centers, hist = harem(f, pop_, points, generations, mutation_rate, 1)
    history.append(hist)

    with open('graphs/log.txt', 'a') as log_file:
        print("-------roulette-------")
        log_file.write("-------roulette-------\n")

    pop_ = deepcopy(pop)
    pop_, centers, hist = roulette(f, pop_, points, generations, mutation_rate, 1)
    history.append(hist)

    with open('graphs/log.txt', 'a') as log_file:
        print("-------tournanment-------")
        log_file.write("-------tournanment-------\n")

    pop_ = deepcopy(pop)
    pop_, centers, hist = tournanment(f, pop_, points, generations, mutation_rate, 1)
    history.append(hist)
    #Plots each method in graph for better comparison
    plt.plot(range(len(history[0])), history[0], label='elitism')
    plt.plot(range(len(history[1])), history[1], label='harem')
    plt.plot(range(len(history[2])), history[2], label='roulette')
    plt.plot(range(len(history[3])), history[3], label='tournanment')
    plt.legend()
    plt.ylabel("Clusters calinskiHarabaz")
    plt.xlabel("Number of generations")
    plt.savefig("graphs/Comparison")
    plt.show()


def plot_clusters(points, centers, name, in_compare):
    '''Plots the points according to the given centers'''
    clusters = [[] for c in centers]

    for _, p in enumerate(points):
        closest = np.argmin([np.linalg.norm(p - c) for c in centers])
        clusters[closest].append(p)

    for c in clusters:
        x = [p[0] for p in c]
        y = [p[1] for p in c]
        plt.plot(x, y, 'o')

    plt.savefig("graphs/" + name)

    #checks if 'plot_clusters' is being used inside the 'compare' function
    if not in_compare:
        plt.show()

    plt.gcf().clear()


''' --------------------------------------------------- '''
def main():
    system("clear")
    print("""\033[92m
     888                              d8b
     888                              Y8P
     888
 .d88888 8888b.  888d888888  888  888 888 88888b.
d88" 888    "88b 888P"  888  888  888 888 888 "88b
888  888.d888888 888    888  888  888 888 888  888
Y88b 888888  888 888    Y88b 888 d88P 888 888  888
 "Y88888"Y888888 888     "Y8888888P"  888 888  888
                                                """)
    print("\033[93mWelcome to Darwin! The Genetic Clustering framework", "If you don't know what to do type help", sep='\n')

    if not os.path.exists("graphs"):
        os.makedirs("graphs")

    open('graphs/log.txt', 'w').close() #cleaning log file

    filename = ""
    f = calinskiHarabaz
    df = None
    points = []
    centers = []
    num_genes = 5
    pop_size = 10
    generations = 10
    mutation_rate = 2
    num_coord = 0
    num_points =  0
    pop = []
    evolve = roulette
    selection = "roulette"
    ignore_rows = 0
    ignore_coord = 0
    std = False

    while(1):
        print("\033[32mdarwin>>\033[96m", end=' ')
        inp = input()

        #initializes population
        if inp.lower() == "init":
            if type(df) == type(None):
                print("\033[91mNo file loaded!")
            else:
                pop = init_pop(points, num_coord, num_genes, pop_size)

        #quits framework
        elif inp.lower() == "quit":
            print("Quiting...")
            break

        #prints data
        elif inp.lower() == "print data":
            if type(df) == type(None):
                print("\033[91mNo file loaded!")
            else:
                print(df)

        #normalizes data
        elif inp.lower() == 'normalize':
            std = not std

        #plots data or clusters
        elif "plot" in inp.lower():
            if type(df) == type(None):
                print("\033[91mNo file loaded!")
            else:
                if "data" in inp.lower():
                    plt.plot([p[0] for p in points], [p[1] for p in points], 'bo')
                    plt.show()
                else:
                    if len(centers) != 0:
                        plot_clusters(points, centers, selection, False)
                    else:
                        print("\033[91mData not cluseterized!")

        #ignores lines
        elif inp[0:12].lower() == "ignore lines":
            ignore_rows = int(inp[13:])
            print(ignore_rows, "lines will be ignored")

        #ignores cloumns
        elif inp[0:14].lower() == "ignore columns":
            ignore_coord = int(inp[15:])
            print(ignore_coord, "columns will be ignored")

        #loads file
        elif inp[0:4].lower() == "load":
            filename = inp[5:]
            try:
                df = pd.read_csv(filename, skiprows=ignore_rows, header=None)
                num_coord = df.shape[1]
                num_points = df.shape[0]

                if ignore_coord > num_coord:
                    ignore_coord = 0
                    print("\033[91mNumber of columns to ignore is greater \
                     than the total ammount of columns!")
                    continue

                for i in range(ignore_coord):
                    df.drop(df.columns[num_coord-i-1], axis=1, inplace=True)

                num_coord = df.shape[1]
                if std:
                    points = preprocessing.scale(df.values)
                else:
                    points = df.values

                ignore_rows = 0
                ignore_coord = 0

                print(df)
                pop = init_pop(points, num_coord, num_genes, pop_size)
                print("loaded " + inp[5:])
            except FileNotFoundError:
                print("\033[91mFile not found!")
            except pd.errors.ParserError:
                print("\033[91mFile not a csv!")
            except pd.errors.EmptyDataError:
                ignore_rows = 0
                print("\033[91mNumber of lines to ignore is\
                 greater than the total ammount of lines!")
            except ValueError:
                print("\033[91mThe header of your dataset is\
                 being considered part of the data. Use 'ignore lines'.")
            except TypeError:
                print("\033[91mThe header of your dataset is\
                 being considered part of the data. Use 'ignore lines'.")

        #sets paramethers
        elif inp[0:3].lower() == "set":
            if "k" in inp or "clusters" in inp:
                try:
                    num_genes = int(inp[5:])
                    print("number of clusters will be ", num_genes)
                except ValueError:
                    print("\033[91mThat's not a number!")
            elif "population" in inp:
                try:
                    pop_size = int(inp[14:])
                    print("New population size is", pop_size)
                except ValueError:
                    print("\033[91mThat's not a number!")
            elif "generations" in inp:
                try:
                    generations = int(inp[15:])
                    print("New number of generations is", generations)
                except ValueError:
                    print("\033[91mThat's not a number!")
            elif "mutation" in inp:
                try:
                    mutation_rate = int(inp[12:])
                    print("New mutation rate is", mutation_rate)
                except ValueError:
                    print("\033[91mThat's not a number!")
            elif "selection" in inp.lower():
                if "harem" in inp.lower():
                    evolve = harem
                    selection = "harem"
                elif "elitism" in inp.lower():
                    evolve = elitism
                    selection = "elitism"
                elif "roulette" in inp.lower():
                    evolve = roulette
                    selection = "roulette"
                else:
                    evolve = tournanment
                    selection = "tournanment"
                print("Selection method set to", selection)
            else:
                print("\033[91mCan't set that!")
            print("Remeber to call init after changing variables")

        #evolves with given paramethers
        elif inp.lower() == "evolve":
            if type(df) == type(None):
                print("\033[91mNo file loaded!")
            else:
                pop, centers, history = evolve(f, pop, points, generations, mutation_rate)

        #compares all the selection methods
        elif inp.lower() == "compare":
            if type(df) == type(None):
                print("\033[91mNo file loaded!")
            else:
                compare(f, pop, points, generations, mutation_rate)

        #show paramethers
        elif inp.lower() == "params":
            print("Number of coordinates = ", num_coord)
            print("Number of points = ", num_points)
            print("Number of clusters = ", num_genes)
            print("Population size = ", pop_size)
            print("Number of generations = ", generations)
            print("Mutation Rate = ", mutation_rate)
            print("Selection Method = ", selection)
            print("Scaled data = ", std)

        #show help menu
        elif inp.lower() == "help":
            helper()

        #prints shape of the loaded data
        elif inp.lower() == "debug":
            if type(df) == type(None):
                print("\033[91mNo file loaded!")
            else:
                print(df.shape)

        else:
            print("\033[91mI don't understand you")


if __name__ == "__main__":
    main()
