import cv2
import numpy as np
import scipy.spatial.distance as ssd
from datasets import DatasetLoader

#
# In der ersten Aufgabe aus dem Bereich Gesichtserkennung sollen Sie sich mit
# Eigenfaces beschaeftigen. Die entsprechenden mathematischen Grundlagen
# sind auf den Folien zu finden.
#
# In dieser Aufgabe werden verschiedene Funktionen der OpenCV Bibliothek
# verwendet. Diese benoetigen Sie aber lediglich zum Anzeigen der Bilder.
# Sie selber muessen keine OpenCV Funktionen verwenden. Bei Interesse koennen
# Sie aber die Dokumentation zu den verwendeten Funktionen unter
# http://docs.opencv.org/ nachschlagen.
#
# Fuer die Eigenfaces wird angenommen, dass jedes Bild die gleiche Groesse hat
# Alle Bilder werden daher auf IMG_SIZE skaliert (heigth, width)
IMG_SIZE = (768, 512)
# Betrachten Sie zunaechst die main Funktion und implementieren Sie die
# entsprechend notwendigen Teile.
#

def main():
    #
    # Um Gesichter zu klassifizieren, benoetigen wir zunaechst eine Menge von
    # Trainingsbildern. In dieser Aufgabe kommen die Bilder aus der FERET
    # Datenbank. Fuer diese Aufgabe wird davon eine kleine Untermenge benutzt.
    # Die entsprechenden Bilder werden ueber den DatasetLoader geladen.
    #
    train_images, train_labels, test_images, test_labels = DatasetLoader.load_feret_small()

    #
    # 3a)
    # Im ersten Schritt sollen fuers Training das durchschnittliche Gesicht
    # und die ersten 15 Eigenfaces berechnet werden. Implementieren Sie
    # hierfuer die entsprechenden Abschnitte in der compute_eigenfaces
    # Funktion weiter unten.
    #
    avg_face, eigen_faces = compute_eigenfaces(train_images, num_eigenfaces=15)

    #
    # Das durchschnittliche Gesicht anzeigen
    #
    avg_face_img = avg_face.reshape(IMG_SIZE).astype(np.uint8)
    cv2.imshow('avg face', avg_face_img)
    cv2.waitKey(0)

    #
    # Die Eigenfaces anzeigen
    #
    display_eigenfaces(eigen_faces)
    cv2.waitKey(0)

    #
    # Im naechsten Abschnitt soll fuer ein uebergebenes Bild getestet werden,
    # ob es ein Gesicht enthaelt oder nicht. Dazu werden alle Testbildern
    # untersucht.
    #
    for i, img in enumerate(test_images):
        #
        # 3b)
        # Um zu bestimmen, ob das aktuelle Testbild img ein Gesicht enthaelt
        # oder nicht, wird das Bild in den Gesichtsraum projiziert und
        # anschliessend im Bildraum wieder rekonstruiert.
        # Implementieren Sie die entsprechenden Abschnitte der
        # reconstruct_img Funktion weiter unten.
        #
        reconstructed_face = reconstruct_img(img, avg_face, eigen_faces)

        #
        # Zeige das Originalbild und die Rekonstruktion an
        #
        cv2.imshow('original face', img)
        cv2.imshow('reconstructed face', reconstructed_face)
        cv2.waitKey(0)

        #
        # 3c)
        # In der Funktion classify_face wird mit Hilfe eines Originalbilds
        # und seiner Rekonstruktion entschieden, ob es sich um ein Gesicht
        # handelt.
        # Implementieren sie die entsprechenden Abschnitt der Funktion
        # classify_face weiter unten. Experimentieren Sie anschliessend mit dem
        # threshold Parameter, um ein gutes Klassifikationsergebnis zu
        # erzielen.
        #
        is_face, dist = classify_face(img, reconstructed_face, threshold=-1)
        print (test_labels[i], "isface:", is_face, " | ", dist)

    #
    # 3d)
    # Im naechsten Schritt soll ein Gesicht der entsprechenden Person
    # zugeordnet werden. Dafuer muss zunaechst eine Repraesentation der
    # Trainingsbilder im Gesichtsraum erzeugt werden.
    # Implementieren Sie die entsprechenden Abschnitte der Funktion
    # get_train_projection weiter unten.
    #
    train_proj = get_train_projection(train_images, avg_face, eigen_faces)

    #
    # 3e)
    # Die Testbilder koennen nun den Trainingsbildern zugeordnet werden, indem
    # auch fuer Sie eine Repraesentation im Gesichtsraum erzeugt wird und diese
    # mit den Trainingsbilderrepraesentationen verglichen wird.
    # Implementieren Sie die entsprechenden Abschnitte der Funktion
    # get_most_similar_face_id weiter unten.
    #
    for i, img in enumerate(test_images):
        idx = get_most_similar_face_id(img, avg_face, eigen_faces, train_proj)

        #
        # Im Folgenden wird ein einfacher Nearest Neighbor Klassifikator
        # benutzt, um das aktuelle Bild einer Person zuzuordnen.
        #
        predicted_label = train_labels[idx]
        real_label = test_labels[i]


        #
        # Ergebnisse anzeigen
        #
        cv2.imshow("pred: " + predicted_label + " | is: " + real_label, img)
        cv2.waitKey(0)

