#!/usr/bin/env python
# Author: Vasil Zlatanov, Nunzio Pucci
# EE4 Pattern Recognition coursework
#
# usage: opt.py [-h] [-t] [-c] [-k] [-m] [-e] [-r] [-a RERANKA]
#               [-b RERANKB] [-l RERANKL] [-n NEIGHBORS] [-v]
#               [-s SHOWRANK] [-1] [-M MULTRANK] [-C] [DATA]
#               [-K KMEAN] [-A] [-P PCA]

import matplotlib.pyplot as plt
import sys
import os
import json
import scipy.io
from sklearn.neighbors import NearestNeighbors
from sklearn.cluster import KMeans
from sklearn.decomposition import PCA
from sklearn.model_selection import train_test_split
from sklearn.metrics import confusion_matrix
from sklearn.metrics import accuracy_score
import argparse
import numpy as np
from numpy import genfromtxt
from numpy import linalg as LA
from timeit import default_timer as timer
from scipy.spatial.distance import cdist
sys.path.append('lib')
from rerank import re_ranking
from kmean import create_kmean_clusters
import logging
from logging import debug
from evaluate import test_model
from evaluate import draw_results 

parser = argparse.ArgumentParser()

parser.add_argument("-t", "--train", help="Use train data instead of query and gallery", action='store_true')
parser.add_argument("-c", "--conf_mat", help="Show visual confusion matrix", action='store_true')
parser.add_argument("-k", "--kmean_alt", help="Perform clustering with generalized labels(not actual kmean)", action='store_true')
parser.add_argument("-m", "--mahalanobis", help="Perform Mahalanobis Distance metric", action='store_true')
parser.add_argument("-e", "--euclidean", help="Use standard euclidean distance", action='store_true')
parser.add_argument("-r", "--rerank", help="Use k-reciprocal rernaking", action='store_true')
parser.add_argument("-a", "--reranka", help="Parameter k1 for rerank", type=int, default = 9)
parser.add_argument("-b", "--rerankb", help="Parameter k2 for rerank", type=int, default = 3)
parser.add_argument("-l", "--rerankl", help="Parameter lambda for rerank", type=float, default = 0.3)
parser.add_argument("-n", "--neighbors", help="Use customized ranklist size NEIGHBORS", type=int, default = 1)
parser.add_argument("-v", "--verbose", help="Use verbose output", action='store_true')
parser.add_argument("-s", "--showrank", help="Save ranklist pics id in a txt file for first SHOWRANK queries", type=int, default = 0)
parser.add_argument("-1", "--normalise", help="Normalise features", action='store_true')
parser.add_argument("-M", "--multrank", help="Run for different ranklist sizes equal to MULTRANK", type=int, default=1)
parser.add_argument("-C", "--comparison", help="Compare baseline and improved metric", action='store_true')
parser.add_argument("--data", help="Folder containing data", default='data')
parser.add_argument("-K", "--kmean", help="Perform Kmean clustering, KMEAN number of clusters", type=int, default=0)
parser.add_argument("-A", "--mAP", help="Display Mean Average Precision", action='store_true')
parser.add_argument("-P", "--PCA", help="Perform pca with PCA eigenvectors", type=int, default=50)

args = parser.parse_args()

if args.verbose:
    logging.basicConfig(level=logging.DEBUG)

def eval(camId, filelist, labels, gallery_idx, train_idx, feature_vectors, args):
    
    if args.train:
        cam = camId[train_idx]
        cam = cam.reshape((cam.shape[0],1))
        labs = labels[train_idx].reshape((labels[train_idx].shape[0],1))
        tt = np.hstack((train_idx, cam)) 
        train, test, train_label, test_label = train_test_split(tt, labs, test_size=0.3, random_state=0)
        del labs
        del cam
        train_data = feature_vectors[train[:,0]]
        test_data = feature_vectors[test[:,0]]
        train_cam = train[:,1]
        test_cam = test[:,1]
        showfiles_train = filelist[train[:,0]]
        showfiles_test = filelist[train[:,0]]
        del train
        del test 
        del tt
    else:
        query_idx = query_idx.reshape(query_idx.shape[0])
        gallery_idx = gallery_idx.reshape(gallery_idx.shape[0])
        camId = camId.reshape(camId.shape[0])
        showfiles_train = filelist[gallery_idx]
        showfiles_test = filelist[query_idx]
        train_data = feature_vectors[gallery_idx]
        test_data = feature_vectors[query_idx]
        train_label = labels[gallery_idx]
        test_label = labels[query_idx]
        train_cam = camId[gallery_idx]
        test_cam = camId[query_idx]
    train_idx = train_idx.reshape(train_idx.shape[0])
    train_model = feature_vectors[train_idx]
    
    if(args.PCA):
        pca=PCA(n_components=args.PCA) #Data variance @100 is 94%
        train_model=pca.fit_transform(train_model)
        train_data=pca.transform(train_data)
        test_data=pca.transform(test_data)


    if args.mAP:
        target_pred, mAP = test_model(train_data, test_data, train_label, test_label, train_cam, test_cam, showfiles_train, showfiles_test, train_model, args)
        return mAP
    else: 
        target_pred = test_model(train_data, test_data, train_label, test_label, train_cam, test_cam, showfiles_train, showfiles_test, train_model, args)
        target_pred = target_pred.reshape(target_pred.shape[1])
        return draw_results(test_label, target_pred)

