#!/usr/bin/python
__author__ = 'ar'

import cv2
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

import skimage as sk
import skimage.io
import skimage.color
from sklearn import cluster
from sklearn.metrics import pairwise as pwd
import sklearn.cluster as skcl
import libKMCUDA as kmcu

import pickle as pkl
import os
import sys
import glob

######################################

class CudaKMeans(object):
    def __init__(self, n_clusters = 5, verbosity=1, metric='L2', seed=None):
        self.kmdata = None
        self.verbosity = verbosity
        self.seed = seed
        self.metric = metric
        self.n_clusters = n_clusters
    def fit(self, dscData):
        if self.seed is not None:
            self.kmdata = kmcu.kmeans_cuda(dscData,
                                           self.n_clusters,
                                           metric=self.metric,
                                           verbosity=self.verbosity,
                                           seed=self.seed)
        else:
            self.kmdata = kmcu.kmeans_cuda(dscData,
                                           self.n_clusters,
                                           metric=self.metric,
                                           verbosity=self.verbosity)
        self.centroids_ = self.kmdata[0]
        tmpNumNAN = np.isnan(self.kmdata[0]).sum()
        if tmpNumNAN > 0:
            print('\t!!!warning!!! found {} NAN in kmean-centroids'.format(tmpNumNAN))
        self.centroids_[np.isnan(self.centroids_)] = -2.
        self.labels_    = self.kmdata[1]
        return self
    def predict(self, dscData):
        tcls = skcl.KMeans(n_clusters=self.n_clusters, n_jobs=6)
        tcls.labels_ = self.labels_
        tcls.cluster_centers_ = self.centroids_
        tret = tcls.predict(dscData)
        # tret = kmcu.knn_cuda(1, dscData, *self.kmdata, metric='L2', verbosity=self.verbosity)
        # tcls = neighbors.KNeighborsClassifier(1, n_jobs=6)
        # tcls.fit(self.centroids_, np.array(range(self.n_clusters)))
        # tret = tcls.predict(dscData)
        return tret


def cropImagePatch(img, pt, angle, radiusPatch, isDebug=False, baseSize=(24,24)):
    px = pt[0]
    py = pt[1]
    #
    imgShape=img.shape
    imgCenter=(int(imgShape[1]/2.0), int(imgShape[0]/2.0))
    #
    matRot=cv2.getRotationMatrix2D(imgCenter, angle, 1.0)
    matShift=np.zeros((2,3))
    matShift[0,0]=1.0
    matShift[0,1]=0.0
    matShift[1,0]=0.0
    matShift[1,1]=1.0
    matShift[0,2]=+imgCenter[0] - px
    matShift[1,2]=+imgCenter[1] - py
    imgShift=cv2.warpAffine(img, matShift, (imgShape[1], imgShape[0]), None, cv2.INTER_CUBIC)
    cropSize=int(np.round(1.0*radiusPatch))
    imgShiftRot=cv2.warpAffine(imgShift, matRot, (imgShape[1], imgShape[0]), None, cv2.INTER_CUBIC)
    imgShiftRotCrop=imgShiftRot[imgCenter[1]-cropSize:imgCenter[1]+cropSize, imgCenter[0]-cropSize:imgCenter[0]+cropSize]
    if isDebug:
        plt.figure()
        plt.subplot(1,3,1)
        plt.imshow(imgShift), plt.title('Shift')
        plt.subplot(1,3,2)
        plt.imshow(imgShiftRot), plt.title('Shift-Rot')
        plt.subplot(1,3,3)
        plt.imshow(imgShiftRotCrop), plt.title('Shift-Rot-Crop')
    brd=1
    baseSizeBrd=(baseSize[0]-brd*2, baseSize[1]-brd*2)
    imgShiftRotCropResize=cv2.resize(imgShiftRotCrop, baseSizeBrd, None,0,0,cv2.INTER_CUBIC)
    imgShiftRotCropResize=cv2.copyMakeBorder(imgShiftRotCropResize,brd,brd,brd,brd,cv2.BORDER_CONSTANT, None, (0,0,0))
    return imgShiftRotCropResize

######################################
def_NumFeaturesPerImage=100
def_NumClustersSqrt=4
def_NumClusters=def_NumClustersSqrt**2
def_numBestPatches= int(def_NumClusters*3/4)
def_sizePatch=32