def compute_eigenfaces(train_faces, num_eigenfaces):
    '''
    Die Methode berechnet zu den uebergebenen train_faces die ersten k Eigenfaces.
    Die Anzahl k wird hier durch num_eigenfaces ausgedrueckt.
    @param train_faces: Liste mit Trainingsbildern
    @param num_eigenfaces: Anzahl der zu berechnenden Eigenfaces
    @return: Das durchschnittliche Gesicht (in Vektorform),
             Die Eigenfaces als 2D Matrix, jede Zeile bildet ein Eigenfaces
    '''
    #
    # 3a)
    # Generieren Sie eine n x m Matrix aus der train_faces List, indem Sie
    # jedes Bild in einen Vektor ueberfuehren. Die Matrix hat anschliessend
    # n Zeilen (Anzahl der Trainingsbilder) und m Spalten (Laenge des Gesichts-
    # vektors).
    # Erzeugen Sie zunaechst die Ausgabematrix mit np.zeros und speichern Sie
    # anschliessend die Bilder in jeweils eine Zeile.
    #

    df = np.zeros((len(train_faces), IMG_SIZE[0] * IMG_SIZE[1]))
    for i in range(0, len(train_faces)):
        df[i] = train_faces[i].flatten()

    #
    # Berechnen Sie das durchschnittliche Gesicht und speichern Sie es in einer
    # Variablen avg_face.
    # Subtrahieren Sie das durchschnittliche Gesicht von allen Gesichtsbildern,
    # um die mittelwertfreien Gesichtsbilder zu erhalten.
    # Benutzen Sie fuer die Subtraktion das in Numpy eingebaute Broadcasting.
    #
    # http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html
    #

    avg_face = np.sum(df, axis=0) / float(len(df))
    df = df - avg_face
    assert(avg_face.shape[0] == IMG_SIZE[0] * IMG_SIZE[1])

    #
    # Benutzen Sie die numpy eigenvalue Funktion (np.linalg.eig),
    # um so die Eigenfaces zu bestimmen. Speichern Sie die Eigenfaces in einer
    # Matrix mit dem Namen e_faces.
    # Achten Sie dabei darauf, dass die Eigenwerte fuer die transformierte Matrix
    # berechnet werden, um das Eigenwertproblem zu vereinfachen (siehe Folien)
    # und Sie daher die Eigenwerte noch Transformieren muessen.
    #
    # Aus welchen der Eigenvektoren werden die k Eigenfaces gebildet?
    #
    # http://docs.scipy.org/doc/numpy/reference/generated/numpy.linalg.eig.html
    #
    #

    # Bilde die 28x28 Kovarainzmatrix
    df_T = np.matmul(df, df.T)

    # Bestimme die Eigenwerte und Eigenvektoren der Kovarianzmatrix
    e_vals, e_vecs = np.linalg.eig(df_T)

    #
    # Kovarianzmatrix im Gesichtsraum bestimmen
    # ACHTUNG: Trick aus den Folien benutzen
    #

    e_faces = np.matmul(df.T, e_vecs).T

    # Normalisiere die Vektoren
    e_faces = e_faces/np.linalg.norm(e_faces, ord=2, axis=1, keepdims=True)

    indices = np.argsort(e_vals)[::-1]
    e_vals = e_vals[indices]

    #
    # Eigenfaces bestimmen
    #
    e_faces = e_faces[indices][:num_eigenfaces]

    assert(e_faces.shape[0] == num_eigenfaces)
    return avg_face, e_faces

def reconstruct_img(img, avg_face, e_faces):
    '''
    Die Funktion projiziert das Eingabebild in den Gesichtsraum und
    rekonstruiert es schliesslich wieder.
    @param img: Eingabebild
    @param avg_face: das durchschnittliche Gesicht als Vektor
    @param e_faces: 2D Matrix von Eigenfaces, jede Zeile ist ein Eigenface
    @return: Die Rekonstruktion von img mit gleicher Bildgroesse
    '''
    #
    # Projiziere das Eingabebild in den Gesichtsraum
    #
    proj_img = get_projection(img, avg_face, e_faces)

    #
    # Benutzen Sie die Projektion, um das Bild aus dem Gesichtsraum
    # zurueck in den Eingaberaum zu transformieren
    #
    print img.shape
    print avg_face.shape
    print e_faces.shape
    print proj_img.shape

    proj_img = np.matmul(proj_img, e_faces)

    proj_img = proj_img.reshape((IMG_SIZE[0],IMG_SIZE[1]))

    #
    # Skalieren Sie das Bild zurueck auf IMG_SIZE
    # und konvertieren Sie die Pixelwerte in den Bereich [0,255]
    # Tipp: Sie koennen hier die matrix_to_img Funktion benutzen
    #

    reconst_img = matrix_to_img(proj_img)

    return reconst_img