def kopt(camId, filelist, labels, gallery_idx, train_idx, feature_vectors, args):
    axis = 0
    search = 0
    steps = 0
    vertical = True
    neg = False
    outofaxis = False
    start = np.array([1,1])
    if args.mAP:
        args.neighbors = 10
    args.train = True
    args.rerank = True
    args.reranka = 1
    args.rerankb = 1
    opt = np.array([1,1])
    checktab = np.zeros((100,100))
    checktab[1][1]=1
    max_acc = eval(camId, filelist, labels, gallery_idx, train_idx, feature_vectors, args)
    print('origin')
    print('vertical')
    while steps<3:
        steps+=1
        while axis<4:
            axis+=1
            p = start[0]
            q = start[1]
            while search <5:
                search+=1
                if vertical:
                    if neg:
                        p = start[0] - 2*search
                        if p < 1:
                            p = 1
                            search = 5
                            outofaxis = True
                    else:
                        p = search*2 + start[0]
                    args.reranka = p
                    if not outofaxis:
                        if checktab[p][q] == 0:
                            checktab[p][q] = 1
                            print('p:',p,' q:',q)
                            acc = eval(camId, filelist, labels, gallery_idx, train_idx, feature_vectors, args)
                    if acc > max_acc:
                        print('new p:',p, ' for accuracy:', acc)
                        max_acc=acc
                        opt[0] = p
                        start[0] = p
                        axis=0
                        steps=0
                        search=6
                else:
                    if neg:
                        q = start[1] - 2*search
                        if q < 1:
                            q = 1
                            search = 5
                            outofaxis = True
                    else:
                        q = search*2 + start[1]
                    args.rerankb = q
                    if not outofaxis:
                        if checktab[p][q] == 0:
                            checktab[p][q]=1
                            print('p:',p,' q:',q)
                            acc = eval(camId, filelist, labels, gallery_idx, train_idx, feature_vectors, args)
                    if acc > max_acc:
                        print('new q:',q, ' for accuracy:', acc)
                        max_acc=acc
                        opt[1] = q
                        start[1] = q
                        axis=0
                        steps=0
                        search=6
            if search==5:
                outofaxis = False
                vertical = not vertical
                print('vertical:',vertical)
            search=0
            if axis==2 or axis == 4:
                neg = not neg
        axis=0
        start[0]+=2
        start[1]+=2
        p=start[0]
        q=start[1]
        args.reranka = start[0]
        args.rerankb = start[1]
        print('p:',p,' q:',q)
        acc = eval(camId, filelist, labels, gallery_idx, train_idx, feature_vectors, args)
        if acc > max_acc:
            print('new p:',p,'new q:',q, ' for accuracy:', acc)
            max_acc=acc
            opt[0] = start[0]
            opt[1] = start[1]
            steps=0
            vertical=True
    print('Maximum Accuracy:',max_acc,' found at p:',opt[0],'|q:',opt[1])
    return max_acc, opt

def main():
    mat = scipy.io.loadmat(os.path.join(args.data,'cuhk03_new_protocol_config_labeled.mat'))
    camId = mat['camId']
    filelist = mat['filelist']
    labels = mat['labels']
    gallery_idx = mat['gallery_idx'] - 1
    query_idx = mat['query_idx'] - 1
    train_idx = mat['train_idx'] - 1
    with open(os.path.join(args.data,'feature_data.json'), 'r') as read_file:
        feature_vectors = np.array(json.load(read_file))
    l=0
    max_acc = np.zeros(11)
    opt = np.zeros((11,2))
    while l < 11:
        args.rerankl = l/10
        print('testing for lambda:',args.rerankl)
        max_acc[l], opt[l] = kopt(camId, filelist, labels, gallery_idx, train_idx, feature_vectors, args)
        l +=1
    print('Max accuracy:',np.max(max_acc),' at p:',opt[np.argmax(max_acc)][0], '| q:',opt[np.argmax(max_acc)][1],'| lambda:',np.argmax(max_acc)/10)
    
if __name__ == "__main__":
    main()