#!/usr/bin/env python3
"""Simple k-means implementation"""
import argparse
import math
import random


def distance(p1, p2) -> float:
    """return the euclidean distance between to points"""

    assert (len(p1) == len(p2))

    if isinstance(p1, tuple):
        p1 = list(p1)
    if isinstance(p2, tuple):
        p2 = list(p2)

    d = 0
    for i, x in enumerate(p1):
        d += (x - p2[i])**2

    return math.sqrt(d)


def gen_points(n, dimensions, max_value, min_value) -> list:
    """generate random points"""
    points = []
    for _ in range(n):
        point = []
        for _ in range(dimensions):
            point.append(random.randint(min_value, max_value))
        points.append(point)
    return points


def init_clusters(n, points) -> list:
    """create random initial clusters
    
    The clusters are chosen randomly between the encountered maximum and
    minimum values in the given points"""

    max_values = []
    min_values = []

    for i in range(len(points[0])):
        max_values.append(max(p[i] for p in points))
        min_values.append(min(p[i] for p in points))

    clusters = []
    for _ in range(n):
        cluster = []
        for i, max_value in enumerate(max_values):
            cluster.append(random.randint(min_values[i], max_value))
        clusters.append(tuple(cluster))

    return clusters


def map_points(clusters, points) -> dict:
    """map points to clusters

    Map each points to the cluster center with the minimal distance."""

    mapping = {p: [] for p in clusters}  # type: ignore
    for p in points:
        min_distance = None
        center = None
        for cluster in clusters:
            d = distance(p, cluster)
            if not min_distance or d < min_distance:
                center = cluster
                min_distance = d
        mapping[center].append(p)

    return mapping


def calc_clusters(mapping) -> list:
    """Calculate and return new centers for each cluster"""
    # return the new centers, use the old ones if no points are in the cluster
    return [
        k_mean(points) if len(points) else center
        for center, points in mapping.items()
    ]


def k_mean(points) -> tuple:
    """return the k-means of the given points"""
    # Start with the values of the first points
    means = list(points[0])

    # Only consider the remaining points
    for p in points[1:]:
        # calculate a running mean for each value
        for i, x in enumerate(p):
            means[i] = (means[i] + x) / 2

    return tuple(means)


def visualize(mapping):
    """visualize a mapping using matplotlib"""
    import matplotlib.pyplot as plt

    for center, points in mapping.items():
        x = [p[0] for p in points]
        y = [p[1] for p in points]
        plt.scatter(x, y, alpha=0.5)
        plt.plot(center[0], center[1], 'x')
    plt.show()


def main():
    """main entry point"""
    parser = argparse.ArgumentParser()
    parser.add_argument("-n", "--clusters", type=int, default=3)
    parser.add_argument("-r", "--random-points", type=int, default=10)
    parser.add_argument("-d", "--point-dimensions", type=int, default=2)
    parser.add_argument("--dont-visualize", action='store_true')
    parser.add_argument("--max", type=int, default=10)
    parser.add_argument("--min", type=int, default=1)
    args = parser.parse_args()

    points = gen_points(args.random_points, args.point_dimensions, args.max,
                        args.min)
    clusters = init_clusters(args.clusters, points)
    while True:
        mapping = map_points(clusters, points)

        if not args.dont_visualize:
            visualize(mapping)

        new_clusters = calc_clusters(mapping)

        if new_clusters == clusters:
            break
        clusters = new_clusters

    print(clusters)


if __name__ == '__main__':
    main()