# -*- coding: utf-8 -*-
"""
Created on Sun May 14 11:20:00 2017
"""

import matplotlib.pyplot as plt
import numpy as np

# inputs
photo = plt.imread("ara.png")
k     = 10

# affichage de la nouvelle image   
plt.figure(1)
plt.imshow(photo)
plt.title('original')   

# extraire les couleurs

R = photo[:,:,0]
G = photo[:,:,1]
B = photo[:,:,2]
[nx,ny] = R.shape

RN = photo[:,:,0]
GN = photo[:,:,1]
BN = photo[:,:,2]

# initilisation de la nouvelle image
Rcenter = np.zeros(k)
Gcenter = np.zeros(k)
Bcenter = np.zeros(k)

dist = np.zeros(k)

# initilisation aléatoire des groupes
Gr0  = np.random.randint(k, size=R.shape)
Gr   = Gr0
Gr1  = np.random.randint(k, size=R.shape)

err = 0

it = 0

while err!=1:
    
    it = it + 1
    # calculer centre de gravite de chaque ensemble
    for i in range(k):
        idtmp = np.where(Gr==i)
        
        if (np.size(idtmp) !=0):
            Rtmp  = R[idtmp]
            Gtmp  = G[idtmp]
            Btmp  = B[idtmp]
            Rcenter[i] = np.mean(Rtmp)
            Gcenter[i] = np.mean(Gtmp)
            Bcenter[i] = np.mean(Btmp)
        
    # affectation de chaque pt au centre le plus proche
    for jx in range(nx):
        for jy in range(ny):
            for i in range(k):
                dist[i]  = (R[jx][jy]-Rcenter[i])**2 + (G[jx][jy]-Gcenter[i])**2 + (B[jx][jy]-Bcenter[i])**2
            min_id = np.argmin(dist)
            Gr1[jx][jy] = min_id
          
    err = np.prod( (Gr==Gr1) )
    
    print ('it=',it)
    
    Gr = Gr1


#reduction: on affecte a chaque ensemble la couleur de son centre 
for i in range(k):
        idtmp = np.where(Gr==i)
        if (np.size(idtmp) != 0):
            for j in range( len(idtmp[0]) ):
                RN[ idtmp[0][j] ][ idtmp[1][j] ]   = Rcenter[i]
                GN[ idtmp[0][j] ][ idtmp[1][j] ]   = Gcenter[i]    
                BN[ idtmp[0][j] ][ idtmp[1][j] ]   = Bcenter[i]   

# construction de la nouvelle image
photoN =  np.zeros((nx, ny,3))           
photoN[:,:,0] = RN
photoN[:,:,1] = GN
photoN[:,:,2] = BN   

# affichage de la nouvelle image      
#plt.imshow(photoN)


plt.figure(2)
plt.imshow(photoN)
plt.title('k-means')