import numpy as np
from PIL import Image
from math import sqrt


def split_img(pic, x, y):
    blocks = []
    bit_map = np.array(list(pic.getdata())).reshape(x, y)
    for r in range(0, x, 4):
        for c in range(0, y, 4):
            blocks.append(bit_map[r:r + 4, c:c + 4])
    return blocks


def calculate_average_vector(blocks):
    average = sum(blocks)/len(blocks)
    return np.array(np.around(average, decimals=0), dtype=np.int64)


def split_code_book(initial, k, e):
    last_code_book = initial
    for i in range(k-1):
        if i == 0:
            new_code_book = np.vstack((last_code_book - e, last_code_book + e))
            new_code_book = np.vsplit(new_code_book, i + 2)
            temp_code_book = last_code_book
            last_code_book = new_code_book
        else:
            if i == 1:
                new_code_book = np.vstack((last_code_book[0] - e, temp_code_book, last_code_book[-1] + e))
            else:
                new_code_book = np.vstack((last_code_book[0] - e))
                for x in range(len(temp_code_book)):
                    new_code_book = np.vstack((new_code_book, temp_code_book[x]))
                new_code_book = np.vstack((new_code_book, last_code_book[-1] + e))
            new_code_book = np.vsplit(new_code_book, i + 2)
            temp_code_book = last_code_book
            last_code_book = new_code_book
    return(new_code_book)


def initiate_distance(blocks, code_book):
    distance_table = []

    for x in range(len(blocks)):  # distance from each vector to each entry codeBook
        for y in range(len(code_book)):
            distance_table.append(euclid_distance(blocks[x], code_book[y]))
    distance_table = (np.array(distance_table).reshape(len(blocks), len(code_book)))
    distance_list = []
    for x in range(len(distance_table)):  # index - vector, value - entry codeBook
        distance_list.append(np.nonzero(distance_table[x] == distance_table[x].min())[0][0])
    return distance_list


def code_book_checker(code_book, distance_list):
    x = 0
    while x in range(len(code_book)):
        if not distance_list.count(x):
            code_book.pop(x)
            for y in range(len(distance_list)):
                if int(distance_list[y]) > x:
                    distance_list[y] -= 1
            x -= 1
        x += 1
    return code_book


def create_clusters(blocks, distance_list, code_book):
    clusters_list = [[] for x in range(len(code_book))]
    for x in range(len(distance_list)):
        clusters_list[distance_list[x]].append(blocks[x])
    return clusters_list


def calculate_distance(blocks, code_book):
    temp_distance = []
    new_distance = []
    for x in range(len(code_book)):
        for y in range(len(blocks[x])):
            temp_distance.append(euclid_distance(code_book[x], blocks[x][y]))
        new_distance.append(temp_distance)
        temp_distance = []
    return new_distance


def modify_code_book(average_distance, clusters, code_book):
    old_average_distance = average_distance
    for i in range(len(code_book)):
        code_book[i] = calculate_average_vector(clusters[i])
    print('кодовая книга\n', code_book)
    print()
    print()
    blocks = []
    for i in range(len(clusters)):
        for j in range(len(clusters[i])):
            blocks.append(clusters[i][j])
    print('блоки\n', blocks)
    distance = initiate_distance(blocks, code_book)

    code_book = code_book_checker(code_book, distance)
    clusters = create_clusters(blocks, distance, code_book)
    new_distance = calculate_distance(clusters, code_book)
    average_distance_list = []
    for i in range(len(np.array(new_distance))):
        average_distance_list.append(np.average(np.array(new_distance)[i]))
    average_distance = np.average(np.array(average_distance_list))
    return(old_average_distance, average_distance, code_book, clusters, distance)


def print_clusters(clusters, code_book):
    for i in range(len(code_book)):
        print('%s кластер:' % i)
        for j in range(len(clusters[i])):
            print(clusters[i][j])
        print()


