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

METHODS = ['forgy', 'random_partition']
SOM_INIT_METHODS = ['random', 'zeros']


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()
    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 method in SOM_INIT_METHODS:
            print(f'Initialization method: {method}')
            errors = []
            for k in range(2, 21, 2):
                som_data = som.init_neurons(data, k, rand, method)
                soms_with_error = som.train_som(som_data, data, algorithm='kohonen')
                error = soms_with_error[-1][1]
                errors.append((k, error))
                soms,_ = zip(*soms_with_error)
                #som.plot_with_data(soms, data, f'_{method}_{k}_data{index}')
            utils.plot_error_data(errors)
            soms_with_errors = []
            for _ in range(100):
                som_data = som.init_neurons(data, k, rand, method)
                soms_with_error = som.train_som(som_data, data, algorithm='kohonen')
                soms_with_errors.append(soms_with_error[-1])
            som.print_som_stats(soms_with_errors, data)
        index += 1

    index = 1
    for data in datas:
        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}')
            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()