def classify_face(img, reconstructed_img, threshold=1000):
    '''
    Klassifiziert ein Bild mit Hilfe seiner Rekonstruktion aus dem Gesichtsraum
    als Gesicht oder kein Gesicht.
    @param img: das Originalbild
    @param reconstructed_img: das ueber die Eigenfaces rekonstruierte Bild
    @param threshold: der Schwellwert, ab dem ein Bild als kein Gesicht klassifiziert wird
    @return: true (Bild enthaelt ein Gesicht) | false (Bild enthaelt kein GEsicht),
             Unterschied zwischen Rekonstruktion und Originalbild
    '''
    #
    # 3c)
    # Berechnen Sie die Differenz zwischen Originalbild und seiner Rekonstruktion
    # und geben Sie (true, differenz) zurueck falls es sich um ein Gesicht handelt.
    # Anderfalls geben Sie (false, differenz) zurueck.
    #
    raise NotImplementedError('Implement me')

def get_train_projection(train_images, avg_face, e_faces):
    '''
    Die Funktion berechnet fuer jedes Bild die Projektion in den Gesichtsraum.
    @param train_images: Liste von Trainingsbildern
    @param avg_face: das durchschnittliche Gesicht
    @param e_faces: die Eigenfaces als 2D Matrix
    @return: 2D Matrix von Projektionsgewichten
    '''
    #
    # 3d)
    # Berechnen Sie fuer jedes Bild aus train_images die Projektion
    # und speichern Sie sie in der Liste train_proj ab.
    # Diese wird weiter unten dann in eine Matrix konvertiert.
    #
    train_proj = []
    raise NotImplementedError('Implement me')

    train_proj_weights = np.array(train_proj)
    return train_proj_weights

def get_most_similar_face_id(img, avg_face, e_faces, train_proj_weights):
    '''
    Die Funktion projiziert das Testbild in den Gesichtsraum und vergleicht
    es dann mit den Trainingsbildern, um so das aehnlichste Gesicht zu erhalten.
    @param img: das Testbild
    @param avg_face: das durchschnittliche Gesicht
    @param e_faces: die Eigenfaces
    @param train_proj_weights: Matrix von Projektionsgewichten fuer die Trainingsbilder
    @return: id des naechsten Trainingsbildes
    '''
    #
    # 3e)
    # Benutzen Sie die ssd.cdist Funktion, um die Distanz
    # des projizierten Testbilds zu allen Trainingsbildern zu erhalten.
    # Verwenden Sie die cityblock Metrik als Distanzmass.
    #
    # Beachten Sie, dass die Funktion zwei 2D Matrizen als Eingabe erwartet.
    # Sie koennen entweder den Vektor der Gewichte in eine Liste mit nur einem
    # Element einbetten oder die numpy Funktion atleast_2d benutzen.
    #
    # Nuetzliche Links:
    # http://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.distance.cdist.html
    # http://docs.scipy.org/doc/numpy/reference/generated/numpy.atleast_2d.html
    #

    raise NotImplementedError('Implement me')
    return idx

def matrix_to_img(mat):
    '''
    Konvertiert die Matrix mat zu einem
    Graustufenbild mit Pixelwerten [0,255]
    @param mat: Eingabematrix mit beliebigen float Typ
    @return: mat konvertiert mit Pixelwerten aus dem Bereich [0, 255]
    '''
    img = np.copy(mat)
    img -= np.min(img)
    img /= np.max(img)
    img *= 255
    img = img.astype(np.uint8)
    return img

def get_projection(img, avg_face, e_faces):
    '''
    Die Funktion berechnet ein durchschnittsbereinigtes Bild aus
    img und projiziert es in den Gesichtsraum.
    @param img: Eingabebild
    @param avg_face: das durchschnittliche Gesicht als Vektor
    @param e_faces: 2D Matrix von Eigenfaces, jede Zeile ist ein Eigenface
    @return: Die Projektionsgewichte als Vektor der Laenge num_eigenfaces
    '''
    mean_free_img_vec = (img.flatten() - avg_face)
    proj_weigths = np.dot(e_faces, mean_free_img_vec)
    return proj_weigths

def display_eigenfaces(eigen_faces):
    '''
    Visualisiert die uebergebenen Eigenfaces
    '''
    e_face_images = []
    for e_face in eigen_faces:
        e_face_img = e_face.reshape(IMG_SIZE)
        e_face_img = matrix_to_img(e_face_img)
        resized_face = cv2.resize(e_face_img, (IMG_SIZE[1]/4, IMG_SIZE[0]/4))
        e_face_images.append(resized_face)

    e_face_images = np.hstack(e_face_images)
    cv2.imshow('eigen faces', e_face_images)
    return

if __name__ == '__main__':
    main()