def euclid_distance(vector, entry_code_book):
    result = 0
    for x in range(len(vector)):
        for y in range(len(vector)):
            result += (vector[y][x] - entry_code_book[y][x])**2
    return sqrt(result)
# codeBookSize = int(input('codeBookSize:'))
# eps = int(input('eps:'))

codeBookSize = 128
eps = 1
D = 0.005
fileName = 'img.bmp'
img = Image.open(fileName)
width, height = img.size[0], img.size[1]
widthVectors, heightVectors = int(width/4), int(height/4)

vectors = split_img(img, width, height)  # split pic
codeBook = split_code_book(calculate_average_vector(vectors), codeBookSize, eps)  # initial codebook
distance = initiate_distance(vectors, codeBook)  # make clusters
print('В инициализированной кодовой книге %s элементов' % len(codeBook))
print('Инициализированная кодовая книга:', codeBook, sep='\n', end='\n\n')
codeBook = code_book_checker(codeBook, distance)  # check codebook
clusters = create_clusters(vectors, distance, codeBook)
newDistance = calculate_distance(clusters, codeBook)

print('Блоки', vectors, sep='\n', end='\n\n')
print('В исправленной кодовой книге %s элементов' % len(codeBook))
print('Исправленная кодовая книга:', codeBook, sep='\n', end='\n\n')
print('Матрица индексов кодовой книги:', np.array(distance).reshape(int(sqrt(len(distance))), int(sqrt(len(distance)))),
      sep='\n', end='\n\n')
np.savetxt('matrix.txt', np.array(distance).reshape(int(sqrt(len(distance))), int(sqrt(len(distance)))), fmt='%u')


averageDistanceList = []
for i in range(len(np.array(newDistance))):
    averageDistanceList.append(np.average(np.array(newDistance)[i]))
averageDistance = np.average(np.array(averageDistanceList))
print('averageDistance', averageDistance)
g = 1
print('Итерация 0:')
lastAverageDistance, averageDistance, codeBook, clusters, distance = \
        modify_code_book(averageDistance, clusters, codeBook)
np.savetxt('matrix0.txt', np.array(distance).reshape(int(sqrt(len(distance))), int(sqrt(len(distance)))), fmt='%u')
print()
exit()

while True:
    lastAverageDistance, averageDistance, codeBook, clusters, distance = \
        modify_code_book(averageDistance, clusters, codeBook)
    print(averageDistance)
    print('В исправленной кодовой книге %s элементов' % len(codeBook))
    print('Исправленная кодовая книга:', codeBook, sep='\n', end='\n\n')
    print('Матрица индексов кодовой книги:',
          np.array(distance).reshape(int(sqrt(len(distance))), int(sqrt(len(distance)))),
          sep='\n', end='\n\n')
    np.savetxt('matrix%s.txt' %g, np.array(distance).reshape(int(sqrt(len(distance))), int(sqrt(len(distance)))), fmt='%u')
    print('%s итерация, значение (D(k-1)-D(k))/D(k) = %s' %
          (g, ((lastAverageDistance - averageDistance) / averageDistance)))
    g += 1
    if (lastAverageDistance - averageDistance)/averageDistance <= D:
        break

newBitMap = []
distance = np.array(distance).reshape(int(sqrt(len(distance))), int(sqrt(len(distance))))
print(len(codeBook))
print(distance)
for i in range(5):
    for j in range(5):
        newBitMap.append(codeBook[distance[i][j]])
newBitMap = np.array(newBitMap)
bitmap = []
for x in range(len(newBitMap)):
    for y in range(len(newBitMap[x])):
        for i in range(len(newBitMap[x][y])):
            bitmap.append(newBitMap[x][y][i])
bitmap = np.array(bitmap)
img3 = Image.fromarray(np.reshape(bitmap, (int(sqrt(len(bitmap))), int(sqrt(len(bitmap))))), mode='L')
img3.save('my.png')
img3.show()

#print_clusters(clusters, codeBook)
# # #if np.any(np.greater(namearray, value)):