import numpy as np
import cv2
from multiprocessing.pool import ThreadPool

SZ = 20 # size of each digit is SZ x SZ
CLASS_N = 10

def load_base(fn):
    print 'loading "%s" ...' % fn
    digits_img = cv2.imread(fn, 0)
    h, w = digits_img.shape
    digits = [np.hsplit(row, w/SZ) for row in np.vsplit(digits_img, h/SZ)]
    digits = np.array(digits).reshape(-1, SZ, SZ)
    digits = np.float32(digits).reshape(-1, SZ*SZ) / 255.0
    labels = np.repeat(np.arange(CLASS_N), len(digits)/CLASS_N)
    return digits, labels

def cross_validate(model_class, params, samples, labels, kfold = 4, pool = None):
    n = len(samples)
    folds = np.array_split(np.arange(n), kfold)
    def f(i):
        model = model_class(**params)
        test_idx = folds[i]
        train_idx = list(folds)
        train_idx.pop(i)
        train_idx = np.hstack(train_idx)
        train_samples, train_labels = samples[train_idx], labels[train_idx]
        test_samples, test_labels = samples[test_idx], labels[test_idx]
        model.train(train_samples, train_labels)
        resp = model.predict(test_samples)
        score = (resp != test_labels).mean()
        print ".",
        return score
    if pool is None:
        scores = map(f, xrange(kfold))
    else:
        scores = pool.map(f, xrange(kfold))
    return np.mean(scores)

class StatModel(object):
    def load(self, fn):
        self.model.load(fn)
    def save(self, fn):
        self.model.save(fn)

class KNearest(StatModel):
    def __init__(self, k = 3):
        self.k = k

    @staticmethod
    def adjust(samples, labels):
        print 'adjusting KNearest ...'
        best_err, best_k = np.inf, -1
        for k in xrange(1, 11):
            err = cross_validate(KNearest, dict(k=k), samples, labels)
            if err < best_err:
                best_err, best_k = err, k
            print 'k = %d, error: %.2f %%' % (k, err*100)
        best_params = dict(k=best_k)
        print 'best params:', best_params
        return best_params

    def train(self, samples, responses):
        self.model = cv2.KNearest()
        self.model.train(samples, responses)

    def predict(self, samples):
        retval, results, neigh_resp, dists = self.model.find_nearest(samples, self.k)
        return results.ravel()

class SVM(StatModel):
    def __init__(self, C = 1, gamma = 0.5):
        self.params = dict( kernel_type = cv2.SVM_RBF, 
                            svm_type = cv2.SVM_C_SVC,
                            C = C,
                            gamma = gamma )

    @staticmethod
    def adjust(samples, labels):
        Cs = np.logspace(0, 5, 10, base=2)
        gammas = np.logspace(-7, -2, 10, base=2)
        scores = np.zeros((len(Cs), len(gammas)))
        scores[:] = np.nan

        print 'adjusting SVM (may take a long time) ...'
        def f(job):
            i, j = job
            params = dict(C = Cs[i], gamma=gammas[j])
            score = cross_validate(SVM, params, samples, labels)
            scores[i, j] = score
            nready = np.isfinite(scores).sum()
            print '%d / %d (best error: %.2f %%, last: %.2f %%)' % (nready, scores.size, np.nanmin(scores)*100, score*100)

        pool = ThreadPool(processes=cv2.getNumberOfCPUs())
        pool.map(f, np.ndindex(*scores.shape))
        print scores

        i, j = np.unravel_index(scores.argmin(), scores.shape)
        best_params = dict(C = Cs[i], gamma=gammas[j])
        print 'best params:', best_params
        print 'best error: %.2f %%' % (scores.min()*100)
        return best_params

    def train(self, samples, responses):
        self.model = cv2.SVM()
        self.model.train(samples, responses, params = self.params)

    def predict(self, samples):
        return self.model.predict_all(samples).ravel()

def main_adjustSVM(samples, labels):
    params = SVM.adjust(samples, labels)
    print 'training SVM on all samples ...'
    model = SVN(**params)
    model.train(samples, labels)
    print 'saving "digits_svm.dat" ...'
    model.save('digits_svm.dat')

def main_adjustKNearest(samples, labels):
    params = KNearest.adjust(samples, labels)

def main_showSVM(samples, labels):
    from common import mosaic

    train_n = int(0.9*len(samples))
    digits_train, digits_test = np.split(samples[shuffle], [train_n])
    labels_train, labels_test = np.split(labels[shuffle], [train_n])

    print 'training SVM ...'
    model = SVM(C=2.16, gamma=0.0536)
    model.train(digits_train, labels_train)

    train_err = (model.predict(digits_train) != labels_train).mean()
    resp_test = model.predict(digits_test)
    test_err = (resp_test != labels_test).mean()
    print 'train errors: %.2f %%' % (train_err*100)
    print 'test errors: %.2f %%' % (test_err*100)

    
    # visualize test results
    vis = []
    for img, flag in zip(digits_test, resp_test == labels_test):
        img = np.uint8(img*255).reshape(SZ, SZ)
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
        if not flag:
            img[...,:2] = 0
        vis.append(img)
    vis = mosaic(25, vis)
    cv2.imshow('test', vis)
    cv2.waitKey()
    


if __name__ == '__main__':
    samples, labels = load_base('digits.png')
    shuffle = np.random.permutation(len(samples))
    samples, labels = samples[shuffle], labels[shuffle]
    
    #main_adjustSVM(samples, labels)
    #main_adjustKNearest(samples, labels)
    main_showSVM(samples, labels)