from pylab import loadtxt, arange, loadtxt, permutation, zeros, sum, \
    plot, subplot, array, scatter, logical_and, figure, savefig, seed
from sys import argv, exit
import classifiers

# Parse input
argc = len(argv)
if argc < 2:
    print 'Usage: python %s METHOD [ K ] [ SEED ]' % argv[0]
    exit()
method = argv[1]
if method == 'knnb':
    if argc < 3:
        print 'Usage: python %s %s K [ SEED ]' % tuple(argv[:2])
        exit()
    k = int(argv[2])
    if argc == 4:
        seed(int(argv[3]))
elif method not in ['nnb', 'mec', 'svm']:
    print 'Unknown classification method "%s"' % argv[1]
    exit()
elif argc == 3:
    seed(int(argv[2]))

# Load dataset
def cnvt(s):
    tab = {'Iris-setosa': 1.0, 'Iris-versicolor': 2.0, 'Iris-virginica': 3.0}
    if tab.has_key(s):
        return tab[s]
    else:
        return -1.0

XC = loadtxt('iris.data', delimiter=',', dtype=float, converters={4: cnvt})
ind = arange(150) # indices into the dataset
ind = permutation(ind) # random permutation
L = ind[0:90] # learning set indices
T = ind[90:] # test set indices

# Learning set
args = [XC[L,0:4].T, XC[L,-1]]
if method == 'knnb':
    args.append(k)
method_class = {'nnb': classifiers.NNb, 'knnb': classifiers.kNNb,
        'mec': classifiers.MEC, 'svm': classifiers.SVM}[method]
classifier = method_class(*args)

# Classification of test set
c = zeros(len(T))
for i in arange(len(T)):
    c[i] = classifier.classify(XC[T[i],0:4])

# Confusion matrix
CM = zeros((3, 3))
for i in range(3):
    for j in range(3):
        CM[i,j] = sum(logical_and(XC[T,4] == i + 1, c == j + 1))
print CM

# Plot test set
figure(1)
color = array([[1,0,0], [0,1,0], [0,0,1]])
for i in range(4):
    for j in range(4):
        subplot(4, 4, 4 * i + j + 1)
        if i == j:
            continue
        scatter(XC[T, i], XC[T, j], s=100, marker='s',
            edgecolor=color[XC[T, 4].astype(int) - 1],
            facecolor=[1,1,1]*len(T))
        scatter(XC[T, i], XC[T, j], s=30, marker='+',
            edgecolor=color[c.astype(int) - 1])
#from pylab import show
#show()
if method == 'knnb':
    filename = 'knnb-%d.pdf' % k
else:
    filename = '%s.pdf' % method
savefig(filename)