######################################
if __name__=='__main__':
    isUseCUDA = True
    numImages = 50
    pathIdx = '/home/ar/data/brodatz/brodatz_dataset_train.csv'
    # paramK = 65536
    # paramN = 180000
    # paramK = 2048
    paramK = 16384
    paramN = 80000
    if not isUseCUDA:
        pathKCls = '{}-kcls-n{}-k{}.pkl'.format(pathIdx, paramN, paramK)
    else:
        pathKCls = '{}-kcls-n{}-k{}-cuda.pkl'.format(pathIdx, paramN, paramK)
    wdir = os.path.dirname(pathIdx)
    dataCSV = pd.read_csv(pathIdx)
    pathImages = np.array([os.path.join(wdir, xx) for xx in dataCSV['path']])
    pathImages = pathImages[np.random.permutation(len(pathImages))[:numImages]]
    #
    # sift=cv2.SIFT(nfeatures=def_NumFeaturesPerImage)
    sift=cv2.xfeatures2d.SIFT_create(nfeatures=def_NumFeaturesPerImage)
    lstK=[]
    lstD=[]
    lstImg=[]
    arrkp=None
    arrkd=None
    cnt=0
    print ('(1) Prepare data:')
    for ii in range(len(pathImages)):
        fimg=pathImages[ii]
        img=cv2.imread(fimg, cv2.IMREAD_GRAYSCALE)
        lstImg.append(img)
        kp, kd=sift.detectAndCompute(img,None)
        kpa=np.zeros((len(kp), 5))
        for kki in range(len(kp)):
            tkp=kp[kki]
            kpa[kki,0]=tkp.pt[0]
            kpa[kki,1]=tkp.pt[1]
            kpa[kki,2]=tkp.size
            kpa[kki,3]=tkp.angle
        kpa[:,4]=ii
        if arrkd is None:
            arrkp=kpa.copy()
            arrkd=kd.copy()
        else:
            kd /= np.tile(np.sum(kd, axis=-1).reshape((-1, 1)), (1, kd.shape[-1]))
            tdscMean = np.tile(np.sum(kd, axis=-1).reshape((-1, 1)), (1, kd.shape[-1]))
            arrkd=np.concatenate([arrkd,kd/tdscMean])
            arrkp=np.concatenate([arrkp,kpa])
        lstK.append(kp)
        lstD.append(kd)
        if not (ii%20):
            print ('%d/%d' % (ii,len(pathImages)))
    #
    print ('(2) Run K-Means:')
    # mdlKMeans=cluster.KMeans(n_clusters=def_NumClusters, max_iter=10, n_jobs=-1)
    print('[**] Loading precalculated KMeans model [{}]'.format(pathKCls))
    with open(pathKCls, 'rb') as f:
        kmeans = pkl.load(f)

    arrCls  = kmeans.predict(arrkd)
    clsUniq = np.unique(arrCls)
    clsHist = np.histogram(arrCls, clsUniq)[0]
    idxMaxs = np.argsort(-clsHist)

    print ('(3) Calculate scores:')
    arrScores_DST = pwd.pairwise_distances(arrkd, kmeans.centroids_, 'l1')
    arrScores_DST = np.sort(arrScores_DST)[:,0]
    # arrScores=np.zeros(len(arrCls))
    # q1 = mdlKMeans.score(arrkd)
    # for ii in xrange(len(arrCls)):
    #     arrScores[ii]=-mdlKMeans.score(arrkd[ii, :].reshape((1, -1)))
    arrScores = arrScores_DST
    #
    print ('(4) Visualize patches:')
    lstMeanCls=[]
    imgBig=np.zeros((def_sizePatch*def_numBestPatches, def_sizePatch*def_NumClusters))
    # arrClsUniq=np.unique(arrCls)
    arrClsUniq = clsUniq[idxMaxs][:def_NumClusters]
    lstMeanPatch=[]
    for cci, cc in enumerate(arrClsUniq):
        idxc=(arrCls==cc)
        print('#idx = {}'.format(np.sum(idxc)))
        tarrkp=arrkp[idxc,:]
        tarrkd=arrkd[idxc,:]
        tscors=arrScores[idxc]
        idxGood=np.argsort(tscors)
        tnumBestPatches=np.min((len(idxGood),def_numBestPatches))
        tMeanPatch=np.zeros((def_sizePatch,def_sizePatch))
        for bbi in range(tnumBestPatches):
            tidx=idxGood[bbi]
            tpts=(tarrkp[tidx,0], tarrkp[tidx,1])
            tsiz=tarrkp[tidx,2]*5.0/2.0
            tang=tarrkp[tidx,3]
            tidxImg=int(tarrkp[tidx,4])
            timgPatch=cropImagePatch(lstImg[tidxImg], tpts,tang,tsiz,baseSize=(def_sizePatch,def_sizePatch))
            imgBig[bbi*def_sizePatch:bbi*def_sizePatch+timgPatch.shape[0], cci*def_sizePatch:cci*def_sizePatch+timgPatch.shape[1]]=timgPatch
            tMeanPatch+=timgPatch.astype(np.float)
        if tnumBestPatches>0:
            tMeanPatch/=tnumBestPatches
        lstMeanPatch.append(tMeanPatch)
        print ('%d/%d' % (cci,len(arrClsUniq)))
    plt.imshow(imgBig, cmap=plt.gray())
    imgBigNorm=cv2.normalize(imgBig, None,0,255,cv2.NORM_MINMAX, cv2.CV_8U)
    cv2.imwrite("VisualDictionary.png", imgBigNorm)
    plt.show()
