import kmeans as km
import som
import numpy as np
import utils
import json

METHODS = ['forgy', 'random_partition']
SOM_INIT_METHODS = ['random', 'zeros']
SOM_ALGORITHMS = ['kohonen', 'neuron gas']
SOM_PARAMETERS_SETS = [(.1, .5), (.1, .5), (.1, 1), (.33, .1), (.33, .5), (.33, 1), (.66, .1), (.66, .5), (.66, 1),
                       (.99, .1), (.99, .5), (.99, 1)]


def get_datas_from_json():
    datas = []
    with open('data1.json', 'r') as d:
        datas.append(json.loads(d.read()))
    with open('data2.json', 'r') as d:
        datas.append(json.loads(d.read()))
    return datas


def get_datas_random():
    datas = []
    for get_data in [utils.get_data1, utils.get_data2]:
        datas.append(get_data())
    return datas


def main():
    datas = get_datas_from_json()
    benchmark_errors = False
    rand = np.random.RandomState(0)
    index = 1
    print("Self-organizing map")
    for data in datas:
        print(f'Data set: {index}')
        utils.plot_data(data)
        for algorithm in SOM_ALGORITHMS:
            print(f'Weights update algorithm: {algorithm}')
            for method in SOM_INIT_METHODS:
                print(f'Initialization method: {method}')
                for param_set in SOM_PARAMETERS_SETS:
                    print(
                        f'Learn rate: {param_set[0]},  Radius square: {param_set[1]}')
                    errors = {}
                    for k in range(2, 21):
                        som_data = som.init_neurons(data, k, rand, method)
                        soms_with_error = som.train_som(som_data, data, learn_rate=param_set[0], radius_sq=param_set[1],
                                                        algorithm=algorithm)
                        error = soms_with_error[-1][1]
                        errors[k] = error
                        soms, _ = zip(*soms_with_error)
                        som.plot_with_data(
                            soms, data, f'_LR{param_set[0]}_RSQ{param_set[1]}_{algorithm}_{method}_neurons{k}_data{index}')
                    if all([i in errors for i in range(2, 21, 2)]):
                        fname = f'som_errors_data{index}_{SOM_PARAMETERS_SETS.index(param_set)}_{algorithm}_{method}.png'
                        utils.plot_error_data([(k, errors[k]) for k in range(2, 21, 2)], fname=fname)
                if benchmark_errors:
                    soms_with_errors = []
                    for _ in range(100):
                        som_data = som.init_neurons(data, 20, rand, method)
                        soms_with_error = som.train_som(
                            som_data, data, algorithm=algorithm)
                        soms_with_errors.append(soms_with_error[-1])
                    som.print_som_stats(soms_with_errors, data)
        index += 1

    index = 1
    for data in datas:
        print(f'Data set {index}')
        utils.plot_data(data)
        for method in METHODS:
            print(f'Method: {method}')
            kmeans_data = {}
            for k in range(2, 21):
                kmeans_with_err = km.kmeans(data, method, k)
                km.print_stats(k, [(iterations[-1], err)
                                   for iterations, err in kmeans_with_err])
                min_err = kmeans_with_err[0][1]
                kmeans = kmeans_with_err[0][0]
                for temp_kmeans, err in kmeans_with_err:
                    if err < min_err:
                        min_err = err
                        kmeans = temp_kmeans
                kmeans_data[k] = (kmeans, min_err)
                km.plot_kmeans(kmeans, k, f'_{method}_{k}_{index}')
                if k in [2, 10]:
                    km.plot_kmeans_change(kmeans, k, f'_{method}_{k}_{index}')
            if all([i in kmeans_data for i in range(2, 21, 2)]):
                error_data = [[i, kmeans_data[i][1]] for i in range(2, 21, 2)]
                utils.plot_error_data(error_data)
            index += 1


if __name__ == '__main__':
    main()